Timestamp in MSE

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java

@@ -106,7 +106,6 @@
 import org.apache.pinot.spi.utils.CommonConstants.Broker;
 import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
 import org.apache.pinot.spi.utils.DataSizeUtils;
-import org.apache.pinot.spi.utils.TimestampIndexUtils;
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
 import org.apache.pinot.sql.FilterKind;
 import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
@@ -935,24 +934,7 @@ private void setTimestampIndexExpressionOverrideHints(@Nullable Expression expre
       return;
     }
     Function function = expression.getFunctionCall();
-    switch (function.getOperator()) {
-      case "datetrunc":
-        String granularString = function.getOperands().get(0).getLiteral().getStringValue().toUpperCase();
-        Expression timeExpression = function.getOperands().get(1);
-        if (((function.getOperandsSize() == 2) || (function.getOperandsSize() == 3 && "MILLISECONDS".equalsIgnoreCase(
-            function.getOperands().get(2).getLiteral().getStringValue()))) && TimestampIndexUtils.isValidGranularity(
-            granularString) && timeExpression.getIdentifier() != null) {
-          String timeColumn = timeExpression.getIdentifier().getName();
-          String timeColumnWithGranularity = TimestampIndexUtils.getColumnWithGranularity(timeColumn, granularString);
-          if (timestampIndexColumns.contains(timeColumnWithGranularity)) {
-            pinotQuery.putToExpressionOverrideHints(expression,
-                RequestUtils.getIdentifierExpression(timeColumnWithGranularity));
-          }
-        }
-        break;
-      default:
-        break;
-    }
+    RequestUtils.applyTimestampIndex(expression, pinotQuery, timestampIndexColumns::contains);
     function.getOperands()
         .forEach(operand -> setTimestampIndexExpressionOverrideHints(operand, timestampIndexColumns, pinotQuery));
   }

pinot-common/src/main/java/org/apache/pinot/common/utils/request/RequestUtils.java

@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Predicate;
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -53,6 +54,7 @@
 import org.apache.pinot.spi.utils.BigDecimalUtils;
 import org.apache.pinot.spi.utils.BytesUtils;
 import org.apache.pinot.spi.utils.CommonConstants.Broker.Request;
+import org.apache.pinot.spi.utils.TimestampIndexUtils;
 import org.apache.pinot.sql.FilterKind;
 import org.apache.pinot.sql.parsers.CalciteSqlParser;
 import org.apache.pinot.sql.parsers.SqlCompilationException;
