From 600c736a9c146f9e1eff985b8b20df16b69f4d3f Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 28 Aug 2026 13:47:25 -0700 Subject: [PATCH] Add an SSE group-key generator provider Downstream Pinot distributions need a narrow extension point to supply specialized group-key generators without replacing the SSE group-by pipeline. Add a LimitedPrivate and Unstable provider hook with compact physical-column metadata. Default, filtered, grouping-set, and StarTree paths retain Pinot built-in generators. Define generator ownership at executor initialization, segment materialization, streaming detach, and combine. Close resource-owning generators once while preserving primary failures. Add real-segment provider and fallback coverage plus focused lifecycle regression tests. --- .../combine/GroupByCombineOperator.java | 72 ++-- .../core/operator/query/GroupByOperator.java | 43 +- .../StreamingGroupByCombineOperator.java | 7 +- .../pinot/core/plan/GroupByPlanNode.java | 11 +- .../plan/maker/InstancePlanMakerImplV2.java | 11 +- .../groupby/AggregationGroupByResult.java | 7 +- .../groupby/DefaultGroupByExecutor.java | 168 ++++++-- .../groupby/GroupKeyGenerator.java | 9 +- .../groupby/GroupKeyGeneratorContext.java | 108 +++++ .../groupby/GroupKeyGeneratorProvider.java | 46 +++ .../GroupByCombineOperatorLifecycleTest.java | 74 ++++ .../query/GroupByOperatorLifecycleTest.java | 73 ++++ .../StreamingGroupByCombineOperatorTest.java | 26 ++ .../DefaultGroupByExecutorProviderTest.java | 94 +++++ .../groupby/GroupKeyGeneratorContextTest.java | 63 +++ .../GroupKeyGeneratorProviderQueriesTest.java | 377 ++++++++++++++++++ 16 files changed, 1103 insertions(+), 86 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContext.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorProvider.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/combine/GroupByCombineOperatorLifecycleTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/query/GroupByOperatorLifecycleTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutorProviderTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContextTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/queries/GroupKeyGeneratorProviderQueriesTest.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java index abc88149a06a..d8b34069d35d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java @@ -106,39 +106,40 @@ protected void processSegments() { ((AcquireReleaseColumnsSegmentOperator) operator).acquire(); } GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock(); - if (_indexedTable == null) { - synchronized (this) { - if (_indexedTable == null) { - _indexedTable = GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, _numTasks, - _executorService); + try (AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult()) { + if (_indexedTable == null) { + synchronized (this) { + if (_indexedTable == null) { + _indexedTable = + GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, _numTasks, + _executorService); + } } } - } - if (resultsBlock.isGroupsTrimmed()) { - _groupsTrimmed = true; - } - // Set groups limit reached flag. - if (resultsBlock.isNumGroupsLimitReached()) { - _numGroupsLimitReached = true; - } - if (resultsBlock.isNumGroupsWarningLimitReached()) { - _numGroupsWarningLimitReached = true; - } + if (resultsBlock.isGroupsTrimmed()) { + _groupsTrimmed = true; + } + // Set groups limit reached flag. + if (resultsBlock.isNumGroupsLimitReached()) { + _numGroupsLimitReached = true; + } + if (resultsBlock.isNumGroupsWarningLimitReached()) { + _numGroupsWarningLimitReached = true; + } - // Merge aggregation group-by result. - // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable - Collection intermediateRecords = resultsBlock.getIntermediateRecords(); - // Count the number of merged keys - int mergedKeys = 0; - // For now, only GroupBy OrderBy query has pre-constructed intermediate records - if (intermediateRecords == null) { // Merge aggregation group-by result. - AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult(); - if (aggregationGroupByResult != null) { - // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable - try { - Iterator dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator(); + // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable + Collection intermediateRecords = resultsBlock.getIntermediateRecords(); + // Count the number of merged keys + int mergedKeys = 0; + // For now, only GroupBy OrderBy query has pre-constructed intermediate records + if (intermediateRecords == null) { + // Merge aggregation group-by result. + if (aggregationGroupByResult != null) { + // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable + Iterator dicGroupKeyIterator = + aggregationGroupByResult.getGroupKeyIterator(); while (dicGroupKeyIterator.hasNext()) { QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME); GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next(); @@ -150,16 +151,13 @@ protected void processSegments() { } _indexedTable.upsert(new Key(keys), new Record(values)); } - } finally { - // Release the resources used by the group key generator - aggregationGroupByResult.closeGroupKeyGenerator(); } - } - } else { - for (IntermediateRecord intermediateResult : intermediateRecords) { - QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME); - //TODO: change upsert api so that it accepts intermediateRecord directly - _indexedTable.upsert(intermediateResult._key, intermediateResult._record); + } else { + for (IntermediateRecord intermediateResult : intermediateRecords) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME); + //TODO: change upsert api so that it accepts intermediateRecord directly + _indexedTable.upsert(intermediateResult._key, intermediateResult._record); + } } } } catch (RuntimeException e) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java index 50afa03d62a7..d4ad9f1824ce 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java @@ -40,6 +40,7 @@ import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils.AggregationInfo; import org.apache.pinot.core.query.aggregation.groupby.DefaultGroupByExecutor; import org.apache.pinot.core.query.aggregation.groupby.GroupByExecutor; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGeneratorProvider; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor; import org.apache.pinot.core.util.GroupByUtils; @@ -62,10 +63,16 @@ public class GroupByOperator extends BaseOperator { private final boolean _useStarTree; private final long _numTotalDocs; private final DataSchema _dataSchema; + private final GroupKeyGeneratorProvider _groupKeyGeneratorProvider; private int _numDocsScanned = 0; public GroupByOperator(QueryContext queryContext, AggregationInfo aggregationInfo, long numTotalDocs) { + this(queryContext, aggregationInfo, numTotalDocs, GroupKeyGeneratorProvider.DEFAULT); + } + + public GroupByOperator(QueryContext queryContext, AggregationInfo aggregationInfo, long numTotalDocs, + GroupKeyGeneratorProvider groupKeyGeneratorProvider) { assert queryContext.getAggregationFunctions() != null && queryContext.getGroupByExpressions() != null; _queryContext = queryContext; _aggregationFunctions = queryContext.getAggregationFunctions(); @@ -73,6 +80,7 @@ public GroupByOperator(QueryContext queryContext, AggregationInfo aggregationInf _projectOperator = aggregationInfo.getProjectOperator(); _useStarTree = aggregationInfo.isUseStarTree(); _numTotalDocs = numTotalDocs; + _groupKeyGeneratorProvider = groupKeyGeneratorProvider; // NOTE: The indexedTable expects that the data schema will have group by columns before aggregation columns int numGroupByExpressions = _groupByExpressions.length; @@ -123,10 +131,33 @@ protected GroupByResultsBlock getNextBlock() { if (_useStarTree) { groupByExecutor = new StarTreeGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator); } else { - groupByExecutor = new DefaultGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator); + groupByExecutor = new DefaultGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator, + _groupKeyGeneratorProvider); } - ValueBlock valueBlock; + GroupByResultsBlock resultsBlock; + boolean closeGroupKeyGenerator; + try { + resultsBlock = buildResultsBlock(groupByExecutor); + closeGroupKeyGenerator = + !_queryContext.isGroupingSets() && resultsBlock.getAggregationGroupByResult() == null; + } catch (RuntimeException | Error e) { + try { + groupByExecutor.getGroupKeyGenerator().close(); + } catch (RuntimeException | Error closeError) { + if (closeError != e) { + e.addSuppressed(closeError); + } + } + throw e; + } + if (closeGroupKeyGenerator) { + groupByExecutor.getGroupKeyGenerator().close(); + } + return resultsBlock; + } + private GroupByResultsBlock buildResultsBlock(GroupByExecutor groupByExecutor) { + ValueBlock valueBlock; while ((valueBlock = _projectOperator.nextBlock()) != null) { _numDocsScanned += valueBlock.getNumDocs(); QueryScanCostContext scanCost = getScanCostContext(); @@ -161,7 +192,6 @@ protected GroupByResultsBlock getNextBlock() { int trimSize = _queryContext.getEffectiveSegmentGroupTrimSize(); boolean unsafeTrim = _queryContext.isUnsafeTrim(); - GroupByResultsBlock resultsBlock; /// Grouping-set queries use a per-set bucketed segment trim (keyed on the $groupingId discriminator) so /// that a global top-K cannot starve low-magnitude sets such as the grand total. The broker still applies /// the final ORDER BY + LIMIT across all sets. @@ -172,15 +202,14 @@ protected GroupByResultsBlock getNextBlock() { groupByExecutor.getNumGroups(), _groupByExpressions.length, numGroupsLimitReached, numGroupsWarningLimitReached); } + + GroupByResultsBlock resultsBlock; // sort and trim segment results if needed if (trimSize > 0 && groupByExecutor.getNumGroups() > trimSize) { TableResizer tableResizer = new TableResizer(_dataSchema, _queryContext); // intermediateRecords is always sorted after trim List intermediateRecords = groupByExecutor.trimGroupByResult(trimSize, tableResizer, !unsafeTrim); - // close groupKeyGenerator after getting intermediateRecords - groupByExecutor.getGroupKeyGenerator().close(); - ServerMetrics.get().addMeteredGlobalValue(ServerMeter.AGGREGATE_TIMES_GROUPS_TRIMMED, 1); resultsBlock = new GroupByResultsBlock(_dataSchema, intermediateRecords, _queryContext); // set trim flag only if it's not safe @@ -199,8 +228,6 @@ protected GroupByResultsBlock getNextBlock() { List intermediateRecords = tableResizer.sortInSegmentResults(groupByExecutor.getGroupKeyGenerator(), groupByExecutor.getGroupByResultHolders(), trimSize); - // close groupKeyGenerator after getting intermediateRecords - groupByExecutor.getGroupKeyGenerator().close(); resultsBlock = new GroupByResultsBlock(_dataSchema, intermediateRecords, _queryContext); } else { // if not sort-aggregate and no trim needed, return segment result as it is diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java index 0ad10b1325d7..66ff2a4be37a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java @@ -164,8 +164,9 @@ protected GroupByResultsBlock detachFromWorkerThreadState(GroupByResultsBlock re if (aggregationGroupByResult == null || resultsBlock.getIntermediateRecords() != null) { return resultsBlock; } - List records = new ArrayList<>(aggregationGroupByResult.getNumGroups()); - try { + List records; + try (aggregationGroupByResult) { + records = new ArrayList<>(aggregationGroupByResult.getNumGroups()); Iterator groupKeyIterator = aggregationGroupByResult.getGroupKeyIterator(); int extractedKeys = 0; while (groupKeyIterator.hasNext()) { @@ -179,8 +180,6 @@ protected GroupByResultsBlock detachFromWorkerThreadState(GroupByResultsBlock re } records.add(IntermediateRecord.withoutOrderByValues(new Key(keys), new Record(values))); } - } finally { - aggregationGroupByResult.closeGroupKeyGenerator(); } GroupByResultsBlock detached = new GroupByResultsBlock(resultsBlock.getDataSchema(), records, _queryContext); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/GroupByPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/GroupByPlanNode.java index e124df80df7f..144ad7c7858b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/GroupByPlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/GroupByPlanNode.java @@ -24,6 +24,7 @@ import org.apache.pinot.core.operator.query.FilteredGroupByOperator; import org.apache.pinot.core.operator.query.GroupByOperator; import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGeneratorProvider; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.SegmentContext; @@ -34,11 +35,18 @@ public class GroupByPlanNode implements PlanNode { private final IndexSegment _indexSegment; private final SegmentContext _segmentContext; private final QueryContext _queryContext; + private final GroupKeyGeneratorProvider _groupKeyGeneratorProvider; public GroupByPlanNode(SegmentContext segmentContext, QueryContext queryContext) { + this(segmentContext, queryContext, GroupKeyGeneratorProvider.DEFAULT); + } + + public GroupByPlanNode(SegmentContext segmentContext, QueryContext queryContext, + GroupKeyGeneratorProvider groupKeyGeneratorProvider) { _indexSegment = segmentContext.getIndexSegment(); _segmentContext = segmentContext; _queryContext = queryContext; + _groupKeyGeneratorProvider = groupKeyGeneratorProvider; } @Override @@ -60,6 +68,7 @@ private GroupByOperator buildNonFilteredGroupByPlan() { AggregationFunctionUtils.buildAggregationInfo(_segmentContext, _queryContext, _queryContext.getAggregationFunctions(), _queryContext.getFilter(), filterOperator, filterPlanNode.getPredicateEvaluators()); - return new GroupByOperator(_queryContext, aggregationInfo, _indexSegment.getSegmentMetadata().getTotalDocs()); + return new GroupByOperator(_queryContext, aggregationInfo, _indexSegment.getSegmentMetadata().getTotalDocs(), + _groupKeyGeneratorProvider); } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java index dc8021857c14..52e43924b618 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java @@ -45,6 +45,7 @@ import org.apache.pinot.core.plan.StreamingInstanceResponsePlanNode; import org.apache.pinot.core.plan.StreamingSelectionPlanNode; import org.apache.pinot.core.query.aggregation.function.AggregationFunction; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGeneratorProvider; import org.apache.pinot.core.query.executor.ResultsBlockStreamer; import org.apache.pinot.core.query.prefetch.FetchPlanner; import org.apache.pinot.core.query.prefetch.FetchPlannerRegistry; @@ -347,7 +348,8 @@ public PlanNode makeSegmentPlanNode(SegmentContext segmentContext, QueryContext List groupByExpressions = queryContext.getGroupByExpressions(); if (groupByExpressions != null) { // Group-by query - return new GroupByPlanNode(segmentContext, queryContext); + return new GroupByPlanNode(segmentContext, queryContext, + getGroupKeyGeneratorProvider(segmentContext, queryContext)); } else { // Aggregation query return new AggregationPlanNode(segmentContext, queryContext); @@ -360,6 +362,13 @@ public PlanNode makeSegmentPlanNode(SegmentContext segmentContext, QueryContext } } + /// Returns the group-key generator provider for a segment group-by query. Filtered aggregations keep Pinot's shared + /// built-in generator and ignore the returned provider. + protected GroupKeyGeneratorProvider getGroupKeyGeneratorProvider(SegmentContext segmentContext, + QueryContext queryContext) { + return GroupKeyGeneratorProvider.DEFAULT; + } + @Override public Plan makeStreamingInstancePlan(List segmentContexts, QueryContext queryContext, ExecutorService executorService, ResultsBlockStreamer streamer) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java index 922385808632..3dfd24908dc8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java @@ -26,7 +26,7 @@ /// It provides an iterator over group-by keys, and provides a method /// to get the aggregation result for the given group-by key. @SuppressWarnings("rawtypes") -public class AggregationGroupByResult { +public class AggregationGroupByResult implements AutoCloseable { private final GroupKeyGenerator _groupKeyGenerator; private final AggregationFunction[] _aggregationFunctions; private final GroupByResultHolder[] _resultHolders; @@ -49,6 +49,11 @@ public Iterator getGroupKeyIterator() { /// Clear and trim DictionaryBasedGroupKeyGenerator after use public void closeGroupKeyGenerator() { + close(); + } + + @Override + public void close() { _groupKeyGenerator.close(); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java index 5bad8a3fc48b..33b6bcce5270 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java @@ -18,9 +18,13 @@ */ package org.apache.pinot.core.query.aggregation.groupby; +import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.OptionalInt; import java.util.Set; import java.util.stream.Collectors; import javax.annotation.Nullable; @@ -38,6 +42,9 @@ import org.apache.pinot.core.query.aggregation.function.AggregationFunction; import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils; import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; /// This class implements group by aggregation. @@ -65,25 +72,62 @@ public class DefaultGroupByExecutor implements GroupByExecutor { public DefaultGroupByExecutor(QueryContext queryContext, ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator) { - this(queryContext, queryContext.getAggregationFunctions(), groupByExpressions, projectOperator, null); + this(queryContext, queryContext.getAggregationFunctions(), groupByExpressions, projectOperator, null, + GroupKeyGeneratorProvider.DEFAULT); + } + + public DefaultGroupByExecutor(QueryContext queryContext, ExpressionContext[] groupByExpressions, + BaseProjectOperator projectOperator, GroupKeyGeneratorProvider groupKeyGeneratorProvider) { + this(queryContext, queryContext.getAggregationFunctions(), groupByExpressions, projectOperator, null, + groupKeyGeneratorProvider); } public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] aggregationFunctions, ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator) { - this(queryContext, aggregationFunctions, groupByExpressions, projectOperator, null); + this(queryContext, aggregationFunctions, groupByExpressions, projectOperator, null, + GroupKeyGeneratorProvider.DEFAULT); } public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] aggregationFunctions, ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator, @Nullable GroupKeyGenerator groupKeyGenerator) { + this(queryContext, aggregationFunctions, groupByExpressions, projectOperator, groupKeyGenerator, + GroupKeyGeneratorProvider.DEFAULT); + } + + private DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] aggregationFunctions, + ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator, + @Nullable GroupKeyGenerator groupKeyGenerator, GroupKeyGeneratorProvider groupKeyGeneratorProvider) { _aggregationFunctions = aggregationFunctions; assert _aggregationFunctions != null; _nullHandlingEnabled = queryContext.isNullHandlingEnabled(); + Objects.requireNonNull(groupKeyGeneratorProvider); boolean hasMVGroupByExpression = false; boolean hasNoDictionaryGroupByExpression = false; + boolean groupingSets = queryContext.isGroupingSets(); + List groupKeys = + groupKeyGenerator == null && !groupingSets + && groupKeyGeneratorProvider != GroupKeyGeneratorProvider.DEFAULT + ? new ArrayList<>(groupByExpressions.length) : null; for (ExpressionContext groupByExpression : groupByExpressions) { ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpression); + if (groupKeys != null) { + DataSource dataSource = columnContext.getDataSource(); + DataType storedType = columnContext.getDataType().getStoredType(); + Optional exactIntegralDomain = Optional.empty(); + OptionalInt cardinalityHint = OptionalInt.empty(); + if (dataSource != null) { + DataSourceMetadata dataSourceMetadata = dataSource.getDataSourceMetadata(); + exactIntegralDomain = getExactIntegralDomain(storedType, dataSourceMetadata); + int cardinality = dataSourceMetadata.getCardinality(); + if (cardinality >= 0) { + cardinalityHint = OptionalInt.of(cardinality); + } + } + groupKeys.add(new GroupKeyGeneratorContext.GroupKeySpec(groupByExpression, storedType, + columnContext.isSingleValue(), columnContext.isDictionaryEncoded(), exactIntegralDomain, cardinalityHint)); + } hasMVGroupByExpression |= !columnContext.isSingleValue(); // A column with EncodingType.RAW + explicit dictionaryIndex has a non-null dictionary but a RAW forward // index that throws on readDictIds; route those through the no-dict GROUP BY generator via the explicit @@ -92,7 +136,6 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a } /// Grouping-set queries expand each row into one group per grouping set, so they always use the /// multi-value (int[][]) executor path even though the union group-by columns are single-valued. - boolean groupingSets = queryContext.isGroupingSets(); _hasMVGroupByExpression = hasMVGroupByExpression || groupingSets; // Initialize group key generator @@ -102,46 +145,107 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a if (queryContext.isOptimizeMaxInitialResultHolderCapacity()) { groupByExpressionSizesFromPredicates = getGroupByExpressionSizesFromPredicates(queryContext); } + GroupKeyGenerator selectedGroupKeyGenerator; if (groupKeyGenerator != null) { - _groupKeyGenerator = groupKeyGenerator; + selectedGroupKeyGenerator = groupKeyGenerator; } else if (groupingSets) { - _groupKeyGenerator = new GroupingSetsGroupKeyGenerator(projectOperator, groupByExpressions, + selectedGroupKeyGenerator = new GroupingSetsGroupKeyGenerator(projectOperator, groupByExpressions, queryContext.getGroupingSets(), numGroupsLimit, _nullHandlingEnabled); + } else if (groupKeyGeneratorProvider == GroupKeyGeneratorProvider.DEFAULT) { + selectedGroupKeyGenerator = createDefaultGroupKeyGenerator(groupByExpressions, projectOperator, + hasNoDictionaryGroupByExpression, numGroupsLimit, maxInitialResultHolderCapacity, + groupByExpressionSizesFromPredicates); } else { - if (hasNoDictionaryGroupByExpression || _nullHandlingEnabled) { - if (groupByExpressions.length == 1) { - // TODO(nhejazi): support MV and dictionary based when null handling is enabled. - _groupKeyGenerator = - new NoDictionarySingleColumnGroupKeyGenerator(projectOperator, groupByExpressions[0], numGroupsLimit, - _nullHandlingEnabled, groupByExpressionSizesFromPredicates); - } else { - _groupKeyGenerator = - new NoDictionaryMultiColumnGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit, - _nullHandlingEnabled, groupByExpressionSizesFromPredicates); - } + GroupKeyGeneratorContext context = createGroupKeyGeneratorContext(Objects.requireNonNull(groupKeys), + groupByExpressionSizesFromPredicates, numGroupsLimit, + maxInitialResultHolderCapacity); + Optional providedGroupKeyGenerator = + Objects.requireNonNull(groupKeyGeneratorProvider.tryCreate(context)); + selectedGroupKeyGenerator = providedGroupKeyGenerator.isPresent() ? providedGroupKeyGenerator.get() + : createDefaultGroupKeyGenerator(groupByExpressions, projectOperator, hasNoDictionaryGroupByExpression, + numGroupsLimit, maxInitialResultHolderCapacity, + groupByExpressionSizesFromPredicates); + } + + try { + // Initialize result holders + int maxNumResults = selectedGroupKeyGenerator.getGlobalGroupKeyUpperBound(); + int initialCapacity = Math.min(maxNumResults, maxInitialResultHolderCapacity); + int numAggregationFunctions = _aggregationFunctions.length; + GroupByResultHolder[] groupByResultHolders = new GroupByResultHolder[numAggregationFunctions]; + for (int i = 0; i < numAggregationFunctions; i++) { + groupByResultHolders[i] = + _aggregationFunctions[i].createGroupByResultHolder(initialCapacity, maxNumResults); + } + + // Initialize map from document Id to group key + int[] svGroupKeys; + int[][] mvGroupKeys; + if (_hasMVGroupByExpression) { + svGroupKeys = null; + mvGroupKeys = THREAD_LOCAL_MV_GROUP_KEYS.get(); } else { - _groupKeyGenerator = new DictionaryBasedGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit, - maxInitialResultHolderCapacity, groupByExpressionSizesFromPredicates); + svGroupKeys = THREAD_LOCAL_SV_GROUP_KEYS.get(); + mvGroupKeys = null; + } + _groupKeyGenerator = selectedGroupKeyGenerator; + _groupByResultHolders = groupByResultHolders; + _svGroupKeys = svGroupKeys; + _mvGroupKeys = mvGroupKeys; + } catch (RuntimeException | Error e) { + if (groupKeyGenerator == null) { + try { + selectedGroupKeyGenerator.close(); + } catch (RuntimeException | Error closeError) { + if (closeError != e) { + e.addSuppressed(closeError); + } + } } + throw e; } + } - // Initialize result holders - int maxNumResults = _groupKeyGenerator.getGlobalGroupKeyUpperBound(); - int initialCapacity = Math.min(maxNumResults, maxInitialResultHolderCapacity); - int numAggregationFunctions = _aggregationFunctions.length; - _groupByResultHolders = new GroupByResultHolder[numAggregationFunctions]; - for (int i = 0; i < numAggregationFunctions; i++) { - _groupByResultHolders[i] = _aggregationFunctions[i].createGroupByResultHolder(initialCapacity, maxNumResults); + private GroupKeyGenerator createDefaultGroupKeyGenerator(ExpressionContext[] groupByExpressions, + BaseProjectOperator projectOperator, boolean hasNoDictionaryGroupByExpression, int numGroupsLimit, + int maxInitialResultHolderCapacity, + @Nullable Map groupByExpressionSizesFromPredicates) { + if (hasNoDictionaryGroupByExpression || _nullHandlingEnabled) { + if (groupByExpressions.length == 1) { + // TODO(nhejazi): support MV and dictionary based when null handling is enabled. + return new NoDictionarySingleColumnGroupKeyGenerator(projectOperator, groupByExpressions[0], numGroupsLimit, + _nullHandlingEnabled, groupByExpressionSizesFromPredicates); + } + return new NoDictionaryMultiColumnGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit, + _nullHandlingEnabled, groupByExpressionSizesFromPredicates); } + return new DictionaryBasedGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit, + maxInitialResultHolderCapacity, groupByExpressionSizesFromPredicates); + } - // Initialize map from document Id to group key - if (_hasMVGroupByExpression) { - _svGroupKeys = null; - _mvGroupKeys = THREAD_LOCAL_MV_GROUP_KEYS.get(); - } else { - _svGroupKeys = THREAD_LOCAL_SV_GROUP_KEYS.get(); - _mvGroupKeys = null; + private GroupKeyGeneratorContext createGroupKeyGeneratorContext(List groupKeys, + @Nullable Map groupByExpressionSizesFromPredicates, int numGroupsLimit, + int maxInitialResultHolderCapacity) { + return new GroupKeyGeneratorContext(groupKeys, List.of(), + groupByExpressionSizesFromPredicates != null ? groupByExpressionSizesFromPredicates : Map.of(), + numGroupsLimit, maxInitialResultHolderCapacity, _nullHandlingEnabled); + } + + private static Optional getExactIntegralDomain(DataType dataType, + DataSourceMetadata dataSourceMetadata) { + Comparable minValue = dataSourceMetadata.getMinValue(); + Comparable maxValue = dataSourceMetadata.getMaxValue(); + if (dataType == DataType.INT && minValue instanceof Integer && maxValue instanceof Integer) { + int min = (Integer) minValue; + int max = (Integer) maxValue; + return min <= max ? Optional.of(new GroupKeyGeneratorContext.IntegralDomain(min, max)) : Optional.empty(); + } + if (dataType == DataType.LONG && minValue instanceof Long && maxValue instanceof Long) { + long min = (Long) minValue; + long max = (Long) maxValue; + return min <= max ? Optional.of(new GroupKeyGeneratorContext.IntegralDomain(min, max)) : Optional.empty(); } + return Optional.empty(); } /// Retrieve the sizes of GroupBy expressions from IN an EQ predicates found in the filter context, if available. diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java index 0e8162d6d462..cea7732d26e1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java @@ -23,7 +23,10 @@ /// Interface for generating group keys. -/// It extends AutoCloseable for thread-local maps to be cleared +/// It extends [AutoCloseable] so generators can deterministically release thread-local or native resources. Generator +/// implementations that own resources must make [#close()] idempotent. Before `close` returns or throws, it must +/// either release owned resources or durably transfer cleanup to an independent retry owner because callers may +/// discard the generator after one close attempt. public interface GroupKeyGenerator extends AutoCloseable { char DELIMITER = '\0'; int INVALID_ID = -1; @@ -58,7 +61,9 @@ public interface GroupKeyGenerator extends AutoCloseable { /// @return current upper bound of the group key. int getCurrentGroupKeyUpperBound(); - /// Returns an iterator of [GroupKey]. Use this interface to iterate through all the group keys. + /// Returns an iterator of [GroupKey]. The iterator may reuse the [GroupKey] wrapper, so callers must read its fields + /// before advancing. Each entry's `_keys` array and elements, however, must be stable, heap-owned values: consumers + /// can retain them after advancing the iterator and after this generator is closed. Iterator getGroupKeys(); /// Return current number of unique keys diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContext.java new file mode 100644 index 000000000000..aa69e5a5862d --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContext.java @@ -0,0 +1,108 @@ +/** + * 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.query.aggregation.groupby; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.OptionalInt; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.spi.annotations.InterfaceAudience; +import org.apache.pinot.spi.annotations.InterfaceStability; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/// Structurally immutable metadata exposed to a [GroupKeyGeneratorProvider]. Collections are defensively copied, but +/// expression objects are borrowed query-lifetime values. The context intentionally excludes operators and segment +/// data sources so providers cannot retain broader query-lifetime objects. +@InterfaceAudience.LimitedPrivate("StarTree") +@InterfaceStability.Unstable +public final class GroupKeyGeneratorContext { + private final List _groupKeys; + private final List> _groupingSets; + private final Map _predicateCardinalityHints; + private final int _numGroupsLimit; + private final int _maxInitialResultHolderCapacity; + private final boolean _nullHandlingEnabled; + + public GroupKeyGeneratorContext(List groupKeys, List> groupingSets, + Map predicateCardinalityHints, int numGroupsLimit, + int maxInitialResultHolderCapacity, boolean nullHandlingEnabled) { + _groupKeys = List.copyOf(Objects.requireNonNull(groupKeys)); + Objects.requireNonNull(groupingSets); + List> copiedGroupingSets = new ArrayList<>(groupingSets.size()); + for (List groupingSet : groupingSets) { + copiedGroupingSets.add(List.copyOf(groupingSet)); + } + _groupingSets = List.copyOf(copiedGroupingSets); + _predicateCardinalityHints = Map.copyOf(Objects.requireNonNull(predicateCardinalityHints)); + _numGroupsLimit = numGroupsLimit; + _maxInitialResultHolderCapacity = maxInitialResultHolderCapacity; + _nullHandlingEnabled = nullHandlingEnabled; + } + + public List getGroupKeys() { + return _groupKeys; + } + + public List> getGroupingSets() { + return _groupingSets; + } + + public Map getPredicateCardinalityHints() { + return _predicateCardinalityHints; + } + + public int getNumGroupsLimit() { + return _numGroupsLimit; + } + + public int getMaxInitialResultHolderCapacity() { + return _maxInitialResultHolderCapacity; + } + + public boolean isNullHandlingEnabled() { + return _nullHandlingEnabled; + } + + /// Metadata for one group-by expression. Integral domains are exact physical-column bounds. Cardinality is a + /// sizing hint because raw-column segment metadata can be approximate. `dictionaryEncoded` describes the forward + /// index, so it is false for a raw forward index with a side dictionary. + public record GroupKeySpec(ExpressionContext expression, DataType storedType, boolean singleValue, + boolean dictionaryEncoded, Optional exactIntegralDomain, + OptionalInt cardinalityHint) { + public GroupKeySpec { + Objects.requireNonNull(expression); + Objects.requireNonNull(storedType); + Objects.requireNonNull(exactIntegralDomain); + Objects.requireNonNull(cardinalityHint); + } + } + + /// Inclusive exact range for an integral physical column. + public record IntegralDomain(long minInclusive, long maxInclusive) { + public IntegralDomain { + if (minInclusive > maxInclusive) { + throw new IllegalArgumentException("minInclusive must not exceed maxInclusive"); + } + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorProvider.java new file mode 100644 index 000000000000..2e3ef444152f --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorProvider.java @@ -0,0 +1,46 @@ +/** + * 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.query.aggregation.groupby; + +import java.util.Optional; +import org.apache.pinot.spi.annotations.InterfaceAudience; +import org.apache.pinot.spi.annotations.InterfaceStability; + + +/// Provides an optional specialized [GroupKeyGenerator] for a single-segment group-by query. +/// +/// [#tryCreate] is called synchronously by the segment execution thread, at most once while each ordinary, +/// non-filtered, non-grouping-set SSE group-by executor is constructed. Callers may reuse one provider across segment +/// plans and queries, so a shared provider can receive concurrent calls and must be thread-safe. +/// +/// A provider must not retain the supplied context and must return an empty [Optional] when it cannot safely handle it; +/// Pinot then uses its built-in generator selection unchanged. Each returned generator must be a fresh, query-owned +/// instance. Pinot owns it until ownership is transferred with the raw group-by result or Pinot makes one close +/// attempt. Resource-owning generators must release their resources or durably transfer cleanup to an independent +/// retry owner before [GroupKeyGenerator#close] returns or throws; Pinot may discard the generator after that attempt. +@FunctionalInterface +@InterfaceAudience.LimitedPrivate("StarTree") +@InterfaceStability.Unstable +public interface GroupKeyGeneratorProvider { + /// The built-in provider sentinel. Callers use identity comparison with this exact instance to preserve Pinot's + /// existing generator-selection hot path without collecting the additional [GroupKeyGeneratorContext] metadata. + GroupKeyGeneratorProvider DEFAULT = context -> Optional.empty(); + + Optional tryCreate(GroupKeyGeneratorContext context); +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/GroupByCombineOperatorLifecycleTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/GroupByCombineOperatorLifecycleTest.java new file mode 100644 index 000000000000..7e743852cfd6 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/GroupByCombineOperatorLifecycleTest.java @@ -0,0 +1,74 @@ +/** + * 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.combine; + +import java.util.List; +import java.util.concurrent.ExecutorService; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; +import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult; +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.query.request.context.utils.QueryContextConverterUtils; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.expectThrows; + + +@SuppressWarnings("rawtypes") +public class GroupByCombineOperatorLifecycleTest { + @Test + public void testClosesRawResultWhenIndexedTableConstructionFails() { + QueryContext queryContext = + QueryContextConverterUtils.getQueryContext("SELECT COUNT(*) FROM testTable GROUP BY intColumn"); + GroupKeyGenerator groupKeyGenerator = mock(GroupKeyGenerator.class); + AggregationGroupByResult rawResult = new AggregationGroupByResult(groupKeyGenerator, + queryContext.getAggregationFunctions(), new GroupByResultHolder[1]); + GroupByResultsBlock resultsBlock = mock(GroupByResultsBlock.class); + when(resultsBlock.getAggregationGroupByResult()).thenReturn(rawResult); + IllegalStateException indexedTableFailure = new IllegalStateException("indexed-table failure"); + when(resultsBlock.getNumGroups()).thenThrow(indexedTableFailure); + Operator operator = mock(Operator.class); + when(operator.nextBlock()).thenReturn(resultsBlock); + ExecutorService executorService = mock(ExecutorService.class); + TestGroupByCombineOperator combineOperator = + new TestGroupByCombineOperator(List.of(operator), queryContext, executorService); + + RuntimeException thrown = expectThrows(RuntimeException.class, combineOperator::process); + + assertSame(thrown.getCause(), indexedTableFailure); + verify(groupKeyGenerator).close(); + } + + private static class TestGroupByCombineOperator extends GroupByCombineOperator { + private TestGroupByCombineOperator(List operators, QueryContext queryContext, + ExecutorService executorService) { + super(operators, queryContext, executorService); + } + + private void process() { + processSegments(); + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/GroupByOperatorLifecycleTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/GroupByOperatorLifecycleTest.java new file mode 100644 index 000000000000..40d936ec80a0 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/GroupByOperatorLifecycleTest.java @@ -0,0 +1,73 @@ +/** + * 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.Optional; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils.AggregationInfo; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.expectThrows; + + +public class GroupByOperatorLifecycleTest { + @Test + public void testClosesProviderGeneratorWhenExecutionFails() { + QueryContext queryContext = + QueryContextConverterUtils.getQueryContext("SELECT COUNT(*) FROM testTable GROUP BY intColumn"); + ExpressionContext expression = queryContext.getGroupByExpressions().get(0); + BaseProjectOperator projectOperator = mock(BaseProjectOperator.class); + ColumnContext columnContext = mock(ColumnContext.class); + when(columnContext.getDataType()).thenReturn(DataType.INT); + when(columnContext.isSingleValue()).thenReturn(true); + when(columnContext.isDictionaryEncoded()).thenReturn(false); + when(projectOperator.getResultColumnContext(expression)).thenReturn(columnContext); + + IllegalStateException executionFailure = new IllegalStateException("project failure"); + IllegalArgumentException closeFailure = new IllegalArgumentException("close failure"); + when(projectOperator.nextBlock()).thenThrow(executionFailure); + GroupKeyGenerator groupKeyGenerator = mock(GroupKeyGenerator.class); + when(groupKeyGenerator.getGlobalGroupKeyUpperBound()).thenReturn(16); + doThrow(closeFailure).when(groupKeyGenerator).close(); + AggregationInfo aggregationInfo = mock(AggregationInfo.class); + doReturn(projectOperator).when(aggregationInfo).getProjectOperator(); + when(aggregationInfo.isUseStarTree()).thenReturn(false); + GroupByOperator operator = new GroupByOperator(queryContext, aggregationInfo, 10, + context -> Optional.of(groupKeyGenerator)); + + IllegalStateException thrown = expectThrows(IllegalStateException.class, operator::nextBlock); + + assertSame(thrown, executionFailure); + assertEquals(thrown.getSuppressed(), new Throwable[]{closeFailure}); + verify(groupKeyGenerator).close(); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java index 7ae970f95cb3..b8ca862d5d29 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java @@ -36,6 +36,7 @@ import org.apache.pinot.core.plan.PlanNode; import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; import org.apache.pinot.core.plan.maker.PlanMaker; +import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; @@ -56,6 +57,9 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.*; @@ -434,6 +438,28 @@ public void testGroupingSetsKeyColumnLayout() { assertEquals(groupSums.get(null), NUM_SEGMENTS * 2550.0, 0.001, "Incorrect rollup total"); } + @Test + public void testDetachPreservesFailureWhenCloseAlsoFails() { + QueryContext queryContext = QueryContextConverterUtils.getQueryContext( + "SELECT groupColumn, COUNT(*) FROM testTable GROUP BY groupColumn"); + IllegalStateException extractionFailure = new IllegalStateException("extraction failure"); + IllegalArgumentException closeFailure = new IllegalArgumentException("close failure"); + AggregationGroupByResult aggregationGroupByResult = mock(AggregationGroupByResult.class); + when(aggregationGroupByResult.getNumGroups()).thenThrow(extractionFailure); + doThrow(closeFailure).when(aggregationGroupByResult).close(); + GroupByResultsBlock resultsBlock = mock(GroupByResultsBlock.class); + when(resultsBlock.getAggregationGroupByResult()).thenReturn(aggregationGroupByResult); + when(resultsBlock.getIntermediateRecords()).thenReturn(null); + + StreamingGroupByCombineOperator combineOperator = + new StreamingGroupByCombineOperator(List.of(), queryContext, EXECUTOR, 10); + IllegalStateException thrown = expectThrows(IllegalStateException.class, + () -> combineOperator.detachFromWorkerThreadState(resultsBlock)); + + assertSame(thrown, extractionFailure); + assertEquals(thrown.getSuppressed(), new Throwable[]{closeFailure}); + } + private List buildOperators(QueryContext queryContext) { List operators = new ArrayList<>(NUM_SEGMENTS); for (IndexSegment indexSegment : _indexSegments) { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutorProviderTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutorProviderTest.java new file mode 100644 index 000000000000..81cfdcae2fbd --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutorProviderTest.java @@ -0,0 +1,94 @@ +/** + * 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.query.aggregation.groupby; + +import java.util.Optional; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +public class DefaultGroupByExecutorProviderTest { + @Test + public void testDefaultProviderPreservesRawSideDictionaryFastPath() { + TestSetup setup = new TestSetup(); + when(setup._columnContext.getDictionary()).thenReturn(mock(Dictionary.class)); + + DefaultGroupByExecutor defaultExecutor = new DefaultGroupByExecutor(setup._queryContext, setup._expressions, + setup._projectOperator, GroupKeyGeneratorProvider.DEFAULT); + + assertTrue(defaultExecutor.getGroupKeyGenerator() instanceof NoDictionarySingleColumnGroupKeyGenerator); + verify(setup._columnContext, never()).getDataSource(); + defaultExecutor.getResult().closeGroupKeyGenerator(); + + DefaultGroupByExecutor fallbackExecutor = new DefaultGroupByExecutor(setup._queryContext, setup._expressions, + setup._projectOperator, context -> Optional.empty()); + assertEquals(fallbackExecutor.getGroupKeyGenerator().getClass(), NoDictionarySingleColumnGroupKeyGenerator.class); + fallbackExecutor.getResult().closeGroupKeyGenerator(); + } + + @Test + public void testClosesProviderGeneratorWhenInitializationFails() { + TestSetup setup = new TestSetup(); + IllegalStateException initializationFailure = new IllegalStateException("initialization failure"); + IllegalArgumentException closeFailure = new IllegalArgumentException("close failure"); + GroupKeyGenerator groupKeyGenerator = mock(GroupKeyGenerator.class); + when(groupKeyGenerator.getGlobalGroupKeyUpperBound()).thenThrow(initializationFailure); + doThrow(closeFailure).when(groupKeyGenerator).close(); + + IllegalStateException thrown = expectThrows(IllegalStateException.class, + () -> new DefaultGroupByExecutor(setup._queryContext, setup._expressions, setup._projectOperator, + context -> Optional.of(groupKeyGenerator))); + + assertSame(thrown, initializationFailure); + assertEquals(thrown.getSuppressed(), new Throwable[]{closeFailure}); + verify(groupKeyGenerator).close(); + } + + private static class TestSetup { + private final QueryContext _queryContext = + QueryContextConverterUtils.getQueryContext("SELECT COUNT(*) FROM testTable GROUP BY intColumn"); + private final ExpressionContext[] _expressions = + _queryContext.getGroupByExpressions().toArray(new ExpressionContext[0]); + private final BaseProjectOperator _projectOperator = mock(BaseProjectOperator.class); + private final ColumnContext _columnContext = mock(ColumnContext.class); + + private TestSetup() { + when(_projectOperator.getResultColumnContext(_expressions[0])).thenReturn(_columnContext); + when(_columnContext.getDataType()).thenReturn(DataType.INT); + when(_columnContext.isSingleValue()).thenReturn(true); + when(_columnContext.isDictionaryEncoded()).thenReturn(false); + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContextTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContextTest.java new file mode 100644 index 000000000000..888accedb593 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGeneratorContextTest.java @@ -0,0 +1,63 @@ +/** + * 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.query.aggregation.groupby; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +public class GroupKeyGeneratorContextTest { + @Test + public void testDefensiveCopies() { + ExpressionContext expression = ExpressionContext.forIdentifier("intColumn"); + List groupKeys = new ArrayList<>(); + groupKeys.add(new GroupKeyGeneratorContext.GroupKeySpec(expression, DataType.INT, true, false, + Optional.of(new GroupKeyGeneratorContext.IntegralDomain(1, 10)), OptionalInt.of(10))); + List groupingSet = new ArrayList<>(List.of(0)); + List> groupingSets = new ArrayList<>(List.of(groupingSet)); + Map hints = new HashMap<>(Map.of(expression, 3)); + + GroupKeyGeneratorContext context = + new GroupKeyGeneratorContext(groupKeys, groupingSets, hints, 100, 20, true); + groupKeys.clear(); + groupingSet.clear(); + groupingSets.clear(); + hints.clear(); + + assertEquals(context.getGroupKeys().size(), 1); + assertEquals(context.getGroupingSets(), List.of(List.of(0))); + assertEquals(context.getPredicateCardinalityHints(), Map.of(expression, 3)); + assertEquals(context.getNumGroupsLimit(), 100); + assertEquals(context.getMaxInitialResultHolderCapacity(), 20); + assertTrue(context.isNullHandlingEnabled()); + expectThrows(UnsupportedOperationException.class, () -> context.getGroupKeys().clear()); + expectThrows(UnsupportedOperationException.class, () -> context.getGroupingSets().get(0).add(1)); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/GroupKeyGeneratorProviderQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/GroupKeyGeneratorProviderQueriesTest.java new file mode 100644 index 000000000000..92f222b383d3 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/queries/GroupKeyGeneratorProviderQueriesTest.java @@ -0,0 +1,377 @@ +/** + * 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.queries; + +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.io.File; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGeneratorContext; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGeneratorProvider; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +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.JsonUtils; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/// Exercises the provider SPI through the real segment plan, projection, group-by, combine, serialization, and broker +/// reduction pipeline. The physical segment includes the three important forward-index shapes: dictionary encoded, +/// raw without a dictionary, and raw with a side dictionary. +public class GroupKeyGeneratorProviderQueriesTest extends BaseQueriesTest { + private static final String TABLE_NAME = "providerTestTable"; + private static final String SEGMENT_NAME = "providerTestSegment"; + private static final String DICTIONARY_INT = "dictionaryInt"; + private static final String RAW_INT = "rawInt"; + private static final String RAW_SIDE_DICTIONARY_INT = "rawSideDictionaryInt"; + private static final String NULLABLE_RAW_INT = "nullableRawInt"; + private static final int PROVIDER_GROUP_LIMIT = 16; + + private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME) + .addSingleValueDimension(DICTIONARY_INT, DataType.INT) + .addSingleValueDimension(RAW_INT, DataType.INT) + .addSingleValueDimension(RAW_SIDE_DICTIONARY_INT, DataType.INT) + .addSingleValueDimension(NULLABLE_RAW_INT, DataType.INT) + .build(); + private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNoDictionaryColumns(List.of(RAW_INT, NULLABLE_RAW_INT)) + .setFieldConfigList(List.of(rawWithDictionary(RAW_SIDE_DICTIONARY_INT))) + .build(); + + private File _indexDir; + private ImmutableSegment _indexSegment; + + @BeforeClass + public void setUp() + throws Exception { + _indexDir = Files.createTempDirectory(getClass().getSimpleName()).toFile(); + + SegmentGeneratorConfig generatorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + generatorConfig.setOutDir(_indexDir.getAbsolutePath()); + generatorConfig.setSegmentName(SEGMENT_NAME); + generatorConfig.setDefaultNullHandlingEnabled(true); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + try (RecordReader recordReader = new GenericRowRecordReader(createRows())) { + driver.init(generatorConfig, recordReader); + driver.build(); + } + + _indexSegment = ImmutableSegmentLoader.load(new File(_indexDir, SEGMENT_NAME), ReadMode.mmap); + assertPhysicalIndexShapes(); + } + + @AfterClass(alwaysRun = true) + public void tearDown() + throws Exception { + if (_indexSegment != null) { + _indexSegment.destroy(); + } + if (_indexDir != null) { + FileUtils.deleteDirectory(_indexDir); + } + } + + @Override + protected String getFilter() { + return ""; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List getIndexSegments() { + return List.of(_indexSegment); + } + + @DataProvider(name = "physicalGroupKeyShapes") + public Object[][] physicalGroupKeyShapes() { + return new Object[][]{ + {RAW_INT, false, false, true, false, 10L, 30L, 3, List.of(10, 20, 30)}, + {RAW_SIDE_DICTIONARY_INT, false, false, true, true, 100L, 300L, 3, List.of(100, 200, 300)}, + {DICTIONARY_INT, true, false, false, false, 1L, 3L, 3, List.of(1, 2, 3)}, + {NULLABLE_RAW_INT, false, true, false, false, (long) Integer.MIN_VALUE, 9L, 3, + Arrays.asList(null, 7, 9)} + }; + } + + @Test(dataProvider = "physicalGroupKeyShapes") + public void testProviderAgainstRealSegment(String column, boolean dictionaryEncoded, boolean nullHandlingEnabled, + boolean expectProviderSelection, boolean materializeSegmentResult, long expectedMin, long expectedMax, + int expectedCardinality, + List expectedKeys) { + String query = (nullHandlingEnabled ? "SET enableNullHandling=true; " : "") + + "SELECT " + column + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + column + + (materializeSegmentResult ? " ORDER BY " + column + " LIMIT 10" : ""); + BrokerResponseNative expectedResponse = getBrokerResponse(query); + + TrackingPlanMaker planMaker = new TrackingPlanMaker(); + BrokerResponseNative actualResponse = getBrokerResponse(query, planMaker); + + assertTrue(expectedResponse.getExceptions().isEmpty(), expectedResponse.getExceptions().toString()); + assertTrue(actualResponse.getExceptions().isEmpty(), actualResponse.getExceptions().toString()); + ResultTable expectedResultTable = expectedResponse.getResultTable(); + ResultTable actualResultTable = actualResponse.getResultTable(); + assertNotNull(expectedResultTable); + assertNotNull(actualResultTable); + assertEquals(actualResultTable.getDataSchema(), expectedResultTable.getDataSchema()); + Map expectedCounts = toCountMap(expectedResultTable); + assertEquals(toCountMap(actualResultTable), expectedCounts); + assertEquals(expectedCounts.keySet(), new HashSet<>(expectedKeys)); + + assertEquals(planMaker._contexts.size(), 1); + GroupKeyGeneratorContext context = planMaker._contexts.get(0); + assertEquals(context.getGroupKeys().size(), 1); + GroupKeyGeneratorContext.GroupKeySpec groupKeySpec = context.getGroupKeys().get(0); + assertEquals(groupKeySpec.expression().getIdentifier(), column); + assertEquals(groupKeySpec.storedType(), DataType.INT); + assertTrue(groupKeySpec.singleValue()); + assertEquals(groupKeySpec.dictionaryEncoded(), dictionaryEncoded); + assertEquals(context.isNullHandlingEnabled(), nullHandlingEnabled); + assertTrue(context.getGroupingSets().isEmpty()); + assertEquals(context.getNumGroupsLimit(), PROVIDER_GROUP_LIMIT); + assertEquals(context.getMaxInitialResultHolderCapacity(), PROVIDER_GROUP_LIMIT); + assertEquals(groupKeySpec.exactIntegralDomain(), + Optional.of(new GroupKeyGeneratorContext.IntegralDomain(expectedMin, expectedMax))); + assertEquals(groupKeySpec.cardinalityHint(), OptionalInt.of(expectedCardinality)); + + if (expectProviderSelection) { + assertEquals(planMaker._generators.size(), 1); + assertEquals(planMaker._generators.get(0)._closeAttempts.get(), 1); + } else { + assertTrue(planMaker._generators.isEmpty()); + } + } + + private void assertPhysicalIndexShapes() { + DataSource dictionaryDataSource = _indexSegment.getDataSource(DICTIONARY_INT); + assertNotNull(dictionaryDataSource.getDictionary()); + assertTrue(dictionaryDataSource.getForwardIndex().isDictionaryEncoded()); + + DataSource rawDataSource = _indexSegment.getDataSource(RAW_INT); + assertNull(rawDataSource.getDictionary()); + assertFalse(rawDataSource.getForwardIndex().isDictionaryEncoded()); + + DataSource rawSideDictionaryDataSource = _indexSegment.getDataSource(RAW_SIDE_DICTIONARY_INT); + assertNotNull(rawSideDictionaryDataSource.getDictionary()); + assertFalse(rawSideDictionaryDataSource.getForwardIndex().isDictionaryEncoded()); + + DataSource nullableRawDataSource = _indexSegment.getDataSource(NULLABLE_RAW_INT); + assertNull(nullableRawDataSource.getDictionary()); + assertFalse(nullableRawDataSource.getForwardIndex().isDictionaryEncoded()); + assertNotNull(nullableRawDataSource.getNullValueVector()); + assertEquals(nullableRawDataSource.getNullValueVector().getNullBitmap().getCardinality(), 2); + } + + private static List createRows() { + List rows = new ArrayList<>(); + rows.add(row(1, 10, 100, 7)); + rows.add(row(2, 20, 200, null)); + rows.add(row(1, 10, 100, 7)); + rows.add(row(3, 30, 300, 9)); + rows.add(row(2, 20, 200, null)); + rows.add(row(3, 30, 300, 9)); + return rows; + } + + private static GenericRow row(int dictionaryValue, int rawValue, int rawSideDictionaryValue, + Integer nullableRawValue) { + GenericRow row = new GenericRow(); + row.putValue(DICTIONARY_INT, dictionaryValue); + row.putValue(RAW_INT, rawValue); + row.putValue(RAW_SIDE_DICTIONARY_INT, rawSideDictionaryValue); + if (nullableRawValue == null) { + row.putDefaultNullValue(NULLABLE_RAW_INT, SCHEMA.getFieldSpecFor(NULLABLE_RAW_INT).getDefaultNullValue()); + } else { + row.putValue(NULLABLE_RAW_INT, nullableRawValue); + } + return row; + } + + private static FieldConfig rawWithDictionary(String column) { + ObjectNode indexes = JsonUtils.newObjectNode(); + ObjectNode forwardIndex = JsonUtils.newObjectNode(); + forwardIndex.put("encodingType", "RAW"); + indexes.set("forward", forwardIndex); + ObjectNode dictionaryIndex = JsonUtils.newObjectNode(); + dictionaryIndex.put("disabled", false); + indexes.set("dictionary", dictionaryIndex); + return new FieldConfig.Builder(column).withEncodingType(FieldConfig.EncodingType.RAW).withIndexes(indexes).build(); + } + + private static Map toCountMap(ResultTable resultTable) { + Map counts = new LinkedHashMap<>(); + for (Object[] row : resultTable.getRows()) { + Object key = row[0]; + assertFalse(counts.containsKey(key), "Duplicate group key: " + key); + counts.put(key, ((Number) row[1]).longValue()); + } + return counts; + } + + private static final class TrackingPlanMaker extends InstancePlanMakerImplV2 { + private final List _contexts = new CopyOnWriteArrayList<>(); + private final List _generators = new CopyOnWriteArrayList<>(); + + private TrackingPlanMaker() { + setNumGroupsLimit(PROVIDER_GROUP_LIMIT); + setMaxInitialResultHolderCapacity(PROVIDER_GROUP_LIMIT); + } + + @Override + protected GroupKeyGeneratorProvider getGroupKeyGeneratorProvider(SegmentContext segmentContext, + QueryContext queryContext) { + return context -> { + _contexts.add(context); + List groupKeySpecs = context.getGroupKeys(); + if (groupKeySpecs.size() != 1 || context.isNullHandlingEnabled() || !context.getGroupingSets().isEmpty()) { + return Optional.empty(); + } + GroupKeyGeneratorContext.GroupKeySpec groupKeySpec = groupKeySpecs.get(0); + if (groupKeySpec.storedType() != DataType.INT || !groupKeySpec.singleValue() + || groupKeySpec.dictionaryEncoded()) { + return Optional.empty(); + } + TrackingIntGroupKeyGenerator generator = + new TrackingIntGroupKeyGenerator(groupKeySpec.expression(), context.getNumGroupsLimit()); + _generators.add(generator); + return Optional.of(generator); + }; + } + } + + private static final class TrackingIntGroupKeyGenerator implements GroupKeyGenerator { + private final ExpressionContext _expression; + private final int _numGroupsLimit; + private final Map _groupIds = new LinkedHashMap<>(); + private final AtomicInteger _closeAttempts = new AtomicInteger(); + + private TrackingIntGroupKeyGenerator(ExpressionContext expression, int numGroupsLimit) { + _expression = expression; + _numGroupsLimit = numGroupsLimit; + } + + @Override + public int getGlobalGroupKeyUpperBound() { + return _numGroupsLimit; + } + + @Override + public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { + int[] values = valueBlock.getBlockValueSet(_expression).getIntValuesSV(); + int numDocs = valueBlock.getNumDocs(); + for (int i = 0; i < numDocs; i++) { + Integer groupId = _groupIds.get(values[i]); + if (groupId == null) { + if (_groupIds.size() >= _numGroupsLimit) { + groupKeys[i] = INVALID_ID; + continue; + } + groupId = _groupIds.size(); + _groupIds.put(values[i], groupId); + } + groupKeys[i] = groupId; + } + } + + @Override + public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { + throw new AssertionError("Single-value provider should not use the multi-value path"); + } + + @Override + public int getCurrentGroupKeyUpperBound() { + return _groupIds.size(); + } + + @Override + public Iterator getGroupKeys() { + Iterator> entries = _groupIds.entrySet().iterator(); + return new Iterator<>() { + @Override + public boolean hasNext() { + return entries.hasNext(); + } + + @Override + public GroupKey next() { + Map.Entry entry = entries.next(); + GroupKey groupKey = new GroupKey(); + groupKey._groupId = entry.getValue(); + groupKey._keys = new Object[]{entry.getKey()}; + return groupKey; + } + }; + } + + @Override + public int getNumKeys() { + return _groupIds.size(); + } + + @Override + public void close() { + _closeAttempts.incrementAndGet(); + } + } +}