Add opt-in off-heap group-by key tables and result holders for SSE

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -586,6 +586,12 @@ public static Integer getNumGroupsWarningLimit(Map<String, String> queryOptions)
     return checkedParseIntPositive(QueryOptionKey.NUM_GROUPS_WARNING_LIMIT, numGroupsWarningLimit);
   }
 
+  @Nullable
+  public static Boolean isGroupByOffHeap(Map<String, String> queryOptions) {
+    String groupByOffHeap = queryOptions.get(QueryOptionKey.GROUP_BY_OFF_HEAP);
+    return groupByOffHeap != null ? Boolean.parseBoolean(groupByOffHeap) : null;
+  }
+
   @Nullable
   public static Integer getMaxInitialResultHolderCapacity(Map<String, String> queryOptions) {
     String maxInitialResultHolderCapacity = queryOptions.get(QueryOptionKey.MAX_INITIAL_RESULT_HOLDER_CAPACITY);

pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java

@@ -329,7 +329,7 @@ private Collection<Record> getUnsortedTopRecords(Map<Key, Record> recordsMap, in
   /// This method is to be called from individual segment if the intermediate results need to be trimmed.
   public List<IntermediateRecord> sortInSegmentResults(GroupKeyGenerator groupKeyGenerator,
       GroupByResultHolder[] groupByResultHolders, int size) {
-    // getNumKeys() does not count nulls
+    // NOTE: getNumKeys() counts every group, including the null group when null handling is enabled
     assert groupKeyGenerator.getNumKeys() <= size;
     Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = groupKeyGenerator.getGroupKeys();
 

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java

@@ -106,38 +106,39 @@ 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);
+        // Hold the group-by result so its group key generator (which may own off-heap resources) is always
+        // released, even when indexed-table creation or the merge below throws
+        AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
+        try {
+          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 {
+          // 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) {
+            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);
@@ -150,16 +151,18 @@ 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);
+        } finally {
+          if (aggregationGroupByResult != null) {
+            // Release the resources used by the group key generator
+            aggregationGroupByResult.closeGroupKeyGenerator();
           }
         }
       } catch (RuntimeException e) {

pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java

@@ -135,6 +135,23 @@ protected GroupByResultsBlock getNextBlock() {
       resultHolderIndexMap.put(_aggregationFunctions[i], i);
     }
 
+    GroupKeyGenerator[] createdGroupKeyGenerator = new GroupKeyGenerator[1];
+    try {
+      return processAndBuildResultsBlock(groupByResultHolders, resultHolderIndexMap, createdGroupKeyGenerator);
+    } catch (Throwable t) {
+      // Release group-by resources (including off-heap key tables and result holders) that would otherwise leak.
+      // Close is idempotent on all generators; on the success path the generator is either closed on the trim/sort
+      // paths below or handed to the combine operator, which closes it after the merge.
+      if (createdGroupKeyGenerator[0] != null) {
+        createdGroupKeyGenerator[0].close();
+      }
+      throw t;
+    }
+  }
+
+  private GroupByResultsBlock processAndBuildResultsBlock(GroupByResultHolder[] groupByResultHolders,
+      IdentityHashMap<AggregationFunction, Integer> resultHolderIndexMap,
+      GroupKeyGenerator[] createdGroupKeyGenerator) {
     GroupKeyGenerator groupKeyGenerator = null;
     for (AggregationInfo aggregationInfo : _aggregationInfos) {
       AggregationFunction[] aggregationFunctions = aggregationInfo.getFunctions();
@@ -160,6 +177,7 @@ protected GroupByResultsBlock getNextBlock() {
       // GroupByExecutor with a pre-existing GroupKeyGenerator so that the GroupKeyGenerator can be shared across
       // loop iterations i.e. across all aggs.
       groupKeyGenerator = groupByExecutor.getGroupKeyGenerator();
+      createdGroupKeyGenerator[0] = groupKeyGenerator;
 
       int numDocsScanned = 0;
       ValueBlock valueBlock;

pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java

@@ -125,6 +125,19 @@ protected GroupByResultsBlock getNextBlock() {
     } else {
       groupByExecutor = new DefaultGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator);
     }
+    try {
+      return processAndBuildResultsBlock(groupByExecutor);
+    } catch (Throwable t) {
+      // Release group-by resources (including off-heap key tables and result holders) that would otherwise leak.
+      // On the success path, ownership either ends inside processAndBuildResultsBlock (trim/sort paths close the
+      // generator there) or moves to the results block consumer (the combine operator closes the generator after
+      // merging the AggregationGroupByResult). Close is idempotent on all generators.
+      groupByExecutor.getGroupKeyGenerator().close();
+      throw t;
+    }
+  }
+
+  private GroupByResultsBlock processAndBuildResultsBlock(GroupByExecutor groupByExecutor) {
     ValueBlock valueBlock;
 
     while ((valueBlock = _projectOperator.nextBlock()) != null) {

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.offheap.OffHeapGroupByBufferPool;
 import org.apache.pinot.core.query.executor.ResultsBlockStreamer;
 import org.apache.pinot.core.query.prefetch.FetchPlanner;
 import org.apache.pinot.core.query.prefetch.FetchPlannerRegistry;
@@ -113,6 +114,8 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
   private int _minSegmentGroupTrimSize = Server.DEFAULT_QUERY_EXECUTOR_MIN_SEGMENT_GROUP_TRIM_SIZE;
   private int _minServerGroupTrimSize = Server.DEFAULT_QUERY_EXECUTOR_MIN_SERVER_GROUP_TRIM_SIZE;
   private int _groupByTrimThreshold = Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_TRIM_THRESHOLD;
+  // Whether to store group-by key tables and fixed-width result holders in off-heap (direct) memory
+  private boolean _groupByOffHeap = Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP;
 
   @Override
   public void init(PinotConfiguration queryExecutorConfig) {
@@ -143,11 +146,16 @@ public void init(PinotConfiguration queryExecutorConfig) {
         Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_TRIM_THRESHOLD);
     Preconditions.checkState(_groupByTrimThreshold > 0,
         "Invalid configurable: groupByTrimThreshold: %d must be positive", _groupByTrimThreshold);
+    _groupByOffHeap =
+        queryExecutorConfig.getProperty(Server.GROUPBY_OFF_HEAP, Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP);
+    OffHeapGroupByBufferPool.setMaxBytesPerThread(
+        queryExecutorConfig.getProperty(Server.GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD,
+            Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD));
     LOGGER.info("Initialized plan maker with maxExecutionThreads: {}, defaultExecutionThreads: {}, "
             + "maxInitialResultHolderCapacity: {}, numGroupsLimit: {}, minSegmentGroupTrimSize: {}, "
-            + "minServerGroupTrimSize: {}, groupByTrimThreshold: {}",
+            + "minServerGroupTrimSize: {}, groupByTrimThreshold: {}, groupByOffHeap: {}",
         _maxExecutionThreads, _defaultExecutionThreads, _maxInitialResultHolderCapacity, _numGroupsLimit,
-        _minSegmentGroupTrimSize, _minServerGroupTrimSize, _groupByTrimThreshold);
+        _minSegmentGroupTrimSize, _minServerGroupTrimSize, _groupByTrimThreshold, _groupByOffHeap);
   }
 
   @VisibleForTesting
@@ -200,6 +208,11 @@ public void setGroupByTrimThreshold(int groupByTrimThreshold) {
     _groupByTrimThreshold = groupByTrimThreshold;
   }
 
+  @VisibleForTesting
+  public void setGroupByOffHeap(boolean groupByOffHeap) {
+    _groupByOffHeap = groupByOffHeap;
+  }
+
   @Override
   public Plan makeInstancePlan(List<SegmentContext> segmentContexts, QueryContext queryContext,
       ExecutorService executorService) {
@@ -292,6 +305,9 @@ void applyQueryOptions(QueryContext queryContext) {
       } else {
         queryContext.setNumGroupsLimit(_numGroupsLimit);
       }
+      // Set groupByOffHeap
+      Boolean groupByOffHeap = QueryOptionsUtils.isGroupByOffHeap(queryOptions);
+      queryContext.setGroupByOffHeap(groupByOffHeap != null ? groupByOffHeap : _groupByOffHeap);
       // Set numGroupsWarningThreshold
       queryContext.setNumGroupsWarningLimit(_numGroupsWarningLimit);
       // Set minSegmentGroupTrimSize

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java

@@ -37,6 +37,10 @@
 import org.apache.pinot.core.plan.DocIdSetPlanNode;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils;
+import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapDoubleGroupByResultHolder;
+import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapIntGroupByResultHolder;
+import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapLongGroupByResultHolder;
+import org.apache.pinot.core.query.aggregation.groupby.offheap.ResourceTrackingGroupKeyGenerator;
 import org.apache.pinot.core.query.request.context.QueryContext;
 
 
@@ -102,36 +106,60 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a
     if (queryContext.isOptimizeMaxInitialResultHolderCapacity()) {
       groupByExpressionSizesFromPredicates = getGroupByExpressionSizesFromPredicates(queryContext);
     }
+    // Off-heap group-by is not enabled for grouping sets yet: GroupingSetsGroupKeyGenerator keeps its key map and
+    // on-the-fly dictionaries on heap, so only the fixed-width result holders could move off-heap, and that
+    // combination is untested. The close plumbing already covers grouping sets (the trim path closes the generator
+    // in GroupByUtils.buildGroupingSetsResultsBlock, and the combine operators close the AggregationGroupByResult's
+    // generator), so enabling it later mainly requires off-heap key storage in that generator plus test coverage.
+    boolean groupByOffHeap = queryContext.isGroupByOffHeap() && !groupingSets;
     if (groupKeyGenerator != null) {
+      // Shared generator (filtered aggregations): if the first executor created it in off-heap mode, it is already
+      // wrapped in a ResourceTrackingGroupKeyGenerator, and this executor registers its holders on the same wrapper
       _groupKeyGenerator = groupKeyGenerator;
-    } else if (groupingSets) {
-      _groupKeyGenerator = new GroupingSetsGroupKeyGenerator(projectOperator, groupByExpressions,
-          queryContext.getGroupingSets(), numGroupsLimit, _nullHandlingEnabled);
     } else {
-      if (hasNoDictionaryGroupByExpression || _nullHandlingEnabled) {
+      GroupKeyGenerator generator;
+      if (groupingSets) {
+        generator = new GroupingSetsGroupKeyGenerator(projectOperator, groupByExpressions,
+            queryContext.getGroupingSets(), numGroupsLimit, _nullHandlingEnabled);
+      } else if (hasNoDictionaryGroupByExpression || _nullHandlingEnabled) {
         if (groupByExpressions.length == 1) {
           // TODO(nhejazi): support MV and dictionary based when null handling is enabled.
-          _groupKeyGenerator =
+          generator =
               new NoDictionarySingleColumnGroupKeyGenerator(projectOperator, groupByExpressions[0], numGroupsLimit,
-                  _nullHandlingEnabled, groupByExpressionSizesFromPredicates);
+                  _nullHandlingEnabled, groupByExpressionSizesFromPredicates, groupByOffHeap);
         } else {
-          _groupKeyGenerator =
+          generator =
               new NoDictionaryMultiColumnGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit,
-                  _nullHandlingEnabled, groupByExpressionSizesFromPredicates);
+                  _nullHandlingEnabled, groupByExpressionSizesFromPredicates, groupByOffHeap);
         }
       } else {
-        _groupKeyGenerator = new DictionaryBasedGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit,
-            maxInitialResultHolderCapacity, groupByExpressionSizesFromPredicates);
+        generator = new DictionaryBasedGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit,
+            maxInitialResultHolderCapacity, groupByExpressionSizesFromPredicates, groupByOffHeap);
       }
