Jackie-Jiang commented on a change in pull request #8172: URL: https://github.com/apache/pinot/pull/8172#discussion_r803114460
########## 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()); + } + queryContext._aggregationFunctions = aggregationFunctions; + queryContext._filteredAggregationFunctions = filteredAggregationFunctions; queryContext._aggregationFunctionIndexMap = aggregationFunctionIndexMap; - queryContext._filteredAggregationsIndexMap = filterExpressionIndexMap; + queryContext._filteredAggregationsIndexMap = filteredAggregationsIndexMap; } } /** - * Helper method to extract AGGREGATION FunctionContexts from the given expression. - * - * NOTE: The left pair of aggregations should be set only for filtered aggregations + * Helper method to extract AGGREGATION FunctionContexts and FILTER FilterContexts from the given expression. */ private static void getAggregations(ExpressionContext expression, - List<Pair<FilterContext, FunctionContext>> aggregations) { + List<Pair<FunctionContext, FilterContext>> filteredAggregations) { FunctionContext function = expression.getFunction(); if (function == null) { return; } if (function.getType() == FunctionContext.Type.AGGREGATION) { // Aggregation - aggregations.add(Pair.of(null, function)); + filteredAggregations.add(Pair.of(function, null)); } else { - List<ExpressionContext> arguments = function.getArguments(); if (function.getFunctionName().equalsIgnoreCase("filter")) { // Filtered aggregation + List<ExpressionContext> arguments = function.getArguments(); Review comment: Reverted -- 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