Add streaming k-way merge to SortedMailboxReceiveOperator

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

@@ -391,6 +391,26 @@ public static Integer getStreamingGroupByFlushThreshold(Map<String, String> quer
     return checkedParseIntNonNegative(QueryOptionKey.STREAMING_GROUP_BY_FLUSH_THRESHOLD, value);
   }
 
+  public static boolean isSortedSelectionMergeEnabled(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_ENABLED));
+  }
+
+  @Nullable
+  public static Integer getSortedSelectionMergeBlockSize(Map<String, String> queryOptions) {
+    String value = queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_BLOCK_SIZE);
+    return checkedParseIntPositive(QueryOptionKey.SORTED_SELECTION_MERGE_BLOCK_SIZE, value);
+  }
+
+  public static boolean isStreamingSortedMailboxReceiveEnabled(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE));
+  }
+
+  @Nullable
+  public static Integer getStreamingSortedMailboxReceiveBlockSize(Map<String, String> queryOptions) {
+    String value = queryOptions.get(QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE);
+    return checkedParseIntPositive(QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, value);
+  }
+
   public static boolean isNullHandlingEnabled(Map<String, String> queryOptions) {
     return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.ENABLE_NULL_HANDLING));
   }

pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java

@@ -38,7 +38,8 @@
 public class QueryOptionsUtilsTest {
   private static final List<String> POSITIVE_INT_KEYS =
       List.of(NUM_REPLICA_GROUPS_TO_QUERY, MAX_EXECUTION_THREADS, NUM_GROUPS_LIMIT, MAX_INITIAL_RESULT_HOLDER_CAPACITY,
-          MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW);
+          MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW, SORTED_SELECTION_MERGE_BLOCK_SIZE,
+          STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE);
   private static final List<String> NON_NEGATIVE_INT_KEYS = List.of(MULTI_STAGE_LEAF_LIMIT);
   private static final List<String> UNBOUNDED_INT_KEYS =
       List.of(MIN_SEGMENT_GROUP_TRIM_SIZE, MIN_SERVER_GROUP_TRIM_SIZE, MIN_BROKER_GROUP_TRIM_SIZE,
@@ -303,6 +304,21 @@ public void testInvertedIndexDistinctCostRatioRejectsNonFiniteValues() {
     }
   }
 
