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

@@ -427,6 +427,11 @@ 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);
+  }
+
   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

@@ -263,9 +263,22 @@ public int size() {
 
   @Override
   public Iterator<Record> iterator() {
+    if (_topRecords == null) {
+      return _lookupMap.values().iterator();
+    }
     return _topRecords.iterator();
   }
 
+
+  /**
+   * 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()) {

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

@@ -0,0 +1,113 @@
+/**
+ * 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.data.table.IndexedTable;
+import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
+import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Non-blocking combine operator for group-by queries.
+ * <p>
+ * Parallelism is bounded by the configured max execution threads via {@link GroupByCombineOperator}.
+ */
+@SuppressWarnings("rawtypes")
+public class NonblockingGroupByCombineOperator extends GroupByCombineOperator {
+  public static final String ALGORITHM = "NON-BLOCKING";
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(NonblockingGroupByCombineOperator.class);
+  private static final String EXPLAIN_NAME = "COMBINE_GROUP_BY";
+
+  public NonblockingGroupByCombineOperator(List<Operator> operators, QueryContext queryContext,
+      ExecutorService executorService) {
+    super(operators, queryContext, executorService);
+    LOGGER.debug("Using {} for group-by combine with {} tasks", ALGORITHM, _numTasks);
+  }
+
+  @Override
+  public String toExplainString() {
+    return EXPLAIN_NAME;
+  }
+
+  /**
+   * Executes query on one segment in a worker thread and merges the results into the indexed table.
+   */
+  @Override
+  protected void processSegments() {
+    int operatorId;
+    IndexedTable indexedTable = null;
+    while (_processingException.get() == null && (operatorId = _nextOperatorId.getAndIncrement()) < _numOperators) {
+      Operator operator = _operators.get(operatorId);
+      try {
+        if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
+          ((AcquireReleaseColumnsSegmentOperator) operator).acquire();
+        }
+        GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock();
+        if (indexedTable == null) {
+          synchronized (this) {
+            if (_indexedTable != null) {
+              indexedTable = _indexedTable;
+              _indexedTable = null;
+            }
+          }
+          if (indexedTable == null) {
+            indexedTable = createIndexedTable(resultsBlock, 1);
+          }
+        }
+        mergeGroupByResultsBlock(indexedTable, resultsBlock, EXPLAIN_NAME);
+      } catch (RuntimeException e) {
+        throw wrapOperatorException(operator, e);
+      } finally {
+        if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
+          ((AcquireReleaseColumnsSegmentOperator) operator).release();
+        }
+      }
+    }
+
+    boolean setGroupByResult = false;
+    while (indexedTable != null && !setGroupByResult) {
+      IndexedTable indexedTableToMerge;
+      synchronized (this) {
+        if (_indexedTable == null) {
+          _indexedTable = indexedTable;
+          setGroupByResult = true;
+          continue;
+        } else {
+          indexedTableToMerge = _indexedTable;
+          _indexedTable = null;
+        }
+      }
+      if (indexedTable.size() > indexedTableToMerge.size()) {
+        indexedTable.merge(indexedTableToMerge);
+        indexedTable.absorbTrimStats(indexedTableToMerge);
+      } else {
+        indexedTableToMerge.merge(indexedTable);
+        indexedTableToMerge.absorbTrimStats(indexedTable);
+        indexedTable = indexedTableToMerge;
+      }
+    }
+  }
+}