gortiz commented on code in PR #19396:
URL: https://github.com/apache/pinot/pull/19396#discussion_r4106395375
##########
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:
RISK (rolling upgrade): with a new broker and old servers, which is the
usual controller→broker→server upgrade window, one query can sort the same rows
up to three times.
1) An old leaf ServerPlanRequestVisitor.visitSort pushes the new SortNode
into V1 as ORDER BY ... LIMIT 2147483647. 2) An old receiver sees isSort=true
and builds the deprecated SortedMailboxReceiveOperator, which does a full sort.
3) The old SortOperator above it does not know its input is sorted and sorts
again. The base plan used (false,false), so it sorted once. The results stay
correct, but queries are slower until all servers are upgraded. A kill switch
would cover this case too.
##########
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:
Drop the Sort created just below (line 158) for the sortOnReceiver=true
branch. It does no work at runtime, and a no-op operator should not stay in the
plan. Removing a Sort whose input already provides the collation is a normal
logical rewrite.
With (sortOnSender=true, sortOnReceiver=true), every receiver already
returns rows in order. A new receiver merges, or falls back to a full sort. An
old receiver uses SortedMailboxReceiveOperator, which does a full sort. The
Sort above has no fetch or offset. On new servers it becomes a
LimitSortOperator with limit MAX_VALUE, which passes every row through. On old
servers it sorts a second time, and that second sort is part of the triple sort
during the upgrade window. Proposal: only wrap the exchange in LogicalSort when
it does not sort on the receiver. Today that is the partitioned (false,false)
branch. Or let a SortRemove-style rule drop the Sort because the exchange's
collation already covers it. Consequences: (1) SortOperator.isInputSorted, the
SortedMultiStageOperator interface, and the LimitSortOperator/SortOperator
Javadoc edits exist only to turn this Sort into a no-op, so they can be
removed. (2) WindowFunctionPlans.json and QueryCompilationTest expectations
lose the Sor
t above the exchange. (3) The comment that cites #19412's reason for the
explicit Sort (fetch/offset handling, deprecated receive op) no longer applies
here, because this Sort has no fetch.
##########
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:
The comparator matches the sender side (FullSortOperator also uses
SortComparator(collations, false)). Null ordering is consistent. OK.
I checked this because a mismatch in null direction between sender and
receiver would silently corrupt the merge.
##########
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:
The equal-head fast path together with the lazy 'dirty' heap adds the most
complexity in this PR. I traced it and it looks correct, but the invariants are
implicit.
Correctness depends on: dirty=true only when active==oldSize>1; every path
that sets _tryEqualHeadMerge=false either heapified already or hits the else
branch at 236; add() is only reached when a cursor was starved, which implies
the heap was heapified. Only one targeted test covers this
(shouldMergeEqualHeadsAcrossRefillIntoBoundedBlocks). A randomized test would
be much cheaper insurance than the reasoning above: N senders, random block
splits, random tie density, random starvation order, and output compared
against a full sort.
##########
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:
MEMORY: when one cursor is starved, readOneBlock pulls from any ready
mailbox into unbounded per-cursor _pending deques. This effectively removes
mailbox backpressure for the fast senders.
The PR describes this, and the worst case equals the old full receiver sort.
The difference is that sender-side FullSortOperators now also hold their full
buffers during the send, so peak cluster memory can be higher than on master
when one sender is slow. Consider a follow-up that bounds pending rows per
cursor, or at least a stat (for example max retained rows) so this can be
observed in production.
##########
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:
On the single-sender path, blocks pass through unsplit, so output blocks can
be larger than 10k rows. That is the same as today's receive, but the class
Javadoc says output blocks are capped at 10,000 rows.
Align the Javadoc or the PR text.
##########
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:
RISK: there is no kill switch. Every global ORDER BY window now takes the
sender-sort + merge path, and the PR's own numbers show +53% on presorted input
and +2-4% on small or single-sender cases. Please add a broker config plus a
query option to turn it off.
Use the same style of mechanism as the normal sort-exchange copy (broker
config CONFIG_OF_SORT_EXCHANGE_COPY_THRESHOLD + query option
sortExchangeCopyThreshold). The trigger must not be that threshold itself,
though. The threshold is keyed on the Sort's fetch, and the window Sort has no
fetch: a query LIMIT sits above the window and is not pushed through it. The
sender sort is therefore always a full sort, and a fetch threshold would never
enable this path. The L x W amplification that motivates the threshold also
does not happen here: a global window already sends all N rows to one receiver,
so sorting on the senders does not change how many rows are sent. The costs
here are different: merge CPU, memory held by both sender buffers and receiver
read-ahead, and the upgrade-window extra sorts. So a boolean (for example a
windowSortOnSender option, default on or off per the author's call) fits
better. A threshold on something meaningful, such as the number of senders,
would also work
. If a LIMIT is ever pushed through ROW_NUMBER/RANK or other prefix-frame
windows, the senders would do a top-L; that is exactly the normal-sort case and
should reuse sortExchangeCopyThreshold. The switch also gives operators a way
back from a production regression without a rollback.
##########
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:
nit: the marker is added to every chunk, including split parts
(waitForMore=true) and EOS.
MailboxContentObserver reads the flag only from the message that completes
the block, so the flag on the other chunks is unused. It costs 26 bytes each,
which is negligible. The flag is set per stream but repeated per message; that
is fine for simplicity.
##########
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:
Inserting a physical SortNode inside the fragmenter works, but it makes the
stage plan differ from the RelNode plan.
EXPLAIN (logical) shows only isSortOnSender=[true]. The SortNode with
fetch=2147483647 appears only in the fragmented or implementation plan. That is
acceptable, but please confirm that EXPLAIN IMPLEMENTATION PLAN renders it
clearly. Also note that the MAX_VALUE fetch reaches old servers' V1 leaf as
LIMIT 2147483647, and that LIMIT replaces any multiStageLeafLimit. On the new
server path it does not, because the sort stays above the leaf boundary.
##########
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:
The marker is keyed on 'input instanceof SortOperator'. That holds for all
three implementations (Limit/TopN/Full), and each one produces a sorted stream
here. OK, but the check is fragile.
If someone later wraps the sender op-chain (for example a stats or tracing
decorator), the marker silently disappears. Every query then takes
fallbackToFullSort, and nothing fails. Consider asking the op for its collation
instead, since SortedMultiStageOperator already exists. Or add a test in
MailboxSendOperatorTest that builds the real op-chain from PlanNodeToOpChain.
##########
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:
API surface: this adds 3 accessors with 'only valid right after the last
read' semantics (getLastReadStream, getFinishedStreamsLastRead view, live
snapshot). They are easy to misuse.
Only the merge receiver uses them. A single read method that returns a small
result (block, source stream, finished streams) would make the protocol
explicit instead of relying on temporal coupling. This is not blocking.
##########
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:
This Preconditions check is correct as a safety net, but the message is
internal-only. If it fires, users see an opaque error.
New senders always mark every message and old senders never do. The check
can therefore only fire on a bug or a truly mixed stream. Fine, but consider a
QueryErrorCode, such as INTERNAL, and a message that tells the user to retry
once the rolling upgrade is complete.
##########
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:
Good catch on hybrid and logical tables: if the sort were pushed into V1 per
physical request, the output would be interleaved sorted runs on one mailbox.
The ServerPlanRequestVisitorTest covers this. The cost is that the leaf no
longer gets V1 ORDER BY pushdown. That pushdown could have used a sorted index
or early segment pruning, but the sort has no limit, so little is lost.
##########
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:
If the Sort above the exchange is dropped (see the comment at
PinotWindowExchangeNodeInsertRule:154), isSort()=true becomes the ONLY thing
that makes a receiver produce ordered rows for the window. Please document that
here, as a contract.
New receivers: isSort && isSortedOnSender selects the merge, and the merge
falls back to a full sort on unmarked data. Old receivers: isSort selects
SortedMailboxReceiveOperator, which does a full sort. Both are needed for
correct results, and nothing above the receiver sorts again. So (a)
SortedMailboxReceiveOperator cannot be removed while old servers may still
receive plans from new brokers, even though it is deprecated; and (b) the
planner must never emit sortedOnSender=true with sort=false. Consider a
Preconditions check in the planner or in PlanFragmenter, and a test that builds
the op-chain for a global ordered window and asserts that the receive operator
orders its output.
--
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]
