Add non-blocking and partitioned group-by combine operators

Introduce two new server-level group-by combine algorithms selectable
via the `groupByAlgorithm` query option:

- NONBLOCKING: Uses a lock-free ConcurrentHashMap to aggregate results
  across segments without holding a global lock, reducing contention
  under high concurrency.

- PARTITIONED: Hash-partitions group keys across independent tables so
  that each partition can be aggregated and trimmed in parallel with
  zero cross-partition synchronization. Includes per-partition
  thread-local accumulation, cached key hashing, and parallel
  finish/trim to minimize overhead for high-cardinality queries.

Also adds:
- `normalizeGroupByAlgorithm` helper in QueryOptionsUtils
- `IndexedTable.transferFrom` for bulk partition merging
- JMH benchmark for comparing combine operator throughput
- Unit and integration tests for both new algorithms

Author: xiangfu0

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

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -427,6 +428,23 @@ public static Integer getMinInitialIndexedTableCapacity(Map<String, String> quer
     return checkedParseIntPositive(QueryOptionKey.MIN_INITIAL_INDEXED_TABLE_CAPACITY, minInitialIndexedTableCapacity);
   }
 
+  @Nullable
+  public static String getGroupByAlgorithm(Map<String, String> queryOptions) {
+    return queryOptions.get(QueryOptionKey.GROUP_BY_ALGORITHM);
+  }
+
+  @Nullable
+  public static String normalizeGroupByAlgorithm(@Nullable String groupByAlgorithm) {
+    String normalizedGroupByAlgorithm = StringUtils.trimToNull(groupByAlgorithm);
+    return normalizedGroupByAlgorithm != null ? normalizedGroupByAlgorithm.toUpperCase(Locale.ROOT) : null;
+  }
+
+  @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-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java

@@ -281,6 +281,13 @@ public void testInvertedIndexDistinctCostRatioRejectsNonFiniteValues() {
     }
   }
 
+  @Test
+  public void shouldNormalizeGroupByAlgorithm() {
+    assertEquals(QueryOptionsUtils.normalizeGroupByAlgorithm(" partitioned "), "PARTITIONED");
+    assertNull(QueryOptionsUtils.normalizeGroupByAlgorithm("   "));
+    assertNull(QueryOptionsUtils.normalizeGroupByAlgorithm(null));
+  }
+
   private static Object getValue(Map<String, String> map, String key) {
     switch (key) {
       // Positive ints

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

@@ -263,9 +263,53 @@ public int size() {
 
   @Override
   public Iterator<Record> iterator() {
+    if (_topRecords == null) {
+      return _lookupMap.values().iterator();
+    }
     return _topRecords.iterator();
   }
 
+
+  /**
+   * Merges all records from the given {@link Table} into this table.
+   * <p>
+   * When {@code table} is an {@link IndexedTable}, this method assumes that the two tables have disjoint key spaces,
+   * i.e. there are no group keys present in both tables. Callers must ensure this precondition holds (for example,
+   * by only merging tables from disjoint partitions).
+   */
+  public void mergePartitionTable(Table table) {
+    _topRecords = null;
+    if (table instanceof IndexedTable) {
+      IndexedTable indexedTable = (IndexedTable) table;
+      for (Map.Entry<Key, Record> entry : indexedTable._lookupMap.entrySet()) {
+        Record existing = _lookupMap.put(entry.getKey(), entry.getValue());
+        Preconditions.checkState(existing == null,
+            "Found overlapping key while merging partition tables; partition key spaces must be disjoint");
+      }
+      _numResizes += indexedTable._numResizes;
+      _resizeTimeNs += indexedTable._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 the resize/trim statistics from another {@link IndexedTable} into this table.
+   * Call this after merging tables to ensure resize metrics are accurately reported.
+   */
+  public void absorbTrimStats(IndexedTable other) {
+    _numResizes += other._numResizes;
+    _resizeTimeNs += other._resizeTimeNs;
+  }
+
   public int getNumResizes() {
     return _numResizes;
   }

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

@@ -53,20 +53,22 @@
  */
 @SuppressWarnings("rawtypes")
 public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<GroupByResultsBlock> {
+  public static final String ALGORITHM = "DEFAULT";
+
   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 +115,11 @@ protected void processSegments() {
         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;
-        }
-
-        // 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();
+              _indexedTable = createIndexedTable(resultsBlock, _numTasks);
             }
           }
-        } 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 +130,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);
@@ -189,13 +193,13 @@ public void onProcessSegmentsFinish() {
   /**
    * {@inheritDoc}
    *
-   * <p>Combines intermediate selection result blocks from underlying operators and returns a merged one.
+   * <p>Combines intermediate group-by results produced by underlying operators and returns a merged results block.
    * <ul>
    *   <li>
-   *     Merges multiple intermediate selection result blocks as a merged one.
+   *     Merges multiple intermediate group-by result blocks into a single merged result.
    *   </li>
    *   <li>
-   *     Set all exceptions encountered during execution into the merged result block
+   *     Sets all exceptions encountered during execution into the merged result block.
    *   </li>
    * </ul>
    */
@@ -205,33 +209,42 @@ 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) {
+    long reportedTimeoutMs = Math.max(timeoutMs, 0L);
+    String userError = "Timed out while combining group-by results after " + reportedTimeoutMs + "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()) {