Partitioned group-by combine with per-partition thread-local tables

Replace the default GroupByCombineOperator (shared ConcurrentHashMap)
with PartitionedGroupByCombineOperator that uses thread-local
per-partition IndexedTables, eliminating cross-thread contention.

Key changes:
- Thread-local partition tables published once under partition locks
- Parallel tree-reduction of partition tables in the final merge
- Cached Key.hashCode() at construction (O(1) thereafter)
- Object2ObjectOpenHashMap replaces HashMap in SimpleIndexedTable
- IndexedTable.merge() reuses existing Keys from source entry set
- mergePartitionTable() propagates trim stats from absorbed tables
- numGroupByPartitions query option controls partitioning
- Sorted combine path preserved for safe-trim with small limits

Benchmark (JMH, 32 segments, 8 threads, GROUP BY ORDER BY LIMIT 500):

10K records/segment:
  500 groups:  127.6ms -> 2.8ms  (45.6x faster)
  5K groups:   132.4ms -> 5.9ms  (22.4x faster)
  50K groups:  144.0ms -> 23.0ms (6.3x faster)
  500K groups: 184.8ms -> 76.9ms (2.4x faster)

100K records/segment (high cardinality):
  5M groups:   1,934ms -> 637ms  (3.0x faster)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: xiangfu0

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

@@ -409,6 +409,12 @@ public static Integer getMinInitialIndexedTableCapacity(Map<String, String> quer
     return checkedParseIntPositive(QueryOptionKey.MIN_INITIAL_INDEXED_TABLE_CAPACITY, minInitialIndexedTableCapacity);
   }
 
+  @Nullable
+  public static Integer getNumGroupByPartitions(Map<String, String> queryOptions) {
+    String numGroupByPartitionsString = queryOptions.get(QueryOptionKey.NUM_GROUP_BY_PARTITIONS);
+    return checkedParseIntPositive(QueryOptionKey.NUM_GROUP_BY_PARTITIONS, numGroupByPartitionsString);
+  }
+
   public static boolean shouldDropResults(Map<String, String> queryOptions) {
     return Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.QueryOptionKey.DROP_RESULTS));
   }

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

@@ -109,6 +109,23 @@ public boolean upsert(Record record) {
     return upsert(new Key(keyValues), record);
   }
 
+  @Override
+  public boolean merge(Table table) {
+    if (table instanceof IndexedTable) {
+      // Optimized path: iterate the source's entry set directly to reuse existing Key objects,
+      // avoiding per-record Arrays.copyOf + new Key + Arrays.hashCode overhead.
+      for (Map.Entry<Key, Record> entry : ((IndexedTable) table)._lookupMap.entrySet()) {
+        upsert(entry.getKey(), entry.getValue());
+      }
+    } else {
+      Iterator<Record> iterator = table.iterator();
+      while (iterator.hasNext()) {
+        upsert(iterator.next());
+      }
+    }
+    return true;
+  }
+
   /**
    * Adds a record with new key or updates a record with existing key.
    */
