apache · xiangfu0 · Aug 28, 2026
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
@@ -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<IntermediateRecord> 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<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
+          // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
+          Collection<IntermediateRecord> 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<GroupKeyGenerator.GroupKey> 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
@@ -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,17 +63,24 @@ public class GroupByOperator extends BaseOperator<GroupByResultsBlock> {
   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();
     _groupByExpressions = queryContext.getGroupByExpressions().toArray(new ExpressionContext[0]);
     _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<IntermediateRecord> 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<IntermediateRecord> 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/...c/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java b/...c/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<IntermediateRecord> records = new ArrayList<>(aggregationGroupByResult.getNumGroups());
-    try {
+    List<IntermediateRecord> records;
+    try (aggregationGroupByResult) {
+      records = new ArrayList<>(aggregationGroupByResult.getNumGroups());
       Iterator<GroupKeyGenerator.GroupKey> 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
@@ -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
@@ -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<ExpressionContext> 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<SegmentContext> segmentContexts, QueryContext queryContext,
       ExecutorService executorService, ResultsBlockStreamer streamer) {

diff --git a/...c/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java b/...c/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<GroupKeyGenerator.GroupKey> getGroupKeyIterator() {
 
   /// Clear and trim DictionaryBasedGroupKeyGenerator after use
   public void closeGroupKeyGenerator() {
+    close();
+  }
+
+  @Override
+  public void close() {
     _groupKeyGenerator.close();
   }