diff --git a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index 63cc18150..d25d3e644 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, ) @@ -214,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 @@ -245,12 +280,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. @@ -334,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( @@ -421,3 +457,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 5947be4c3..e1cac56f2 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,68 @@ //! 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. +//! 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 +//! build side and `UnspecifiedDistribution` on the probe side, so a broadcast +//! leaf is never a co-partitioned sibling. +//! +//! 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 +//! 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 +//! +//! **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 +//! upstream partition therefore still becomes a single downstream task, +//! however large it is. //! //! # Default off //! @@ -68,27 +102,32 @@ //! //! # 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` or `K <= 1`, the rewrite is degenerate 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. Classify each leaf. Broadcast and pass-through leaves drop out; +//! unresolved, statistics-free, and inconsistent leaves make their group +//! undecidable. +//! 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`. 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 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. 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 //! @@ -99,15 +138,16 @@ //! 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 //! -//! 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 +156,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, @@ -124,6 +164,280 @@ 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, + /// 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, 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 an explicit partitioning, and anything feeding a + /// `Partitioned` join gets a hash `RepartitionExec` from + /// `EnforceDistribution`, which `DistributedExchangeRule` turns into a + /// repartitioning 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`. + 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), +} + +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::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)] +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, +} + +/// 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, + }; + } + // Must run after the broadcast check above: `new_broadcast` and + // `to_broadcast` also leave `partitioning` unset, so a broadcast leaf would + // misclassify as `PassThrough` if this check ran first. + if ex.preserves_child_ordering() { + return ClassifiedLeaf { + m, + kind: LeafKind::PassThrough, + }; + } + // 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 LeafKind::UnknownBytes, + } + } + sizes.push(total); + } + LeafKind::Sizes(sizes) + }); + ClassifiedLeaf { + m, + kind: kind.unwrap_or(LeafKind::Unresolved), + } +} + +/// Partition classified leaves into alignment groups keyed by upstream +/// partition count `M`. +/// +/// 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(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. +/// +/// 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( + leaves: &[ClassifiedLeaf], +) -> 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(sizes) => { + if let Some(members) = entry { + members.push(GroupMember { idx, sizes }); + } + } + _ => *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. +/// +/// 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::() { + Some(( + exchange.input().clone(), + format!( + "ExchangeExec plan_id={} stage_id={:?}", + exchange.plan_id, + exchange.stage_id() + ), + )) + } else if let Some(adaptive) = plan.downcast_ref::() { + Some(( + adaptive.input().clone(), + format!("AdaptiveDatafusionExec stage_id={:?}", adaptive.stage_id()), + )) + } 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 +/// 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(); + 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; + } + 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. @@ -155,154 +469,83 @@ 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, root)) = 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); }; + debug!("[coalesce-rule] root={root}"); - // 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 { + if leaves.is_empty() { + debug!("[coalesce-rule] no leaves; bail"); + return Ok(plan); + } + + let classified: Vec = leaves + .iter() + .map(|arc| classify_leaf(as_exchange(arc))) + .collect(); + for (arc, leaf) in leaves.iter().zip(&classified) { let ex = as_exchange(arc); debug!( - "[coalesce-rule] leaf: plan_id={} stage_id={:?} partitioning={} M={} resolved={} existing_coalesce={:?}", + "[coalesce-rule] leaf: plan_id={} stage_id={:?} partitioning={} M={} kind={}", 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.m, + leaf.kind.label(), ); + 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 + ); + } } - // Leaf-scan stage with no upstream Exchanges → nothing to coalesce. - if leaves.is_empty() { - debug!("[coalesce-rule] no leaves; bail"); - 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) { - debug!("[coalesce-rule] broadcast leaf present; bail entire group"); - return Ok(plan); - } + // 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()]; - // 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) - { - return Ok(plan); - } + 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(|member| member.sizes).collect(); + let summed = sum_sizes(&sizes, m); + debug!("[coalesce-rule] group M={m} summed bytes: {summed:?}"); - // 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]; - 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); + let Some(cp) = decide(&summed, target, small, merged) else { + continue; }; - for (i, locs) in parts.iter().enumerate() { - summed[i] += locs - .iter() - .filter_map(|l| l.partition_stats.num_bytes()) - .sum::(); + // 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 member in &members { + decisions[member.idx] = Some(cp.clone()); } } - 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" - ); - 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), - }); - for arc in &leaves { + for (arc, decision) in leaves.iter().zip(decisions) { let ex = as_exchange(arc); debug!( - "[coalesce-rule] set_coalesce(K={k}) on plan_id={} (was {:?})", + "[coalesce-rule] set_coalesce({:?}) on plan_id={}", + decision.as_ref().map(|cp| cp.groups.len()), ex.plan_id, - ex.coalesce().as_ref().map(|cp| cp.groups.len()), ); - ex.set_coalesce(cp.clone()); + ex.set_coalesce(decision); } Ok(plan) } @@ -315,3 +558,313 @@ 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] + ); + } + + 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_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]))); + + 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); + } + + 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, + } + } + + fn pass_through(m: usize) -> ClassifiedLeaf { + ClassifiedLeaf { + m, + kind: LeafKind::PassThrough, + } + } + + /// 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![ + 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!(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] + 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!(member_indices(&groups, 8), Some(Some(vec![1, 2]))); + 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!(member_indices(&groups, 8), Some(Some(vec![1, 2]))); + } + + #[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!(member_indices(&groups, 8), Some(None)); + assert_eq!(member_indices(&groups, 4), Some(Some(vec![2]))); + } + + #[test] + 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/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 0ac47ae72..e329488cc 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -25,15 +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}; +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 ballista_core::serde::scheduler::{ - ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, - PartitionId, PartitionLocation, PartitionStats, -}; 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, displayable}; use datafusion::prelude::{SessionConfig, SessionContext}; use std::sync::Arc; @@ -83,40 +86,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; @@ -441,3 +410,350 @@ 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(()) +} + +/// 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. +// +// 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. +/// +/// `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)), + plan_id, + ); + 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(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) +} + +/// 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(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())?; + + assert_eq!(decision(&left), Some((2, 8))); + 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(()) +} + +/// #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(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())?; + + 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(1, 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(()) +} diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs b/ballista/scheduler/src/state/aqe/test/mod.rs index 52f173d19..6af69b910 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![