From 232d76110bfbaba270a0495fb4295891b2224786 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 08:02:14 -0600 Subject: [PATCH 01/10] refactor(scheduler): let the AQE coalesce slot be cleared set_coalesce now takes an Option so a rule pass can unset a decision, and to_broadcast no longer carries a coalesce plan into a node that cannot use one. --- .../src/state/aqe/execution_plan/exchange.rs | 71 +++++++++++++++++-- .../aqe/optimizer_rule/coalesce_partitions.rs | 2 +- 2 files changed, 65 insertions(+), 8 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index 63cc181502..ea72c5f4ee 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs @@ -132,7 +132,10 @@ impl ExchangeExec { plan_id, self.stage_id.clone(), self.shuffle_partitions.clone(), - self.coalesce.clone(), + // A broadcast exchange never coalesces: its reader flattens every + // upstream location into one partition. Start with an empty slot + // rather than inheriting a decision the new node cannot use. + Arc::new(Mutex::new(None)), true, self.inactive_stage, ) @@ -245,12 +248,16 @@ impl ExchangeExec { &self.input } - /// Attaches a `CoalescePlan` to this Exchange. The adapter consumes the - /// plan when converting Exchange → ShuffleReader: a Some value triggers - /// `try_new_coalesced` (K-partition reader); None uses `try_new` - /// (M-partition reader). Idempotent overwrite. - pub fn set_coalesce(&self, cp: Arc) { - self.coalesce.lock().replace(cp); + /// Attaches or clears the `CoalescePlan` on this Exchange. The adapter + /// consumes the plan when converting Exchange → ShuffleReader: a `Some` + /// value triggers `try_new_coalesced` (K-partition reader); `None` uses + /// `try_new` (M-partition reader). + /// + /// `CoalescePartitionsRule` clears every leaf it collected before deciding + /// anything, so passing `None` is a normal part of a rule pass and not an + /// error path. + pub fn set_coalesce(&self, cp: Option>) { + *self.coalesce.lock() = cp; } /// Returns the attached `CoalescePlan`, if `set_coalesce` was called. @@ -421,3 +428,53 @@ impl ExecutionPlan for ExchangeExec { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::aqe::test::mock_schema; + use ballista_core::execution_plans::PartitionGroup; + use datafusion::physical_plan::empty::EmptyExec; + + fn a_plan() -> Arc { + Arc::new(CoalescePlan { + upstream_partition_count: 4, + groups: vec![PartitionGroup { + upstream_indices: vec![0, 1, 2, 3], + }], + }) + } + + fn an_exchange() -> ExchangeExec { + ExchangeExec::new( + Arc::new(EmptyExec::new(mock_schema())), + Some(Partitioning::UnknownPartitioning(4)), + 0, + ) + } + + #[test] + fn set_coalesce_none_clears_a_previous_decision() { + let exchange = an_exchange(); + exchange.set_coalesce(Some(a_plan())); + assert!(exchange.coalesce().is_some()); + + exchange.set_coalesce(None); + assert!(exchange.coalesce().is_none()); + } + + #[test] + fn to_broadcast_does_not_carry_the_coalesce_decision_forward() { + // A broadcast reader flattens every upstream location into one + // partition, so a decision made while this exchange was a shuffle is + // meaningless afterwards and must not follow it across. + let exchange = an_exchange(); + exchange.set_coalesce(Some(a_plan())); + + let broadcast = exchange.to_broadcast(1); + + assert!(broadcast.coalesce().is_none()); + // The original is untouched: `to_broadcast` builds a new node. + assert!(exchange.coalesce().is_some()); + } +} diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 5947be4c3d..6572d161f6 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -302,7 +302,7 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { ex.plan_id, ex.coalesce().as_ref().map(|cp| cp.groups.len()), ); - ex.set_coalesce(cp.clone()); + ex.set_coalesce(Some(cp.clone())); } Ok(plan) } From dc460ea7204dd8dd89dde616e3970b6b0cabe7e6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 08:07:03 -0600 Subject: [PATCH 02/10] refactor(scheduler): extract sum_sizes and decide from the coalesce rule Makes the summing and bin-packing testable without a planner, and replaces the K<=1 and K>=M guards with the single K>=M test, which allows a tiny stage to collapse to one downstream partition. --- .../aqe/optimizer_rule/coalesce_partitions.rs | 187 ++++++++++++++---- 1 file changed, 149 insertions(+), 38 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 6572d161f6..25f9f5b93b 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -76,7 +76,7 @@ //! 4. Bin-pack the summed sizes into `K` buckets near //! `target_partition_bytes` (Spark's `advisoryPartitionSizeInBytes`, //! 64 MB by default) using `split_size_list_by_target_size`. -//! 5. If `K >= M` or `K <= 1`, the rewrite is degenerate and is skipped. +//! 5. If `K >= M`, the rewrite is not a reduction and is skipped. //! 6. Otherwise, attach a shared [`CoalescePlan`] (with `K` partition //! groups) to every leaf `ExchangeExec` via `set_coalesce(..)`. The //! adapter consumes that decision when it builds the downstream @@ -124,6 +124,51 @@ use crate::state::aqe::coalesce::{ use crate::state::aqe::execution_plan::AdaptiveDatafusionExec; use crate::state::aqe::execution_plan::ExchangeExec; +/// Sum an alignment group's per-upstream-partition byte counts element-wise. +/// +/// `summed[i]` is the total downstream work for upstream index `i` across the +/// whole group: for a partitioned join, the task reading group `i` reads that +/// index from both sides, so their sizes add. +/// +/// `m` is the group's upstream partition count. Slices are zipped against the +/// accumulator rather than indexed, so a shorter or longer member cannot index +/// out of bounds. Addition saturates. +fn sum_sizes(sizes: &[&[u64]], m: usize) -> Vec { + let mut summed = vec![0u64; m]; + for leaf in sizes { + for (acc, &size) in summed.iter_mut().zip(leaf.iter()) { + *acc = acc.saturating_add(size); + } + } + summed +} + +/// Bin-pack a summed size list into a coalesce decision. +/// +/// Returns `None` when the rewrite would not reduce the partition count +/// (`K >= M`). That single test covers every degenerate case: an empty list and +/// a one-partition input both pack to `K = 1`, which is not below their `M`. +/// A `K` of 1 over a larger `M` *is* a reduction and is returned. +fn decide( + summed: &[u64], + target: u64, + small_factor: f64, + merged_factor: f64, +) -> Option { + let m = summed.len(); + let starts = + split_size_list_by_target_size(summed, target, small_factor, merged_factor); + let k = starts.len(); + if k >= m { + debug!("[coalesce-rule] K={k} is not a reduction from M={m}; no decision"); + return None; + } + Some(CoalescePlan { + upstream_partition_count: m as u32, + groups: start_indices_to_partition_groups(&starts, m), + }) +} + /// AQE rule that attaches a coalesce decision to every leaf `ExchangeExec` /// feeding the current stage, so the downstream reader exposes `K < M` /// partitions. @@ -245,13 +290,11 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { return Ok(plan); } - // Sum byte sizes element-wise across the alignment group. Upstream - // partition `i` is the same logical hash bucket on every leaf, so - // `summed[i]` is the total downstream work for that bucket. If any - // leaf is still unresolved we bail — early `replan_stages()` passes - // run before all upstream stages finalize, and the rule reruns on - // every later pass anyway, so the no-op is free. - let mut summed = vec![0u64; m]; + // Collect each leaf's per-partition byte sizes. If any leaf is still + // unresolved we bail — early `replan_stages()` passes run before all + // upstream stages finalize, and the rule reruns on every later pass + // anyway, so the no-op is free. + let mut sizes: Vec> = Vec::with_capacity(leaves.len()); for arc in &leaves { let ex = as_exchange(arc); let Some(parts) = ex.shuffle_partitions() else { @@ -261,44 +304,31 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { ); return Ok(plan); }; - for (i, locs) in parts.iter().enumerate() { - summed[i] += locs + sizes.push( + parts .iter() - .filter_map(|l| l.partition_stats.num_bytes()) - .sum::(); - } + .map(|locs| { + locs.iter() + .filter_map(|l| l.partition_stats.num_bytes()) + .fold(0u64, |acc, b| acc.saturating_add(b)) + }) + .collect(), + ); } + + let borrowed: Vec<&[u64]> = sizes.iter().map(|s| s.as_slice()).collect(); + let summed = sum_sizes(&borrowed, m); debug!("[coalesce-rule] summed bytes per upstream partition: {summed:?}"); - // One bin-pack decision for the whole alignment group, packing toward - // `target_partition_bytes` (Spark's `advisoryPartitionSizeInBytes`). - // The rule is opt-in (`coalesce.enabled=false` by default), so users - // get parallelism preservation unless they explicitly trade it for - // larger tasks. This corresponds to Spark's - // `parallelismFirst=false` mode — direct advisory-driven packing. - let starts = split_size_list_by_target_size(&summed, target, small, merged); - let k = starts.len(); - debug!("[coalesce-rule] bin-pack result: K={k} M={m}"); - if k >= m || k <= 1 { - debug!( - "[coalesce-rule] K degenerate (K>=M or K<=1); bail without setting coalesce" - ); + let Some(cp) = decide(&summed, target, small, merged) else { return Ok(plan); - } - - // Attach the same `CoalescePlan` to every member of the alignment - // group. Sharing the plan (not just the K value) keeps the upstream - // index → group mapping identical across leaves — hash buckets that - // were aligned at M stay aligned at K, and the join's - // partition-count requirement still holds after the rewrite. - let cp = Arc::new(CoalescePlan { - upstream_partition_count: m as u32, - groups: start_indices_to_partition_groups(&starts, m), - }); + }; + let cp = Arc::new(cp); for arc in &leaves { let ex = as_exchange(arc); debug!( - "[coalesce-rule] set_coalesce(K={k}) on plan_id={} (was {:?})", + "[coalesce-rule] set_coalesce(K={}) on plan_id={} (was {:?})", + cp.groups.len(), ex.plan_id, ex.coalesce().as_ref().map(|cp| cp.groups.len()), ); @@ -315,3 +345,84 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { false } } + +#[cfg(test)] +mod tests { + use super::*; + + /// The production defaults, at the byte scale the existing snapshot tests + /// use, so a trace here lines up with a trace there. + const TARGET: u64 = 200; + const SMALL: f64 = 0.2; + const MERGED: f64 = 1.2; + + fn decide_at_defaults(summed: &[u64]) -> Option { + decide(summed, TARGET, SMALL, MERGED) + } + + #[test] + fn sum_sizes_adds_element_wise() { + let a: &[u64] = &[1, 2, 3]; + let b: &[u64] = &[10, 20, 30]; + assert_eq!(sum_sizes(&[a, b], 3), vec![11, 22, 33]); + } + + #[test] + fn sum_sizes_of_no_leaves_is_all_zero() { + assert_eq!(sum_sizes(&[], 3), vec![0, 0, 0]); + } + + #[test] + fn sum_sizes_saturates_instead_of_overflowing() { + // Byte counts come off the wire; a corrupt pair must not panic a + // release-mode scheduler differently from a debug one. + let a: &[u64] = &[u64::MAX]; + let b: &[u64] = &[1]; + assert_eq!(sum_sizes(&[a, b], 1), vec![u64::MAX]); + } + + #[test] + fn decide_declines_when_every_partition_is_already_full() { + // 300 > 200 so each partition flushes on its own and the post-flush + // merge is rejected: K = M = 8, no reduction, nothing to do. + assert!(decide_at_defaults(&[300; 8]).is_none()); + } + + #[test] + fn decide_declines_for_a_single_upstream_partition() { + // M = 1 always packs to K = 1, which is not a reduction. + assert!(decide_at_defaults(&[10]).is_none()); + } + + #[test] + fn decide_declines_for_an_empty_size_list() { + assert!(decide_at_defaults(&[]).is_none()); + } + + #[test] + fn decide_packs_eight_fiftys_into_two_partitions() { + // 4 x 50 fills a bucket to exactly 200; the fifth overshoots and + // flushes; the remaining 4 fill the second bucket; the post-loop merge + // is rejected (400 >= 200 * 1.2). This is the trace the existing + // end-to-end snapshot asserts. + let plan = decide_at_defaults(&[50; 8]).expect("K=2 is a reduction from M=8"); + assert_eq!(plan.upstream_partition_count, 8); + assert_eq!(plan.groups.len(), 2); + assert_eq!(plan.groups[0].upstream_indices, vec![0, 1, 2, 3]); + assert_eq!(plan.groups[1].upstream_indices, vec![4, 5, 6, 7]); + } + + #[test] + fn decide_packs_a_tiny_stage_into_a_single_partition() { + // 8 x 10 = 80 never reaches the 200 target, so the whole stage becomes + // one downstream task. Previously refused by a `K <= 1` guard, which is + // exactly the case where per-task overhead dominates. + let plan = decide_at_defaults(&[10; 8]).expect("K=1 is a reduction from M=8"); + assert_eq!(plan.upstream_partition_count, 8); + assert_eq!(plan.groups.len(), 1); + assert_eq!( + plan.groups[0].upstream_indices, + vec![0, 1, 2, 3, 4, 5, 6, 7] + ); + } +} From c2ca76e2783a6bc3c29d8979df02788567588253 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 08:13:53 -0600 Subject: [PATCH 03/10] fix(scheduler): read coalesce M from the resolved shuffle shape The rule read M from leaf 0's declared partitioning but indexed the summed byte vector by the resolved location vector, which is how a broadcast leaf (declared 1, resolved M) drove the TPC-H Q22 out-of-bounds. Classification now checks the two against each other and treats missing byte statistics as unusable rather than as zero. --- .../aqe/optimizer_rule/coalesce_partitions.rs | 235 ++++++++++++++---- .../src/state/aqe/test/coalesce_rule.rs | 40 +-- ballista/scheduler/src/state/aqe/test/mod.rs | 45 ++++ 3 files changed, 236 insertions(+), 84 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 25f9f5b93b..95c8043594 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -124,6 +124,85 @@ use crate::state::aqe::coalesce::{ use crate::state::aqe::execution_plan::AdaptiveDatafusionExec; use crate::state::aqe::execution_plan::ExchangeExec; +/// What one leaf `ExchangeExec` contributes to its alignment group. +#[derive(Debug, PartialEq, Eq)] +enum LeafKind { + /// A broadcast exchange. Its reader flattens every upstream location into a + /// single output partition, so there is no partition structure to coalesce, + /// and a `CollectLeft` join never requires it to be co-partitioned with + /// anything. Excluded from every alignment group. + Broadcast, + /// The upstream stage has not finalized. The group defers to a later pass. + Unresolved, + /// At least one location reports no `num_bytes`. + UnknownBytes, + /// The resolved location vector's length disagrees with the declared + /// partition count. Should be unreachable for a non-broadcast exchange. + Inconsistent { declared: usize, resolved: usize }, + /// Per-upstream-partition byte counts. Length equals the declared count. + Sizes(Vec), +} + +/// A leaf `ExchangeExec` reduced to what the rule needs: its alignment-group +/// key and its contribution to that group. +#[derive(Debug)] +struct ClassifiedLeaf { + /// Upstream partition count `M`, the alignment-group key. + m: usize, + kind: LeafKind, +} + +/// Reduce one leaf `ExchangeExec` to a [`ClassifiedLeaf`]. +/// +/// `m` is the declared partition count. For a `Sizes` leaf the resolved shape +/// is checked to match it, so `m == sizes.len()` holds for every leaf that +/// reaches the summing step and the summed vector can never be indexed out of +/// bounds. +fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { + let m = ex.properties().partitioning.partition_count(); + if ex.broadcast { + return ClassifiedLeaf { + m, + kind: LeafKind::Broadcast, + }; + } + let Some(parts) = ex.shuffle_partitions() else { + return ClassifiedLeaf { + m, + kind: LeafKind::Unresolved, + }; + }; + if parts.len() != m { + return ClassifiedLeaf { + m, + kind: LeafKind::Inconsistent { + declared: m, + resolved: parts.len(), + }, + }; + } + let mut sizes = Vec::with_capacity(parts.len()); + for locations in &parts { + let mut total = 0u64; + for location in locations { + match location.partition_stats.num_bytes() { + Some(bytes) => total = total.saturating_add(bytes), + None => { + return ClassifiedLeaf { + m, + kind: LeafKind::UnknownBytes, + }; + } + } + } + sizes.push(total); + } + ClassifiedLeaf { + m, + kind: LeafKind::Sizes(sizes), + } +} + /// Sum an alignment group's per-upstream-partition byte counts element-wise. /// /// `summed[i]` is the total downstream work for upstream index `i` across the @@ -266,58 +345,36 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { return Ok(plan); } - // 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) { + let classified: Vec = leaves + .iter() + .map(|arc| classify_leaf(as_exchange(arc))) + .collect(); + + // Still one alignment group in this task; Task 4 splits it by M. + if classified.iter().any(|c| c.kind == LeafKind::Broadcast) { debug!("[coalesce-rule] broadcast leaf present; bail entire group"); return Ok(plan); } - - // The alignment-group invariant assumes a shared `M`. In every plan - // shape we currently produce, all leaves of one stage subtree are - // hash-partitioned by the same target_partitions setting upstream, - // so reading `M` from leaf 0 is sufficient. - 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) - { + let m = classified[0].m; + if classified.iter().any(|c| c.m != m) { + debug!("[coalesce-rule] heterogeneous M; bail entire group"); return Ok(plan); } - - // Collect each leaf's per-partition byte sizes. If any leaf is still - // unresolved we bail — early `replan_stages()` passes run before all - // upstream stages finalize, and the rule reruns on every later pass - // anyway, so the no-op is free. - let mut sizes: Vec> = Vec::with_capacity(leaves.len()); - for arc in &leaves { - let ex = as_exchange(arc); - let Some(parts) = ex.shuffle_partitions() else { - debug!( - "[coalesce-rule] leaf plan_id={} unresolved; bail entire group", - ex.plan_id - ); - return Ok(plan); - }; - sizes.push( - parts - .iter() - .map(|locs| { - locs.iter() - .filter_map(|l| l.partition_stats.num_bytes()) - .fold(0u64, |acc, b| acc.saturating_add(b)) - }) - .collect(), - ); + let mut sizes: Vec<&[u64]> = Vec::with_capacity(classified.len()); + for (leaf, arc) in classified.iter().zip(&leaves) { + match &leaf.kind { + LeafKind::Sizes(s) => sizes.push(s.as_slice()), + other => { + debug!( + "[coalesce-rule] leaf plan_id={} is {other:?}; bail entire group", + as_exchange(arc).plan_id + ); + return Ok(plan); + } + } } - let borrowed: Vec<&[u64]> = sizes.iter().map(|s| s.as_slice()).collect(); - let summed = sum_sizes(&borrowed, m); + let summed = sum_sizes(&sizes, m); debug!("[coalesce-rule] summed bytes per upstream partition: {summed:?}"); let Some(cp) = decide(&summed, target, small, merged) else { @@ -425,4 +482,92 @@ mod tests { vec![0, 1, 2, 3, 4, 5, 6, 7] ); } + + use crate::state::aqe::test::{ + mock_schema, partitions_with_byte_sizes, partitions_with_optional_byte_sizes, + }; + use ballista_core::serde::scheduler::PartitionLocation; + use datafusion::physical_plan::Partitioning; + use datafusion::physical_plan::empty::EmptyExec; + + /// A shuffle exchange declaring `m` partitions, optionally resolved. + fn shuffle_exchange( + m: usize, + resolved: Option>>, + ) -> ExchangeExec { + let exchange = ExchangeExec::new( + Arc::new(EmptyExec::new(mock_schema())), + Some(Partitioning::UnknownPartitioning(m)), + 0, + ); + if let Some(parts) = resolved { + exchange.resolve_shuffle_partitions(parts); + } + exchange + } + + #[test] + fn classify_reads_sizes_from_the_resolved_shuffle_shape() { + let exchange = + shuffle_exchange(3, Some(partitions_with_byte_sizes(&[10, 20, 30]))); + + let leaf = classify_leaf(&exchange); + + assert_eq!(leaf.m, 3); + assert_eq!(leaf.kind, LeafKind::Sizes(vec![10, 20, 30])); + } + + #[test] + fn classify_reports_broadcast_regardless_of_resolved_shape() { + // The Q22 regression guard. A broadcast exchange declares + // `UnknownPartitioning(1)` but resolves to one entry per upstream + // partition. Reading M from the declaration and indexing by the + // resolution is what panicked; classification never lets a broadcast + // leaf reach the summing step at all. + let exchange = + ExchangeExec::new_broadcast(Arc::new(EmptyExec::new(mock_schema())), None, 0); + exchange.resolve_shuffle_partitions(partitions_with_byte_sizes(&[50; 8])); + + let leaf = classify_leaf(&exchange); + + assert_eq!(leaf.kind, LeafKind::Broadcast); + } + + #[test] + fn classify_reports_inconsistent_when_the_resolved_shape_disagrees() { + let exchange = shuffle_exchange(8, Some(partitions_with_byte_sizes(&[50; 4]))); + + let leaf = classify_leaf(&exchange); + + assert_eq!( + leaf.kind, + LeafKind::Inconsistent { + declared: 8, + resolved: 4 + } + ); + } + + #[test] + fn classify_reports_unresolved_before_the_upstream_stage_finalizes() { + let exchange = shuffle_exchange(8, None); + + assert_eq!(classify_leaf(&exchange).kind, LeafKind::Unresolved); + } + + #[test] + fn classify_reports_unknown_bytes_when_a_location_has_no_size() { + // One missing value makes the leaf's totals untrustworthy: counting it + // as zero would let real partitions pack into an oversized task. + let exchange = shuffle_exchange( + 3, + Some(partitions_with_optional_byte_sizes(&[ + Some(10), + None, + Some(30), + ])), + ); + + assert_eq!(classify_leaf(&exchange).kind, LeafKind::UnknownBytes); + } } diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 0ac47ae721..58b810a975 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -26,12 +26,8 @@ use crate::assert_plan; use crate::state::aqe::planner::AdaptivePlanner; -use crate::state::aqe::test::{mock_batch, mock_schema}; +use crate::state::aqe::test::{mock_batch, mock_schema, partitions_with_byte_sizes}; use ballista_core::extension::SessionConfigExt; -use ballista_core::serde::scheduler::{ - ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, - PartitionId, PartitionLocation, PartitionStats, -}; use datafusion::datasource::MemTable; use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{SessionConfig, SessionContext}; @@ -83,40 +79,6 @@ fn register_partitioned_table( Ok(()) } -/// Build a `Vec>` of length `per_partition_bytes.len()` -/// where each upstream partition reports the given byte size. The rule sums -/// `partition_stats.num_bytes` across leaves before bin-packing — that's the -/// only field these tests need to vary. -fn partitions_with_byte_sizes( - per_partition_bytes: &[u64], -) -> Vec> { - per_partition_bytes - .iter() - .enumerate() - .map(|(idx, &bytes)| { - vec![PartitionLocation { - map_partition_id: 0, - partition_id: PartitionId { - job_id: "".into(), - stage_id: 0, - partition_id: idx, - }, - executor_meta: ExecutorMetadata { - id: "".to_string(), - host: "".to_string(), - port: 0, - grpc_port: 0, - specification: ExecutorSpecification::default().with_vcores(0), - os_info: ExecutorOperatingSystemSpecification::default(), - }, - partition_stats: PartitionStats::new(Some(1), None, Some(bytes)), - file_id: None, - is_sort_shuffle: false, - }] - }) - .collect() -} - /// Happy path: M=8 upstream partitions @ 50 bytes each, target=200. /// Bin-pack trace (small_factor=0.2 → 40, merged_factor=1.2 → 240): /// i=0..3 accumulate into bucket=200; i=4 overshoots, flush, start new; diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs b/ballista/scheduler/src/state/aqe/test/mod.rs index 52f173d191..6af69b910e 100644 --- a/ballista/scheduler/src/state/aqe/test/mod.rs +++ b/ballista/scheduler/src/state/aqe/test/mod.rs @@ -92,6 +92,51 @@ pub(crate) fn mock_partitions_with_statistics_no_data() -> Vec>` of length `per_partition_bytes.len()` +/// with one location per upstream partition. `Some(n)` reports `n` bytes; +/// `None` reports no size at all, which is what the coalesce rule treats as +/// unusable statistics. +pub(crate) fn partitions_with_optional_byte_sizes( + per_partition_bytes: &[Option], +) -> Vec> { + per_partition_bytes + .iter() + .enumerate() + .map(|(idx, &bytes)| { + vec![PartitionLocation { + map_partition_id: 0, + partition_id: PartitionId { + job_id: "".into(), + stage_id: 0, + partition_id: idx, + }, + executor_meta: ExecutorMetadata { + id: "".to_string(), + host: "".to_string(), + port: 0, + grpc_port: 0, + specification: ExecutorSpecification::default().with_vcores(0), + os_info: ExecutorOperatingSystemSpecification::default(), + }, + partition_stats: PartitionStats::new(Some(1), None, bytes), + file_id: None, + is_sort_shuffle: false, + }] + }) + .collect() +} + +/// Build a `Vec>` where every upstream partition reports +/// the given byte size. The coalesce rule sums `partition_stats.num_bytes` +/// across leaves before bin-packing, so that is the only field these tests vary. +pub(crate) fn partitions_with_byte_sizes( + per_partition_bytes: &[u64], +) -> Vec> { + let optional: Vec> = + per_partition_bytes.iter().copied().map(Some).collect(); + partitions_with_optional_byte_sizes(&optional) +} + /// Returns schema with three columns (a,b,c) all of [DataType::Int32] type pub(crate) fn mock_schema() -> SchemaRef { Arc::new(Schema::new(vec![ From cfc6eab3b03739d5d1f3c962f3663c08fc6b8420 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 08:22:26 -0600 Subject: [PATCH 04/10] feat(scheduler): coalesce AQE shuffle partitions per alignment group Groups leaf exchanges by upstream partition count and decides each group independently, instead of bailing the whole stage when any leaf is a broadcast exchange or when leaves disagree on partition count. Closes #2166 Closes #2167 --- .../aqe/optimizer_rule/coalesce_partitions.rs | 381 +++++++++++------- 1 file changed, 245 insertions(+), 136 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 95c8043594..0f11a6b2e4 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -23,34 +23,37 @@ //! subtree, collects every leaf [`ExchangeExec`] — the resolved upstream //! shuffles feeding this stage — and decides whether to coalesce. //! -//! # The alignment group +//! # Alignment groups //! -//! Every leaf `ExchangeExec` in a single stage subtree forms one **alignment -//! group**. Why one group, not one decision per leaf? +//! Leaf Exchanges are grouped by upstream partition count `M`, and each group +//! is decided independently. //! -//! - Hash-partitioned joins (`HashJoinExec(Partitioned)`, `SortMergeJoinExec`) -//! require their two inputs to have the *same partition count* and to be -//! hash-partitioned on the join key. If we coalesced left to `K=4` and -//! right to `K=2`, DataFusion's `EnforceDistribution` would either reject -//! the plan or insert remediation repartitions that undo the optimization. -//! - Both join legs read shuffle output from upstream stages that wrote -//! `M` partitions using the *same* hash function on the *same* key -//! (that's what made them joinable in the first place). So upstream -//! partition `i` of the left and upstream partition `i` of the right -//! hold rows that must meet at downstream partition `f(i)`. Coalescing -//! them with the *same* mapping `i → group(i)` keeps that meeting point -//! consistent; coalescing them with different mappings scatters it. +//! Grouping exists because hash-partitioned joins (`HashJoinExec(Partitioned)`, +//! `SortMergeJoinExec`) require their two inputs to have the same partition +//! count and the same hash mapping. Both legs read shuffle output written with +//! the same hash function on the same key, so upstream partition `i` of the +//! left and upstream partition `i` of the right hold rows that must meet at one +//! downstream partition. Coalescing them with the same `i → group(i)` mapping +//! keeps that meeting point; coalescing them differently scatters it. Every +//! member of a group therefore gets the *same* `CoalescePlan`, not merely the +//! same `K`. //! -//! Practically: we treat all leaf Exchanges as a single workload, sum their -//! per-partition byte counts element-wise, bin-pack the summed sizes once, -//! and attach the *same* `CoalescePlan` to every leaf. Joins with two leaves -//! and chains of joins with three or more leaves all go through the same -//! code path — there is no per-leaf decision. +//! Grouping by `M` is sufficient, not merely convenient. DataFusion's +//! `EnforceDistribution` guarantees both sides of a `Partitioned` join have +//! equal partition counts, so two leaves feeding one partitioned join always +//! share an `M` and always land in the same group. Two leaves with different +//! `M` provably are not co-partitioned siblings of one join. //! -//! Concretely for `[25; 8]` bytes per partition on both sides of a join: -//! summed `[50; 8]`, bin-pack at target `200` produces `K=2` (4 upstream -//! partitions per group), both leaves get `coalesce=2 of 8`, the downstream -//! join runs with 2 partitions on each side, hash buckets stay aligned. +//! Broadcast leaves are excluded from every group. A broadcast reader flattens +//! all upstream locations into a single output partition, so there is nothing +//! to coalesce, and a `CollectLeft` join requires `SinglePartition` on the +//! build side and `UnspecifiedDistribution` on the probe side, so a broadcast +//! leaf is never a co-partitioned sibling. +//! +//! The grouping is conservative in the other direction: leaves that share an +//! `M` without sharing any co-partitioning requirement, such as the two arms of +//! a union, still land in one group and have their sizes summed. That inflates +//! the per-index totals and under-coalesces, which is safe. //! //! # Default off //! @@ -68,27 +71,27 @@ //! //! # Algorithm //! -//! 1. Find leaf `ExchangeExec`s — the alignment group. If empty, this -//! stage reads from scans and has nothing to coalesce. -//! 2. All leaves share the upstream partition count `M` (the writer side). -//! 3. Sum per-partition byte sizes element-wise across the group to get -//! combined work per upstream index. -//! 4. Bin-pack the summed sizes into `K` buckets near -//! `target_partition_bytes` (Spark's `advisoryPartitionSizeInBytes`, -//! 64 MB by default) using `split_size_list_by_target_size`. -//! 5. If `K >= M`, the rewrite is not a reduction and is skipped. -//! 6. Otherwise, attach a shared [`CoalescePlan`] (with `K` partition -//! groups) to every leaf `ExchangeExec` via `set_coalesce(..)`. The -//! adapter consumes that decision when it builds the downstream -//! `ShuffleReaderExec`s. +//! 1. Collect leaf `ExchangeExec`s. If none, this stage reads from scans and +//! has nothing to coalesce. +//! 2. Clear every leaf's coalesce slot, so the pass recomputes wholesale and +//! no decision survives from an earlier pass. +//! 3. Classify each leaf. Broadcast leaves drop out; unresolved, statistics- +//! free, and inconsistent leaves make their group undecidable. +//! 4. Group the rest by `M`. +//! 5. Per group: sum per-partition byte sizes element-wise, then bin-pack the +//! sums toward `target_partition_bytes` (Spark's +//! `advisoryPartitionSizeInBytes`, 64 MB by default) with +//! `split_size_list_by_target_size`. +//! 6. If `K >= M` the rewrite is not a reduction and the group is left alone. +//! Otherwise attach one shared [`CoalescePlan`] to every member. `K == 1` +//! is a legitimate outcome for a stage that fits in one partition. //! //! # Carrier semantics //! -//! The `CoalescePlan` lives on the upstream `ExchangeExec`; the rule does -//! not rewrite the plan tree. Idempotency is structural — `set_coalesce` -//! overwrites the slot with an equivalent plan on re-entry, and the -//! bin-pack is a pure function of the resolved byte sizes, so the second -//! pass produces the same decision. +//! The `CoalescePlan` lives on the upstream `ExchangeExec`; the rule does not +//! rewrite the plan tree. Idempotency comes from the clear-then-set contract in +//! step 2 combined with the bin-pack being a pure function of resolved byte +//! sizes, which do not change once a stage has finalized. //! //! # Grouping discipline //! @@ -103,11 +106,12 @@ //! //! # Behavior preservation //! -//! When `ballista.planner.coalesce.enabled=false`, when the subtree has no -//! leaf Exchanges, or when the bin-pack returns a degenerate `K`, the rule -//! is a no-op and returns the input `Arc` verbatim (preserving -//! `Arc::ptr_eq`). +//! The rule never rewrites the plan tree; decisions travel on each leaf's +//! interior-mutable slot and the input `Arc` is returned verbatim, preserving +//! `Arc::ptr_eq`. When `ballista.planner.coalesce.enabled=false` the rule +//! short-circuits before touching any slot. +use std::collections::BTreeMap; use std::sync::Arc; use ballista_core::config::BallistaConfig; @@ -116,7 +120,7 @@ use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::config::ConfigOptions; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_plan::ExecutionPlan; -use log::debug; +use log::{debug, warn}; use crate::state::aqe::coalesce::{ split_size_list_by_target_size, start_indices_to_partition_groups, @@ -203,6 +207,82 @@ fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { } } +/// Partition classified leaves into alignment groups keyed by upstream +/// partition count `M`. +/// +/// Broadcast leaves are dropped: they carry no partition structure to coalesce +/// and are never a co-partitioned join sibling, so excluding them cannot break +/// an alignment invariant. +/// +/// A group maps to `Some(indices)` when every member carries usable sizes, and +/// to `None` when any member is unresolved, missing byte statistics, or +/// inconsistent. Skipping is per group: one undecidable leaf no longer +/// suppresses coalescing for leaves it has no relationship with. +/// +/// `BTreeMap` rather than `HashMap` so iteration order, and therefore the debug +/// log and the order decisions are applied, is stable across passes. +fn group_by_upstream_count( + leaves: &[ClassifiedLeaf], +) -> BTreeMap>> { + let mut groups: BTreeMap>> = BTreeMap::new(); + for (idx, leaf) in leaves.iter().enumerate() { + if leaf.kind == LeafKind::Broadcast { + continue; + } + let entry = groups.entry(leaf.m).or_insert_with(|| Some(Vec::new())); + match &leaf.kind { + LeafKind::Sizes(_) => { + if let Some(members) = entry { + members.push(idx); + } + } + _ => *entry = None, + } + } + groups +} + +/// Unwrap a stage root to the subtree the rule inspects. +/// +/// The root is either an `ExchangeExec` (intermediate stage) or an +/// `AdaptiveDatafusionExec` (final stage). Anything else means the adapter is +/// about to fail anyway, so the rule declines rather than guessing. +fn stage_input(plan: &Arc) -> Option> { + if let Some(exchange) = plan.downcast_ref::() { + Some(exchange.input().clone()) + } else if let Some(adaptive) = plan.downcast_ref::() { + Some(adaptive.input().clone()) + } else { + None + } +} + +/// Collect every leaf `ExchangeExec` feeding this stage. +/// +/// `Jump` after each hit stops the walk from descending into the upstream +/// stage's compute: those nodes belong to whatever stage wrote them, not to +/// this one. +fn collect_leaf_exchanges( + input: &Arc, +) -> datafusion::common::Result>> { + let mut leaves: Vec> = Vec::new(); + input.apply(|node| { + if node.is::() { + leaves.push(node.clone()); + Ok(TreeNodeRecursion::Jump) + } else { + Ok(TreeNodeRecursion::Continue) + } + })?; + Ok(leaves) +} + +/// Downcast a collected leaf back to `&ExchangeExec`. +fn as_exchange(arc: &Arc) -> &ExchangeExec { + arc.downcast_ref::() + .expect("collect_leaf_exchanges filters to ExchangeExec") +} + /// Sum an alignment group's per-upstream-partition byte counts element-wise. /// /// `summed[i]` is the total downstream work for upstream index `i` across the @@ -279,117 +359,81 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { ); // Get the subtree below the root. Two root kinds, same outcome. - let input = if let Some(ex) = plan.downcast_ref::() { - debug!( - "[coalesce-rule] root=ExchangeExec plan_id={} stage_id={:?} stage_resolved={}", - ex.plan_id, - ex.stage_id(), - ex.shuffle_partitions().is_some(), - ); - ex.input().clone() - } else if let Some(adp) = plan.downcast_ref::() { - debug!( - "[coalesce-rule] root=AdaptiveDatafusionExec stage_id={:?}", - adp.stage_id(), - ); - adp.input().clone() - } else { + let Some(input) = stage_input(&plan) else { debug!( "[coalesce-rule] root is neither ExchangeExec nor AdaptiveDatafusionExec; bail" ); - return Ok(plan); // unexpected root — adapter will fail anyway, just bail + return Ok(plan); }; - // Collect the alignment group: every leaf `ExchangeExec` feeding - // this stage. `Jump` after each hit stops the walk from descending - // into the upstream stage's compute — those nodes aren't part of - // *this* stage's group, they belong to whatever stage wrote them. - let mut leaves: Vec> = Vec::new(); - input.apply(|node| { - if node.is::() { - leaves.push(node.clone()); - Ok(TreeNodeRecursion::Jump) - } else { - Ok(TreeNodeRecursion::Continue) - } - })?; - - // Helper: downcast each Arc back to &ExchangeExec. - fn as_exchange(arc: &Arc) -> &ExchangeExec { - arc.downcast_ref::() - .expect("filtered to ExchangeExec above") - } - + let leaves = collect_leaf_exchanges(&input)?; debug!( "[coalesce-rule] collected {} leaf ExchangeExec(s)", leaves.len() ); - for arc in &leaves { - let ex = as_exchange(arc); - debug!( - "[coalesce-rule] leaf: plan_id={} stage_id={:?} partitioning={} M={} resolved={} existing_coalesce={:?}", - ex.plan_id, - ex.stage_id(), - ex.properties().partitioning, - ex.properties().partitioning.partition_count(), - ex.shuffle_partitions().is_some(), - ex.coalesce() - .as_ref() - .map(|cp| (cp.groups.len(), cp.upstream_partition_count)), - ); - } - - // Leaf-scan stage with no upstream Exchanges → nothing to coalesce. if leaves.is_empty() { debug!("[coalesce-rule] no leaves; bail"); return Ok(plan); } + // Wholesale recompute: clear every leaf before deciding anything, so a + // group can never be left with one member carrying a decision from an + // earlier pass and another carrying none. + for arc in &leaves { + as_exchange(arc).set_coalesce(None); + } + let classified: Vec = leaves .iter() .map(|arc| classify_leaf(as_exchange(arc))) .collect(); - - // Still one alignment group in this task; Task 4 splits it by M. - if classified.iter().any(|c| c.kind == LeafKind::Broadcast) { - debug!("[coalesce-rule] broadcast leaf present; bail entire group"); - return Ok(plan); - } - let m = classified[0].m; - if classified.iter().any(|c| c.m != m) { - debug!("[coalesce-rule] heterogeneous M; bail entire group"); - return Ok(plan); - } - let mut sizes: Vec<&[u64]> = Vec::with_capacity(classified.len()); - for (leaf, arc) in classified.iter().zip(&leaves) { - match &leaf.kind { - LeafKind::Sizes(s) => sizes.push(s.as_slice()), - other => { - debug!( - "[coalesce-rule] leaf plan_id={} is {other:?}; bail entire group", - as_exchange(arc).plan_id - ); - return Ok(plan); - } - } - } - - let summed = sum_sizes(&sizes, m); - debug!("[coalesce-rule] summed bytes per upstream partition: {summed:?}"); - - let Some(cp) = decide(&summed, target, small, merged) else { - return Ok(plan); - }; - let cp = Arc::new(cp); - for arc in &leaves { + for (arc, leaf) in leaves.iter().zip(&classified) { let ex = as_exchange(arc); debug!( - "[coalesce-rule] set_coalesce(K={}) on plan_id={} (was {:?})", - cp.groups.len(), + "[coalesce-rule] leaf: plan_id={} stage_id={:?} partitioning={} M={} kind={:?}", ex.plan_id, - ex.coalesce().as_ref().map(|cp| cp.groups.len()), + ex.stage_id(), + ex.properties().partitioning, + leaf.m, + leaf.kind, ); - ex.set_coalesce(Some(cp.clone())); + if let LeafKind::Inconsistent { declared, resolved } = &leaf.kind { + warn!( + "[coalesce-rule] leaf plan_id={} declares {declared} partitions but \ + resolved {resolved}; skipping its alignment group", + ex.plan_id + ); + } + } + + for (m, members) in group_by_upstream_count(&classified) { + let Some(members) = members else { + debug!("[coalesce-rule] group M={m} has an unusable leaf; skipping"); + continue; + }; + let sizes: Vec<&[u64]> = members + .iter() + .filter_map(|&idx| match &classified[idx].kind { + LeafKind::Sizes(sizes) => Some(sizes.as_slice()), + _ => None, + }) + .collect(); + let summed = sum_sizes(&sizes, m); + debug!("[coalesce-rule] group M={m} summed bytes: {summed:?}"); + + let Some(cp) = decide(&summed, target, small, merged) else { + continue; + }; + let cp = Arc::new(cp); + for &idx in &members { + let ex = as_exchange(&leaves[idx]); + debug!( + "[coalesce-rule] set_coalesce(K={}) on plan_id={}", + cp.groups.len(), + ex.plan_id, + ); + ex.set_coalesce(Some(cp.clone())); + } } Ok(plan) } @@ -570,4 +614,69 @@ mod tests { assert_eq!(classify_leaf(&exchange).kind, LeafKind::UnknownBytes); } + + fn sized(m: usize, sizes: Vec) -> ClassifiedLeaf { + ClassifiedLeaf { + m, + kind: LeafKind::Sizes(sizes), + } + } + + fn broadcast(m: usize) -> ClassifiedLeaf { + ClassifiedLeaf { + m, + kind: LeafKind::Broadcast, + } + } + + fn unresolved(m: usize) -> ClassifiedLeaf { + ClassifiedLeaf { + m, + kind: LeafKind::Unresolved, + } + } + + #[test] + fn grouping_splits_leaves_by_upstream_partition_count() { + let leaves = vec![ + sized(8, vec![50; 8]), + sized(8, vec![50; 8]), + sized(4, vec![50; 4]), + sized(4, vec![50; 4]), + ]; + + let groups = group_by_upstream_count(&leaves); + + assert_eq!(groups.len(), 2); + assert_eq!(groups.get(&8), Some(&Some(vec![0, 1]))); + assert_eq!(groups.get(&4), Some(&Some(vec![2, 3]))); + } + + #[test] + fn grouping_drops_broadcast_leaves_without_disturbing_their_siblings() { + // The #2166 case: a broadcast leaf must not take the shuffle leaves in + // the same stage down with it. + let leaves = vec![broadcast(1), sized(8, vec![50; 8]), sized(8, vec![50; 8])]; + + let groups = group_by_upstream_count(&leaves); + + assert_eq!(groups.len(), 1); + assert_eq!(groups.get(&8), Some(&Some(vec![1, 2]))); + assert!(!groups.contains_key(&1)); + } + + #[test] + fn grouping_skips_only_the_group_holding_an_unusable_leaf() { + let leaves = vec![sized(8, vec![50; 8]), unresolved(8), sized(4, vec![50; 4])]; + + let groups = group_by_upstream_count(&leaves); + + assert_eq!(groups.get(&8), Some(&None)); + assert_eq!(groups.get(&4), Some(&Some(vec![2]))); + } + + #[test] + fn grouping_of_only_broadcast_leaves_is_empty() { + assert!(group_by_upstream_count(&[broadcast(1)]).is_empty()); + } } From bb8da564a8191a4166a500486626f845f931329c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 08:31:13 -0600 Subject: [PATCH 05/10] test(scheduler): cover mixed-broadcast and heterogeneous-M coalesce stages --- .../src/state/aqe/test/coalesce_rule.rs | 153 ++++++++++++++++++ 1 file changed, 153 insertions(+) diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 58b810a975..e81a5c87d8 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -25,11 +25,18 @@ //! against `split_size_list_by_target_size`. use crate::assert_plan; +use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; +use crate::state::aqe::optimizer_rule::CoalescePartitionsRule; use crate::state::aqe::planner::AdaptivePlanner; use crate::state::aqe::test::{mock_batch, mock_schema, partitions_with_byte_sizes}; +use ballista_core::execution_plans::{CoalescePlan, PartitionGroup}; use ballista_core::extension::SessionConfigExt; use datafusion::datasource::MemTable; use datafusion::execution::SessionStateBuilder; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion::prelude::{SessionConfig, SessionContext}; use std::sync::Arc; @@ -403,3 +410,149 @@ async fn shuffle_reader_uses_coalesced_k_when_rule_fires() -> datafusion::error: Ok(()) } + +// --------------------------------------------------------------------------- +// Rule-level tests. +// +// The cases below build the stage subtree directly rather than planning SQL, +// because their subject *is* the shape of the leaf set: a broadcast leaf beside +// shuffle leaves, and leaves that disagree on partition count. Both are awkward +// to coax out of a SQL planner and trivial to state directly. +// --------------------------------------------------------------------------- + +/// The configuration the rule reads, at the byte scale the snapshot tests use. +fn rule_config() -> SessionConfig { + SessionConfig::new_with_ballista() + .with_ballista_coalesce_enabled(true) + .with_ballista_coalesce_target_partition_bytes(200) +} + +/// A resolved shuffle leaf declaring `m` partitions of `bytes` bytes each. +fn resolved_shuffle_leaf(m: usize, bytes: u64) -> Arc { + let exchange = ExchangeExec::new( + Arc::new(EmptyExec::new(mock_schema())), + Some(Partitioning::UnknownPartitioning(m)), + 0, + ); + exchange.resolve_shuffle_partitions(partitions_with_byte_sizes(&vec![bytes; m])); + Arc::new(exchange) +} + +/// A resolved broadcast leaf whose upstream wrote `m` partitions. Note the +/// declared partition count is 1 while the resolved shape is `m`: that gap is +/// what the rule used to index off the end of. +fn resolved_broadcast_leaf(m: usize, bytes: u64) -> Arc { + let exchange = ExchangeExec::new_broadcast( + Arc::new(EmptyExec::new(mock_schema())), + None, + 1, + ); + exchange.resolve_shuffle_partitions(partitions_with_byte_sizes(&vec![bytes; m])); + Arc::new(exchange) +} + +/// Wrap leaves in a final-stage root so the rule sees them as one stage's +/// alignment set. Takes `Arc` so callers can keep a typed handle +/// on each leaf and read its decision back after the rule runs. +fn stage_over(leaves: Vec>) -> datafusion::error::Result> { + let mut inputs: Vec> = Vec::with_capacity(leaves.len()); + for leaf in leaves { + // Push, rather than cast: `as` cannot perform the unsizing coercion + // from `Arc` to `Arc`. + inputs.push(leaf); + } + let input: Arc = if inputs.len() == 1 { + inputs.pop().expect("one leaf") + } else { + UnionExec::try_new(inputs)? + }; + Ok(Arc::new(AdaptiveDatafusionExec::new(99, input))) +} + +/// `(K, M)` of a leaf's decision, or `None` when it has none. +fn decision(leaf: &ExchangeExec) -> Option<(usize, u32)> { + leaf.coalesce() + .map(|cp| (cp.groups.len(), cp.upstream_partition_count)) +} + +/// #2166: a broadcast leaf beside shuffle leaves. The broadcast leaf is +/// excluded from the alignment group rather than suppressing it, so the shuffle +/// leaves still coalesce, and the broadcast leaf itself is left alone. +/// +/// Byte trace: two shuffle leaves at `[25; 8]` sum to `[50; 8]`; at target 200 +/// that packs to K=2, the same trace as the hash-join snapshot above. +#[test] +fn should_coalesce_shuffle_leaves_beside_a_broadcast_leaf() +-> datafusion::error::Result<()> { + let broadcast = resolved_broadcast_leaf(8, 25); + let left = resolved_shuffle_leaf(8, 25); + let right = resolved_shuffle_leaf(8, 25); + let plan = stage_over(vec![broadcast.clone(), left.clone(), right.clone()])?; + + CoalescePartitionsRule.optimize(plan, rule_config().options())?; + + assert_eq!(decision(&left), Some((2, 8))); + assert_eq!(decision(&right), Some((2, 8))); + assert_eq!(decision(&broadcast), None); + + Ok(()) +} + +/// #2167: leaves with differing upstream partition counts. Each group packs +/// against its own M instead of the whole stage bailing. +/// +/// Byte trace: the M=8 group is one leaf at `[50; 8]`, which packs 4 partitions +/// per 200-byte bucket into K=2. The M=4 group is one leaf at `[50; 4]`, which +/// sums to exactly 200 and never overshoots, so it packs into K=1. +#[test] +fn should_coalesce_each_partition_count_group_against_its_own_m() +-> datafusion::error::Result<()> { + let eight = resolved_shuffle_leaf(8, 50); + let four = resolved_shuffle_leaf(4, 50); + let plan = stage_over(vec![eight.clone(), four.clone()])?; + + CoalescePartitionsRule.optimize(plan, rule_config().options())?; + + assert_eq!(decision(&eight), Some((2, 8))); + assert_eq!(decision(&four), Some((1, 4))); + + Ok(()) +} + +/// A leaf whose group cannot be decided must not keep a decision an earlier +/// pass left on it. Without the clear, the unresolved leaf's sibling would read +/// as coalesced while the leaf itself read as not, and a join across them would +/// see mismatched partition counts. +#[test] +fn should_clear_a_stale_decision_when_the_group_becomes_undecidable() +-> datafusion::error::Result<()> { + let resolved = resolved_shuffle_leaf(8, 50); + let unresolved = Arc::new(ExchangeExec::new( + Arc::new(EmptyExec::new(mock_schema())), + Some(Partitioning::UnknownPartitioning(8)), + 2, + )); + + // Stand in for a decision an earlier pass attached. + let stale = Arc::new(CoalescePlan { + upstream_partition_count: 8, + groups: vec![ + PartitionGroup { + upstream_indices: vec![0, 1, 2, 3], + }, + PartitionGroup { + upstream_indices: vec![4, 5, 6, 7], + }, + ], + }); + resolved.set_coalesce(Some(stale)); + + let plan = stage_over(vec![resolved.clone(), unresolved.clone()])?; + + CoalescePartitionsRule.optimize(plan, rule_config().options())?; + + assert_eq!(decision(&resolved), None); + assert_eq!(decision(&unresolved), None); + + Ok(()) +} From a2ed9c8711c2e5637aeca4cfa0bd7f31380a21fa Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 08:37:21 -0600 Subject: [PATCH 06/10] test(scheduler): fix formatting and assert shared coalesce plan across leaves --- .../src/state/aqe/test/coalesce_rule.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index e81a5c87d8..23f31c1db5 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -442,11 +442,8 @@ fn resolved_shuffle_leaf(m: usize, bytes: u64) -> Arc { /// declared partition count is 1 while the resolved shape is `m`: that gap is /// what the rule used to index off the end of. fn resolved_broadcast_leaf(m: usize, bytes: u64) -> Arc { - let exchange = ExchangeExec::new_broadcast( - Arc::new(EmptyExec::new(mock_schema())), - None, - 1, - ); + let exchange = + ExchangeExec::new_broadcast(Arc::new(EmptyExec::new(mock_schema())), None, 1); exchange.resolve_shuffle_partitions(partitions_with_byte_sizes(&vec![bytes; m])); Arc::new(exchange) } @@ -454,7 +451,9 @@ fn resolved_broadcast_leaf(m: usize, bytes: u64) -> Arc { /// Wrap leaves in a final-stage root so the rule sees them as one stage's /// alignment set. Takes `Arc` so callers can keep a typed handle /// on each leaf and read its decision back after the rule runs. -fn stage_over(leaves: Vec>) -> datafusion::error::Result> { +fn stage_over( + leaves: Vec>, +) -> datafusion::error::Result> { let mut inputs: Vec> = Vec::with_capacity(leaves.len()); for leaf in leaves { // Push, rather than cast: `as` cannot perform the unsizing coercion @@ -495,6 +494,14 @@ fn should_coalesce_shuffle_leaves_beside_a_broadcast_leaf() assert_eq!(decision(&right), Some((2, 8))); assert_eq!(decision(&broadcast), None); + // Same K is not enough: a hash join needs identical `i -> group(i)` + // boundaries on both sides, which is why the rule hands every member of a + // group the same plan rather than packing each leaf separately. + assert!(Arc::ptr_eq( + &left.coalesce().expect("left is coalesced"), + &right.coalesce().expect("right is coalesced"), + )); + Ok(()) } From e0b34f1374803d6aea0ea230d514eb504500d007 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 09:07:10 -0600 Subject: [PATCH 07/10] fix(scheduler): address review of the AQE coalesce rewrite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Document what the coalesce rule does not reason about — per-partition ordering does not survive a coalesced read, and the bin-pack never splits a skewed partition — so the alignment argument is not over-trusted. Record why grouping by M stays checkable by inspection, and why the wholesale clear is only safe while the rule and the adapter share one per-stage closure. Restore the bin-pack K/M log line on both paths and the root-identity line that attributes leaf and group lines to their stage, and summarise a Sizes leaf by length rather than dumping every element. Make the unreachable non-Sizes arm of the per-group collection panic rather than silently under-coalesce, and give the grouping fixture's leaves distinct plan ids. Add end-to-end coverage for the two cases only the decision layer had: a stage holding a broadcast leaf beside a shuffle leaf, driven through AdaptivePlanner to the ShuffleReaderExec, and a stage packing to K = 1. --- .../aqe/optimizer_rule/coalesce_partitions.rs | 87 +++++++++- ballista/scheduler/src/state/aqe/planner.rs | 8 + .../src/state/aqe/test/coalesce_rule.rs | 154 ++++++++++++++++-- 3 files changed, 231 insertions(+), 18 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 0f11a6b2e4..8025b1a72e 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -44,6 +44,14 @@ //! share an `M` and always land in the same group. Two leaves with different //! `M` provably are not co-partitioned siblings of one join. //! +//! That argument stays checkable by inspection because the rule never touches +//! a leaf's *declared* partitioning: the count `EnforceDistribution` equalised +//! is still `M` on every replan pass, so the grouping key is the same value the +//! distribution guarantee was made about. The substitution of `K` for `M` +//! happens later and elsewhere, in `BallistaAdapter`, and it happens +//! identically for every member of a group because they all carry the same +//! [`CoalescePlan`]. +//! //! Broadcast leaves are excluded from every group. A broadcast reader flattens //! all upstream locations into a single output partition, so there is nothing //! to coalesce, and a `CollectLeft` join requires `SinglePartition` on the @@ -55,6 +63,25 @@ //! a union, still land in one group and have their sizes summed. That inflates //! the per-index totals and under-coalesces, which is safe. //! +//! # What this rule does not reason about +//! +//! The alignment argument above is about hash co-partitioning only. It says +//! nothing about the two properties below, and neither does the rule. +//! +//! **Ordering.** A coalesced reader concatenates several upstream partitions +//! into one output partition and does not merge them, so per-partition ordering +//! does not survive the rewrite. A stage whose consumer depends on the reader's +//! ordering — a `SortPreservingMergeExec` sitting above the pass-through +//! `ExchangeExec` that `DistributedExchangeRule` inserts beneath it, for +//! instance — can therefore produce wrongly ordered output when this rule +//! fires. The rule does not currently detect that shape. This is not a +//! consequence of per-group coalescing; it is a property of coalescing at all, +//! and it predates the grouping rewrite. Untracked at the time of writing. +//! +//! **Skew.** The bin-pack only merges neighbouring partitions; it never splits +//! an oversized one. A single hot upstream partition therefore still becomes a +//! single downstream task, however large it is. +//! //! # Default off //! //! `ballista.planner.coalesce.enabled` is `false` by default. The rule is an opt-in @@ -93,6 +120,15 @@ //! step 2 combined with the bin-pack being a pure function of resolved byte //! sizes, which do not change once a stage has finalized. //! +//! The wholesale clear is only safe because `AdaptivePlanner::actionable_stages` +//! runs this rule and `BallistaAdapter::adapt_to_ballista` inside the *same* +//! per-stage closure, so each stage's decision is consumed before the next +//! stage is optimized. An `ExchangeExec` is a shared `Arc` — at once the root +//! of the stage that writes it and a leaf of every stage that reads it — so +//! splitting that loop into "optimize every stage, then adapt every stage" +//! would let a later stage's clear wipe an earlier stage's decision before the +//! adapter ever read it. +//! //! # Grouping discipline //! //! The bin-pack groups **neighboring** upstream partitions only — each @@ -102,7 +138,7 @@ //! Spark's `CoalesceShufflePartitions` and is what keeps hash //! co-partitioning intact across the rewrite: a hash bucket that used to //! live at index `i` still lives in the single output group that covers -//! `i`, on every leaf of the alignment group. +//! `i`, on every leaf of its alignment group. //! //! # Behavior preservation //! @@ -147,12 +183,33 @@ enum LeafKind { Sizes(Vec), } +impl LeafKind { + /// A short label for the debug log. + /// + /// `Sizes` is summarised by its length: a stage with `M` in the thousands + /// would otherwise print every element, once per leaf, once per pass, and + /// the group's summed vector is logged on its own line anyway. + fn label(&self) -> String { + match self { + LeafKind::Broadcast => "Broadcast".to_string(), + LeafKind::Unresolved => "Unresolved".to_string(), + LeafKind::UnknownBytes => "UnknownBytes".to_string(), + LeafKind::Inconsistent { declared, resolved } => { + format!("Inconsistent(declared={declared}, resolved={resolved})") + } + LeafKind::Sizes(sizes) => format!("Sizes(len={})", sizes.len()), + } + } +} + /// A leaf `ExchangeExec` reduced to what the rule needs: its alignment-group /// key and its contribution to that group. #[derive(Debug)] struct ClassifiedLeaf { /// Upstream partition count `M`, the alignment-group key. m: usize, + /// What the leaf contributes to that group: its per-partition byte counts, + /// or the reason it has none to contribute. kind: LeafKind, } @@ -247,10 +304,23 @@ fn group_by_upstream_count( /// The root is either an `ExchangeExec` (intermediate stage) or an /// `AdaptiveDatafusionExec` (final stage). Anything else means the adapter is /// about to fail anyway, so the rule declines rather than guessing. +/// +/// The root's identity is logged here because the rule runs once per runnable +/// stage within a single pass, so without it the leaf and group lines of +/// several stages interleave with nothing saying which stage each belongs to. fn stage_input(plan: &Arc) -> Option> { if let Some(exchange) = plan.downcast_ref::() { + debug!( + "[coalesce-rule] root=ExchangeExec plan_id={} stage_id={:?}", + exchange.plan_id, + exchange.stage_id(), + ); Some(exchange.input().clone()) } else if let Some(adaptive) = plan.downcast_ref::() { + debug!( + "[coalesce-rule] root=AdaptiveDatafusionExec stage_id={:?}", + adaptive.stage_id(), + ); Some(adaptive.input().clone()) } else { None @@ -318,6 +388,7 @@ fn decide( let starts = split_size_list_by_target_size(summed, target, small_factor, merged_factor); let k = starts.len(); + debug!("[coalesce-rule] bin-pack result: K={k} M={m}"); if k >= m { debug!("[coalesce-rule] K={k} is not a reduction from M={m}; no decision"); return None; @@ -390,12 +461,12 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { for (arc, leaf) in leaves.iter().zip(&classified) { let ex = as_exchange(arc); debug!( - "[coalesce-rule] leaf: plan_id={} stage_id={:?} partitioning={} M={} kind={:?}", + "[coalesce-rule] leaf: plan_id={} stage_id={:?} partitioning={} M={} kind={}", ex.plan_id, ex.stage_id(), ex.properties().partitioning, leaf.m, - leaf.kind, + leaf.kind.label(), ); if let LeafKind::Inconsistent { declared, resolved } = &leaf.kind { warn!( @@ -413,9 +484,13 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { }; let sizes: Vec<&[u64]> = members .iter() - .filter_map(|&idx| match &classified[idx].kind { - LeafKind::Sizes(sizes) => Some(sizes.as_slice()), - _ => None, + .map(|&idx| match &classified[idx].kind { + LeafKind::Sizes(sizes) => sizes.as_slice(), + other => unreachable!( + "group_by_upstream_count only puts LeafKind::Sizes leaves in a \ + decidable group, but member {idx} of group M={m} is {}", + other.label() + ), }) .collect(); let summed = sum_sizes(&sizes, m); diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 67d303761a..44cab94552 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -369,6 +369,14 @@ impl AdaptivePlanner { // group. This avoids cross-stage gluing and stale state // that would arise if the rule walked the entire residual // plan in `default_optimizers()`. + // + // `adapt_to_ballista` must stay in this same closure. + // The rule clears the coalesce slot of every leaf it + // collects before deciding anything, and an + // `ExchangeExec` is a shared `Arc` seen by every stage + // that reads it, so optimizing all stages first and + // adapting them afterwards would let a later stage's + // clear wipe an earlier stage's decision unread. let plan = CoalescePartitionsRule.optimize(plan, config)?; // adapt_to_ballista takes an job_id, we are passing a job_name. Need to transform to fix compiler. let job_id = self.job_name.clone().into(); diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 23f31c1db5..70022cb106 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -36,7 +36,7 @@ use datafusion::execution::SessionStateBuilder; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::union::UnionExec; -use datafusion::physical_plan::{ExecutionPlan, Partitioning}; +use datafusion::physical_plan::{ExecutionPlan, Partitioning, displayable}; use datafusion::prelude::{SessionConfig, SessionContext}; use std::sync::Arc; @@ -411,6 +411,130 @@ async fn shuffle_reader_uses_coalesced_k_when_rule_fires() -> datafusion::error: Ok(()) } +/// `K == 1` through the adapter. Allowing a whole stage to collapse onto one +/// downstream task is this branch's one behaviour change, and the decision +/// layer alone cannot show that `ShuffleReaderExec::try_new_coalesced` accepts +/// a single group and a 1-partition `Partitioning::Hash`. +/// +/// Byte trace: 8 partitions × 10 bytes = 80, which never reaches the 200-byte +/// target, so the bin-pack flushes once at the end: `starts = [0]`, K=1 < M=8. +#[tokio::test] +async fn shuffle_reader_collapses_to_one_partition_when_the_stage_is_tiny() +-> datafusion::error::Result<()> { + let ctx = coalesce_context(8, true); + ctx.register_batch("t", mock_batch()?)?; + + let plan = ctx + .sql("select min(a) as c0, max(b) as c1, c as c2 from t group by c") + .await? + .create_physical_plan() + .await?; + let mut planner = + AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?; + + let stages = planner.runnable_stages()?.unwrap(); + assert_eq!(1, stages.len()); + + planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[10; 8]))?; + + // The reader declares one partition, not eight: the `K <= 1` guard the old + // rule carried would have left this at `partitioning: Hash([c@0], 8)` with + // no `coalesce` field at all. + let stages = planner.runnable_stages()?.unwrap(); + assert_eq!(1, stages.len()); + assert_plan!(stages[0].plan.as_ref(), @ r" + ShuffleWriterExec: partitioning: None + ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2] + AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)] + ShuffleReaderExec: upstream_stage: 0, partitioning: Hash([c@0], 1), coalesce: 1 of 8 + "); + + Ok(()) +} + +/// #2166 through the adapter. The rule-level test below states the same case +/// against hand-built exchanges; this one proves the planner reaches it, since +/// `BallistaAdapter::adapt_to_ballista` is where the original TPC-H Q22 panic +/// surfaced and where the old whole-stage bail cost real coalescing. +/// +/// The join is planned as a `DynamicJoinSelectionExec` over two hash exchanges +/// — DataFusion's own collect-left promotion is disabled by the zeroed +/// `hash_join_single_partition_threshold*`, so the strategy is AQE's to pick. +/// Once both upstream stages finalize, stage 0's measured 8 bytes fall under +/// the 100-byte broadcast threshold and stage 1's 400 do not, so `SelectJoinRule` +/// promotes stage 0's exchange to a broadcast and leaves stage 1's a shuffle. +/// The consuming stage therefore holds one leaf of each kind. +/// +/// Byte trace for the surviving alignment group: one leaf at `[50; 8]`, target +/// 200 → K=2, the same trace as the happy-path test. +#[tokio::test] +async fn shuffle_leaf_still_coalesces_beside_a_broadcast_leaf_end_to_end() +-> datafusion::error::Result<()> { + let config = SessionConfig::new_with_ballista() + .with_target_partitions(8) + .with_round_robin_repartition(false) + .with_ballista_coalesce_enabled(true) + .with_ballista_coalesce_target_partition_bytes(200) + // Between stage 0's measured 8 bytes and stage 1's 400, so exactly one + // side of the join is broadcast. + .with_ballista_broadcast_join_threshold_bytes(100) + .set_u64( + "datafusion.optimizer.hash_join_single_partition_threshold", + 0, + ) + .set_u64( + "datafusion.optimizer.hash_join_single_partition_threshold_rows", + 0, + ); + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + register_partitioned_table(&ctx, "t1", 8)?; + register_partitioned_table(&ctx, "t2", 8)?; + + // `try_new` rather than `try_from_plan`: the broadcast only becomes + // available through `DelayJoinSelectionRule`, which runs in the + // logical-plan preparation pass. + let lp = ctx + .sql("select t1.a, t2.b from t1 join t2 on t1.c = t2.c") + .await? + .into_optimized_plan()?; + let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?; + + let stages = planner.runnable_stages()?.unwrap(); + assert_eq!(2, stages.len()); + + planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[1; 8]))?; + planner.finalise_stage_internal(1, partitions_with_byte_sizes(&[50; 8]))?; + + let stages = planner.runnable_stages()?.unwrap(); + assert_eq!(1, stages.len()); + let stage = displayable(stages[0].plan.as_ref()) + .indent(true) + .to_string(); + + // Under the pre-#2166 rule the broadcast leaf suppressed the whole stage and + // the second reader read `Hash([c@1], 8)` with no coalesce field at all. + assert_plan!(stages[0].plan.as_ref(), @ r" + ShuffleWriterExec: partitioning: None + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(c@1, c@1)], projection=[a@0, b@2] + ShuffleReaderExec: upstream_stage: 0, broadcast: true, upstream_partition_count: 8 + ShuffleReaderExec: upstream_stage: 1, partitioning: Hash([c@1], 2), coalesce: 2 of 8 + "); + assert!( + stage.contains("broadcast: true"), + "the stage must still hold a broadcast reader; plan was:\n{stage}" + ); + assert!( + stage.contains("coalesce: 2 of 8"), + "the shuffle leaf beside it must still coalesce; plan was:\n{stage}" + ); + + Ok(()) +} + // --------------------------------------------------------------------------- // Rule-level tests. // @@ -428,11 +552,14 @@ fn rule_config() -> SessionConfig { } /// A resolved shuffle leaf declaring `m` partitions of `bytes` bytes each. -fn resolved_shuffle_leaf(m: usize, bytes: u64) -> Arc { +/// +/// `plan_id` is per-leaf because the planner never issues two leaves the same +/// one, and the rule's per-leaf debug output identifies leaves by it. +fn resolved_shuffle_leaf(plan_id: usize, m: usize, bytes: u64) -> Arc { let exchange = ExchangeExec::new( Arc::new(EmptyExec::new(mock_schema())), Some(Partitioning::UnknownPartitioning(m)), - 0, + plan_id, ); exchange.resolve_shuffle_partitions(partitions_with_byte_sizes(&vec![bytes; m])); Arc::new(exchange) @@ -441,9 +568,12 @@ fn resolved_shuffle_leaf(m: usize, bytes: u64) -> Arc { /// A resolved broadcast leaf whose upstream wrote `m` partitions. Note the /// declared partition count is 1 while the resolved shape is `m`: that gap is /// what the rule used to index off the end of. -fn resolved_broadcast_leaf(m: usize, bytes: u64) -> Arc { - let exchange = - ExchangeExec::new_broadcast(Arc::new(EmptyExec::new(mock_schema())), None, 1); +fn resolved_broadcast_leaf(plan_id: usize, m: usize, bytes: u64) -> Arc { + let exchange = ExchangeExec::new_broadcast( + Arc::new(EmptyExec::new(mock_schema())), + None, + plan_id, + ); exchange.resolve_shuffle_partitions(partitions_with_byte_sizes(&vec![bytes; m])); Arc::new(exchange) } @@ -483,9 +613,9 @@ fn decision(leaf: &ExchangeExec) -> Option<(usize, u32)> { #[test] fn should_coalesce_shuffle_leaves_beside_a_broadcast_leaf() -> datafusion::error::Result<()> { - let broadcast = resolved_broadcast_leaf(8, 25); - let left = resolved_shuffle_leaf(8, 25); - let right = resolved_shuffle_leaf(8, 25); + let broadcast = resolved_broadcast_leaf(0, 8, 25); + let left = resolved_shuffle_leaf(1, 8, 25); + let right = resolved_shuffle_leaf(2, 8, 25); let plan = stage_over(vec![broadcast.clone(), left.clone(), right.clone()])?; CoalescePartitionsRule.optimize(plan, rule_config().options())?; @@ -514,8 +644,8 @@ fn should_coalesce_shuffle_leaves_beside_a_broadcast_leaf() #[test] fn should_coalesce_each_partition_count_group_against_its_own_m() -> datafusion::error::Result<()> { - let eight = resolved_shuffle_leaf(8, 50); - let four = resolved_shuffle_leaf(4, 50); + let eight = resolved_shuffle_leaf(0, 8, 50); + let four = resolved_shuffle_leaf(1, 4, 50); let plan = stage_over(vec![eight.clone(), four.clone()])?; CoalescePartitionsRule.optimize(plan, rule_config().options())?; @@ -533,7 +663,7 @@ fn should_coalesce_each_partition_count_group_against_its_own_m() #[test] fn should_clear_a_stale_decision_when_the_group_becomes_undecidable() -> datafusion::error::Result<()> { - let resolved = resolved_shuffle_leaf(8, 50); + let resolved = resolved_shuffle_leaf(1, 8, 50); let unresolved = Arc::new(ExchangeExec::new( Arc::new(EmptyExec::new(mock_schema())), Some(Partitioning::UnknownPartitioning(8)), From 132d84c42128a25169a324f815cb1a8d7cf52aaa Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 09:25:56 -0600 Subject: [PATCH 08/10] fix(scheduler): exclude pass-through exchanges from coalesce alignment groups CoalescePartitionsRule could attach a coalesce decision to a pass-through ExchangeExec (partitioning: None) sitting directly beneath a SortPreservingMergeExec. DistributedExchangeRule inserts exactly this shape to carry a stage's ordering across the boundary; a coalesced ShuffleReaderExec shuffles its concatenated partition locations, destroying that ordering and letting the SortPreservingMergeExec above silently emit wrongly ordered ORDER BY results. Exclude pass-through leaves from alignment groups the same way broadcast leaves already are: classify_leaf reports a new PassThrough LeafKind, and group_by_upstream_count drops it via the same non-poisoning continue path so it does not suppress coalescing for unrelated same-M siblings. --- .../aqe/optimizer_rule/coalesce_partitions.rs | 128 +++++++++++++++--- .../src/state/aqe/test/coalesce_rule.rs | 64 +++++++++ 2 files changed, 173 insertions(+), 19 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 8025b1a72e..6c362a0257 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -58,6 +58,16 @@ //! build side and `UnspecifiedDistribution` on the probe side, so a broadcast //! leaf is never a co-partitioned sibling. //! +//! Pass-through leaves — `ExchangeExec`s with `partitioning: None` — are +//! excluded too. `DistributedExchangeRule` creates exactly two of these: one +//! beneath a `CoalescePartitionsExec`, one beneath a `SortPreservingMergeExec`. +//! Neither is a co-partitioned join sibling, for the same reason a broadcast +//! leaf isn't: `SelectJoinRule` always builds join-leg exchanges with +//! `Some(partitioning)`, and anything feeding a `Partitioned` join gets a hash +//! `RepartitionExec` from `EnforceDistribution`, which `DistributedExchangeRule` +//! turns into a `Some(partitioning)` exchange, never a pass-through one. See +//! below for why excluding them also costs nothing. +//! //! The grouping is conservative in the other direction: leaves that share an //! `M` without sharing any co-partitioning requirement, such as the two arms of //! a union, still land in one group and have their sizes summed. That inflates @@ -65,22 +75,26 @@ //! //! # What this rule does not reason about //! -//! The alignment argument above is about hash co-partitioning only. It says -//! nothing about the two properties below, and neither does the rule. -//! -//! **Ordering.** A coalesced reader concatenates several upstream partitions -//! into one output partition and does not merge them, so per-partition ordering -//! does not survive the rewrite. A stage whose consumer depends on the reader's -//! ordering — a `SortPreservingMergeExec` sitting above the pass-through -//! `ExchangeExec` that `DistributedExchangeRule` inserts beneath it, for -//! instance — can therefore produce wrongly ordered output when this rule -//! fires. The rule does not currently detect that shape. This is not a -//! consequence of per-group coalescing; it is a property of coalescing at all, -//! and it predates the grouping rewrite. Untracked at the time of writing. +//! **Ordering** used to be an unhandled gap here: a coalesced reader +//! concatenates several upstream partitions into one output partition without +//! merging them, so per-partition ordering does not survive the rewrite, and a +//! `SortPreservingMergeExec` reading a coalesced partition would merge streams +//! that are no longer sorted and emit wrongly ordered output. The rule now +//! closes it structurally instead of detecting the shape: the *only* thing +//! that carries a child's ordering across a stage boundary is a pass-through +//! `ExchangeExec` (`partitioning: None`), and those are excluded from every +//! alignment group above. `DistributedExchangeRule` creates a pass-through +//! exchange at exactly two sites — beneath a `CoalescePartitionsExec` and +//! beneath a `SortPreservingMergeExec` — and only the second is dangerous: a +//! `CoalescePartitionsExec` merges everything into one unordered partition +//! immediately above it, so coalescing underneath was never worth anything +//! there either. Declining both therefore costs nothing and needs no +//! ordering-specific reasoning in the rule itself. //! -//! **Skew.** The bin-pack only merges neighbouring partitions; it never splits -//! an oversized one. A single hot upstream partition therefore still becomes a -//! single downstream task, however large it is. +//! **Skew** is still a real, untracked gap. The bin-pack only merges +//! neighbouring partitions; it never splits an oversized one. A single hot +//! upstream partition therefore still becomes a single downstream task, +//! however large it is. //! //! # Default off //! @@ -172,6 +186,19 @@ enum LeafKind { /// and a `CollectLeft` join never requires it to be co-partitioned with /// anything. Excluded from every alignment group. Broadcast, + /// A pass-through exchange (`partitioning: None`): `DistributedExchangeRule` + /// inserts these directly beneath a `SortPreservingMergeExec` or a + /// `CoalescePartitionsExec` to mark a stage boundary without re-partitioning + /// the child. Its only job is to carry the child's partitioning *and* + /// ordering across that boundary unchanged. A coalesced reader concatenates + /// several upstream partitions into one without merging them, which destroys + /// both, so this leaf is excluded from every alignment group. It is also + /// never a co-partitioned join sibling: `SelectJoinRule` always builds + /// join-leg exchanges with `Some(partitioning)`, and anything feeding a + /// `Partitioned` join gets a hash `RepartitionExec` from + /// `EnforceDistribution`, which `DistributedExchangeRule` turns into a + /// `Some(partitioning)` exchange, not this one. + PassThrough, /// The upstream stage has not finalized. The group defers to a later pass. Unresolved, /// At least one location reports no `num_bytes`. @@ -192,6 +219,7 @@ impl LeafKind { fn label(&self) -> String { match self { LeafKind::Broadcast => "Broadcast".to_string(), + LeafKind::PassThrough => "PassThrough".to_string(), LeafKind::Unresolved => "Unresolved".to_string(), LeafKind::UnknownBytes => "UnknownBytes".to_string(), LeafKind::Inconsistent { declared, resolved } => { @@ -227,6 +255,12 @@ fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { kind: LeafKind::Broadcast, }; } + if ex.partitioning.is_none() { + return ClassifiedLeaf { + m, + kind: LeafKind::PassThrough, + }; + } let Some(parts) = ex.shuffle_partitions() else { return ClassifiedLeaf { m, @@ -267,15 +301,25 @@ fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { /// Partition classified leaves into alignment groups keyed by upstream /// partition count `M`. /// -/// Broadcast leaves are dropped: they carry no partition structure to coalesce -/// and are never a co-partitioned join sibling, so excluding them cannot break -/// an alignment invariant. +/// Broadcast and pass-through leaves are dropped: neither carries partition +/// structure to coalesce, and neither is ever a co-partitioned join sibling, +/// so excluding them cannot break an alignment invariant. /// /// A group maps to `Some(indices)` when every member carries usable sizes, and /// to `None` when any member is unresolved, missing byte statistics, or /// inconsistent. Skipping is per group: one undecidable leaf no longer /// suppresses coalescing for leaves it has no relationship with. /// +/// Broadcast and pass-through leaves take the early `continue` below rather +/// than falling into the `_ => *entry = None` poisoning arm. That distinction +/// matters: poisoning is for leaves whose *group* is undecidable (unresolved, +/// missing stats, inconsistent shape), which must suppress coalescing for +/// every `Sizes` sibling at the same `M` because the rule cannot tell whether +/// they were meant to align with it. A broadcast or pass-through leaf is never +/// such a sibling — it is excluded from alignment entirely — so its presence +/// must leave an otherwise-decidable group of `Sizes` leaves at the same `M` +/// alone. +/// /// `BTreeMap` rather than `HashMap` so iteration order, and therefore the debug /// log and the order decisions are applied, is stable across passes. fn group_by_upstream_count( @@ -283,7 +327,7 @@ fn group_by_upstream_count( ) -> BTreeMap>> { let mut groups: BTreeMap>> = BTreeMap::new(); for (idx, leaf) in leaves.iter().enumerate() { - if leaf.kind == LeafKind::Broadcast { + if matches!(leaf.kind, LeafKind::Broadcast | LeafKind::PassThrough) { continue; } let entry = groups.entry(leaf.m).or_insert_with(|| Some(Vec::new())); @@ -652,6 +696,21 @@ mod tests { assert_eq!(leaf.kind, LeafKind::Broadcast); } + #[test] + fn classify_reports_pass_through_even_once_resolved_with_byte_sizes() { + // A pass-through exchange (`partitioning: None`) exists to carry a + // child's partitioning and ordering across a stage boundary, not to be + // coalesced. It must classify as `PassThrough` and never reach the + // summing step, however fully it has resolved. + let exchange = + ExchangeExec::new(Arc::new(EmptyExec::new(mock_schema())), None, 0); + exchange.resolve_shuffle_partitions(partitions_with_byte_sizes(&[50; 8])); + + let leaf = classify_leaf(&exchange); + + assert_eq!(leaf.kind, LeafKind::PassThrough); + } + #[test] fn classify_reports_inconsistent_when_the_resolved_shape_disagrees() { let exchange = shuffle_exchange(8, Some(partitions_with_byte_sizes(&[50; 4]))); @@ -711,6 +770,13 @@ mod tests { } } + fn pass_through(m: usize) -> ClassifiedLeaf { + ClassifiedLeaf { + m, + kind: LeafKind::PassThrough, + } + } + #[test] fn grouping_splits_leaves_by_upstream_partition_count() { let leaves = vec![ @@ -740,6 +806,30 @@ mod tests { assert!(!groups.contains_key(&1)); } + #[test] + fn grouping_drops_pass_through_leaves_without_poisoning_their_siblings() { + // A pass-through leaf shares M=8 with two shuffle leaves. If it took + // the `_ => *entry = None` poisoning arm instead of the early + // `continue`, it would wipe the group and suppress coalescing for both + // `Sizes` siblings — the same mistake the broadcast case guards + // against above, but for a different `LeafKind`. + let leaves = vec![ + pass_through(8), + sized(8, vec![50; 8]), + sized(8, vec![50; 8]), + ]; + + let groups = group_by_upstream_count(&leaves); + + assert_eq!(groups.len(), 1); + assert_eq!(groups.get(&8), Some(&Some(vec![1, 2]))); + } + + #[test] + fn grouping_of_only_pass_through_leaves_is_empty() { + assert!(group_by_upstream_count(&[pass_through(1)]).is_empty()); + } + #[test] fn grouping_skips_only_the_group_holding_an_unusable_leaf() { let leaves = vec![sized(8, vec![50; 8]), unresolved(8), sized(4, vec![50; 4])]; diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 70022cb106..e329488cc9 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -535,6 +535,70 @@ async fn shuffle_leaf_still_coalesces_beside_a_broadcast_leaf_end_to_end() Ok(()) } +/// Guard against the ordering bug: a leaf `ExchangeExec` that +/// `DistributedExchangeRule` inserted directly beneath a +/// `SortPreservingMergeExec` is a pass-through exchange (`partitioning: None`) +/// — it exists only to carry the upstream `SortExec`'s per-partition ordering +/// across the stage boundary. A coalesced reader concatenates several upstream +/// partitions into one without merging them, which would destroy that +/// ordering and the `SortPreservingMergeExec` above would silently emit +/// wrongly ordered rows. `classify_leaf` must classify this leaf as +/// `PassThrough`, not `Sizes`, so `CoalescePartitionsRule` declines it. +/// +/// Byte trace: stage 0 finalizes with 8 partitions x 50 bytes, which packs to +/// K=2 at target=200 (the same trace as the happy-path test) *if the rule were +/// allowed to coalesce this leaf*. The assertion is that it is not: the +/// resulting `ShuffleReaderExec` carries no `coalesce:` field and keeps all 8 +/// upstream partitions, proving the guard fired rather than the bin-pack +/// merely declining on its own. +#[tokio::test] +async fn should_not_coalesce_a_pass_through_exchange_beneath_sort_preserving_merge() +-> datafusion::error::Result<()> { + let ctx = coalesce_context(8, true); + register_partitioned_table(&ctx, "t", 8)?; + + let plan = ctx + .sql("select a from t order by a") + .await? + .create_physical_plan() + .await?; + let mut planner = + AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?; + + // Stage 0 is the upstream writer: each of the 8 source partitions is + // locally sorted and written as-is (`ShuffleWriterExec: partitioning: None` + // — no hash repartitioning, since a plain ORDER BY has no partitioning + // requirement of its own). + let stages = planner.runnable_stages()?.unwrap(); + assert_eq!(1, stages.len()); + assert_plan!(stages[0].plan.as_ref(), @ r" + ShuffleWriterExec: partitioning: None + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1, 1, 1, 1, 1] + "); + + // Finalize stage 0 with byte sizes that would comfortably coalesce + // (8 x 50 = 400, target = 200, same trace as `should_attach_coalesce_when_ + // partitions_pack_below_m`) if the pass-through guard did not exclude this + // leaf from its alignment group. + planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[50; 8]))?; + + // Stage 1 is the final stage. Its `ShuffleReaderExec` still declares all 8 + // partitions and carries no `coalesce:` field: the guard suppressed the + // decision that would otherwise have collapsed it to K=2, which is exactly + // what protects the `SortPreservingMergeExec` above from merging streams + // that are no longer sorted. + let stages = planner.runnable_stages()?.unwrap(); + assert_eq!(1, stages.len()); + assert_plan!(stages[0].plan.as_ref(), @ r" + ShuffleWriterExec: partitioning: None + SortPreservingMergeExec: [a@0 ASC NULLS LAST] + ShuffleReaderExec: upstream_stage: 0, partitioning: UnknownPartitioning(8) + "); + + Ok(()) +} + // --------------------------------------------------------------------------- // Rule-level tests. // From 3722f2d3b7fe38a7aa64a7e64fdea78103c079b9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 09:30:59 -0600 Subject: [PATCH 09/10] docs(scheduler): fix stale algorithm summary and note broadcast-check ordering The Algorithm section's step 3 still said only broadcast leaves drop out during classification, contradicting the pass-through exclusion added above it. Also note why the broadcast check in classify_leaf must run before the partitioning::is_none() check: both new_broadcast and to_broadcast set partitioning to None, so reordering the checks would silently misclassify every broadcast leaf as pass-through. --- .../src/state/aqe/optimizer_rule/coalesce_partitions.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 6c362a0257..9318c830f7 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -116,8 +116,9 @@ //! has nothing to coalesce. //! 2. Clear every leaf's coalesce slot, so the pass recomputes wholesale and //! no decision survives from an earlier pass. -//! 3. Classify each leaf. Broadcast leaves drop out; unresolved, statistics- -//! free, and inconsistent leaves make their group undecidable. +//! 3. Classify each leaf. Broadcast and pass-through leaves drop out; +//! unresolved, statistics-free, and inconsistent leaves make their group +//! undecidable. //! 4. Group the rest by `M`. //! 5. Per group: sum per-partition byte sizes element-wise, then bin-pack the //! sums toward `target_partition_bytes` (Spark's @@ -255,6 +256,9 @@ fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { kind: LeafKind::Broadcast, }; } + // Must run after the broadcast check above: `new_broadcast` and + // `to_broadcast` also set `partitioning: None`, so a broadcast leaf would + // misclassify as `PassThrough` if this check ran first. if ex.partitioning.is_none() { return ClassifiedLeaf { m, From b839f036843a6ce61b7540fe372d0a740a31ba71 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 10:02:06 -0600 Subject: [PATCH 10/10] refactor(scheduler): tighten the AQE coalesce rule after review Grouping now hands each member its own size slice instead of an index into a parallel vector, so the caller sums a group without re-matching on LeafKind and the unreachable! that guarded that match is gone. The rule also decides first and applies second, writing every leaf's slot exactly once per pass with its final value. That removes the window in which a leaf held a cleared slot, and with it the cross-module requirement that actionable_stages keep optimize and adapt in one closure. Also: name the pass-through condition as ExchangeExec::preserves_child_ordering so the rule and maintains_input_order share one definition, read byte counts under the lock rather than deep-cloning every PartitionLocation, and drop the hand-spelled LeafKind label arms that Debug already produces. --- .../src/state/aqe/execution_plan/exchange.rs | 37 ++- .../aqe/optimizer_rule/coalesce_partitions.rs | 307 ++++++++++-------- ballista/scheduler/src/state/aqe/planner.rs | 8 - 3 files changed, 196 insertions(+), 156 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index ea72c5f4ee..d25d3e6441 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs @@ -217,6 +217,38 @@ impl ExchangeExec { self.shuffle_partitions.lock().clone() } + /// Runs `f` against the resolved shuffle partitions in place, returning + /// `None` if they have not been resolved yet. + /// + /// Prefer this over [`Self::shuffle_partitions`] when only a summary is + /// needed. That method deep-clones the whole vector, and every + /// `PartitionLocation` in it carries several `String`s, so reading a byte + /// count per partition would otherwise allocate proportionally to the + /// upstream partition count on every call. + pub fn with_shuffle_partitions( + &self, + f: impl FnOnce(&[Vec]) -> R, + ) -> Option { + self.shuffle_partitions.lock().as_deref().map(f) + } + + /// Whether this exchange carries its child's ordering across the stage + /// boundary unchanged. + /// + /// True exactly for a pass-through exchange, which `DistributedExchangeRule` + /// inserts beneath a `SortPreservingMergeExec` or `CoalescePartitionsExec` + /// to mark a boundary without re-partitioning. A repartitioning exchange + /// makes no such promise. + /// + /// NOTE: the `true` answer is only sound because `CoalescePartitionsRule` + /// declines to coalesce these leaves. A coalesced `ShuffleReaderExec` + /// concatenates several upstream partitions into one output partition and + /// randomises the order it reads their locations in, which would destroy + /// the ordering this reports as preserved. + pub fn preserves_child_ordering(&self) -> bool { + self.partitioning.is_none() + } + /// Flattens partition locations into single vector, /// this method is usually used when we want to collect partitions /// to form a broadcast join @@ -341,10 +373,7 @@ impl ExecutionPlan for ExchangeExec { } fn maintains_input_order(&self) -> Vec { - match self.partitioning { - Some(_) => vec![false; self.children().len()], - None => vec![true; self.children().len()], - } + vec![self.preserves_child_ordering(); self.children().len()] } fn with_new_children( diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 9318c830f7..e1cac56f21 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -58,15 +58,9 @@ //! build side and `UnspecifiedDistribution` on the probe side, so a broadcast //! leaf is never a co-partitioned sibling. //! -//! Pass-through leaves — `ExchangeExec`s with `partitioning: None` — are -//! excluded too. `DistributedExchangeRule` creates exactly two of these: one -//! beneath a `CoalescePartitionsExec`, one beneath a `SortPreservingMergeExec`. -//! Neither is a co-partitioned join sibling, for the same reason a broadcast -//! leaf isn't: `SelectJoinRule` always builds join-leg exchanges with -//! `Some(partitioning)`, and anything feeding a `Partitioned` join gets a hash -//! `RepartitionExec` from `EnforceDistribution`, which `DistributedExchangeRule` -//! turns into a `Some(partitioning)` exchange, never a pass-through one. See -//! below for why excluding them also costs nothing. +//! Pass-through leaves are excluded too, for the reason spelled out on +//! [`LeafKind::PassThrough`]: they are never co-partitioned join siblings +//! either, and coalescing them would destroy the ordering they exist to carry. //! //! The grouping is conservative in the other direction: leaves that share an //! `M` without sharing any co-partitioning requirement, such as the two arms of @@ -75,21 +69,17 @@ //! //! # What this rule does not reason about //! -//! **Ordering** used to be an unhandled gap here: a coalesced reader -//! concatenates several upstream partitions into one output partition without -//! merging them, so per-partition ordering does not survive the rewrite, and a -//! `SortPreservingMergeExec` reading a coalesced partition would merge streams -//! that are no longer sorted and emit wrongly ordered output. The rule now -//! closes it structurally instead of detecting the shape: the *only* thing -//! that carries a child's ordering across a stage boundary is a pass-through -//! `ExchangeExec` (`partitioning: None`), and those are excluded from every -//! alignment group above. `DistributedExchangeRule` creates a pass-through -//! exchange at exactly two sites — beneath a `CoalescePartitionsExec` and -//! beneath a `SortPreservingMergeExec` — and only the second is dangerous: a -//! `CoalescePartitionsExec` merges everything into one unordered partition -//! immediately above it, so coalescing underneath was never worth anything -//! there either. Declining both therefore costs nothing and needs no -//! ordering-specific reasoning in the rule itself. +//! **Ordering** is handled structurally rather than by detecting the shapes +//! that depend on it. A coalesced reader concatenates several upstream +//! partitions into one output partition without merging them, so per-partition +//! ordering does not survive the rewrite. The only exchange that carries a +//! child's ordering across a stage boundary is a pass-through one — see +//! [`ExchangeExec::preserves_child_ordering`] — and those are excluded from +//! every alignment group, so the rule needs no ordering-specific reasoning of +//! its own. Excluding them costs nothing: of the two sites that create one, +//! only the `SortPreservingMergeExec` site was ever at risk, and the +//! `CoalescePartitionsExec` site merges everything into a single unordered +//! partition immediately above, where coalescing bought nothing anyway. //! //! **Skew** is still a real, untracked gap. The bin-pack only merges //! neighbouring partitions; it never splits an oversized one. A single hot @@ -114,35 +104,30 @@ //! //! 1. Collect leaf `ExchangeExec`s. If none, this stage reads from scans and //! has nothing to coalesce. -//! 2. Clear every leaf's coalesce slot, so the pass recomputes wholesale and -//! no decision survives from an earlier pass. -//! 3. Classify each leaf. Broadcast and pass-through leaves drop out; +//! 2. Classify each leaf. Broadcast and pass-through leaves drop out; //! unresolved, statistics-free, and inconsistent leaves make their group //! undecidable. -//! 4. Group the rest by `M`. -//! 5. Per group: sum per-partition byte sizes element-wise, then bin-pack the +//! 3. Group the rest by `M`. +//! 4. Per group: sum per-partition byte sizes element-wise, then bin-pack the //! sums toward `target_partition_bytes` (Spark's //! `advisoryPartitionSizeInBytes`, 64 MB by default) with -//! `split_size_list_by_target_size`. -//! 6. If `K >= M` the rewrite is not a reduction and the group is left alone. -//! Otherwise attach one shared [`CoalescePlan`] to every member. `K == 1` -//! is a legitimate outcome for a stage that fits in one partition. +//! `split_size_list_by_target_size`. If `K >= M` the rewrite is not a +//! reduction and the group is left alone. `K == 1` is a legitimate outcome +//! for a stage that fits in one partition. +//! 5. Write every leaf's decision in one pass: the group's shared +//! [`CoalescePlan`] for a member of a decided group, `None` for every +//! other leaf. //! //! # Carrier semantics //! //! The `CoalescePlan` lives on the upstream `ExchangeExec`; the rule does not -//! rewrite the plan tree. Idempotency comes from the clear-then-set contract in -//! step 2 combined with the bin-pack being a pure function of resolved byte -//! sizes, which do not change once a stage has finalized. -//! -//! The wholesale clear is only safe because `AdaptivePlanner::actionable_stages` -//! runs this rule and `BallistaAdapter::adapt_to_ballista` inside the *same* -//! per-stage closure, so each stage's decision is consumed before the next -//! stage is optimized. An `ExchangeExec` is a shared `Arc` — at once the root -//! of the stage that writes it and a leaf of every stage that reads it — so -//! splitting that loop into "optimize every stage, then adapt every stage" -//! would let a later stage's clear wipe an earlier stage's decision before the -//! adapter ever read it. +//! rewrite the plan tree. Every collected leaf's slot is written exactly once +//! per pass, with its final value — the group's shared plan, or `None` for a +//! leaf no group decided to coalesce. A leaf therefore never carries a decision +//! from an earlier pass, and never holds a transient cleared slot mid-pass +//! either. Idempotency follows from that plus the bin-pack being a pure +//! function of resolved byte sizes, which do not change once a stage has +//! finalized. //! //! # Grouping discipline //! @@ -187,18 +172,19 @@ enum LeafKind { /// and a `CollectLeft` join never requires it to be co-partitioned with /// anything. Excluded from every alignment group. Broadcast, - /// A pass-through exchange (`partitioning: None`): `DistributedExchangeRule` + /// A pass-through exchange, in the sense of + /// [`ExchangeExec::preserves_child_ordering`]. `DistributedExchangeRule` /// inserts these directly beneath a `SortPreservingMergeExec` or a /// `CoalescePartitionsExec` to mark a stage boundary without re-partitioning - /// the child. Its only job is to carry the child's partitioning *and* + /// the child, so their whole job is to carry the child's partitioning *and* /// ordering across that boundary unchanged. A coalesced reader concatenates /// several upstream partitions into one without merging them, which destroys /// both, so this leaf is excluded from every alignment group. It is also /// never a co-partitioned join sibling: `SelectJoinRule` always builds - /// join-leg exchanges with `Some(partitioning)`, and anything feeding a + /// join-leg exchanges with an explicit partitioning, and anything feeding a /// `Partitioned` join gets a hash `RepartitionExec` from /// `EnforceDistribution`, which `DistributedExchangeRule` turns into a - /// `Some(partitioning)` exchange, not this one. + /// repartitioning exchange, not this one. PassThrough, /// The upstream stage has not finalized. The group defers to a later pass. Unresolved, @@ -219,18 +205,27 @@ impl LeafKind { /// the group's summed vector is logged on its own line anyway. fn label(&self) -> String { match self { - LeafKind::Broadcast => "Broadcast".to_string(), - LeafKind::PassThrough => "PassThrough".to_string(), - LeafKind::Unresolved => "Unresolved".to_string(), - LeafKind::UnknownBytes => "UnknownBytes".to_string(), - LeafKind::Inconsistent { declared, resolved } => { - format!("Inconsistent(declared={declared}, resolved={resolved})") - } LeafKind::Sizes(sizes) => format!("Sizes(len={})", sizes.len()), + other => format!("{other:?}"), } } } +/// One member of a decidable alignment group: which leaf it is, and the +/// per-partition byte counts it contributes. +/// +/// Carrying the sizes rather than a bare index is what lets the caller sum a +/// group without re-matching on [`LeafKind`]. The borrow checker then enforces +/// what [`group_by_upstream_count`] already established — that a decidable +/// group holds only `Sizes` leaves — instead of a runtime assertion. +#[derive(Debug, PartialEq, Eq)] +struct GroupMember<'a> { + /// Index into the leaf vector the classification was built from. + idx: usize, + /// This leaf's per-upstream-partition byte counts. Length is the group's `M`. + sizes: &'a [u64], +} + /// A leaf `ExchangeExec` reduced to what the rule needs: its alignment-group /// key and its contribution to that group. #[derive(Debug)] @@ -257,48 +252,40 @@ fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { }; } // Must run after the broadcast check above: `new_broadcast` and - // `to_broadcast` also set `partitioning: None`, so a broadcast leaf would + // `to_broadcast` also leave `partitioning` unset, so a broadcast leaf would // misclassify as `PassThrough` if this check ran first. - if ex.partitioning.is_none() { + if ex.preserves_child_ordering() { return ClassifiedLeaf { m, kind: LeafKind::PassThrough, }; } - let Some(parts) = ex.shuffle_partitions() else { - return ClassifiedLeaf { - m, - kind: LeafKind::Unresolved, - }; - }; - if parts.len() != m { - return ClassifiedLeaf { - m, - kind: LeafKind::Inconsistent { + // Read the byte counts in place: the locations carry several `String`s + // each, and cloning them out just to sum a `u64` per partition would + // allocate proportionally to `M` on every pass. + let kind = ex.with_shuffle_partitions(|parts| { + if parts.len() != m { + return LeafKind::Inconsistent { declared: m, resolved: parts.len(), - }, - }; - } - let mut sizes = Vec::with_capacity(parts.len()); - for locations in &parts { - let mut total = 0u64; - for location in locations { - match location.partition_stats.num_bytes() { - Some(bytes) => total = total.saturating_add(bytes), - None => { - return ClassifiedLeaf { - m, - kind: LeafKind::UnknownBytes, - }; + }; + } + let mut sizes = Vec::with_capacity(parts.len()); + for locations in parts { + let mut total = 0u64; + for location in locations { + match location.partition_stats.num_bytes() { + Some(bytes) => total = total.saturating_add(bytes), + None => return LeafKind::UnknownBytes, } } + sizes.push(total); } - sizes.push(total); - } + LeafKind::Sizes(sizes) + }); ClassifiedLeaf { m, - kind: LeafKind::Sizes(sizes), + kind: kind.unwrap_or(LeafKind::Unresolved), } } @@ -309,7 +296,7 @@ fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { /// structure to coalesce, and neither is ever a co-partitioned join sibling, /// so excluding them cannot break an alignment invariant. /// -/// A group maps to `Some(indices)` when every member carries usable sizes, and +/// A group maps to `Some(members)` when every member carries usable sizes, and /// to `None` when any member is unresolved, missing byte statistics, or /// inconsistent. Skipping is per group: one undecidable leaf no longer /// suppresses coalescing for leaves it has no relationship with. @@ -328,17 +315,17 @@ fn classify_leaf(ex: &ExchangeExec) -> ClassifiedLeaf { /// log and the order decisions are applied, is stable across passes. fn group_by_upstream_count( leaves: &[ClassifiedLeaf], -) -> BTreeMap>> { - let mut groups: BTreeMap>> = BTreeMap::new(); +) -> BTreeMap>>> { + let mut groups: BTreeMap>>> = BTreeMap::new(); for (idx, leaf) in leaves.iter().enumerate() { if matches!(leaf.kind, LeafKind::Broadcast | LeafKind::PassThrough) { continue; } let entry = groups.entry(leaf.m).or_insert_with(|| Some(Vec::new())); match &leaf.kind { - LeafKind::Sizes(_) => { + LeafKind::Sizes(sizes) => { if let Some(members) = entry { - members.push(idx); + members.push(GroupMember { idx, sizes }); } } _ => *entry = None, @@ -353,23 +340,27 @@ fn group_by_upstream_count( /// `AdaptiveDatafusionExec` (final stage). Anything else means the adapter is /// about to fail anyway, so the rule declines rather than guessing. /// -/// The root's identity is logged here because the rule runs once per runnable -/// stage within a single pass, so without it the leaf and group lines of -/// several stages interleave with nothing saying which stage each belongs to. -fn stage_input(plan: &Arc) -> Option> { +/// The returned string identifies the root for the debug log. The rule runs +/// once per runnable stage within a single pass, so without it the leaf and +/// group lines of several stages interleave with nothing saying which stage +/// each belongs to. +fn stage_input( + plan: &Arc, +) -> Option<(Arc, String)> { if let Some(exchange) = plan.downcast_ref::() { - debug!( - "[coalesce-rule] root=ExchangeExec plan_id={} stage_id={:?}", - exchange.plan_id, - exchange.stage_id(), - ); - Some(exchange.input().clone()) + Some(( + exchange.input().clone(), + format!( + "ExchangeExec plan_id={} stage_id={:?}", + exchange.plan_id, + exchange.stage_id() + ), + )) } else if let Some(adaptive) = plan.downcast_ref::() { - debug!( - "[coalesce-rule] root=AdaptiveDatafusionExec stage_id={:?}", - adaptive.stage_id(), - ); - Some(adaptive.input().clone()) + Some(( + adaptive.input().clone(), + format!("AdaptiveDatafusionExec stage_id={:?}", adaptive.stage_id()), + )) } else { None } @@ -478,12 +469,13 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { ); // Get the subtree below the root. Two root kinds, same outcome. - let Some(input) = stage_input(&plan) else { + let Some((input, root)) = stage_input(&plan) else { debug!( "[coalesce-rule] root is neither ExchangeExec nor AdaptiveDatafusionExec; bail" ); return Ok(plan); }; + debug!("[coalesce-rule] root={root}"); let leaves = collect_leaf_exchanges(&input)?; debug!( @@ -495,13 +487,6 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { return Ok(plan); } - // Wholesale recompute: clear every leaf before deciding anything, so a - // group can never be left with one member carrying a decision from an - // earlier pass and another carrying none. - for arc in &leaves { - as_exchange(arc).set_coalesce(None); - } - let classified: Vec = leaves .iter() .map(|arc| classify_leaf(as_exchange(arc))) @@ -525,39 +510,43 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { } } + // Decide first, apply second. Every leaf's slot is then written exactly + // once per pass, with its final value: a leaf whose group decided to + // coalesce gets that group's shared plan, and every other leaf is + // actively reset rather than left carrying a decision from an earlier + // pass. There is no window in which a leaf holds a cleared slot the + // adapter could read. + let mut decisions: Vec>> = vec![None; leaves.len()]; + for (m, members) in group_by_upstream_count(&classified) { let Some(members) = members else { debug!("[coalesce-rule] group M={m} has an unusable leaf; skipping"); continue; }; - let sizes: Vec<&[u64]> = members - .iter() - .map(|&idx| match &classified[idx].kind { - LeafKind::Sizes(sizes) => sizes.as_slice(), - other => unreachable!( - "group_by_upstream_count only puts LeafKind::Sizes leaves in a \ - decidable group, but member {idx} of group M={m} is {}", - other.label() - ), - }) - .collect(); + let sizes: Vec<&[u64]> = members.iter().map(|member| member.sizes).collect(); let summed = sum_sizes(&sizes, m); debug!("[coalesce-rule] group M={m} summed bytes: {summed:?}"); let Some(cp) = decide(&summed, target, small, merged) else { continue; }; + // The same `Arc` on every member, so the upstream-index → group + // mapping is identical across the group, not merely the same `K`. let cp = Arc::new(cp); - for &idx in &members { - let ex = as_exchange(&leaves[idx]); - debug!( - "[coalesce-rule] set_coalesce(K={}) on plan_id={}", - cp.groups.len(), - ex.plan_id, - ); - ex.set_coalesce(Some(cp.clone())); + for member in &members { + decisions[member.idx] = Some(cp.clone()); } } + + for (arc, decision) in leaves.iter().zip(decisions) { + let ex = as_exchange(arc); + debug!( + "[coalesce-rule] set_coalesce({:?}) on plan_id={}", + decision.as_ref().map(|cp| cp.groups.len()), + ex.plan_id, + ); + ex.set_coalesce(decision); + } Ok(plan) } @@ -781,6 +770,21 @@ mod tests { } } + /// The leaf indices of group `m`: `None` if there is no such group, + /// `Some(None)` if the group exists but is undecidable, `Some(Some(..))` + /// with its members otherwise. The tests below are about *membership*, so + /// this keeps them from restating each member's sizes. + fn member_indices( + groups: &BTreeMap>>>, + m: usize, + ) -> Option>> { + groups.get(&m).map(|members| { + members + .as_ref() + .map(|members| members.iter().map(|member| member.idx).collect()) + }) + } + #[test] fn grouping_splits_leaves_by_upstream_partition_count() { let leaves = vec![ @@ -793,8 +797,25 @@ mod tests { let groups = group_by_upstream_count(&leaves); assert_eq!(groups.len(), 2); - assert_eq!(groups.get(&8), Some(&Some(vec![0, 1]))); - assert_eq!(groups.get(&4), Some(&Some(vec![2, 3]))); + assert_eq!(member_indices(&groups, 8), Some(Some(vec![0, 1]))); + assert_eq!(member_indices(&groups, 4), Some(Some(vec![2, 3]))); + } + + #[test] + fn grouping_carries_each_members_sizes_alongside_its_index() { + // The members carry their own sizes so the caller can sum a group + // without re-matching on `LeafKind`. + let leaves = vec![sized(2, vec![10, 20]), sized(2, vec![30, 40])]; + + let groups = group_by_upstream_count(&leaves); + + let members = groups + .get(&2) + .expect("group M=2") + .as_ref() + .expect("decidable"); + assert_eq!(members[0].sizes, &[10, 20]); + assert_eq!(members[1].sizes, &[30, 40]); } #[test] @@ -806,7 +827,7 @@ mod tests { let groups = group_by_upstream_count(&leaves); assert_eq!(groups.len(), 1); - assert_eq!(groups.get(&8), Some(&Some(vec![1, 2]))); + assert_eq!(member_indices(&groups, 8), Some(Some(vec![1, 2]))); assert!(!groups.contains_key(&1)); } @@ -826,12 +847,7 @@ mod tests { let groups = group_by_upstream_count(&leaves); assert_eq!(groups.len(), 1); - assert_eq!(groups.get(&8), Some(&Some(vec![1, 2]))); - } - - #[test] - fn grouping_of_only_pass_through_leaves_is_empty() { - assert!(group_by_upstream_count(&[pass_through(1)]).is_empty()); + assert_eq!(member_indices(&groups, 8), Some(Some(vec![1, 2]))); } #[test] @@ -840,12 +856,15 @@ mod tests { let groups = group_by_upstream_count(&leaves); - assert_eq!(groups.get(&8), Some(&None)); - assert_eq!(groups.get(&4), Some(&Some(vec![2]))); + assert_eq!(member_indices(&groups, 8), Some(None)); + assert_eq!(member_indices(&groups, 4), Some(Some(vec![2]))); } #[test] - fn grouping_of_only_broadcast_leaves_is_empty() { + fn grouping_of_only_excluded_leaves_is_empty() { + // Neither excluded kind creates a group entry at all, so a stage made + // up of nothing else has nothing to decide. assert!(group_by_upstream_count(&[broadcast(1)]).is_empty()); + assert!(group_by_upstream_count(&[pass_through(1)]).is_empty()); } } diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 44cab94552..67d303761a 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -369,14 +369,6 @@ impl AdaptivePlanner { // group. This avoids cross-stage gluing and stale state // that would arise if the rule walked the entire residual // plan in `default_optimizers()`. - // - // `adapt_to_ballista` must stay in this same closure. - // The rule clears the coalesce slot of every leaf it - // collects before deciding anything, and an - // `ExchangeExec` is a shared `Arc` seen by every stage - // that reads it, so optimizing all stages first and - // adapting them afterwards would let a later stage's - // clear wipe an earlier stage's decision unread. let plan = CoalescePartitionsRule.optimize(plan, config)?; // adapt_to_ballista takes an job_id, we are passing a job_name. Need to transform to fix compiler. let job_id = self.job_name.clone().into();