fix: coalesce AQE shuffle partitions per alignment group - #2181
Draft
andygrove wants to merge 10 commits into
Draft
fix: coalesce AQE shuffle partitions per alignment group#2181andygrove wants to merge 10 commits into
andygrove wants to merge 10 commits into
Conversation
set_coalesce now takes an Option so a rule pass can unset a decision, and to_broadcast no longer carries a coalesce plan into a node that cannot use one.
Makes the summing and bin-packing testable without a planner, and replaces the K<=1 and K>=M guards with the single K>=M test, which allows a tiny stage to collapse to one downstream partition.
The rule read M from leaf 0's declared partitioning but indexed the summed byte vector by the resolved location vector, which is how a broadcast leaf (declared 1, resolved M) drove the TPC-H Q22 out-of-bounds. Classification now checks the two against each other and treats missing byte statistics as unusable rather than as zero.
Groups leaf exchanges by upstream partition count and decides each group independently, instead of bailing the whole stage when any leaf is a broadcast exchange or when leaves disagree on partition count. Closes apache#2166 Closes apache#2167
Document what the coalesce rule does not reason about — per-partition ordering does not survive a coalesced read, and the bin-pack never splits a skewed partition — so the alignment argument is not over-trusted. Record why grouping by M stays checkable by inspection, and why the wholesale clear is only safe while the rule and the adapter share one per-stage closure. Restore the bin-pack K/M log line on both paths and the root-identity line that attributes leaf and group lines to their stage, and summarise a Sizes leaf by length rather than dumping every element. Make the unreachable non-Sizes arm of the per-group collection panic rather than silently under-coalesce, and give the grouping fixture's leaves distinct plan ids. Add end-to-end coverage for the two cases only the decision layer had: a stage holding a broadcast leaf beside a shuffle leaf, driven through AdaptivePlanner to the ShuffleReaderExec, and a stage packing to K = 1.
…t groups CoalescePartitionsRule could attach a coalesce decision to a pass-through ExchangeExec (partitioning: None) sitting directly beneath a SortPreservingMergeExec. DistributedExchangeRule inserts exactly this shape to carry a stage's ordering across the boundary; a coalesced ShuffleReaderExec shuffles its concatenated partition locations, destroying that ordering and letting the SortPreservingMergeExec above silently emit wrongly ordered ORDER BY results. Exclude pass-through leaves from alignment groups the same way broadcast leaves already are: classify_leaf reports a new PassThrough LeafKind, and group_by_upstream_count drops it via the same non-poisoning continue path so it does not suppress coalescing for unrelated same-M siblings.
… ordering The Algorithm section's step 3 still said only broadcast leaves drop out during classification, contradicting the pass-through exclusion added above it. Also note why the broadcast check in classify_leaf must run before the partitioning::is_none() check: both new_broadcast and to_broadcast set partitioning to None, so reordering the checks would silently misclassify every broadcast leaf as pass-through.
Grouping now hands each member its own size slice instead of an index into a parallel vector, so the caller sums a group without re-matching on LeafKind and the unreachable! that guarded that match is gone. The rule also decides first and applies second, writing every leaf's slot exactly once per pass with its final value. That removes the window in which a leaf held a cleared slot, and with it the cross-module requirement that actionable_stages keep optimize and adapt in one closure. Also: name the pass-through condition as ExchangeExec::preserves_child_ordering so the rule and maintains_input_order share one definition, read byte counts under the lock rather than deep-cloning every PartitionLocation, and drop the hand-spelled LeafKind label arms that Debug already produces.
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
1 participant
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.
Which issue does this PR close?
Closes #2166.
Closes #2167.
Rationale for this change
CoalescePartitionsRulebailed out of an entire stage under two conditions: when any leafExchangeExecwas a broadcast exchange (#2166), and when the leaves disagreed on upstream partition countM(#2167). Both guards were added in response to an out-of-bounds panic on TPC-H Q22, and both were documented in the code as temporary.They share one root cause. The rule read
Mfrom leaf 0's declared partitioning but indexed the summed byte vector by the resolvedshuffle_partitions()vector. Those two disagree for a broadcast exchange:ExchangeExec::new_with_detailsforcesPartitioning::UnknownPartitioning(1)whenbroadcastis set, whileStageOutput::partition_locations_broadcast()returns one entry per upstream partition. A broadcast leaf at index 0 therefore yieldedm = 1, and a sibling with 8 resolved partitions wrote past the end ofvec![0u64; 1].Fixing the cause makes both guards unnecessary. Broadcast leaves have no partition structure to coalesce and are never a co-partitioned join sibling, so they can be excluded rather than allowed to disable the stage. Leaves with different
Mprovably cannot be sides of the same partitioned join, sinceEnforceDistributionequalises the partition counts of aPartitionedjoin's inputs, so they can be grouped and decided independently.The practical effect of the guards was that any stage mixing a broadcast leg with shuffle legs, or holding leaves at different partition counts, kept one task per upstream partition no matter how small the per-partition output. That is exactly the small-task overhead the rule exists to remove.
What changes are included in this PR?
The rule (
ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs) is restructured into a short driver over pure functions: classify each leaf, group the survivors byM, sum each group's per-partition bytes, bin-pack each group independently, and attach one sharedCoalescePlanto every member of a decided group.Mnow comes from the resolved shuffle shape, checked against the declared partition count. A mismatch is reported and its group skipped, so the Q22 out-of-bounds is unrepresentable rather than guarded against.Mand decided per group.partitioning: None) are also excluded.DistributedExchangeRuleinserts these beneathSortPreservingMergeExecandCoalescePartitionsExec. A coalesced reader concatenates several upstream partitions into one output partition andShuffleReaderExecrandomises location order for executor load balancing, so per-partition ordering does not survive coalescing and aSortPreservingMergeExecabove would merge unsorted streams. This one predates the change (the old rule coalesced the same single-leaf shape), but per-group coalescing widens where it can fire, so it is fixed here rather than left in place.num_byteson a partition location now makes a leaf unusable and skips its group. Previously it counted as zero bytes, which under-counted and over-coalesced into oversized tasks.K == 1is allowed. The oldK <= 1guard refused to act on a stage whose whole output fits in one target-sized partition, which is where per-task overhead dominates most. The only degenerate case now isK >= M, which subsumesM == 1correctly.ExchangeExec::set_coalescetakes anOption, and the rule clears every leaf it collects before deciding anything. A group can therefore never be left with one member carrying a decision from an earlier pass and a sibling carrying none. Relatedly,ExchangeExec::to_broadcastno longer clones the coalesce slot into the promoted node, which also removes the aliasing that made the two nodes share one slot.Module documentation now records the alignment-group argument, why grouping by
Mis sufficient rather than merely convenient, why the wholesale clear is safe only whileactionable_stagesruns the rule and the adapter in the same per-stage closure, and what the rule does not reason about (skew).Tests. Unit tests for the extracted pure functions and for leaf classification, including a regression guard for the Q22 shape. Rule-level tests for a stage mixing a broadcast leaf with shuffle leaves and for a stage with heterogeneous
M, asserting viaArc::ptr_eqthat group members share one plan rather than merely agreeing onK. End-to-end tests throughAdaptivePlannerfor a broadcast join stage, for a stage collapsing toK == 1, and for theORDER BYshape that must not coalesce. The six pre-existing snapshots are unchanged.Are there any user-facing changes?
No public API changes;
ExchangeExecand the rule are private to the scheduler crate. No configuration keys added or changed.Behaviour changes only when
ballista.planner.coalesce.enabled=true, which is off by default, as is the adaptive planner itself. Under that flag, stages that previously kept one task per upstream partition may now run fewer, larger tasks, and a stage whose output fits in one target-sized partition may collapse to a single task.Benchmarks
Plan shape changes when
coalesce.enabled=true, so TPC-H numbers under that flag are wanted before merge, Q22 in particular since it is the query that originally panicked.