andygrove opened a new pull request, #2181: URL: https://github.com/apache/datafusion-ballista/pull/2181
# Which issue does this PR close? Closes #2166. Closes #2167. # Rationale for this change `CoalescePartitionsRule` bailed out of an entire stage under two conditions: when any leaf `ExchangeExec` was a broadcast exchange (#2166), and when the leaves disagreed on upstream partition count `M` (#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 `M` from leaf 0's *declared* partitioning but indexed the summed byte vector by the *resolved* `shuffle_partitions()` vector. Those two disagree for a broadcast exchange: `ExchangeExec::new_with_details` forces `Partitioning::UnknownPartitioning(1)` when `broadcast` is set, while `StageOutput::partition_locations_broadcast()` returns one entry per upstream partition. A broadcast leaf at index 0 therefore yielded `m = 1`, and a sibling with 8 resolved partitions wrote past the end of `vec![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 `M` provably cannot be sides of the same partitioned join, since `EnforceDistribution` equalises the partition counts of a `Partitioned` join'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 by `M`, sum each group's per-partition bytes, bin-pack each group independently, and attach one shared `CoalescePlan` to every member of a decided group. - `M` now 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. - Broadcast leaves are excluded from grouping instead of disabling it. - Remaining leaves are grouped by `M` and decided per group. - Pass-through exchanges (`partitioning: None`) are also excluded. `DistributedExchangeRule` inserts these beneath `SortPreservingMergeExec` and `CoalescePartitionsExec`. A coalesced reader concatenates several upstream partitions into one output partition and `ShuffleReaderExec` randomises location order for executor load balancing, so per-partition ordering does not survive coalescing and a `SortPreservingMergeExec` above 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. - Missing `num_bytes` on 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 == 1` is allowed. The old `K <= 1` guard 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 is `K >= M`, which subsumes `M == 1` correctly. - `ExchangeExec::set_coalesce` takes an `Option`, 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_broadcast` no 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 `M` is sufficient rather than merely convenient, why the wholesale clear is safe only while `actionable_stages` runs 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 via `Arc::ptr_eq` that group members share one plan rather than merely agreeing on `K`. End-to-end tests through `AdaptivePlanner` for a broadcast join stage, for a stage collapsing to `K == 1`, and for the `ORDER BY` shape that must not coalesce. The six pre-existing snapshots are unchanged. # Are there any user-facing changes? No public API changes; `ExchangeExec` and 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. -- 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]
