andygrove opened a new issue, #2167: URL: https://github.com/apache/datafusion-ballista/issues/2167
## Problem The AQE coalesce-shuffle-partitions rule (`CoalescePartitionsRule`) assumes every leaf `ExchangeExec` in a stage's alignment group shares the same upstream partition count `M`. It reads `M` from leaf 0 and, if any other leaf disagrees, bails the whole group: https://github.com/apache/datafusion-ballista/blob/main/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs#L231-L246 ```rust let m = as_exchange(&leaves[0]) .properties() .partitioning .partition_count(); // TODO: per-M subgrouping; for now bail on heterogeneous M (Q22 panic guard). if leaves .iter() .any(|arc| as_exchange(arc).properties().partitioning.partition_count() != m) { return Ok(plan); } ``` The bail was added as a guard after a panic on TPC-H Q22 (the element-wise byte-sum over `summed[i]` indexes into a `vec![0u64; m]` and would go out of bounds when a leaf has more than `M` partitions). It is tracked only as an inline `TODO`; there is no issue for it and it is not on the "enable AQE by default" checklist (#2092). ## Impact Any stage whose leaf exchanges do not all share the same partition count is skipped entirely by the coalesce rule, so it keeps one task per upstream partition. The guard is correct (it prevents the panic) but it is a coverage gap: heterogeneous-`M` stages never benefit from coalescing. ## Proposed direction Implement per-`M` subgrouping: partition the leaves into alignment subgroups by their upstream partition count, and run the sum → bin-pack → `set_coalesce` path independently per subgroup, so each subgroup coalesces against its own `M`. Leaves that must stay co-partitioned with each other (same hash key, same `M`) land in the same subgroup, preserving the join partition-count invariant within the subgroup. ## Acceptance criteria - [ ] Coalescing runs on stages with heterogeneous leaf partition counts instead of bailing. - [ ] No out-of-bounds / panic on the byte-sum step (regression guard for the original Q22 panic). - [ ] Hash co-partitioning is preserved within each per-`M` subgroup. - [ ] Test covering a stage with leaves of differing `M`. 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]
