Unify behavior of both query engines to follow standard SQL behavior; allow overriding behavior for performance

Author: yashmayya

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -234,12 +234,8 @@ public static boolean isServerReturnFinalResultKeyUnpartitioned(Map<String, Stri
     return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SERVER_RETURN_FINAL_RESULT_KEY_UNPARTITIONED));
   }
 
-  public static Optional<Boolean> isFilteredAggregationsComputeAllGroups(Map<String, String> queryOptions) {
-    String value = queryOptions.get(QueryOptionKey.FILTERED_AGGREGATIONS_COMPUTE_ALL_GROUPS);
-    if (value == null) {
-      return Optional.empty();
-    }
-    return Optional.of(Boolean.parseBoolean(value));
+  public static boolean isFilteredAggregationsSkipEmptyGroups(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.FILTERED_AGGREGATIONS_SKIP_EMPTY_GROUPS));
   }
 
   @Nullable

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java

@@ -385,12 +385,14 @@ public static List<AggregationInfo> buildFilteredAggregationInfos(SegmentContext
       }
     }
 
-    if (!nonFilteredFunctions.isEmpty() || (QueryOptionsUtils.isFilteredAggregationsComputeAllGroups(
-        queryContext.getQueryOptions()).orElse(false))) {
-      // If there are no non-filtered aggregation functions, but the query option to compute all groups is set, we
-      // add a new AggregationInfo with an empty AggregationFunction array and the main query filter so that the
-      // GroupByExecutor will compute all the groups (from the result of applying the main query filter) but no
-      // unnecessary additional aggregation will be done since the AggregationFunction array is empty.
+    if (!nonFilteredFunctions.isEmpty() || ((queryContext.getGroupByExpressions() != null)
+        && !QueryOptionsUtils.isFilteredAggregationsSkipEmptyGroups(queryContext.getQueryOptions()))) {
+      // If there are no non-filtered aggregation functions for a group by query, we still add a new AggregationInfo
+      // with an empty AggregationFunction array and the main query filter so that the GroupByExecutor will compute all
+      // the groups (from the result of applying the main query filter) but no unnecessary additional aggregation will
+      // be done since the AggregationFunction array is empty. However, if the query option to skip empty groups is
+      // enabled, we don't do this in order to avoid unnecessary computation of empty groups (which can be very
+      // expensive if the main filter has high selectivity).
       AggregationFunction[] aggregationFunctions = nonFilteredFunctions.toArray(new AggregationFunction[0]);
       aggregationInfos.add(
           buildAggregationInfo(segmentContext, queryContext, aggregationFunctions, mainFilter, mainFilterOperator,

pinot-core/src/test/java/org/apache/pinot/queries/FilteredAggregationsTest.java

@@ -41,6 +41,7 @@
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.GenericRow;
 import org.apache.pinot.spi.data.readers.RecordReader;
+import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -394,9 +395,9 @@ public void testGroupBy() {
 
   @Test
   public void testGroupByMultipleColumns() {
-    String filterQuery =
-        "SELECT SUM(INT_COL) FILTER(WHERE INT_COL > 25000) testSum FROM MyTable GROUP BY BOOLEAN_COL, STRING_COL "
-            + "ORDER BY BOOLEAN_COL, STRING_COL";
+    String filterQuery = "SET " + CommonConstants.Broker.Request.QueryOptionKey.FILTERED_AGGREGATIONS_SKIP_EMPTY_GROUPS
+        + "=true; SELECT SUM(INT_COL) FILTER(WHERE INT_COL > 25000) testSum FROM MyTable GROUP BY BOOLEAN_COL, "
+        + "STRING_COL ORDER BY BOOLEAN_COL, STRING_COL";
     String nonFilterQuery =
         "SELECT SUM(INT_COL) testSum FROM MyTable WHERE INT_COL > 25000 GROUP BY BOOLEAN_COL, STRING_COL "
             + "ORDER BY BOOLEAN_COL, STRING_COL";

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

@@ -1040,16 +1040,16 @@ public void testFilteredAggregationWithNoValueMatchingAggregationFilterWithOptio
       throws Exception {
     // Use a hint to ensure that the aggregation will not be pushed to the leaf stage, so that we can test the
     // MultistageGroupByExecutor
-    String sqlQuery = "SET " + CommonConstants.Broker.Request.QueryOptionKey.FILTERED_AGGREGATIONS_COMPUTE_ALL_GROUPS
-        + "=false; SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */"
+    String sqlQuery = "SET " + CommonConstants.Broker.Request.QueryOptionKey.FILTERED_AGGREGATIONS_SKIP_EMPTY_GROUPS
+        + "=true; SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */"
         + "AirlineID, COUNT(*) FILTER (WHERE Origin = 'garbage') FROM mytable WHERE AirlineID > 20000 GROUP BY "
         + "AirlineID";
 
     JsonNode result = postQuery(sqlQuery);
     assertNoError(result);
 
     // Result set will be empty since the aggregation filter does not match any rows, and we've set the query option to
-    // not compute all groups
+    // skip empty groups
     assertEquals(result.get("numRowsResultSet").asInt(), 0);
   }
 

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

@@ -3730,12 +3730,10 @@ public void testSkipIndexes(boolean useMultiStageQueryEngine)
     reloadAllSegments(TEST_UPDATED_RANGE_INDEX_QUERY, true, numTotalDocs);
   }
 
-  @Test
-  public void testFilteredAggregationWithNoValueMatchingAggregationFilterDefault()
+  @Test(dataProvider = "useBothQueryEngines")
+  public void testFilteredAggregationWithNoValueMatchingAggregationFilterDefault(boolean useMultiStageQueryEngine)
       throws Exception {
-    // The multi-stage query engine computes all groups by default (even without the query option) for group by queries
-    // with only filtered aggregates since that is the SQL standard behavior.
-    setUseMultiStageQueryEngine(false);
+    setUseMultiStageQueryEngine(useMultiStageQueryEngine);
 
     String sqlQuery =
         "SELECT AirlineID, COUNT(*) FILTER (WHERE Origin = 'garbage') FROM mytable WHERE AirlineID > 20000 GROUP BY "
@@ -3744,30 +3742,31 @@ public void testFilteredAggregationWithNoValueMatchingAggregationFilterDefault()
     JsonNode result = postQuery(sqlQuery);
     assertNoError(result);
 
-    // Result set will be empty by default since the aggregation filter does not match any rows
-    assertEquals(result.get("numRowsResultSet").asInt(), 0);
+    // Ensure that result set is not empty since all groups should be computed by default
+    assertTrue(result.get("numRowsResultSet").asInt() > 0);
+
+    // Ensure that the count is 0 for all groups (because the aggregation filter does not match any rows)
+    JsonNode rows = result.get("resultTable").get("rows");
+    for (int i = 0; i < rows.size(); i++) {
+      assertEquals(rows.get(i).get(1).asInt(), 0);
+      // Ensure that the main filter was applied
+      assertTrue(rows.get(i).get(0).asInt() > 20000);
+    }
   }
 
   @Test(dataProvider = "useBothQueryEngines")
   public void testFilteredAggregationWithNoValueMatchingAggregationFilterWithOption(boolean useMultiStageQueryEngine)
       throws Exception {
     setUseMultiStageQueryEngine(useMultiStageQueryEngine);
     String sqlQuery =
-        "SET " + CommonConstants.Broker.Request.QueryOptionKey.FILTERED_AGGREGATIONS_COMPUTE_ALL_GROUPS + "=true; "
+        "SET " + CommonConstants.Broker.Request.QueryOptionKey.FILTERED_AGGREGATIONS_SKIP_EMPTY_GROUPS + "=true; "
             + "SELECT AirlineID, COUNT(*) FILTER (WHERE Origin = 'garbage') FROM mytable WHERE AirlineID > 20000 "
             + "GROUP BY AirlineID";
     JsonNode result = postQuery(sqlQuery);
     assertNoError(result);
 
-    // Ensure that result set is not empty since all groups should be computed now
-    assertTrue(result.get("numRowsResultSet").asInt() > 0);
-
-    // Ensure that the count is 0 for all groups (because the aggregation filter does not match any rows)
-    JsonNode rows = result.get("resultTable").get("rows");
-    for (int i = 0; i < rows.size(); i++) {
-      assertEquals(rows.get(i).get(1).asInt(), 0);
-      // Ensure that the main filter was applied
-      assertTrue(rows.get(i).get(0).asInt() > 20000);
-    }
+    // Result set will be empty since the aggregation filter does not match any rows, and we've set the option to skip
+    // empty groups
+    assertEquals(result.get("numRowsResultSet").asInt(), 0);
   }
 }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java

@@ -58,7 +58,7 @@ public class MultistageGroupByExecutor {
   private final AggType _aggType;
   private final DataSchema _resultSchema;
   private final int _numGroupsLimit;
-  private final boolean _filteredAggregationsComputeAllGroups;
+  private final boolean _filteredAggregationsSkipEmptyGroups;
 
   // Group By Result holders for each mode
   private final GroupByResultHolder[] _aggregateResultHolders;
@@ -82,8 +82,7 @@ public MultistageGroupByExecutor(int[] groupKeyIds, AggregationFunction[] aggFun
 
     // By default, we compute all groups for SQL compliant results. However, we allow overriding this behavior via
     // query option for improved performance.
-    _filteredAggregationsComputeAllGroups =
-        QueryOptionsUtils.isFilteredAggregationsComputeAllGroups(opChainMetadata).orElse(true);
+    _filteredAggregationsSkipEmptyGroups = QueryOptionsUtils.isFilteredAggregationsSkipEmptyGroups(opChainMetadata);
 
     int numFunctions = aggFunctions.length;
     if (!aggType.isInputIntermediateFormat()) {
@@ -247,9 +246,10 @@ private void processAggregate(TransferableBlock block) {
           aggFunction.aggregateGroupBySV(numMatchedRows, filteredIntKeys, groupByResultHolder, blockValSetMap);
         }
       }
-      if (intKeys == null && _filteredAggregationsComputeAllGroups) {
+      if (intKeys == null && !_filteredAggregationsSkipEmptyGroups) {
         // _groupIdGenerator should still have all the groups even if there are only filtered aggregates for SQL
-        // compliant results.
+        // compliant results. However, if the query option to skip empty groups is set, we avoid this step for
+        // improved performance.
         generateGroupByKeys(block);
       }
     }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java

@@ -227,10 +227,6 @@ private static void updateQueryOptions(PinotQuery pinotQuery, OpChainExecutionCo
     Map<String, String> queryOptions = new HashMap<>(executionContext.getOpChainMetadata());
     queryOptions.put(CommonConstants.Broker.Request.QueryOptionKey.TIMEOUT_MS,
         Long.toString(executionContext.getDeadlineMs() - System.currentTimeMillis()));
-    // All groups should be computed by default for group by queries with only filtered aggregations in the
-    // multi-stage query engine since that is the standard SQL behavior.
-    queryOptions.putIfAbsent(CommonConstants.Broker.Request.QueryOptionKey.FILTERED_AGGREGATIONS_COMPUTE_ALL_GROUPS,
-        "true");
     pinotQuery.setQueryOptions(queryOptions);
   }
 

pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java

@@ -459,14 +459,14 @@ public static class QueryOptionKey {
         // fashion with limited compute.
         public static final String IS_SECONDARY_WORKLOAD = "isSecondaryWorkload";
 
-        // For group by queries with only filtered aggregations (and no non-filtered aggregations), the v1 query engine
-        // does not compute all groups by default - instead, it will only compute the groups from the filtered result
-        // set (i.e., union of the main query filter and all the individual aggregation filters). This is good for
-        // performance, since indexes can be used, but the result won't match standard SQL behavior (where all groups
-        // are expected to be returned). If this option is set to true, the v1 query engine will compute all groups for
-        // group by queries with only filtered aggregations. This could require a full scan if the main query does not
-        // have a filter and performance could be much worse, but the result will be SQL compliant.
-        public static final String FILTERED_AGGREGATIONS_COMPUTE_ALL_GROUPS = "filteredAggregationsComputeAllGroups";
+        // For group by queries with only filtered aggregations (and no non-filtered aggregations), the default behavior
+        // is to compute all groups over the rows matching the main query filter. This ensures SQL compliant results,
+        // since empty groups are also expected to be returned in such queries. However, this could be quite inefficient
+        // if the main query does not have a filter (since a scan would be required to compute all groups). In case
+        // users are okay with skipping empty groups - i.e., only the groups matching at least one aggregation filter
+        // will be returned - this query option can be set. This is useful for performance, since indexes can be used
+        // for the aggregation filters and a full scan can be avoided.
+        public static final String FILTERED_AGGREGATIONS_SKIP_EMPTY_GROUPS = "filteredAggregationsSkipEmptyGroups";
       }
 
       public static class QueryOptionValue {