+      _groupKeyGenerator = groupByOffHeap ? new ResourceTrackingGroupKeyGenerator(generator) : generator;
     }
 
-    // Initialize result holders
+    // Initialize result holders. In off-heap mode, fixed-width holders are mirrored off-heap and registered on the
+    // resource-tracking generator so the existing generator close() call sites release them.
+    ResourceTrackingGroupKeyGenerator offHeapResourceTracker =
+        _groupKeyGenerator instanceof ResourceTrackingGroupKeyGenerator
+            ? (ResourceTrackingGroupKeyGenerator) _groupKeyGenerator : null;
     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);
+    try {
+      for (int i = 0; i < numAggregationFunctions; i++) {
+        _groupByResultHolders[i] = offHeapResourceTracker != null
+            ? createOffHeapCapableResultHolder(_aggregationFunctions[i], initialCapacity, maxNumResults,
+                offHeapResourceTracker)
+            : _aggregationFunctions[i].createGroupByResultHolder(initialCapacity, maxNumResults);
+      }
+    } catch (Throwable t) {
+      // Holder creation failed midway: release the generator (and any off-heap holders already registered on it)
+      // because the caller never gets an executor reference to clean up. Close is idempotent.
+      _groupKeyGenerator.close();
+      throw t;
     }
 
     // Initialize map from document Id to group key