@@ -631,4 +633,32 @@ public static Map<String, String> getOptionsFromJson(JsonNode request, String op
   public static Map<String, String> getOptionsFromString(String optionStr) {
     return Splitter.on(';').omitEmptyStrings().trimResults().withKeyValueSeparator('=').split(optionStr);
   }
+
+  public static void applyTimestampIndex(Expression expression, PinotQuery query) {
+    applyTimestampIndex(expression, query, timeColumnWithGranularity -> true);
+  }
+
+  public static void applyTimestampIndex(
+      Expression expression, PinotQuery query, Predicate<String> timeColumnWithGranularityPredicate
+  ) {
+    if (!expression.isSetFunctionCall()) {
+      return;
+    }
+    Function function = expression.getFunctionCall();
+    if (!function.getOperator().equalsIgnoreCase("datetrunc")) {
+      return;
+    }
+    String granularString = function.getOperands().get(0).getLiteral().getStringValue().toUpperCase();
+    Expression timeExpression = function.getOperands().get(1);
+    if (((function.getOperandsSize() == 2) || (function.getOperandsSize() == 3 && "MILLISECONDS".equalsIgnoreCase(
+        function.getOperands().get(2).getLiteral().getStringValue()))) && TimestampIndexUtils.isValidGranularity(
+        granularString) && timeExpression.getIdentifier() != null) {
+      String timeColumn = timeExpression.getIdentifier().getName();
+      String timeColumnWithGranularity = TimestampIndexUtils.getColumnWithGranularity(timeColumn, granularString);
+
+      if (timeColumnWithGranularityPredicate.test(timeColumnWithGranularity)) {
+        query.putToExpressionOverrideHints(expression, getIdentifierExpression(timeColumnWithGranularity));
+      }
+    }
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.ExecutorService;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.collections4.MapUtils;
+import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pinot.common.metrics.ServerMetrics;
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.common.request.context.FilterContext;
@@ -46,6 +47,7 @@
 import org.apache.pinot.core.plan.StreamingInstanceResponsePlanNode;
 import org.apache.pinot.core.plan.StreamingSelectionPlanNode;
 import org.apache.pinot.core.plan.TimeSeriesPlanNode;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
 import org.apache.pinot.core.query.executor.ResultsBlockStreamer;
 import org.apache.pinot.core.query.prefetch.FetchPlanner;
 import org.apache.pinot.core.query.prefetch.FetchPlannerRegistry;
@@ -321,6 +323,7 @@ public Plan makeStreamingInstancePlan(List<SegmentContext> segmentContexts, Quer
   public PlanNode makeStreamingSegmentPlanNode(SegmentContext segmentContext, QueryContext queryContext) {
     if (QueryContextUtils.isSelectionOnlyQuery(queryContext) && queryContext.getLimit() != 0) {
       // Use streaming operator only for non-empty selection-only query
+      rewriteQueryContextWithHints(queryContext, segmentContext.getIndexSegment());
       return new StreamingSelectionPlanNode(segmentContext, queryContext);
     } else {
       return makeSegmentPlanNode(segmentContext, queryContext);
@@ -344,6 +347,17 @@ public static void rewriteQueryContextWithHints(QueryContext queryContext, Index
     selectExpressions.replaceAll(
         expression -> overrideWithExpressionHints(expression, indexSegment, expressionOverrideHints));
 
+    List<Pair<AggregationFunction, FilterContext>> filtAggrFuns = queryContext.getFilteredAggregationFunctions();
+    if (filtAggrFuns != null) {
+      for (Pair<AggregationFunction, FilterContext> filteredAggregationFunction : filtAggrFuns) {
+        FilterContext right = filteredAggregationFunction.getRight();
+        if (right != null) {
+          Predicate predicate = right.getPredicate();
+          predicate.setLhs(overrideWithExpressionHints(predicate.getLhs(), indexSegment, expressionOverrideHints));
+        }
+      }
+    }
+
     List<ExpressionContext> groupByExpressions = queryContext.getGroupByExpressions();
     if (CollectionUtils.isNotEmpty(groupByExpressions)) {
       groupByExpressions.replaceAll(

pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java

@@ -19,6 +19,7 @@
 package org.apache.pinot.integration.tests;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.google.common.collect.Lists;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
@@ -28,7 +29,6 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -91,9 +91,10 @@ public abstract class BaseClusterIntegrationTest extends ClusterTest {
   protected static final List<String> DEFAULT_NO_DICTIONARY_COLUMNS =
       Arrays.asList("ActualElapsedTime", "ArrDelay", "DepDelay", "CRSDepTime");
   protected static final String DEFAULT_SORTED_COLUMN = "Carrier";
-  protected static final List<String> DEFAULT_INVERTED_INDEX_COLUMNS = Arrays.asList("FlightNum", "Origin", "Quarter");
-  private static final List<String> DEFAULT_BLOOM_FILTER_COLUMNS = Arrays.asList("FlightNum", "Origin");
-  private static final List<String> DEFAULT_RANGE_INDEX_COLUMNS = Collections.singletonList("Origin");
+  protected static final List<String> DEFAULT_INVERTED_INDEX_COLUMNS
+      = Lists.newArrayList("FlightNum", "Origin", "Quarter");
+  private static final List<String> DEFAULT_BLOOM_FILTER_COLUMNS = Lists.newArrayList("FlightNum", "Origin");
+  private static final List<String> DEFAULT_RANGE_INDEX_COLUMNS = Lists.newArrayList("Origin");
   protected static final int DEFAULT_NUM_REPLICAS = 1;
   protected static final boolean DEFAULT_NULL_HANDLING_ENABLED = false;
 

pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java

@@ -509,7 +509,7 @@ protected JsonNode getDebugInfo(final String uri)
   /**
    * Queries the broker's sql query endpoint (/query/sql)
    */
-  protected JsonNode postQuery(String query)
+  public JsonNode postQuery(String query)
       throws Exception {
     return postQuery(query, getBrokerQueryApiUrl(getBrokerBaseApiUrl(), useMultiStageQueryEngine()), null,
         getExtraQueryProperties());

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ExplainIntegrationTestTrait.java

@@ -0,0 +1,75 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.integration.tests;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import org.intellij.lang.annotations.Language;
+import org.testng.Assert;
+
+
+public interface ExplainIntegrationTestTrait {
+
+  JsonNode postQuery(@Language("sql") String query)
+      throws Exception;
+
+  default void explainLogical(@Language("sql") String query, String expected) {
+    try {
+      JsonNode jsonNode = postQuery("explain plan without implementation for " + query);
+      JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);
+
+      Assert.assertEquals(plan.asText(), expected);
+    } catch (RuntimeException e) {
+      throw e;
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  default void explain(@Language("sql") String query, String expected) {
+    try {
+      JsonNode jsonNode = postQuery("explain plan for " + query);
+      JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);
+
+      Assert.assertEquals(plan.asText(), expected);
+    } catch (RuntimeException e) {
+      throw e;
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  default void explainVerbose(@Language("sql") String query, String expected) {
+    try {
+      JsonNode jsonNode = postQuery("set explainPlanVerbose=true; explain plan for " + query);
+      JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);
+
+      String actual = plan.asText()
+          .replaceAll("numDocs=\\[[^\\]]*]", "numDocs=[any]")
+          .replaceAll("segment=\\[[^\\]]*]", "segment=[any]")
+          .replaceAll("totalDocs=\\[[^\\]]*]", "totalDocs=[any]");
+
+
+      Assert.assertEquals(actual, expected);
+    } catch (RuntimeException e) {
+      throw e;
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+}

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineExplainIntegrationTest.java

@@ -18,24 +18,22 @@
  */
 package org.apache.pinot.integration.tests;
 
-import com.fasterxml.jackson.databind.JsonNode;
 import java.io.File;
 import java.util.List;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.util.TestUtils;
-import org.intellij.lang.annotations.Language;
 import org.testcontainers.shaded.org.apache.commons.io.FileUtils;
-import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
 
-public class MultiStageEngineExplainIntegrationTest extends BaseClusterIntegrationTest {
+public class MultiStageEngineExplainIntegrationTest extends BaseClusterIntegrationTest
+    implements ExplainIntegrationTestTrait {
 
   @BeforeClass
   public void setUp()
@@ -161,7 +159,7 @@ public void simpleQueryVerbose() {
             + "            Project(columns=[[]])\n"
             + "              DocIdSet(maxDocs=[10000])\n"
             + "                FilterMatchEntireSegment(numDocs=[any])\n");
-    //@formatter:on
+        //@formatter:on
   }
 
   @Test
@@ -171,7 +169,7 @@ public void simpleQueryLogical() {
         "Execution Plan\n"
             + "LogicalProject(EXPR$0=[1])\n"
             + "  LogicalTableScan(table=[[default, mytable]])\n");
-    //@formatter:on
+        //@formatter:on
   }
 
   @AfterClass
@@ -186,49 +184,4 @@ public void tearDown()
 
     FileUtils.deleteDirectory(_tempDir);
   }
-
-  private void explainVerbose(@Language("sql") String query, String expected) {
-    try {
-      JsonNode jsonNode = postQuery("set explainPlanVerbose=true; explain plan for " + query);
-      JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);
-
-      String actual = plan.asText()
-          .replaceAll("numDocs=\\[[^\\]]*]", "numDocs=[any]")
-          .replaceAll("segment=\\[[^\\]]*]", "segment=[any]")
-          .replaceAll("totalDocs=\\[[^\\]]*]", "totalDocs=[any]");
-
-
-      Assert.assertEquals(actual, expected);
-    } catch (RuntimeException e) {
-      throw e;
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
-  }
-
-  private void explain(@Language("sql") String query, String expected) {
-    try {
-      JsonNode jsonNode = postQuery("explain plan for " + query);
-      JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);
-
-      Assert.assertEquals(plan.asText(), expected);
-    } catch (RuntimeException e) {
-      throw e;
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
-  }
-
-  private void explainLogical(@Language("sql") String query, String expected) {
-    try {
-      JsonNode jsonNode = postQuery("set explainAskingServers=false; explain plan for " + query);
-      JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);
-
-      Assert.assertEquals(plan.asText(), expected);
-    } catch (RuntimeException e) {
-      throw e;
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
-  }
 }