Jackie-Jiang commented on a change in pull request #8172:
URL: https://github.com/apache/pinot/pull/8172#discussion_r806164875



##########
File path: 
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
##########
@@ -459,104 +457,87 @@ public QueryContext build() {
      * Helper method to generate the aggregation functions for the query.
      */
     private void generateAggregationFunctions(QueryContext queryContext) {
-      List<AggregationFunction> aggregationFunctions = new ArrayList<>();
-      List<Pair<AggregationFunction, FilterContext>> filteredAggregations = 
new ArrayList<>();
-      Map<FunctionContext, Integer> aggregationFunctionIndexMap = new 
HashMap<>();
-      Map<Pair<FunctionContext, FilterContext>, Integer> 
filterExpressionIndexMap = new HashMap<>();
+      List<Pair<AggregationFunction, FilterContext>> 
filteredAggregationFunctions = new ArrayList<>();
+      Map<Pair<FunctionContext, FilterContext>, Integer> 
filteredAggregationsIndexMap = new HashMap<>();
 
       // Add aggregation functions in the SELECT clause
       // NOTE: DO NOT deduplicate the aggregation functions in the SELECT 
clause because that involves protocol change.
-      List<Pair<FilterContext, FunctionContext>> aggregationsInSelect = new 
ArrayList<>();
+      List<Pair<FunctionContext, FilterContext>> filteredAggregations = new 
ArrayList<>();
       for (ExpressionContext selectExpression : 
queryContext._selectExpressions) {
-        getAggregations(selectExpression, aggregationsInSelect);
+        getAggregations(selectExpression, filteredAggregations);
       }
-      for (Pair<FilterContext, FunctionContext> pair : aggregationsInSelect) {
-        FunctionContext function = pair.getRight();
-        int functionIndex = filteredAggregations.size();
-        AggregationFunction aggregationFunction =
-            AggregationFunctionFactory.getAggregationFunction(function, 
queryContext);
-
-        FilterContext filterContext = null;
-        // If the left pair is not null, implies a filtered aggregation
-        if (pair.getLeft() != null) {
+      for (Pair<FunctionContext, FilterContext> pair : filteredAggregations) {
+        FunctionContext aggregation = pair.getLeft();
+        FilterContext filter = pair.getRight();
+        if (filter != null) {
+          // Filtered aggregation
           if (_groupByExpressions != null) {
             throw new IllegalStateException("GROUP BY with FILTER clauses is 
not supported");
           }
           queryContext._hasFilteredAggregations = true;
-          filterContext = pair.getLeft();
-          Pair<FunctionContext, FilterContext> filterContextPair =
-              Pair.of(function, filterContext);
-          if (!filterExpressionIndexMap.containsKey(filterContextPair)) {
-            int filterMapIndex = filterExpressionIndexMap.size();
-            filterExpressionIndexMap.put(filterContextPair, filterMapIndex);
-          }
         }
-        filteredAggregations.add(Pair.of(aggregationFunction, filterContext));
-        aggregationFunctionIndexMap.put(function, functionIndex);
+        int functionIndex = filteredAggregationFunctions.size();
+        AggregationFunction aggregationFunction =
+            AggregationFunctionFactory.getAggregationFunction(aggregation, 
queryContext);
+        filteredAggregationFunctions.add(Pair.of(aggregationFunction, filter));
+        filteredAggregationsIndexMap.put(Pair.of(aggregation, filter), 
functionIndex);
       }
 
-      // Add aggregation functions in the HAVING clause but not in the SELECT 
clause
+      // Add aggregation functions in the HAVING and ORDER-BY clause but not 
in the SELECT clause
+      filteredAggregations.clear();
       if (queryContext._havingFilter != null) {
-        List<Pair<FilterContext, FunctionContext>> aggregationsInHaving = new 
ArrayList<>();
-        getAggregations(queryContext._havingFilter, aggregationsInHaving);
-        for (Pair<FilterContext, FunctionContext> pair : aggregationsInHaving) 
{
-          FunctionContext function = pair.getRight();
-          if (!aggregationFunctionIndexMap.containsKey(function)) {
-            int functionIndex = filteredAggregations.size();
-            filteredAggregations.add(Pair.of(
-                AggregationFunctionFactory.getAggregationFunction(function, 
queryContext), null));
-            aggregationFunctionIndexMap.put(function, functionIndex);
-          }
-        }
+        getAggregations(queryContext._havingFilter, filteredAggregations);
       }
-
-      // Add aggregation functions in the ORDER-BY clause but not in the 
SELECT or HAVING clause
       if (queryContext._orderByExpressions != null) {
-        List<Pair<FilterContext, FunctionContext>> aggregationsInOrderBy = new 
ArrayList<>();
         for (OrderByExpressionContext orderByExpression : 
queryContext._orderByExpressions) {
-          getAggregations(orderByExpression.getExpression(), 
aggregationsInOrderBy);
+          getAggregations(orderByExpression.getExpression(), 
filteredAggregations);
         }
-        for (Pair<FilterContext, FunctionContext> pair : 
aggregationsInOrderBy) {
-          FunctionContext function = pair.getRight();
-          if (!aggregationFunctionIndexMap.containsKey(function)) {
-            int functionIndex = filteredAggregations.size();
-            filteredAggregations.add(Pair.of(
-                AggregationFunctionFactory.getAggregationFunction(function, 
queryContext), null));
-            aggregationFunctionIndexMap.put(function, functionIndex);
-          }
+      }
+      for (Pair<FunctionContext, FilterContext> pair : filteredAggregations) {
+        if (!filteredAggregationsIndexMap.containsKey(pair)) {
+          FunctionContext aggregation = pair.getLeft();
+          FilterContext filter = pair.getRight();
+          int functionIndex = filteredAggregationFunctions.size();
+          AggregationFunction aggregationFunction =
+              AggregationFunctionFactory.getAggregationFunction(aggregation, 
queryContext);
+          filteredAggregationFunctions.add(Pair.of(aggregationFunction, 
filter));
+          filteredAggregationsIndexMap.put(Pair.of(aggregation, filter), 
functionIndex);
         }
       }
 
-      if (!filteredAggregations.isEmpty()) {
-        for (Pair<AggregationFunction, FilterContext> pair : 
filteredAggregations) {
-          aggregationFunctions.add(pair.getLeft());
+      if (!filteredAggregationFunctions.isEmpty()) {
+        int numAggregations = filteredAggregationFunctions.size();
+        AggregationFunction[] aggregationFunctions = new 
AggregationFunction[numAggregations];
+        for (int i = 0; i < numAggregations; i++) {
+          aggregationFunctions[i] = 
filteredAggregationFunctions.get(i).getLeft();
         }
-
-        queryContext._aggregationFunctions = aggregationFunctions.toArray(new 
AggregationFunction[0]);
-        queryContext._filteredAggregations = filteredAggregations;
+        Map<FunctionContext, Integer> aggregationFunctionIndexMap = new 
HashMap<>();
+        for (Map.Entry<Pair<FunctionContext, FilterContext>, Integer> entry : 
filteredAggregationsIndexMap.entrySet()) {
+          aggregationFunctionIndexMap.put(entry.getKey().getLeft(), 
entry.getValue());

Review comment:
       I'd suggest doing that when adding the FILTER support for group-by 
queries because it does some overhead (looking up pair vs function context)




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org

Reply via email to