@@ -263,9 +280,49 @@ public int size() {
 
   @Override
   public Iterator<Record> iterator() {
+    if (_topRecords == null) {
+      return _lookupMap.values().iterator();
+    }
     return _topRecords.iterator();
   }
 
+
+  /**
+   * Merges records from a disjoint partition table into this table using {@code putAll}.
+   * <p><b>Precondition:</b> the source table must have no overlapping keys with this table
+   * (i.e. they come from disjoint hash partitions). If keys overlap, {@code putAll} will
+   * silently overwrite existing entries instead of aggregating them. Use {@link #merge(Table)}
+   * for tables that may share keys.
+   */
+  public void mergePartitionTable(Table table) {
+    _topRecords = null;
+    if (table instanceof IndexedTable) {
+      IndexedTable other = (IndexedTable) table;
+      _lookupMap.putAll(other._lookupMap);
+      _numResizes += other._numResizes;
+      _resizeTimeNs += other._resizeTimeNs;
+    } else {
+      Iterator<Record> iterator = table.iterator();
+      while (iterator.hasNext()) {
+        // NOTE: For performance concern, does not check the return value of the upsert(). Override this method if
+        //       upsert() can return false.
+        upsert(iterator.next());
+      }
+    }
+    if (_lookupMap.size() >= _trimThreshold) {
+      resize();
+    }
+  }
+
+  /**
+   * Absorbs trim/resize statistics from another table that was merged into this one.
+   * Call after {@link #merge(Table)} to preserve the absorbed table's trim history.
+   */
+  public void absorbResizeStats(IndexedTable other) {
+    _numResizes += other._numResizes;
+    _resizeTimeNs += other._resizeTimeNs;
+  }
+
   public int getNumResizes() {
     return _numResizes;
   }

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

@@ -38,9 +38,11 @@
  */
 public class Key implements Comparable<Key> {
   private final Object[] _values;
+  private final int _hashCode;
 
   public Key(Object[] values) {
     _values = values;
+    _hashCode = Arrays.hashCode(values);
   }
 
   public Object[] getValues() {
@@ -56,7 +58,7 @@ public boolean equals(Object o) {
 
   @Override
   public int hashCode() {
-    return Arrays.hashCode(_values);
+    return _hashCode;
   }
 
   @Override

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

@@ -18,23 +18,26 @@
  */
 package org.apache.pinot.core.data.table;
 
-import java.util.HashMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import java.util.concurrent.ExecutorService;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.core.query.request.context.QueryContext;
 
 
 /**
- * {@link Table} implementation for aggregating TableRecords based on combination of keys
+ * {@link Table} implementation for aggregating TableRecords based on combination of keys.
+ * Uses fastutil's {@link Object2ObjectOpenHashMap} for open-addressing, which provides better
+ * cache locality and avoids per-entry {@code Map.Entry} object allocations compared to
+ * {@code java.util.HashMap}.
  */
 @NotThreadSafe
 public class SimpleIndexedTable extends IndexedTable {
 
   public SimpleIndexedTable(DataSchema dataSchema, boolean hasFinalInput, QueryContext queryContext, int resultSize,
       int trimSize, int trimThreshold, int initialCapacity, ExecutorService executorService) {
-    super(dataSchema, hasFinalInput, queryContext, resultSize, trimSize, trimThreshold, new HashMap<>(initialCapacity),
-        executorService);
+    super(dataSchema, hasFinalInput, queryContext, resultSize, trimSize, trimThreshold,
+        new Object2ObjectOpenHashMap<>(initialCapacity), executorService);
   }
 
   /**

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

@@ -56,17 +56,17 @@ public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<Group
   private static final Logger LOGGER = LoggerFactory.getLogger(GroupByCombineOperator.class);
   private static final String EXPLAIN_NAME = "COMBINE_GROUP_BY";
 
-  private final int _numAggregationFunctions;
-  private final int _numGroupByExpressions;
-  private final int _numColumns;
+  protected final int _numAggregationFunctions;
+  protected final int _numGroupByExpressions;
+  protected final int _numColumns;
   // We use a CountDownLatch to track if all Futures are finished by the query timeout, and cancel the unfinished
   // _futures (try to interrupt the execution if it already started).
-  private final CountDownLatch _operatorLatch;
+  protected final CountDownLatch _operatorLatch;
 
-  private volatile IndexedTable _indexedTable;
-  private volatile boolean _groupsTrimmed;
-  private volatile boolean _numGroupsLimitReached;
-  private volatile boolean _numGroupsWarningLimitReached;
+  protected volatile IndexedTable _indexedTable;
+  protected volatile boolean _groupsTrimmed;
+  protected volatile boolean _numGroupsLimitReached;
+  protected volatile boolean _numGroupsWarningLimitReached;
 
   public GroupByCombineOperator(List<Operator> operators, QueryContext queryContext, ExecutorService executorService) {
     super(null, operators, overrideMaxExecutionThreads(queryContext, operators.size()), executorService);
@@ -113,59 +113,11 @@ protected void processSegments() {
         if (_indexedTable == null) {
           synchronized (this) {
             if (_indexedTable == null) {
-              _indexedTable = GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, _numTasks,
-                  _executorService);
+              _indexedTable = createIndexedTable(resultsBlock, _numTasks);
             }
           }
         }
-
-        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();
-              while (dicGroupKeyIterator.hasNext()) {
-                QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME);
-                GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
-                Object[] keys = groupKey._keys;
-                Object[] values = Arrays.copyOf(keys, _numColumns);
-                int groupId = groupKey._groupId;
-                for (int i = 0; i < _numAggregationFunctions; i++) {
-                  values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
-                }
-                _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);
-          }
-        }
+        mergeGroupByResultsBlock(_indexedTable, resultsBlock, EXPLAIN_NAME);
       } catch (RuntimeException e) {
         throw wrapOperatorException(operator, e);
       } finally {
@@ -176,6 +128,56 @@ protected void processSegments() {
     }
   }
 
+  protected IndexedTable createIndexedTable(GroupByResultsBlock resultsBlock, int numThreads) {
+    return GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, numThreads, _executorService);
+  }
+
+  protected void updateCombineResultsStats(GroupByResultsBlock resultsBlock) {
+    if (resultsBlock.isGroupsTrimmed()) {
+      _groupsTrimmed = true;
+    }
+    if (resultsBlock.isNumGroupsLimitReached()) {
+      _numGroupsLimitReached = true;
+    }
+    if (resultsBlock.isNumGroupsWarningLimitReached()) {
+      _numGroupsWarningLimitReached = true;
+    }
+  }
+
+  protected void mergeGroupByResultsBlock(IndexedTable indexedTable, GroupByResultsBlock resultsBlock,
+      String explainName) {
+    updateCombineResultsStats(resultsBlock);
+
+    Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
+    int mergedKeys = 0;
+    if (intermediateRecords == null) {
+      AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
+      if (aggregationGroupByResult != null) {
+        try {
+          Iterator<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
+          while (dicGroupKeyIterator.hasNext()) {
+            QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, explainName);
+            GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
+            Object[] keys = groupKey._keys;
+            Object[] values = Arrays.copyOf(keys, _numColumns);
+            int groupId = groupKey._groupId;
+            for (int i = 0; i < _numAggregationFunctions; i++) {
+              values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
+            }
+            indexedTable.upsert(new Key(keys), new Record(values));
+          }
+        } finally {
+          aggregationGroupByResult.closeGroupKeyGenerator();
+        }
+      }
+    } else {
+      for (IntermediateRecord intermediateResult : intermediateRecords) {
+        QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, explainName);
+        indexedTable.upsert(intermediateResult._key, intermediateResult._record);
+      }
+    }
+  }
+
   @Override
   public void onProcessSegmentsException(Throwable t) {
     _processingException.compareAndSet(null, t);
@@ -205,33 +207,41 @@ public BaseResultsBlock mergeResults()
     long timeoutMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
     boolean opCompleted = _operatorLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
     if (!opCompleted) {
-      // If this happens, the broker side should already timed out, just log the error and return
-      String userError = "Timed out while combining group-by order-by results after " + timeoutMs + "ms";
-      String logMsg = userError + ", queryContext = " + _queryContext;
-      LOGGER.error(logMsg);
-      return new ExceptionResultsBlock(new QueryErrorMessage(QueryErrorCode.EXECUTION_TIMEOUT, userError, logMsg));
+      return getTimeoutResultsBlock(timeoutMs);
     }
 
     Throwable ex = _processingException.get();
     if (ex != null) {
-      String userError = "Caught exception while processing group-by order-by query";
-      String devError = userError + ": " + ex.getMessage();
-      QueryErrorMessage errMsg;
-      if (ex instanceof QueryException) {
-        // If the exception is a QueryException, use the error code from the exception and trust the error message
-        errMsg = new QueryErrorMessage(((QueryException) ex).getErrorCode(), devError, devError);
-      } else {
-        // If the exception is not a QueryException, use the generic error code and don't expose the exception message
-        errMsg = new QueryErrorMessage(QueryErrorCode.QUERY_EXECUTION, userError, devError);
-      }
-      return new ExceptionResultsBlock(errMsg);
+      return getExceptionResultsBlock(ex);
     }
 
-    if (_indexedTable.isTrimmed() && _queryContext.isUnsafeTrim()) {
+    return getMergedResultsBlock(_indexedTable);
+  }
+
+  protected ExceptionResultsBlock getTimeoutResultsBlock(long timeoutMs) {
+    String userError = "Timed out while combining group-by results after " + Math.max(0, timeoutMs) + "ms";
+    String logMsg = userError + ", queryContext = " + _queryContext;
+    LOGGER.error(logMsg);
+    return new ExceptionResultsBlock(new QueryErrorMessage(QueryErrorCode.EXECUTION_TIMEOUT, userError, logMsg));
+  }
+
+  protected ExceptionResultsBlock getExceptionResultsBlock(Throwable ex) {
+    String userError = "Caught exception while processing group-by query";
+    String devError = userError + ": " + ex.getMessage();
+    QueryErrorMessage errMsg;
+    if (ex instanceof QueryException) {
+      errMsg = new QueryErrorMessage(((QueryException) ex).getErrorCode(), devError, devError);
+    } else {
+      errMsg = new QueryErrorMessage(QueryErrorCode.QUERY_EXECUTION, userError, devError);
+    }
+    return new ExceptionResultsBlock(errMsg);
+  }
+
+  protected GroupByResultsBlock getMergedResultsBlock(IndexedTable indexedTable) {
+    if (indexedTable.isTrimmed() && _queryContext.isUnsafeTrim()) {
       _groupsTrimmed = true;
     }
 
-    IndexedTable indexedTable = _indexedTable;
     if (_queryContext.isServerReturnFinalResult()) {
       indexedTable.finish(true, true);
     } else if (_queryContext.isServerReturnFinalResultKeyUnpartitioned()) {