AQE coalesce rule bails on heterogeneous leaf partition counts (per-M subgrouping TODO / Q22 panic guard)

[andygrove]

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

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)