egalpin commented on code in PR #10000: URL: https://github.com/apache/pinot/pull/10000#discussion_r1060977130
########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java: ########## @@ -0,0 +1,221 @@ +/** + * 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.core.operator.query; + +import java.util.Collection; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.data.table.IntermediateRecord; +import org.apache.pinot.core.data.table.TableResizer; +import org.apache.pinot.core.operator.BaseOperator; +import org.apache.pinot.core.operator.ExecutionStatistics; +import org.apache.pinot.core.operator.blocks.TransformBlock; +import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; +import org.apache.pinot.core.operator.transform.TransformOperator; +import org.apache.pinot.core.query.aggregation.function.AggregationFunction; +import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult; +import org.apache.pinot.core.query.aggregation.groupby.DefaultGroupByExecutor; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.util.GroupByUtils; +import org.apache.pinot.spi.trace.Tracing; + + +/** + * The <code>FilteredGroupByOperator</code> class provides the operator for group-by query on a single segment when + * there are 1 or more filter expressions on aggregations. + */ +@SuppressWarnings("rawtypes") +public class FilteredGroupByOperator extends BaseOperator<GroupByResultsBlock> { + private static final String EXPLAIN_NAME = "GROUP_BY_FILTERED"; + + @Nullable + private final AggregationFunction[] _aggregationFunctions; + private final List<Pair<AggregationFunction[], TransformOperator>> _aggFunctionsWithTransformOperator; + private final ExpressionContext[] _groupByExpressions; + private final long _numTotalDocs; + private long _numDocsScanned; + private long _numEntriesScannedInFilter; + private long _numEntriesScannedPostFilter; + private final DataSchema _dataSchema; + private final QueryContext _queryContext; + + public FilteredGroupByOperator( + AggregationFunction[] aggregationFunctions, + List<Pair<AggregationFunction[], TransformOperator>> aggFunctionsWithTransformOperator, + ExpressionContext[] groupByExpressions, + long numTotalDocs, + QueryContext queryContext) { + _aggregationFunctions = aggregationFunctions; + _aggFunctionsWithTransformOperator = aggFunctionsWithTransformOperator; + _groupByExpressions = groupByExpressions; + _numTotalDocs = numTotalDocs; + _queryContext = queryContext; + + // NOTE: The indexedTable expects that the data schema will have group by columns before aggregation columns + int numGroupByExpressions = groupByExpressions.length; + int numAggregationFunctions = aggregationFunctions.length; + int numColumns = numGroupByExpressions + numAggregationFunctions; + String[] columnNames = new String[numColumns]; + DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[numColumns]; + + // Extract column names and data types for group-by columns + for (int i = 0; i < numGroupByExpressions; i++) { + ExpressionContext groupByExpression = groupByExpressions[i]; + columnNames[i] = groupByExpression.toString(); + columnDataTypes[i] = DataSchema.ColumnDataType.fromDataTypeSV( + aggFunctionsWithTransformOperator.get(i).getRight().getResultMetadata(groupByExpression).getDataType()); + } + + // Extract column names and data types for aggregation functions + for (int i = 0; i < numAggregationFunctions; i++) { + AggregationFunction aggregationFunction = aggregationFunctions[i]; + int index = numGroupByExpressions + i; + columnNames[index] = aggregationFunction.getResultColumnName(); + columnDataTypes[index] = aggregationFunction.getIntermediateResultColumnType(); + } + + _dataSchema = new DataSchema(columnNames, columnDataTypes); + } + + @Override + protected GroupByResultsBlock getNextBlock() { + GroupKeyGenerator groupKeyGenerator = null; + // TODO(egalpin): Support Startree query resolution when possible, even with FILTER expressions + assert _aggregationFunctions != null; + boolean numGroupsLimitReached = false; + int numAggregations = _aggregationFunctions.length; + + GroupByResultHolder[] groupByResultHolders = new GroupByResultHolder[numAggregations]; + IdentityHashMap<AggregationFunction, Integer> resultHolderIndexMap = new IdentityHashMap<>(numAggregations); + for (int i = 0; i < numAggregations; i++) { + resultHolderIndexMap.put(_aggregationFunctions[i], i); + } + + for (Pair<AggregationFunction[], TransformOperator> filteredAggregation : _aggFunctionsWithTransformOperator) { + TransformOperator transformOperator = filteredAggregation.getRight(); + AggregationFunction[] filteredAggFunctions = filteredAggregation.getLeft(); + + // Perform aggregation group-by on all the blocks + DefaultGroupByExecutor groupByExecutor; + if (groupKeyGenerator == null) { + // The group key generator should be shared across all AggregationFunctions so that agg results can be Review Comment: @Jackie-Jiang please let me know where I may have misstated the reasoning -- 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