@@ -144,6 +172,31 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a
     }
   }
 
+  /// Mirrors fixed-width result holders off-heap. The holder type and default value are discovered through a
+  /// zero-capacity probe (aggregation functions choose both — createGroupByResultHolder must stay side-effect-free
+  /// for the probe to be safe), and any non-fixed-width holder (object holders, dummy
+  /// holders, custom implementations) is recreated on-heap with the real initial capacity. Off-heap holders are
+  /// registered on the resource tracker, which releases them when the group key generator is closed.
+  private static GroupByResultHolder createOffHeapCapableResultHolder(AggregationFunction<?, ?> function,
+      int initialCapacity, int maxCapacity, ResourceTrackingGroupKeyGenerator resourceTracker) {
+    GroupByResultHolder probe = function.createGroupByResultHolder(0, maxCapacity);
+    GroupByResultHolder holder;
+    if (probe.getClass() == DoubleGroupByResultHolder.class) {
+      holder = new OffHeapDoubleGroupByResultHolder(initialCapacity, maxCapacity,
+          ((DoubleGroupByResultHolder) probe).getDefaultValue());
+    } else if (probe.getClass() == LongGroupByResultHolder.class) {
+      holder = new OffHeapLongGroupByResultHolder(initialCapacity, maxCapacity,
+          ((LongGroupByResultHolder) probe).getDefaultValue());
+    } else if (probe.getClass() == IntGroupByResultHolder.class) {
+      holder = new OffHeapIntGroupByResultHolder(initialCapacity, maxCapacity,
+          ((IntGroupByResultHolder) probe).getDefaultValue());
+    } else {
+      return function.createGroupByResultHolder(initialCapacity, maxCapacity);
+    }
+    resourceTracker.register((AutoCloseable) holder);
+    return holder;
+  }
+
   /// Retrieve the sizes of GroupBy expressions from IN an EQ predicates found in the filter context, if available.
   /// 1. If the filter context is null or lacks GroupBy expressions, return null.
   /// 2. Ensure the top-level filter context consists solely of AND-type filters; other types for example OR we cannot