apache / apache/datafusion-ballista

AQE coalesce rule bails the entire alignment group when any leaf is a broadcast exchange

Open
#2,166 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.