Replace default group-by combine with non-blocking per-thread operator (CONCURRENT fallback)

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

@@ -450,6 +450,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

@@ -155,6 +155,16 @@ protected void resize() {
     _resizeTimeNs += resizeTimeNs;
   }
 
+  /**
+   * Trims this table to at most {@code trimSize} entries if it has an ORDER BY clause. No-op otherwise.
+   * Intended for coordinated trim across multiple tables (e.g., in a partitioned combine).
+   */
+  public void trim() {
+    if (_hasOrderBy && !_lookupMap.isEmpty()) {
+      resize();
+    }
+  }
+
   @Override
   public void finish(boolean sort, boolean storeFinalResult) {
     if (_hasOrderBy) {
@@ -263,11 +273,34 @@ public int size() {
     return _topRecords != null ? _topRecords.size() : _lookupMap.size();
   }
 
+  /**
+   * Returns an iterator over all records in this table.
+   * <p>
+   * If {@link #finish} has already been called, iterates over the trimmed/sorted top records. Otherwise iterates
+   * directly over the internal lookup map. In the pre-finish case the caller is responsible for ensuring no
+   * concurrent modifications are made to this table while iterating.
+   */
   @Override
   public Iterator<Record> iterator() {
+    if (_topRecords == null) {
+      return _lookupMap.values().iterator();
+    }
     return _topRecords.iterator();
   }
 
+  /**
+   * Accumulates the resize/trim statistics from {@code other} into this table.
+   * <p>
+   * This method is <em>not</em> thread-safe. It must only be called when both tables are exclusively owned by the
+   * calling thread (i.e., neither table is concurrently modified or iterated).
+   *
+   * @param other the table whose stats are absorbed; must be exclusively owned by the caller
+   */
+  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/data/table/Key.java

@@ -38,6 +38,11 @@
  */
 public class Key implements Comparable<Key> {
   private final Object[] _values;
+  // Lazily cached hash code. Declared volatile so that threads reading a Key from a shared map (e.g.
+  // ConcurrentHashMap in the CONCURRENT algorithm) observe the cached value written by the inserting thread.
+  // 0 means not yet computed; if Arrays.hashCode returns 0 the field stays 0 and is recomputed on every
+  // call (extremely rare, not a correctness concern).
+  private volatile int _hashCode;
 
   public Key(Object[] values) {
     _values = values;
@@ -56,7 +61,12 @@ public boolean equals(Object o) {
 
   @Override
   public int hashCode() {
-    return Arrays.hashCode(_values);
+    int h = _hashCode;
+    if (h == 0) {
+      h = Arrays.hashCode(_values);
+      _hashCode = h;
+    }
+    return h;
   }
 
   @Override

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

@@ -53,23 +53,57 @@
  */
 @SuppressWarnings("rawtypes")
 public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<GroupByResultsBlock> {
+  public static final String ALGORITHM = "CONCURRENT";
+
   private static final Logger LOGGER = LoggerFactory.getLogger(GroupByCombineOperator.class);
-  private static final String EXPLAIN_NAME = "COMBINE_GROUP_BY";
+  protected static final String EXPLAIN_NAME = "COMBINE_GROUP_BY";
 
-  private final int _numAggregationFunctions;
+  protected final int _numAggregationFunctions;
   /// Number of key columns: union group-by columns plus the synthetic $groupingId column for grouping sets.
   /// Key columns precede the aggregation columns in the record layout.
-  private final int _numKeyColumns;
-  private final int _numColumns;
+  protected final int _numKeyColumns;
+  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;
 
+  // Shared indexed table for the single-slot tournament exchange. Written exactly once (null → non-null) under
+  // synchronized(this) via DCL. Subsequent reads are unsynchronized volatile reads, which is safe because the field
+  // never reverts to null in this class. Subclasses that call stealSharedTable() / depositSharedTable() must do so
+  // under synchronized(this).
   private volatile IndexedTable _indexedTable;
+  // Written by multiple worker threads via updateCombineResultsStats(). Only transitions false→true,
+  // so concurrent volatile writes are safe without additional synchronization.
   private volatile boolean _groupsTrimmed;
   private volatile boolean _numGroupsLimitReached;
   private volatile boolean _numGroupsWarningLimitReached;
 
+  /**
+   * Atomically steals the shared indexed table (setting it to null) and returns it.
+   * Returns null if no shared table is available. Must be called under {@code synchronized(this)}.
+   */
+  protected IndexedTable stealSharedTable() {
+    IndexedTable table = _indexedTable;
+    _indexedTable = null;
+    return table;
+  }
+
+  /**
+   * Deposits the given indexed table as the shared result.
+   * Must be called under {@code synchronized(this)}.
+   */
+  protected void depositSharedTable(IndexedTable table) {
+    _indexedTable = table;
+  }
+
+  /**
+   * Returns true if the shared indexed table is non-null.
+   * Must be called under {@code synchronized(this)}.
+   */
+  protected boolean hasSharedTable() {
+    return _indexedTable != null;
+  }
+
   public GroupByCombineOperator(List<Operator> operators, QueryContext queryContext, ExecutorService executorService) {
     super(null, operators, overrideMaxExecutionThreads(queryContext, operators.size()), executorService);
 
@@ -115,59 +149,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[_numKeyColumns + 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 {
@@ -178,6 +164,67 @@ protected void processSegments() {
     }
   }
 
+  protected IndexedTable createIndexedTable(GroupByResultsBlock resultsBlock, int numThreads) {
+    return GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, numThreads, _executorService);
+  }
+
+  protected IndexedTable createIndexedTable(GroupByResultsBlock resultsBlock, int numThreads, int numGroupsHint) {
+    return GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, numThreads, _executorService,
+        numGroupsHint);
+  }
+
+  protected IndexedTable createIndexedTable(GroupByResultsBlock resultsBlock, int numThreads, int numGroupsHint,
+      int groupTrimThresholdOverride) {
+    return GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, numThreads, _executorService,
+        numGroupsHint, groupTrimThresholdOverride);
+  }
+
+  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[_numKeyColumns + 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);
@@ -188,52 +235,78 @@ public void onProcessSegmentsFinish() {
     _operatorLatch.countDown();
   }
 
+  /**
+   * Waits for all worker tasks to finish, up to the given timeout. Returns {@code true} if all tasks completed before
+   * the timeout expired, {@code false} otherwise.
+   */
+  protected boolean awaitOperatorLatch(long timeoutMs)
+      throws InterruptedException {
+    return _operatorLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
+  }
+
   /**
    * {@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>
    */
   @Override
   public BaseResultsBlock mergeResults()
       throws Exception {
     long timeoutMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
-    boolean opCompleted = _operatorLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
+    boolean opCompleted = awaitOperatorLatch(timeoutMs);
     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 == null) {
+      // No segments were processed (e.g. _numOperators == 0 or all blocks contained no groups).
+      // This should not happen in normal execution but is guarded defensively.
+      String msg = "No groups produced by combine operator (indexedTable is null)";
+      LOGGER.warn(msg);
+      return new ExceptionResultsBlock(new QueryErrorMessage(QueryErrorCode.QUERY_EXECUTION, msg, msg));
+    }
+
+    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);
+  }
 
-    if (_indexedTable.isTrimmed() && _queryContext.isUnsafeTrim()) {
+  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()) {