+  @Test
+  public void testStreamingSortedMailboxReceiveEnabled() {
+    // Unset and any non-"true" value are equivalent: the k-way merge stays off.
+    assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(Map.of()));
+    assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(new HashMap<>()));
+    assertTrue(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(
+        Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "true")));
+    assertTrue(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(
+        Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "TRUE")));
+    assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(
+        Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "false")));
+    assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(
+        Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "1")));
+  }
+
   private static Object getValue(Map<String, String> map, String key) {
     switch (key) {
       // Positive ints
@@ -320,6 +336,10 @@ private static Object getValue(Map<String, String> map, String key) {
         return QueryOptionsUtils.getMaxRowsInJoin(map);
       case MAX_ROWS_IN_WINDOW:
         return QueryOptionsUtils.getMaxRowsInWindow(map);
+      case SORTED_SELECTION_MERGE_BLOCK_SIZE:
+        return QueryOptionsUtils.getSortedSelectionMergeBlockSize(map);
+      case STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE:
+        return QueryOptionsUtils.getStreamingSortedMailboxReceiveBlockSize(map);
       // Non-negative ints
       case MULTI_STAGE_LEAF_LIMIT:
         return QueryOptionsUtils.getMultiStageLeafLimit(map);

pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java

@@ -34,6 +34,7 @@
 import org.apache.pinot.core.operator.combine.SelectionOrderByCombineOperator;
 import org.apache.pinot.core.operator.combine.SequentialSortedGroupByCombineOperator;
 import org.apache.pinot.core.operator.combine.SortedGroupByCombineOperator;
+import org.apache.pinot.core.operator.combine.StreamingSelectionOrderByCombineOperator;
 import org.apache.pinot.core.operator.streaming.StreamingGroupByCombineOperator;
 import org.apache.pinot.core.operator.streaming.StreamingSelectionOnlyCombineOperator;
 import org.apache.pinot.core.query.executor.ResultsBlockStreamer;
@@ -134,6 +135,17 @@ private BaseCombineOperator getCombineOperator() {
         // Use streaming operator only for non-empty selection-only query
         return new StreamingSelectionOnlyCombineOperator(operators, _queryContext, _executorService);
       }
+      // Streaming selection order-by (opt-in via the sortedSelectionMergeEnabled hint). Selection-only already
+      // returned above, so reaching here with a non-empty limit and an order-by present implies selection order-by.
+      if (_queryContext.isSortedSelectionMergeEnabled() && QueryContextUtils.isSelectionQuery(_queryContext)
+          && _queryContext.getLimit() != 0) {
+        List<OrderByExpressionContext> orderByExpressions = _queryContext.getOrderByExpressions();
+        if (orderByExpressions != null
+            && orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
+          return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService,
+              true /* streaming */);
+        }
+      }
       int flushThreshold = _queryContext.getStreamingGroupByFlushThreshold();
       if (flushThreshold > 0 && QueryContextUtils.isAggregationQuery(_queryContext)
           && _queryContext.getGroupByExpressions() != null) {
@@ -165,6 +177,10 @@ private BaseCombineOperator getCombineOperator() {
         List<OrderByExpressionContext> orderByExpressions = _queryContext.getOrderByExpressions();
         assert orderByExpressions != null;
         if (orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
+          if (_queryContext.isSortedSelectionMergeEnabled()) {
+            return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService,
+                false /* streaming */);
+          }
           return new MinMaxValueBasedSelectionOrderByCombineOperator(operators, _queryContext, _executorService);
         } else {
           return new SelectionOrderByCombineOperator(operators, _queryContext, _executorService);

pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java

@@ -32,6 +32,7 @@
 import org.apache.pinot.core.operator.query.SelectionOrderByOperator;
 import org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByDescOperation;
 import org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByLinearOperator;
+import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator;
 import org.apache.pinot.core.query.request.context.QueryContext;
 import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
 import org.apache.pinot.segment.spi.IndexSegment;
@@ -90,11 +91,36 @@ public Operator<SelectionResultsBlock> run() {
         maxDocsPerCall = Math.min(limit + _queryContext.getOffset(), DocIdSetPlanNode.MAX_DOC_PER_CALL);
       }
 
-      BaseProjectOperator<?> projectOperator = getSortedByProject(expressions, maxDocsPerCall, orderByExpressions);
       boolean asc = orderByExpressions.get(0).isAsc();
       // Remember that we cannot use asc == projectOperator.isAscending() because empty operators are considered
       // both ascending and descending
       DocIdOrderedOperator.DocIdOrder queryOrder = DocIdOrderedOperator.DocIdOrder.fromAsc(asc);
+
+      // Opt-in streaming path: emit one globally-sorted block at a time so a downstream k-way-merge combine can pull
+      // lazily. Only build it when the first order-by column is an identifier (kept consistent with the combine-side
+      // gate) and the forward-scan project is order-compatible; the DESC-incompatible sorted case still falls back to
+      // the materialized SelectionPartiallyOrderedByDescOperation below so global order stays correct.
+      if (_queryContext.isSortedSelectionMergeEnabled()
+          && orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
+        // When there are non-order-by output expressions, only fetch the order-by expressions during the forward scan
+        // (the streaming operator fetches the rest in a second pass); otherwise fetch all expressions.
+        List<ExpressionContext> projectExpressions = expressions;
+        if (expressions.size() > numOrderByExpressions) {
+          projectExpressions = new ArrayList<>(numOrderByExpressions);
+          for (OrderByExpressionContext orderByExpression : orderByExpressions) {
+            projectExpressions.add(orderByExpression.getExpression());
+          }
+        }
+        BaseProjectOperator<?> streamingProjectOperator =
+            getSortedByProject(projectExpressions, maxDocsPerCall, orderByExpressions);
+        if (streamingProjectOperator.isCompatibleWith(queryOrder)) {
+          return new StreamingSelectionOrderByOperator(_indexSegment, _queryContext, expressions,
+              streamingProjectOperator, sortedColumnsPrefixSize);
+        }
+        // DESC-incompatible: fall through to the materialized fallback (rebuilds the project over all expressions).
+      }
+
+      BaseProjectOperator<?> projectOperator = getSortedByProject(expressions, maxDocsPerCall, orderByExpressions);
       if (projectOperator.isCompatibleWith(queryOrder)) {
         return new SelectionPartiallyOrderedByLinearOperator(_indexSegment, _queryContext, expressions, projectOperator,
             sortedColumnsPrefixSize);

pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java

@@ -271,6 +271,17 @@ void applyQueryOptions(QueryContext queryContext) {
     }
     queryContext.setMaxExecutionThreads(maxExecutionThreads);
 
+    // Set streaming selection order-by options (opt-in; gated on selection queries to prevent accidental routing
+    // if a downstream guard is ever missed)
+    if (QueryContextUtils.isSelectionQuery(queryContext)) {
+      queryContext.setSortedSelectionMergeEnabled(QueryOptionsUtils.isSortedSelectionMergeEnabled(queryOptions));
+      Integer sortedSelectionMergeBlockSize =
+          QueryOptionsUtils.getSortedSelectionMergeBlockSize(queryOptions);
+      if (sortedSelectionMergeBlockSize != null) {
+        queryContext.setSortedSelectionMergeBlockSize(sortedSelectionMergeBlockSize);
+      }
+    }
+
     // Set group-by query options
     if (QueryContextUtils.isAggregationQuery(queryContext) && queryContext.getGroupByExpressions() != null) {
       // Set maxInitialResultHolderCapacity

pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java

@@ -44,6 +44,7 @@
 import org.apache.pinot.segment.spi.datasource.DataSource;
 import org.apache.pinot.spi.config.table.FieldConfig;
 import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.CommonConstants.Broker;
 import org.apache.pinot.spi.utils.CommonConstants.Server;
 
 
@@ -143,6 +144,11 @@ public class QueryContext {
   private int _effectiveSegmentGroupTrimSize;
   // Flush threshold for streaming group-by (0 = disabled)
   private int _streamingGroupByFlushThreshold;
+  /// Opt-in: use the streaming k-way-merge selection ORDER BY combine over sorted segments
+  private boolean _sortedSelectionMergeEnabled;
+  /// Output block size (rows) for the streaming selection ORDER BY combine
+  private int _sortedSelectionMergeBlockSize = Broker.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE;
+
   // Whether null handling is enabled
   private boolean _nullHandlingEnabled;
   // Whether server returns the final result
@@ -547,6 +553,22 @@ public void setStreamingGroupByFlushThreshold(int streamingGroupByFlushThreshold
     _streamingGroupByFlushThreshold = streamingGroupByFlushThreshold;
   }
 
+  public boolean isSortedSelectionMergeEnabled() {
+    return _sortedSelectionMergeEnabled;
+  }
+
+  public void setSortedSelectionMergeEnabled(boolean sortedSelectionMergeEnabled) {
+    _sortedSelectionMergeEnabled = sortedSelectionMergeEnabled;
+  }
+
+  public int getSortedSelectionMergeBlockSize() {
+    return _sortedSelectionMergeBlockSize;
+  }
+
+  public void setSortedSelectionMergeBlockSize(int sortedSelectionMergeBlockSize) {
+    _sortedSelectionMergeBlockSize = sortedSelectionMergeBlockSize;
+  }
+
   public boolean isNullHandlingEnabled() {
     return _nullHandlingEnabled;
   }