xiangfu0 commented on code in PR #19396: URL: https://github.com/apache/pinot/pull/19396#discussion_r4110236587
########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxMergeReceiveOperator.java: ########## @@ -0,0 +1,668 @@ +/** + * 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.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.operator.utils.AsyncStream; +import org.apache.pinot.query.runtime.operator.utils.SortUtils; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Receives streams that the plan declares sorted on the sender and merges them by the exchange collation. +/// +/// An explicit sender [SortOperator] establishes the row ordering; [MailboxSendOperator] only transports that +/// ordering. The transport marker confirms rollout compatibility and is not itself a sorting mechanism. +/// +/// The plan declaration alone is not trusted during a rolling upgrade. Every data block must carry the transport's +/// sender-sort confirmation. Before this operator emits its first row it obtains a head row, or EOS, from every live +/// sender. If any sender's first data is unconfirmed, all tentatively buffered rows are folded into a full receiver +/// sort. Once output starts, losing the confirmation is a protocol violation because already emitted rows cannot be +/// recovered into that fallback. +/// +/// The merge reads whichever mailbox is ready instead of blocking on one sender. This prevents a sender that is +/// backpressured by another receiver from creating a cross-receiver wait cycle. Rows are emitted in blocks of at most +/// 10,000 while cursor state carries the ordering frontier across calls. A fast sender can be read ahead while another +/// sender is starved, so retained input is workload-dependent and can approach the legacy full receiver sort in the +/// worst case. +/// +/// This operator is driven by a single consumer thread and is not thread-safe. +public class SortedMailboxMergeReceiveOperator extends BaseMailboxReceiveOperator implements SortedMultiStageOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxMergeReceiveOperator.class); + + private static final String EXPLAIN_NAME = "SORTED_MAILBOX_MERGE_RECEIVE"; + private static final String MERGE_SCOPE = "SortedMailboxMergeReceiveOperator"; + private final DataSchema _dataSchema; + private final List<RelFieldCollation> _collations; + private final Comparator<Object[]> _comparator; + private final SenderCursorHeap _readyCursors; + private final boolean _singleSortedSender; + /// Senders that have not finished but do not currently have a row ready. Nothing can be emitted while this is + /// non-empty because any one of these senders may hold the next row. + private final Set<SenderCursor> _starvedCursors = Collections.newSetFromMap(new IdentityHashMap<>()); + private final Map<AsyncStream<ReceivingMailbox.MseBlockWithStats>, SenderCursor> _cursorsByStream = + new IdentityHashMap<>(); + private boolean _mergeOutputStarted; + private boolean _fallbackToSort; + private boolean _tryEqualHeadMerge; + private int _fallbackOutputIndex = -1; + + /// Rows buffered only for the mixed-version fallback. The sorted list is handed downstream as-is, so cleanup must + /// drop this reference rather than clear it. + @Nullable + private List<Object[]> _rows; + + @Nullable + private MseBlock _eosBlock; + + public SortedMailboxMergeReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { + super(context, node); + Preconditions.checkState(node.isSort(), "Receiver-side sorting must be enabled"); + Preconditions.checkState(node.isSortedOnSender(), "Sender-side sorting must be enabled"); + Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); + _dataSchema = node.getDataSchema(); + _collations = List.copyOf(node.getCollations()); + _comparator = new SortUtils.SortComparator(_collations, false); Review Comment: Confirmed. The merge uses the same comparator construction and null-direction handling as the regular sort path, so ascending, descending, and null ordering remain consistent. No change was needed. ########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java: ########## @@ -141,18 +141,17 @@ public void onMatch(RelOptRuleCall call) { exchange = PinotLogicalExchange.create(input, RelDistributions.hash(windowGroup.keys.toList()), prePartitioned); } else { // PARTITION BY and ORDER BY on different key(s) - // Add a LogicalSortExchange hashed on the partition by keys and collation based on order by keys. - // The ordering itself is established by the Sort placed over the exchange below, not by the receive - // operator - see the comment at the transformTo call. + // Keep the receiver full-sort path for a partitioned exchange. Sorting before the hash exchange compares rows + // routed to different receivers and blocks streaming; the explicit Sort retained above the exchange establishes + // the required ordering after partitioning. exchange = PinotLogicalSortExchange.create(input, RelDistributions.hash(windowGroup.keys.toList()), windowGroup.orderKeys, false, false, prePartitioned); } } // WindowAggregateOperator requires its input ordered on the ORDER BY keys and does no ordering of its own, so - // where the exchange carries a collation the ordering has to be established above it. Place an explicit Sort - // rather than asking the receive operator to sort: SortOperator is the operator that knows fetch/offset, and - // SortedMailboxReceiveOperator is deprecated. The Sort carries no fetch, so it keeps every row - the same - // semantics as the unbounded list the receive operator used. + // where the exchange carries a collation the ordering has to be established above it. Keep an explicit Sort as + // the semantic boundary. A confirmed merge receiver advertises the exact collation, allowing SortOperator to + // stream through it; a legacy receiver still performs the full sort. Review Comment: Done. I removed the redundant Sort above a receiver-sorted exchange and deleted SortedMultiStageOperator, SortOperator.isInputSorted, and their associated tests and documentation. The old plan-fixture changes are reverted as well. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java: ########## @@ -193,7 +194,9 @@ void record(PlanNode node, MultiStageOperator operator) { @Override public MultiStageOperator visitMailboxReceive(MailboxReceiveNode node, OpChainExecutionContext context) { try { - if (node.isSort()) { + if (node.isSort() && node.isSortedOnSender()) { Review Comment: Addressed by documenting the isSortOnSender/isSortOnReceiver ordering contract and enforcing sortOnSender implies sortOnReceiver in PlanFragmenter. Plan-shape, real op-chain selection, and randomized merge tests now cover the contract end to end. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java: ########## @@ -185,6 +186,16 @@ public PlanNode visitExchange(ExchangeNode node, Context context) { // Create a new context for the next PlanFragment with MailboxSendNode as the root node. PlanNode nextPlanFragmentRoot = node.getInputs().get(0).visit(this, new Context(senderPlanFragmentId)); + if (node.isSortOnSender()) { + Preconditions.checkState(!node.getCollations().isEmpty(), + "Sender sorting requires a non-empty exchange collation"); + // Ordering belongs to an explicit operator in the sender fragment. MailboxSendOperator only preserves and + // transports this output; ServerPlanRequestVisitor keeps this SortNode above the V1 leaf boundary. + // SortOperator applies the broker response limit when fetch is absent. This internal sort must retain every + // sender row, so use the largest representable fetch with a zero effective offset. + nextPlanFragmentRoot = new SortNode(senderPlanFragmentId, nextPlanFragmentRoot.getDataSchema(), null, + List.of(nextPlanFragmentRoot), node.getCollations(), Integer.MAX_VALUE, -1); Review Comment: Added implementation-EXPLAIN coverage. The enabled plan asserts both MAIL_SEND [SORTED] and the sender-side SORT LIMIT 2147483647, in addition to checking that no redundant Sort remains above the merge receiver. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java: ########## @@ -64,13 +67,17 @@ public class MailboxSendOperator extends MultiStageOperator { private final BlockExchange _exchange; private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class); - // TODO: Support sort on sender public MailboxSendOperator(OpChainExecutionContext context, MultiStageOperator input, MailboxSendNode node) { - this(context, input, statMap -> getBlockExchange(context, node, statMap)); + this(context, input, statMap -> getBlockExchange(context, node, statMap, isSortedOnSender(input, node))); _statMap.merge(StatKey.STAGE, context.getStageId()); _statMap.merge(StatKey.PARALLELISM, 1); } + @VisibleForTesting + static boolean isSortedOnSender(MultiStageOperator input, MailboxSendNode node) { + return input instanceof SortOperator && node.hasExplicitSortInput(); Review Comment: I kept the structural check as a fail-safe and added a real op-chain conversion test that verifies an explicit SortNode selects the sorted MailboxService overload. If a future wrapper hides the sort, the sender omits confirmation and the receiver safely falls back instead of risking incorrect ordering. ########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java: ########## @@ -120,13 +120,13 @@ public void onMatch(RelOptRuleCall call) { exchange = PinotLogicalExchange.create(input, RelDistributions.hash(List.of())); } else { // Only ORDER BY - // Add a LogicalSortExchange with collation on the order by key(s) and an empty hash partition key. - // The ordering itself is established by the Sort placed over the exchange below, not by the receive - // operator - see the comment at the transformTo call. + // Sort each sender explicitly and merge the sorted mailbox streams at the receiver. The Sort retained above + // the exchange is the semantic ordering boundary and becomes a streaming limit when the merge receiver + // advertises this exact collation. // TODO: Revisit whether we should use hash distribution exchange = - PinotLogicalSortExchange.create(input, RelDistributions.hash(List.of()), windowGroup.orderKeys, false, - false); + PinotLogicalSortExchange.create(input, RelDistributions.hash(List.of()), windowGroup.orderKeys, true, Review Comment: Addressed. The merge receiver now establishes final ordering directly, so the redundant post-exchange Sort is gone. The optimization is also default-off with query and broker controls: disabled mode preserves the master plan, while enabled mixed-version execution is limited to sender plus legacy receiver sorting rather than three sorts. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxMergeReceiveOperator.java: ########## @@ -0,0 +1,668 @@ +/** + * 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.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.operator.utils.AsyncStream; +import org.apache.pinot.query.runtime.operator.utils.SortUtils; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Receives streams that the plan declares sorted on the sender and merges them by the exchange collation. +/// +/// An explicit sender [SortOperator] establishes the row ordering; [MailboxSendOperator] only transports that +/// ordering. The transport marker confirms rollout compatibility and is not itself a sorting mechanism. +/// +/// The plan declaration alone is not trusted during a rolling upgrade. Every data block must carry the transport's +/// sender-sort confirmation. Before this operator emits its first row it obtains a head row, or EOS, from every live +/// sender. If any sender's first data is unconfirmed, all tentatively buffered rows are folded into a full receiver +/// sort. Once output starts, losing the confirmation is a protocol violation because already emitted rows cannot be +/// recovered into that fallback. +/// +/// The merge reads whichever mailbox is ready instead of blocking on one sender. This prevents a sender that is +/// backpressured by another receiver from creating a cross-receiver wait cycle. Rows are emitted in blocks of at most +/// 10,000 while cursor state carries the ordering frontier across calls. A fast sender can be read ahead while another +/// sender is starved, so retained input is workload-dependent and can approach the legacy full receiver sort in the +/// worst case. +/// +/// This operator is driven by a single consumer thread and is not thread-safe. +public class SortedMailboxMergeReceiveOperator extends BaseMailboxReceiveOperator implements SortedMultiStageOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxMergeReceiveOperator.class); + + private static final String EXPLAIN_NAME = "SORTED_MAILBOX_MERGE_RECEIVE"; + private static final String MERGE_SCOPE = "SortedMailboxMergeReceiveOperator"; + private final DataSchema _dataSchema; + private final List<RelFieldCollation> _collations; + private final Comparator<Object[]> _comparator; + private final SenderCursorHeap _readyCursors; + private final boolean _singleSortedSender; + /// Senders that have not finished but do not currently have a row ready. Nothing can be emitted while this is + /// non-empty because any one of these senders may hold the next row. + private final Set<SenderCursor> _starvedCursors = Collections.newSetFromMap(new IdentityHashMap<>()); + private final Map<AsyncStream<ReceivingMailbox.MseBlockWithStats>, SenderCursor> _cursorsByStream = + new IdentityHashMap<>(); + private boolean _mergeOutputStarted; + private boolean _fallbackToSort; + private boolean _tryEqualHeadMerge; + private int _fallbackOutputIndex = -1; + + /// Rows buffered only for the mixed-version fallback. The sorted list is handed downstream as-is, so cleanup must + /// drop this reference rather than clear it. + @Nullable + private List<Object[]> _rows; + + @Nullable + private MseBlock _eosBlock; + + public SortedMailboxMergeReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { + super(context, node); + Preconditions.checkState(node.isSort(), "Receiver-side sorting must be enabled"); + Preconditions.checkState(node.isSortedOnSender(), "Sender-side sorting must be enabled"); + Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); + _dataSchema = node.getDataSchema(); + _collations = List.copyOf(node.getCollations()); + _comparator = new SortUtils.SortComparator(_collations, false); + List<AsyncStream<ReceivingMailbox.MseBlockWithStats>> streams = _multiConsumer.getLiveStreamsSnapshot(); + _readyCursors = new SenderCursorHeap(streams.size(), _comparator); + _singleSortedSender = streams.size() == 1; + if (!_singleSortedSender) { + for (AsyncStream<ReceivingMailbox.MseBlockWithStats> stream : streams) { + SenderCursor cursor = new SenderCursor(stream); + _cursorsByStream.put(stream, cursor); + _starvedCursors.add(cursor); + } + } + } + + @Override + protected Logger logger() { + return LOGGER; + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + @Override + public List<RelFieldCollation> getCollations() { + return _collations; + } + + @Override + protected MseBlock getNextBlock() { + if (_fallbackOutputIndex >= 0 && _rows != null) { + return emitFallbackBlock(); + } + if (_eosBlock != null) { + return _eosBlock; + } + if (_isEarlyTerminated) { + return readUntilEos(); + } + return _singleSortedSender ? readSingleSortedSender() : mergeNextBlock(); + } + + /// Passes through one confirmed sorted sender without copying its rows through the merge heap. + private MseBlock readSingleSortedSender() { + while (true) { + MseBlock block = _multiConsumer.readMseBlockBlocking(); + if (block.isEos()) { + return terminate(block); + } + MseBlock.Data dataBlock = (MseBlock.Data) block; + checkActiveTerminationAndSampleUsage(); + if (!_multiConsumer.isLastBlockSortedOnSender()) { Review Comment: Corrected the class and PR documentation. The 10,000-row cap applies to multi-sender merge and fallback output; a confirmed single sender is passed through with its original block boundaries. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java: ########## @@ -259,8 +259,17 @@ public Void visitMailboxReceive(MailboxReceiveNode node, ServerPlanRequestContex @Override public Void visitMailboxSend(MailboxSendNode node, ServerPlanRequestContext context) { - if (visit(node.getInputs().get(0), context)) { - context.setLeafStageBoundaryNode(node.getInputs().get(0)); + PlanNode input = node.getInputs().get(0); + if (node.hasExplicitSortInput()) { Review Comment: Confirmed. ServerPlanRequestVisitorTest covers the leaf-boundary invariant needed for hybrid and logical-table physical requests. The conservative loss of unbounded-sort pushdown is accepted here; no extra hybrid-specific code change was needed. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxMergeReceiveOperator.java: ########## @@ -0,0 +1,668 @@ +/** + * 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.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.operator.utils.AsyncStream; +import org.apache.pinot.query.runtime.operator.utils.SortUtils; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Receives streams that the plan declares sorted on the sender and merges them by the exchange collation. +/// +/// An explicit sender [SortOperator] establishes the row ordering; [MailboxSendOperator] only transports that +/// ordering. The transport marker confirms rollout compatibility and is not itself a sorting mechanism. +/// +/// The plan declaration alone is not trusted during a rolling upgrade. Every data block must carry the transport's +/// sender-sort confirmation. Before this operator emits its first row it obtains a head row, or EOS, from every live +/// sender. If any sender's first data is unconfirmed, all tentatively buffered rows are folded into a full receiver +/// sort. Once output starts, losing the confirmation is a protocol violation because already emitted rows cannot be +/// recovered into that fallback. +/// +/// The merge reads whichever mailbox is ready instead of blocking on one sender. This prevents a sender that is +/// backpressured by another receiver from creating a cross-receiver wait cycle. Rows are emitted in blocks of at most +/// 10,000 while cursor state carries the ordering frontier across calls. A fast sender can be read ahead while another +/// sender is starved, so retained input is workload-dependent and can approach the legacy full receiver sort in the +/// worst case. +/// +/// This operator is driven by a single consumer thread and is not thread-safe. +public class SortedMailboxMergeReceiveOperator extends BaseMailboxReceiveOperator implements SortedMultiStageOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxMergeReceiveOperator.class); + + private static final String EXPLAIN_NAME = "SORTED_MAILBOX_MERGE_RECEIVE"; + private static final String MERGE_SCOPE = "SortedMailboxMergeReceiveOperator"; + private final DataSchema _dataSchema; + private final List<RelFieldCollation> _collations; + private final Comparator<Object[]> _comparator; + private final SenderCursorHeap _readyCursors; + private final boolean _singleSortedSender; + /// Senders that have not finished but do not currently have a row ready. Nothing can be emitted while this is + /// non-empty because any one of these senders may hold the next row. + private final Set<SenderCursor> _starvedCursors = Collections.newSetFromMap(new IdentityHashMap<>()); + private final Map<AsyncStream<ReceivingMailbox.MseBlockWithStats>, SenderCursor> _cursorsByStream = + new IdentityHashMap<>(); + private boolean _mergeOutputStarted; + private boolean _fallbackToSort; + private boolean _tryEqualHeadMerge; + private int _fallbackOutputIndex = -1; + + /// Rows buffered only for the mixed-version fallback. The sorted list is handed downstream as-is, so cleanup must + /// drop this reference rather than clear it. + @Nullable + private List<Object[]> _rows; + + @Nullable + private MseBlock _eosBlock; + + public SortedMailboxMergeReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { + super(context, node); + Preconditions.checkState(node.isSort(), "Receiver-side sorting must be enabled"); + Preconditions.checkState(node.isSortedOnSender(), "Sender-side sorting must be enabled"); + Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); + _dataSchema = node.getDataSchema(); + _collations = List.copyOf(node.getCollations()); + _comparator = new SortUtils.SortComparator(_collations, false); + List<AsyncStream<ReceivingMailbox.MseBlockWithStats>> streams = _multiConsumer.getLiveStreamsSnapshot(); + _readyCursors = new SenderCursorHeap(streams.size(), _comparator); + _singleSortedSender = streams.size() == 1; + if (!_singleSortedSender) { + for (AsyncStream<ReceivingMailbox.MseBlockWithStats> stream : streams) { + SenderCursor cursor = new SenderCursor(stream); + _cursorsByStream.put(stream, cursor); + _starvedCursors.add(cursor); + } + } + } + + @Override + protected Logger logger() { + return LOGGER; + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + @Override + public List<RelFieldCollation> getCollations() { + return _collations; + } + + @Override + protected MseBlock getNextBlock() { + if (_fallbackOutputIndex >= 0 && _rows != null) { + return emitFallbackBlock(); + } + if (_eosBlock != null) { + return _eosBlock; + } + if (_isEarlyTerminated) { + return readUntilEos(); + } + return _singleSortedSender ? readSingleSortedSender() : mergeNextBlock(); + } + + /// Passes through one confirmed sorted sender without copying its rows through the merge heap. + private MseBlock readSingleSortedSender() { + while (true) { + MseBlock block = _multiConsumer.readMseBlockBlocking(); + if (block.isEos()) { + return terminate(block); + } + MseBlock.Data dataBlock = (MseBlock.Data) block; + checkActiveTerminationAndSampleUsage(); + if (!_multiConsumer.isLastBlockSortedOnSender()) { + fallbackToFullSort(dataBlock.asRowHeap().getRows()); + return sortAllRows(); + } + if (dataBlock.getNumRows() > 0) { + _mergeOutputStarted = true; + return dataBlock; + } + } + } + + /// Merges the sorted senders, emitting at most [SortOperator#DEFAULT_MAX_ROWS_PER_BLOCK] rows per call. + private MseBlock mergeNextBlock() { + ArrayList<Object[]> rows = new ArrayList<>(0); + while (rows.size() < SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK) { + if (!_starvedCursors.isEmpty()) { + MseBlock.Eos error; + boolean receivedMoreRows = false; + if (rows.isEmpty()) { + error = readOneBlock(); + } else { + // Rows already removed from the heap are a globally ordered prefix. Consume any immediately available + // cursor progress so blocks can still be coalesced, but return that safe prefix instead of waiting merely + // to fill the output block. + MseBlock block = _multiConsumer.pollMseBlockOrStreamCompletion(); + if (block == null && _multiConsumer.getFinishedStreamsLastRead().isEmpty()) { + break; + } + error = processReadBlock(block); + receivedMoreRows = block != null && block.isData() && ((MseBlock.Data) block).getNumRows() > 0; + } + if (error != null) { + return terminate(error); + } + if (_fallbackToSort) { + // These rows were already removed from cursors while building this not-yet-emitted block. + _rows.addAll(rows); + return sortAllRows(); + } + if (receivedMoreRows) { + // A refill proves this is not a one-block result. Restore the established full-block capacity once so + // fragmented input cannot trigger repeated growth while this output block is assembled. + rows.ensureCapacity(SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK); + } + continue; + } + if (_readyCursors.isEmpty()) { + break; + } + if (rows.isEmpty()) { + // Keep empty and tiny results cheap without sacrificing the one-allocation path for full output blocks. + // Sum all currently buffered rows once; later output blocks retain the established full-block capacity. + int initialCapacity = _mergeOutputStarted ? SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK + : _readyCursors.getCappedAvailableRowCount(SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK); + // ArrayList grows by 1.5x. Round near-full first blocks up now so a small refill cannot allocate an oversized + // replacement in addition to the nearly full initial array. + if (initialCapacity + (initialCapacity >> 1) >= SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK) { + initialCapacity = SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK; + } + rows.ensureCapacity(initialCapacity); + } + if (_tryEqualHeadMerge && _readyCursors.size() > 1) { + int remaining = SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK - rows.size(); + if (remaining >= _readyCursors.size()) { + if (_readyCursors.allHeadsEqual()) { + int previousRowCount = rows.size(); + List<SenderCursor> exhausted = _readyCursors.advanceEqualHeads(rows); + if (exhausted != null) { + for (SenderCursor cursor : exhausted) { + if (!cursor._finished) { + _starvedCursors.add(cursor); + } + } + } + if (!_starvedCursors.isEmpty()) { + _tryEqualHeadMerge = false; + } + checkActiveTerminationAndSampleUsageAfterBatch(previousRowCount, rows.size()); + continue; + } else { + _tryEqualHeadMerge = false; + } + } else { + _readyCursors.ensureOrdered(); + } + } + SenderCursor cursor = _readyCursors.peek(); + rows.add(cursor.next()); + if (cursor.hasRow()) { + // The cursor's key can only move forward, so restoring the heap from the root takes one sift-down. A generic + // PriorityQueue poll followed by add performs two independent heap repairs for every emitted row. + _readyCursors.updateTop(); + } else if (!cursor._finished) { + _readyCursors.removeTop(); + _starvedCursors.add(cursor); + _tryEqualHeadMerge = false; + } else { + _readyCursors.removeTop(); + _tryEqualHeadMerge = true; + } + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(rows.size(), MERGE_SCOPE, + _context.getActiveDeadlineMs()); + } + if (rows.isEmpty()) { + return terminate(SuccessMseBlock.INSTANCE); + } + _mergeOutputStarted = true; + return new RowHeapDataBlock(rows, _dataSchema); + } + + /// Reads one block from whichever sender is ready and updates only the cursor that produced it. + /// + /// @return the error that ended the read, or `null` when the read succeeded + @Nullable + private MseBlock.Eos readOneBlock() { + return processReadBlock(_multiConsumer.readMseBlockOrStreamCompletionBlocking()); + } + + @Nullable + private MseBlock.Eos processReadBlock(@Nullable MseBlock block) { + if (block == null) { + updateFinishedCursors(); + if (_starvedCursors.isEmpty()) { + _tryEqualHeadMerge = true; + } + return null; + } + if (block.isEos()) { + updateFinishedCursors(); + MseBlock.Eos eos = (MseBlock.Eos) block; + if (eos.isError()) { + return eos; + } + // Aggregate success is returned only after every sender has emitted EOS. + _starvedCursors.clear(); + _tryEqualHeadMerge = true; + return null; + } + AsyncStream<ReceivingMailbox.MseBlockWithStats> stream = _multiConsumer.getLastReadStream(); + Preconditions.checkState(stream != null, "Read a data block from no mailbox on stage: %s", _context.getStageId()); + SenderCursor cursor = _cursorsByStream.get(stream); + Preconditions.checkState(cursor != null, "Read a data block from unknown mailbox: %s", stream.getId()); + List<Object[]> rows = ((MseBlock.Data) block).asRowHeap().getRows(); + checkActiveTerminationAndSampleUsage(); + if (!_multiConsumer.isLastBlockSortedOnSender()) { + fallbackToFullSort(rows); + return null; + } + cursor.offer(rows); + updateFinishedCursors(); + if (cursor.hasRow() && _starvedCursors.remove(cursor)) { + _readyCursors.add(cursor); + } + if (_starvedCursors.isEmpty()) { + _tryEqualHeadMerge = true; + } + return null; + } + + /// Removes only the starved cursors whose EOS was consumed by the last read. + private void updateFinishedCursors() { + for (AsyncStream<ReceivingMailbox.MseBlockWithStats> stream : _multiConsumer.getFinishedStreamsLastRead()) { + SenderCursor cursor = _cursorsByStream.get(stream); + if (cursor != null) { + cursor._finished = true; + if (!cursor.hasRow()) { + _starvedCursors.remove(cursor); + _cursorsByStream.remove(stream); + } + } + } + } + + /// Switches to a full receiver sort when a legacy sender omits the transport confirmation. + private void fallbackToFullSort(List<Object[]> unconfirmedRows) { + Preconditions.checkState(!_mergeOutputStarted, Review Comment: Fixed. A confirmation change after output begins now returns an INTERNAL query error with an actionable message to retry after the rolling upgrade or disable windowSortOnSender. The regression test asserts both the error code and the message. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java: ########## @@ -432,12 +435,14 @@ protected void sendContent(ByteString byteString, boolean waitForMore, boolean b if (!bypassReady && isTerminated()) { return; } - MailboxContent content = MailboxContent.newBuilder() + MailboxContent.Builder contentBuilder = MailboxContent.newBuilder() .setMailboxId(_id) .setPayload(byteString) - .setWaitForMore(waitForMore) - .build(); - _contentObserver.onNext(content); + .setWaitForMore(waitForMore); + if (_sortedOnSender) { + contentBuilder.putMetadata(ChannelUtils.MAILBOX_METADATA_SORTED_ON_SENDER, Boolean.TRUE.toString()); Review Comment: Confirmed. The ordering marker intentionally remains on every mailbox message so each message is self-describing; this is a fixed 26-byte per-message cost rather than a per-row cost. No code change was needed here. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java: ########## @@ -269,6 +315,37 @@ private E readDroppingSuccessEos() { return block; } + /// Returns the stream that produced the element the last blocking read returned. + /// + /// This is only meaningful right after that call returned a data or an error element. The element that ends all + /// the streams is not produced by any of them, and a stream that emitted its EOS is dropped from the ones this + /// consumer tracks, so `null` is returned in both cases. + /// + /// Consumers that need to keep the elements of each stream apart, like a merge of already sorted streams, use + /// this to tell which stream the element they just read belongs to. + /// + /// This method is called by the consumer thread. + @Nullable + public AsyncStream<E> getLastReadStream() { Review Comment: Agreed that a combined ReadResult API would make the temporal coupling clearer. I am deferring that shared mailbox API refactor because this operator is the single consumer here and the comment is nonblocking; it can be handled separately without expanding this PR. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxMergeReceiveOperator.java: ########## @@ -0,0 +1,668 @@ +/** + * 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.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.operator.utils.AsyncStream; +import org.apache.pinot.query.runtime.operator.utils.SortUtils; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Receives streams that the plan declares sorted on the sender and merges them by the exchange collation. +/// +/// An explicit sender [SortOperator] establishes the row ordering; [MailboxSendOperator] only transports that +/// ordering. The transport marker confirms rollout compatibility and is not itself a sorting mechanism. +/// +/// The plan declaration alone is not trusted during a rolling upgrade. Every data block must carry the transport's +/// sender-sort confirmation. Before this operator emits its first row it obtains a head row, or EOS, from every live +/// sender. If any sender's first data is unconfirmed, all tentatively buffered rows are folded into a full receiver +/// sort. Once output starts, losing the confirmation is a protocol violation because already emitted rows cannot be +/// recovered into that fallback. +/// +/// The merge reads whichever mailbox is ready instead of blocking on one sender. This prevents a sender that is +/// backpressured by another receiver from creating a cross-receiver wait cycle. Rows are emitted in blocks of at most +/// 10,000 while cursor state carries the ordering frontier across calls. A fast sender can be read ahead while another +/// sender is starved, so retained input is workload-dependent and can approach the legacy full receiver sort in the +/// worst case. +/// +/// This operator is driven by a single consumer thread and is not thread-safe. +public class SortedMailboxMergeReceiveOperator extends BaseMailboxReceiveOperator implements SortedMultiStageOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxMergeReceiveOperator.class); + + private static final String EXPLAIN_NAME = "SORTED_MAILBOX_MERGE_RECEIVE"; + private static final String MERGE_SCOPE = "SortedMailboxMergeReceiveOperator"; + private final DataSchema _dataSchema; + private final List<RelFieldCollation> _collations; + private final Comparator<Object[]> _comparator; + private final SenderCursorHeap _readyCursors; + private final boolean _singleSortedSender; + /// Senders that have not finished but do not currently have a row ready. Nothing can be emitted while this is + /// non-empty because any one of these senders may hold the next row. + private final Set<SenderCursor> _starvedCursors = Collections.newSetFromMap(new IdentityHashMap<>()); + private final Map<AsyncStream<ReceivingMailbox.MseBlockWithStats>, SenderCursor> _cursorsByStream = + new IdentityHashMap<>(); + private boolean _mergeOutputStarted; + private boolean _fallbackToSort; + private boolean _tryEqualHeadMerge; + private int _fallbackOutputIndex = -1; + + /// Rows buffered only for the mixed-version fallback. The sorted list is handed downstream as-is, so cleanup must + /// drop this reference rather than clear it. + @Nullable + private List<Object[]> _rows; + + @Nullable + private MseBlock _eosBlock; + + public SortedMailboxMergeReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { + super(context, node); + Preconditions.checkState(node.isSort(), "Receiver-side sorting must be enabled"); + Preconditions.checkState(node.isSortedOnSender(), "Sender-side sorting must be enabled"); + Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); + _dataSchema = node.getDataSchema(); + _collations = List.copyOf(node.getCollations()); + _comparator = new SortUtils.SortComparator(_collations, false); + List<AsyncStream<ReceivingMailbox.MseBlockWithStats>> streams = _multiConsumer.getLiveStreamsSnapshot(); + _readyCursors = new SenderCursorHeap(streams.size(), _comparator); + _singleSortedSender = streams.size() == 1; + if (!_singleSortedSender) { + for (AsyncStream<ReceivingMailbox.MseBlockWithStats> stream : streams) { + SenderCursor cursor = new SenderCursor(stream); + _cursorsByStream.put(stream, cursor); + _starvedCursors.add(cursor); + } + } + } + + @Override + protected Logger logger() { + return LOGGER; + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + @Override + public List<RelFieldCollation> getCollations() { + return _collations; + } + + @Override + protected MseBlock getNextBlock() { + if (_fallbackOutputIndex >= 0 && _rows != null) { + return emitFallbackBlock(); + } + if (_eosBlock != null) { + return _eosBlock; + } + if (_isEarlyTerminated) { + return readUntilEos(); + } + return _singleSortedSender ? readSingleSortedSender() : mergeNextBlock(); + } + + /// Passes through one confirmed sorted sender without copying its rows through the merge heap. + private MseBlock readSingleSortedSender() { + while (true) { + MseBlock block = _multiConsumer.readMseBlockBlocking(); + if (block.isEos()) { + return terminate(block); + } + MseBlock.Data dataBlock = (MseBlock.Data) block; + checkActiveTerminationAndSampleUsage(); + if (!_multiConsumer.isLastBlockSortedOnSender()) { + fallbackToFullSort(dataBlock.asRowHeap().getRows()); + return sortAllRows(); + } + if (dataBlock.getNumRows() > 0) { + _mergeOutputStarted = true; + return dataBlock; + } + } + } + + /// Merges the sorted senders, emitting at most [SortOperator#DEFAULT_MAX_ROWS_PER_BLOCK] rows per call. + private MseBlock mergeNextBlock() { + ArrayList<Object[]> rows = new ArrayList<>(0); + while (rows.size() < SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK) { + if (!_starvedCursors.isEmpty()) { Review Comment: Agreed: pending per-sender blocks remain unbounded in the worst case. I kept bounded queues and retained-row telemetry out of this minimal patch; issue #19395 remains open to track the broader strict-memory-bound work. The PR description now states this limitation explicitly. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxMergeReceiveOperator.java: ########## @@ -0,0 +1,668 @@ +/** + * 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.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.operator.utils.AsyncStream; +import org.apache.pinot.query.runtime.operator.utils.SortUtils; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Receives streams that the plan declares sorted on the sender and merges them by the exchange collation. +/// +/// An explicit sender [SortOperator] establishes the row ordering; [MailboxSendOperator] only transports that +/// ordering. The transport marker confirms rollout compatibility and is not itself a sorting mechanism. +/// +/// The plan declaration alone is not trusted during a rolling upgrade. Every data block must carry the transport's +/// sender-sort confirmation. Before this operator emits its first row it obtains a head row, or EOS, from every live +/// sender. If any sender's first data is unconfirmed, all tentatively buffered rows are folded into a full receiver +/// sort. Once output starts, losing the confirmation is a protocol violation because already emitted rows cannot be +/// recovered into that fallback. +/// +/// The merge reads whichever mailbox is ready instead of blocking on one sender. This prevents a sender that is +/// backpressured by another receiver from creating a cross-receiver wait cycle. Rows are emitted in blocks of at most +/// 10,000 while cursor state carries the ordering frontier across calls. A fast sender can be read ahead while another +/// sender is starved, so retained input is workload-dependent and can approach the legacy full receiver sort in the +/// worst case. +/// +/// This operator is driven by a single consumer thread and is not thread-safe. +public class SortedMailboxMergeReceiveOperator extends BaseMailboxReceiveOperator implements SortedMultiStageOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxMergeReceiveOperator.class); + + private static final String EXPLAIN_NAME = "SORTED_MAILBOX_MERGE_RECEIVE"; + private static final String MERGE_SCOPE = "SortedMailboxMergeReceiveOperator"; + private final DataSchema _dataSchema; + private final List<RelFieldCollation> _collations; + private final Comparator<Object[]> _comparator; + private final SenderCursorHeap _readyCursors; + private final boolean _singleSortedSender; + /// Senders that have not finished but do not currently have a row ready. Nothing can be emitted while this is + /// non-empty because any one of these senders may hold the next row. + private final Set<SenderCursor> _starvedCursors = Collections.newSetFromMap(new IdentityHashMap<>()); + private final Map<AsyncStream<ReceivingMailbox.MseBlockWithStats>, SenderCursor> _cursorsByStream = + new IdentityHashMap<>(); + private boolean _mergeOutputStarted; + private boolean _fallbackToSort; + private boolean _tryEqualHeadMerge; + private int _fallbackOutputIndex = -1; + + /// Rows buffered only for the mixed-version fallback. The sorted list is handed downstream as-is, so cleanup must + /// drop this reference rather than clear it. + @Nullable + private List<Object[]> _rows; + + @Nullable + private MseBlock _eosBlock; + + public SortedMailboxMergeReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { + super(context, node); + Preconditions.checkState(node.isSort(), "Receiver-side sorting must be enabled"); + Preconditions.checkState(node.isSortedOnSender(), "Sender-side sorting must be enabled"); + Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); + _dataSchema = node.getDataSchema(); + _collations = List.copyOf(node.getCollations()); + _comparator = new SortUtils.SortComparator(_collations, false); + List<AsyncStream<ReceivingMailbox.MseBlockWithStats>> streams = _multiConsumer.getLiveStreamsSnapshot(); + _readyCursors = new SenderCursorHeap(streams.size(), _comparator); + _singleSortedSender = streams.size() == 1; + if (!_singleSortedSender) { + for (AsyncStream<ReceivingMailbox.MseBlockWithStats> stream : streams) { + SenderCursor cursor = new SenderCursor(stream); + _cursorsByStream.put(stream, cursor); + _starvedCursors.add(cursor); + } + } + } + + @Override + protected Logger logger() { + return LOGGER; + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + @Override + public List<RelFieldCollation> getCollations() { + return _collations; + } + + @Override + protected MseBlock getNextBlock() { + if (_fallbackOutputIndex >= 0 && _rows != null) { + return emitFallbackBlock(); + } + if (_eosBlock != null) { + return _eosBlock; + } + if (_isEarlyTerminated) { + return readUntilEos(); + } + return _singleSortedSender ? readSingleSortedSender() : mergeNextBlock(); + } + + /// Passes through one confirmed sorted sender without copying its rows through the merge heap. + private MseBlock readSingleSortedSender() { + while (true) { + MseBlock block = _multiConsumer.readMseBlockBlocking(); + if (block.isEos()) { + return terminate(block); + } + MseBlock.Data dataBlock = (MseBlock.Data) block; + checkActiveTerminationAndSampleUsage(); + if (!_multiConsumer.isLastBlockSortedOnSender()) { + fallbackToFullSort(dataBlock.asRowHeap().getRows()); + return sortAllRows(); + } + if (dataBlock.getNumRows() > 0) { + _mergeOutputStarted = true; + return dataBlock; + } + } + } + + /// Merges the sorted senders, emitting at most [SortOperator#DEFAULT_MAX_ROWS_PER_BLOCK] rows per call. + private MseBlock mergeNextBlock() { + ArrayList<Object[]> rows = new ArrayList<>(0); + while (rows.size() < SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK) { + if (!_starvedCursors.isEmpty()) { + MseBlock.Eos error; + boolean receivedMoreRows = false; + if (rows.isEmpty()) { + error = readOneBlock(); + } else { + // Rows already removed from the heap are a globally ordered prefix. Consume any immediately available + // cursor progress so blocks can still be coalesced, but return that safe prefix instead of waiting merely + // to fill the output block. + MseBlock block = _multiConsumer.pollMseBlockOrStreamCompletion(); + if (block == null && _multiConsumer.getFinishedStreamsLastRead().isEmpty()) { + break; + } + error = processReadBlock(block); + receivedMoreRows = block != null && block.isData() && ((MseBlock.Data) block).getNumRows() > 0; + } + if (error != null) { + return terminate(error); + } + if (_fallbackToSort) { + // These rows were already removed from cursors while building this not-yet-emitted block. + _rows.addAll(rows); + return sortAllRows(); + } + if (receivedMoreRows) { + // A refill proves this is not a one-block result. Restore the established full-block capacity once so + // fragmented input cannot trigger repeated growth while this output block is assembled. + rows.ensureCapacity(SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK); + } + continue; + } + if (_readyCursors.isEmpty()) { + break; + } + if (rows.isEmpty()) { + // Keep empty and tiny results cheap without sacrificing the one-allocation path for full output blocks. + // Sum all currently buffered rows once; later output blocks retain the established full-block capacity. + int initialCapacity = _mergeOutputStarted ? SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK + : _readyCursors.getCappedAvailableRowCount(SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK); + // ArrayList grows by 1.5x. Round near-full first blocks up now so a small refill cannot allocate an oversized + // replacement in addition to the nearly full initial array. + if (initialCapacity + (initialCapacity >> 1) >= SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK) { + initialCapacity = SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK; + } + rows.ensureCapacity(initialCapacity); + } + if (_tryEqualHeadMerge && _readyCursors.size() > 1) { Review Comment: Added a deterministic randomized comparison against an independent full sort. It covers 50 fixed-seed scenarios, 2-8 senders, key cardinalities 1-32, empty and duplicate-heavy inputs, randomized block boundaries, and temporary mailbox starvation, while also checking for lost or duplicated rows. ########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java: ########## @@ -120,13 +120,13 @@ public void onMatch(RelOptRuleCall call) { exchange = PinotLogicalExchange.create(input, RelDistributions.hash(List.of())); } else { // Only ORDER BY - // Add a LogicalSortExchange with collation on the order by key(s) and an empty hash partition key. - // The ordering itself is established by the Sort placed over the exchange below, not by the receive - // operator - see the comment at the transformTo call. + // Sort each sender explicitly and merge the sorted mailbox streams at the receiver. The Sort retained above + // the exchange is the semantic ordering boundary and becomes a streaming limit when the merge receiver + // advertises this exact collation. // TODO: Revisit whether we should use hash distribution exchange = - PinotLogicalSortExchange.create(input, RelDistributions.hash(List.of()), windowGroup.orderKeys, false, - false); + PinotLogicalSortExchange.create(input, RelDistributions.hash(List.of()), windowGroup.orderKeys, true, Review Comment: Added both controls: SET windowSortOnSender=true|false and broker config pinot.broker.multistage.window.sort.on.sender. The query option overrides the broker default, and the default is false because the expanded benchmarks found a stable presorted-input regression. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
