Fix test + Init TimeSeriesBuilderFactoryProvider in the Broker

apache · Dec 19, 2024 · 9b6e5a0 · 9b6e5a0
1 parent d668f3a
commit 9b6e5a0
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/...broker/src/main/java/org/apache/pinot/broker/requesthandler/TimeSeriesRequestHandler.java b/...broker/src/main/java/org/apache/pinot/broker/requesthandler/TimeSeriesRequestHandler.java
@@ -53,6 +53,7 @@
 import org.apache.pinot.tsdb.planner.physical.TimeSeriesDispatchablePlan;
 import org.apache.pinot.tsdb.spi.RangeTimeSeriesRequest;
 import org.apache.pinot.tsdb.spi.TimeSeriesLogicalPlanResult;
+import org.apache.pinot.tsdb.spi.series.TimeSeriesBuilderFactoryProvider;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -70,6 +71,7 @@ public TimeSeriesRequestHandler(PinotConfiguration config, String brokerId, Brok
     _queryEnvironment = new TimeSeriesQueryEnvironment(config, routingManager, tableCache);
     _queryEnvironment.init(config);
     _queryDispatcher = queryDispatcher;
+    TimeSeriesBuilderFactoryProvider.init(config);
   }
 
   @Override

diff --git a/...va/org/apache/pinot/query/runtime/timeseries/PhysicalTimeSeriesServerPlanVisitorTest.java b/...va/org/apache/pinot/query/runtime/timeseries/PhysicalTimeSeriesServerPlanVisitorTest.java
@@ -29,6 +29,9 @@
 import org.apache.pinot.tsdb.spi.AggInfo;
 import org.apache.pinot.tsdb.spi.TimeBuckets;
 import org.apache.pinot.tsdb.spi.plan.LeafTimeSeriesPlanNode;
+import org.apache.pinot.tsdb.spi.series.SimpleTimeSeriesBuilderFactory;
+import org.apache.pinot.tsdb.spi.series.TimeSeriesBuilderFactoryProvider;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.mock;
@@ -38,8 +41,14 @@
 
 
 public class PhysicalTimeSeriesServerPlanVisitorTest {
+  private static final String LANGUAGE = "m3ql";
   private static final int DUMMY_DEADLINE_MS = 10_000;
 
+  @BeforeClass
+  public void setUp() {
+    TimeSeriesBuilderFactoryProvider.registerSeriesBuilderFactory(LANGUAGE, new SimpleTimeSeriesBuilderFactory());
+  }
+
   @Test
   public void testCompileQueryContext() {
     final String planId = "id";
@@ -52,14 +61,14 @@ public void testCompileQueryContext() {
     // Case-1: Without offset, simple column based group-by expression, simple column based value, and non-empty filter.
     {
       TimeSeriesExecutionContext context =
-          new TimeSeriesExecutionContext("m3ql", TimeBuckets.ofSeconds(1000L, Duration.ofSeconds(10), 100),
+          new TimeSeriesExecutionContext(LANGUAGE, TimeBuckets.ofSeconds(1000L, Duration.ofSeconds(10), 100),
               DUMMY_DEADLINE_MS, Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
       LeafTimeSeriesPlanNode leafNode =
           new LeafTimeSeriesPlanNode(planId, Collections.emptyList(), tableName, timeColumn, TimeUnit.SECONDS, 0L,
               filterExpr, "orderCount", aggInfo, Collections.singletonList("cityName"));
       QueryContext queryContext = serverPlanVisitor.compileQueryContext(leafNode, context);
       assertNotNull(queryContext.getTimeSeriesContext());
-      assertEquals(queryContext.getTimeSeriesContext().getLanguage(), "m3ql");
+      assertEquals(queryContext.getTimeSeriesContext().getLanguage(), LANGUAGE);
       assertEquals(queryContext.getTimeSeriesContext().getOffsetSeconds(), 0L);
       assertEquals(queryContext.getTimeSeriesContext().getTimeColumn(), timeColumn);
       assertEquals(queryContext.getTimeSeriesContext().getValueExpression().getIdentifier(), "orderCount");
@@ -70,7 +79,7 @@ public void testCompileQueryContext() {
     // Case-2: With offset, complex group-by expression, complex value, and non-empty filter
     {
       TimeSeriesExecutionContext context =
-          new TimeSeriesExecutionContext("m3ql", TimeBuckets.ofSeconds(1000L, Duration.ofSeconds(10), 100),
+          new TimeSeriesExecutionContext(LANGUAGE, TimeBuckets.ofSeconds(1000L, Duration.ofSeconds(10), 100),
               DUMMY_DEADLINE_MS, Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
       LeafTimeSeriesPlanNode leafNode =
           new LeafTimeSeriesPlanNode(planId, Collections.emptyList(), tableName, timeColumn, TimeUnit.SECONDS, 10L,
@@ -80,7 +89,7 @@ public void testCompileQueryContext() {
       assertNotNull(queryContext.getGroupByExpressions());
       assertEquals("concat(cityName,stateName,'-')", queryContext.getGroupByExpressions().get(0).toString());
       assertNotNull(queryContext.getTimeSeriesContext());
-      assertEquals(queryContext.getTimeSeriesContext().getLanguage(), "m3ql");
+      assertEquals(queryContext.getTimeSeriesContext().getLanguage(), LANGUAGE);
       assertEquals(queryContext.getTimeSeriesContext().getOffsetSeconds(), 10L);
       assertEquals(queryContext.getTimeSeriesContext().getTimeColumn(), timeColumn);
       assertEquals(queryContext.getTimeSeriesContext().getValueExpression().toString(), "times(orderCount,'2')");

diff --git a/...-spi/src/main/java/org/apache/pinot/tsdb/spi/series/TimeSeriesBuilderFactoryProvider.java b/...-spi/src/main/java/org/apache/pinot/tsdb/spi/series/TimeSeriesBuilderFactoryProvider.java
@@ -51,7 +51,7 @@ public static void init(PinotConfiguration pinotConfiguration) {
         TimeSeriesBuilderFactory seriesBuilderFactory = (TimeSeriesBuilderFactory) untypedSeriesBuilderFactory;
         seriesBuilderFactory.init(pinotConfiguration.subset(
             PinotTimeSeriesConfiguration.CONFIG_PREFIX + "." + language));
-        FACTORY_MAP.put(language, seriesBuilderFactory);
+        FACTORY_MAP.putIfAbsent(language, seriesBuilderFactory);
       } catch (Exception e) {
         throw new RuntimeException(e);
       }