Add streaming k-way merge to SortedMailboxReceiveOperator - #19121
Draft
rohityadav1993 wants to merge 3 commits into
Draft
Add streaming k-way merge to SortedMailboxReceiveOperator#19121rohityadav1993 wants to merge 3 commits into
rohityadav1993 wants to merge 3 commits into
Conversation
An unbounded leaf-stage ORDER BY (as injected for sorted merge join inputs) routes to MinMaxValueBasedSelectionOrderByCombineOperator, which merges every segment's rows into a single block before returning anything. At large data volumes this exceeds the leaf stage's CPU budget and ThreadAccountant raises EarlyTerminationException inside SelectionOperatorUtils.mergeWithOrdering(), surfacing to the broker as a spurious "Cancelled by sender". This adds a streaming alternative, opt-in via the `streamingSelectionOrderBy` query option: - StreamingSelectionOrderByOperator emits sorted blocks incrementally for a segment that is physically sorted on the leading ORDER BY column, reading the sorted forward index in order instead of building a priority queue. - StreamingSelectionOrderByCombineOperator performs a k-way heap merge across segment operators and emits bounded blocks (`streamingSelectionOrderByBlockSize`, default 10000) rather than one materialized result. - SelectionPlanNode and CombinePlanNode select these operators when the option is set and the sortedness precondition holds; otherwise behaviour is unchanged. Part of apache#18667.
Correctness and safety fixes on top of the streaming selection ORDER BY combine operator: - StreamingSelectionOrderByCombineOperator now checks the query deadline and termination state inside the merge loop. The merge drains its heap on the caller thread and, for single-block cursors, may never re-enter a child operator, so nothing else observed cancellation, pause, or timeout for up to (limit + offset) iterations. This restored parity with MinMaxValueBasedSelectionOrderByCombineOperator, which this operator replaces on the same query shape. - Child blocks whose data schema disagrees with the merge schema are dropped and reported as MERGE_RESPONSE errors instead of being merged blindly, mirroring SelectionOrderByResultsBlockMerger. Segments on a server can disagree on schema mid-reload, and merging rows of differing width under a single schema corrupted the result rather than failing. - Segments are released when an Error (not just an Exception) escapes the merge, so an OutOfMemoryError while accumulating output rows can no longer leave segments acquired for the lifetime of the server. In single-stage mode there is no start()/stop() backstop around this operator. - StreamingSelectionOrderByOperator reports docs and entries scanned to the shared QueryScanCostContext, as every other selection operator does. Without it, scan-based query killing was blind on this path. Also renames the query option from `streamingSelectionOrderBy` to `sortedSelectionMergeEnabled` (and `...BlockSize` to `sortedSelectionMergeBlockSize`). The option gates the classic single-stage combine path as well as the MSE leaf path, where nothing streams, so the old name misdescribed it; query option strings are effectively permanent public API. The default block size moves to a named constant. Part of apache#18667.
SortedMailboxReceiveOperator accumulates every row from all senders and then sorts, even when each sender's stream is already sorted on the collation keys (noted as a TODO in the class). That defeats streaming and makes the receive stage's peak memory proportional to the full input. When the sending stage is known to produce sorted output, the receiver now performs a k-way heap merge across mailboxes and emits bounded sorted blocks as data arrives: - BlockingMultiStreamConsumer gains a StreamHandle<T> abstraction with awaitDataOrTerminal()/poll(), so a consumer can peek the head of each stream without draining it — the primitive the heap merge needs. - PlanFragmenter/PinotLogicalQueryPlanner propagate sender-side collation so the receiver can tell whether its inputs are individually sorted. - Opt-in via the `streamingSortedMailboxReceive` query option, with `streamingSortedMailboxReceiveBlockSize` controlling emitted block size. Without the option, the existing accumulate-then-sort path is used unchanged. Part of apache#18667.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #19121 +/- ##
============================================
+ Coverage 65.49% 65.60% +0.10%
Complexity 1423 1423
============================================
Files 3430 3434 +4
Lines 218010 218881 +871
Branches 34648 34853 +205
============================================
+ Hits 142784 143591 +807
- Misses 63666 63691 +25
- Partials 11560 11599 +39
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
SortedMailboxReceiveOperatoraccumulates every row from all senders and only then sorts, even wheneach sender's stream is already sorted on the collation keys — flagged as a TODO in the class today.
That defeats streaming and makes the receive stage's peak memory proportional to the whole input.
This is the second bullet of Challenge 1 in #18667.
Approach
When the planner can prove each sender emits sorted output, the receiver performs a k-way heap merge
across mailboxes and emits bounded sorted blocks as data arrives.
BlockingMultiStreamConsumergains aStreamHandle<T>abstraction(
awaitDataOrTerminal()/poll()) so a consumer can hold the head of each stream withoutdraining it — the primitive the heap merge needs. A consumer must use either the handles or
readMseBlockBlocking(), never both; the mode is latched.relieveSiblingsOnce()does one non-blocking poll per sibling stream per refill, so a fast senderwith a disjoint key range cannot sit blocked on a full mailbox holding an MSE worker thread.
SortedMailboxReceiveOperatorimplements the merge with a per-mailbox cursor that caches itshead row, and verifies the precondition at runtime: each sender's rows must be non-decreasing, and
a violation surfaces as an error block naming the offending mailbox instead of silently wrong
output.
PlanFragmenter/PinotLogicalQueryPlanner/QueryEnvironmentmark a leaf selectionORDER BYsender fragment as sorted-on-sender when the plan shape and collation match.This gate fails closed:
resolvesToSinglePhysicalTable()rejects hybrid tables, logicaltables, unknown tables, and a null
TableCache. This matters — a scan over a hybrid table compilesinto two
ServerQueryRequests thatLeafOperatorruns concurrently into one mailbox, producingtwo independently sorted runs concatenated, not a sorted stream. Empty collations are also
rejected, so a
LIMITwithoutORDER BY(which compiles to a collation-lessLogicalSort) isnever marked sorted.
BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED(new) reports per receive operator, inthe response
stageStats, whether the merge actually ran. Reporting is presence-based —StatMapdropsfalsebooleans, sokWayMergeUsed: trueappears on the merge path and nothingappears on the accumulate-then-sort path. This is the only runtime discriminator between the two
paths; both render as
type: MAILBOX_RECEIVE.Opt-in query options
streamingSortedMailboxReceivefalsestreamingSortedMailboxReceiveBlockSize10000sortedSelectionMergeEnabled(from the parent PR)falseORDER BYshape, where the rel plan contains no sort exchangeisSortedOnSenderis set either by a rel-level sort exchange that already declares sender-sidesorting, or — for a plain leaf selection
ORDER BY— byPlanFragmenterundersortedSelectionMergeEnabled. Only the latter case needs both options.Rolling-upgrade note
K_WAY_MERGE_USEDis appended as the lastStatKeyconstant.StatMapserializes keys byordinal and deserializes without a bounds check, so a peer on a build predating the key cannot
decode stage stats containing it. Because BOOLEAN keys are written only when
true, and the key isonly ever
truewhenstreamingSortedMailboxReceiveis on, leaving the option off (the default)keeps a mixed-version cluster safe. This is documented on the option constant. Bounds-checking the
ordinal in
StatMap.merge(DataInput)is a worthwhile separate fix.No behaviour change when the option is off
With
streamingSortedMailboxReceiveunset, the existing accumulate-then-sort path runs unchanged andno new stat is emitted. The planner-side marking is additionally a no-op under
usePhysicalOptimizer, since the v2 path does not go throughPlanFragmenter.Tests
SortedMailboxReceiveOperatorTest(extended)PlanFragmenterTest(new)SortOperatorTest(extended)BlockingMultiStreamConsumerTest(extended)pinot-query-plannersuiteStreamingSortedMailboxReceiveTest(new integration test)The integration test asserts
kWayMergeUsedinstageStatson the merge arm and its absence on thebaseline arm, so it cannot pass with the merge path deleted.
What this buys, honestly
On a local cluster over an 87.8M-doc, 56-segment table, the merge is confirmed engaged
(
kWayMergeUsed: trueon both join-input receives) and results are identical to theaccumulate-then-sort path. It is not a throughput win for that workload — enabling the mailbox
merge alone was the slowest sorted variant measured. The benefit is bounded peak memory in the
receive stage and latency to first row, not wall clock. Row order among equal collation keys is
unspecified and may differ between the two paths, as with any SQL
ORDER BY; the output multiset isidentical.
Stacking
Part of #18667.