andygrove opened a new issue, #2166: URL: https://github.com/apache/datafusion-ballista/issues/2166
## 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) -- 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]
