apache / apache/datafusion-ballista
AQE coalesce rule bails the entire alignment group when any leaf is a broadcast exchange
- Dominant language
- Rust
- Stars
- 2.1k
- Forks
- 320
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 66
Description
## Problem
The AQE coalesce-shuffle-partitions rule (`CoalescePartitionsRule`) bails out of the **entire alignment group** as soon as any leaf `ExchangeExec` is a broadcast exchange:
https://github.com/apache/datafusion-ballista/blob/main/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs#L224-L229
```rust
// this is temporary fix until we figure it out how to
// make this work with broadcast
if leaves.iter().any(|arc| as_exchange(arc).broadcast) {
debug!("[coalesce-rule] broadcast leaf present; bail entire group");
return Ok(plan);
}
```
This is documented in the code as a temporary fix. The comment is the only place the limitation is tracked; there is no issue for it, and it is not on the "enable AQE by default" checklist (#2092).
## Impact
When a stage reads from a mix of shuffle exchanges and at least one broadcast exchange (e.g. a hash/shuffle join whose plan also contains a broadcast leg), coalescing is skipped for **all** leaves in that stage — including the non-broadcast shuffle inputs that could safely be coalesced. Downstream stages keep one task per upstream partition even when the per-partition output is tiny, which is exactly the small/empty-task overhead the coalesce rule exists to remove.
Broadcast leaves have a single partition and are not part of the hash co-partitioning invariant that ties the alignment group together, so in principle they can be excluded from the alignment group and the remaining shuffle leaves coalesced normally — but that interaction has not been worked out, hence the current all-or-nothing bail.
## Proposed direction
- Exclude broadcast leaves from the alignment group rather than bailing the whole group.
- Coalesce the remaining shuffle leaves (which share `M` and the hash mapping) as usual.
- Add coverage for a stage subtree that mixes a broadcast leaf with shuffle leaves, asserting the shuffle leaves still get a `CoalescePlan` and the broadcast leaf is left untouched.
## Acceptance criteria
- [ ] A broadcast leaf no longer disables coalescing for the sibling shuffle leaves in the same stage.
- [ ] Hash co-partitioning across the (non-broadcast) alignment group is preserved after the rewrite.
- [ ] Test covering the mixed broadcast + shuffle leaf case.
Parent epic: #2092 (Enable AQE by default) / #1359 (AQE design)
Contributor guide
Research direction
Start in ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs around the broadcast-leaf guard, then trace how alignment groups are built and rewritten. Add coverage for a mixed broadcast and shuffle subtree; done means shuffle leaves receive CoalescePlan, the broadcast leaf is unchanged, and hash co-partitioning remains intact.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100