diff --git a/.github/workflows/h2o.yml b/.github/workflows/h2o.yml index 997935510e..4cc17f5506 100644 --- a/.github/workflows/h2o.yml +++ b/.github/workflows/h2o.yml @@ -171,6 +171,7 @@ jobs: --partitions 4 \ --verify \ -c ballista.planner.adaptive.enabled=true \ + -c ballista.planner.parallel_window.enabled=true \ -c ballista.scheduler.max_partitions_per_task=0 echo "::endgroup::" done diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index b2d7ed6624..2995e94eab 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -136,6 +136,11 @@ pub const BALLISTA_COALESCE_ENABLED: &str = "ballista.planner.coalesce.enabled"; /// This could benefit the workload by injecting EmptyExec in the plan (i.e during joins) pub const BALLISTA_PROPAGATE_EMPTY_ENABLED: &str = "ballista.planner.propagate_empty.enabled"; +/// Configuration key to enable the AQE `ParallelWindowRule`, which rewrites +/// bounded-RANGE-frame windows into a distributed range-shuffle so BWAG's +/// single-partition constraint isn't a serial bottleneck. Opt-in. +pub const BALLISTA_PARALLEL_WINDOW_ENABLED: &str = + "ballista.planner.parallel_window.enabled"; /// Configuration key for the target post-coalesce partition byte size (bytes). /// Mirrors Spark's `spark.sql.adaptive.advisoryPartitionSizeInBytes`. pub const BALLISTA_COALESCE_TARGET_PARTITION_BYTES: &str = @@ -323,6 +328,14 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| of a join, allowing downstream work to be skipped.".to_string(), DataType::Boolean, Some(true.to_string())), + ConfigEntry::new(BALLISTA_PARALLEL_WINDOW_ENABLED.to_string(), + "Enables the AQE parallel-window rule (ParallelWindowRule), which \ + rewrites bounded-RANGE-frame windows into a distributed range-shuffle \ + so BoundedWindowAggExec's single-partition constraint is not a serial \ + bottleneck. Disabled by default — opt in when the workload contains \ + matching window shapes.".to_string(), + DataType::Boolean, + Some(false.to_string())), ConfigEntry::new( BALLISTA_COALESCE_TARGET_PARTITION_BYTES.to_string(), "Target post-coalesce partition size in bytes. Mirrors Spark's \ @@ -706,6 +719,11 @@ impl BallistaConfig { self.get_bool_setting(BALLISTA_PROPAGATE_EMPTY_ENABLED) } + /// Returns whether the AQE parallel-window rule is enabled. + pub fn parallel_window_enabled(&self) -> bool { + self.get_bool_setting(BALLISTA_PARALLEL_WINDOW_ENABLED) + } + /// Returns compression codec that will be used during write stage of shuffle pub fn shuffle_compression_codec( &self, diff --git a/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs index d8f044442e..2b01987740 100644 --- a/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs +++ b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs @@ -74,11 +74,11 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; -// The rule's `as_candidate` gates guarantee no PARTITION BY + single Column -// ORDER BY over a sorted source, so `BWAG::try_new` is always invoked with -// `InputOrderMode::Sorted` and `can_repartition=false` (partition_keys() is -// empty either way when there's no PARTITION BY). Hardcode both to keep the -// wire and the type small. +// `maybe_rewrite_bwag`'s shape gates guarantee no PARTITION BY + single +// Column ORDER BY over a sorted source, so `BWAG::try_new` is always invoked +// with `InputOrderMode::Sorted` and `can_repartition=false` (partition_keys() +// is empty either way when there's no PARTITION BY). Hardcode both to keep +// the wire and the type small. const BWAG_INPUT_ORDER_MODE: InputOrderMode = InputOrderMode::Sorted; const BWAG_CAN_REPARTITION: bool = false; diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index eff45c6f17..312041f7d5 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -827,12 +827,23 @@ pub struct TaskRuntimeStats { pub report: RuntimeStatsReport, } -/// Walk `plan` for the first `UnorderedRangeRepartitionExec` or -/// `OrderedRangeRepartitionExec` and return its routing expression -/// (`order_by[0].expr`). `Ok(None)` means no range-repartition operator -/// in the plan; `Err(_)` means one was found but its `order_by` was +/// Walk the partition-preserving spine of `plan` for the +/// `UnorderedRangeRepartitionExec` or `OrderedRangeRepartitionExec` that +/// drives this stage's output partitioning, and return its routing +/// expression (`order_by[0].expr`). +/// +/// The spine is the chain of partition-preserving ops (see +/// [`super::preserves_partitioning`]) between the stage root and the barrier +/// that sets the stage's output partitioning. Descent stops at any +/// non-preserving op (join, union, hash-agg, unknown node) — an RRE +/// below such a barrier drives a different logical partitioning that +/// this stage's output no longer carries. +/// +/// `Ok(None)` means no range-repartition op drives this stage's +/// partitioning; `Err(_)` means one was found but its `order_by` was /// empty (invariant break — a range repartition without a routing key -/// can't route anything). +/// can't route anything), or the spine hit a partition-preserving node +/// with more than one child (shape bug in the whitelist). pub fn repartition_routing_expr( plan: &dyn ExecutionPlan, ) -> Result>> { @@ -848,36 +859,65 @@ pub fn repartition_routing_expr( [] => internal_err!("OrderedRangeRepartitionExec has empty ORDER BY"), }; } - for child in plan.children() { - if let Some(expr) = repartition_routing_expr(child.as_ref())? { - return Ok(Some(expr)); - } + if !super::preserves_partitioning(plan) { + return Ok(None); + } + let children = plan.children(); + match children.as_slice() { + [] => Ok(None), + [child] => repartition_routing_expr(child.as_ref()), + _ => internal_err!( + "partition-preserving op `{}` has {} children — the whitelist \ + assumes single-child; expand the algorithm if this fires", + plan.name(), + children.len() + ), } - Ok(None) } /// Rebuild a stage's `Vec>` under range-repartition /// overlap semantics: for each producer file in `original_partitions`, /// find its sketch (from `reports`), and route the file into every -/// downstream partition whose global cut range overlaps +/// downstream partition whose *halo-widened* range overlaps /// `[sketch.min(), sketch.max()]`. /// -/// Downstream partition ranges follow the half-open convention: -/// - `k = 0` → `(-∞, cuts[0])` -/// - `0 < k < K - 1` → `[cuts[k-1], cuts[k])` -/// - `k = K - 1` → `[cuts[K-2], +∞)` +/// Downstream partition ranges follow the half-open convention, widened +/// by the downstream `RangeFilterExec`'s halos on each side: +/// - `k = 0` → `(-∞, cuts[0] + halo_hi)` +/// - `0 < k < K - 1` → `[cuts[k-1] - halo_lo, cuts[k] + halo_hi)` +/// - `k = K - 1` → `[cuts[K-2] - halo_lo, +∞)` /// /// `[min, max]` overlaps `[lower, upper)` iff `max >= lower AND min < upper`. /// +/// `halo_lo`/`halo_hi` are `0.0` when the downstream stage has no halo +/// consumer (hash-agg, no-window range-repartition) — the check collapses +/// to raw cuts. When the downstream stage has a `RangeFilterExec` with +/// non-zero halo (bounded RANGE-frame windows), the caller passes the +/// widened halos so files straddling the halo band route to both sides. +/// Skipping this widening loses boundary rows from downstream window sums. +/// /// Files without a corresponding sketch (missing entirely, or present /// with `count == 0`) are safe to skip only when `partition_stats.num_rows` /// confirms the file is empty (`Some(0)`). If the file has rows or the /// row count is unknown (`None`), silently skipping would lose data — /// error out instead. +/// +/// # Arguments +/// +/// * `original_partitions` — passthrough shuffle output, `partitions[k]` +/// holds every file the writer produced for global partition `k`. +/// * `reports` — one per completed producer task; each carries the +/// per-sub-part sketches used for overlap lookup. +/// * `global_cuts` — K-1 monotone quantile cuts derived from merged +/// sketches; produce K downstream buckets. +/// * `halo_lo` / `halo_hi` — downstream `RangeFilterExec`'s halo widths +/// in the routing expression's units. `0.0` for non-halo consumers. pub fn cut_partitions( original_partitions: Vec>, reports: &[TaskRuntimeStats], global_cuts: &[f64], + halo_lo: f64, + halo_hi: f64, ) -> Result>> { use std::collections::HashMap; @@ -927,13 +967,15 @@ pub fn cut_partitions( } continue; }; - // Bucket i has (lower, upper) = (cuts[i-1], cuts[i]) with ±∞ at the - // ends, and matches iff `sketch_max >= lower && sketch_min < upper`. - // Monotone cuts → the set of matching buckets is a contiguous range - // [b_lo, b_hi], found by two partition_points over `global_cuts`. + // Bucket i has (lower, upper) = (cuts[i-1] - halo_lo, cuts[i] + + // halo_hi) with ±∞ at the ends, and matches iff `sketch_max + + // halo_lo >= cuts[i-1] && sketch_min - halo_hi < cuts[i]`. + // Monotone cuts → the set of matching buckets is a contiguous + // range [b_lo, b_hi], found by two partition_points over + // `global_cuts` with the sketch shifted by the halos. let (sketch_min, sketch_max) = (sketch.min(), sketch.max()); - let b_lo = global_cuts.partition_point(|&c| c <= sketch_min); - let b_hi = global_cuts.partition_point(|&c| c <= sketch_max); + let b_lo = global_cuts.partition_point(|&c| c <= sketch_min - halo_hi); + let b_hi = global_cuts.partition_point(|&c| c <= sketch_max + halo_lo); for bucket in &mut remapped[b_lo..=b_hi] { bucket.push(file.clone()); } @@ -1743,7 +1785,8 @@ mod overlap_remap_tests { // Passthrough map: both producers wrote to sub_part_id=0. let original_partitions = vec![vec![location(0, 100), location(0, 200)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2, "K = cuts.len() + 1"); // Partition 0: only producer 100. assert_eq!(remapped[0].len(), 1); @@ -1755,7 +1798,7 @@ mod overlap_remap_tests { /// A straddling sub-part — one whose sketched [min, max] spans the cut /// — appears in BOTH downstream partitions' lists. This is the case - /// PerPartitionFilterExec exists to clean up. + /// RangeFilterExec exists to clean up. #[test] fn overlap_remap_straddling_producer_appears_in_both_partitions() { // Producer 300 covers [5, 25) — straddles the cut at 15. @@ -1763,7 +1806,8 @@ mod overlap_remap_tests { let cuts = vec![15.0]; let original_partitions = vec![vec![location(0, 300)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert_eq!(remapped[0].len(), 1, "straddler in partition 0"); assert_eq!(remapped[0][0].file_id, Some(300)); @@ -1783,7 +1827,7 @@ mod overlap_remap_tests { bad.file_id = None; let original_partitions = vec![vec![bad]]; - let err = cut_partitions(original_partitions, &reports, &cuts) + let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) .expect_err("missing file_id must surface as an error"); assert!( err.to_string().contains("missing file_id"), @@ -1805,7 +1849,8 @@ mod overlap_remap_tests { let cuts = vec![10.0]; let original_partitions = vec![vec![location(0, 200)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1819,7 +1864,8 @@ mod overlap_remap_tests { let cuts = vec![10.0]; let original_partitions = vec![vec![location(0, 100)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1836,7 +1882,7 @@ mod overlap_remap_tests { orphan.partition_stats = PartitionStats::new(Some(5), None, None); let original_partitions = vec![vec![orphan]]; - let err = cut_partitions(original_partitions, &reports, &cuts) + let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) .expect_err("file with rows but no sketch must error"); let msg = err.to_string(); assert!( @@ -1876,7 +1922,8 @@ mod overlap_remap_tests { location(0, 6), ]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 4); let ids = |b: &[PartitionLocation]| { let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); @@ -1900,7 +1947,7 @@ mod overlap_remap_tests { orphan.partition_stats = PartitionStats::default(); // num_rows = None let original_partitions = vec![vec![orphan]]; - let err = cut_partitions(original_partitions, &reports, &cuts) + let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) .expect_err("file with unknown rows but no sketch must error"); let msg = err.to_string(); assert!( @@ -1908,4 +1955,82 @@ mod overlap_remap_tests { "unexpected error: {msg}" ); } + + /// Halo widening lets each partition see files that sit within + /// `[halo_lo, halo_hi]` of its raw cut range — the downstream + /// `RangeFilterExec`'s frame-context rows come from those files, and + /// missing any of them causes RANGE-frame window sums to drop rows + /// at boundaries. + /// + /// The K=5 layout with `halo_lo != halo_hi` proves three things at + /// once: (a) `halo_lo` widens downward, (b) `halo_hi` widens upward, + /// (c) the halo band stays *local* — it does not bleed across two + /// cut hops to far-away partitions. + #[test] + fn overlap_remap_halo_band_widens_both_sides_without_bleeding_to_far_partitions() { + // K=5, asymmetric halos so we can tell halo_lo and halo_hi apart. + let cuts = vec![10.0, 20.0, 30.0, 40.0]; + let halo_lo = 1.0; + let halo_hi = 2.0; + // Effective partition ranges: + // P0: (-∞, 12) P1: [9, 22) P2: [19, 32) P3: [29, 42) P4: [39, +∞) + let reports = vec![ + // 100 sits deep inside P0 — far from P1's halo, stays P0-only. + sketch_report(100, vec![vec![5.0, 6.0]]), + // 200 is entirely below cut 20 but within halo_lo=1 of it — + // routes to P1 (own bucket) AND P2 (halo band from below). + // Must NOT reach P0 (two cut hops away). + sketch_report(200, vec![vec![18.0, 19.0]]), + // 300 sits cleanly inside P2 — no halo participation. + sketch_report(300, vec![vec![25.0, 26.0]]), + // 400 is entirely above cut 30 but within halo_hi=2 of it — + // routes to P2 (halo band from above) AND P3 (own bucket). + // Must NOT reach P4 (two cut hops away). + sketch_report(400, vec![vec![31.0, 32.0]]), + // 500 sits deep inside P4 — far from P3's halo, stays P4-only. + sketch_report(500, vec![vec![45.0, 46.0]]), + ]; + let original_partitions = vec![vec![ + location(0, 100), + location(0, 200), + location(0, 300), + location(0, 400), + location(0, 500), + ]]; + + let remapped = + cut_partitions(original_partitions, &reports, &cuts, halo_lo, halo_hi) + .unwrap(); + let ids = |b: &[PartitionLocation]| { + let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); + v.sort(); + v + }; + assert_eq!(remapped.len(), 5); + assert_eq!( + ids(&remapped[0]), + vec![100u64], + "P0 sees only its own bucket — 200's halo band belongs to P1/P2, not here", + ); + assert_eq!( + ids(&remapped[1]), + vec![200u64], + "P1 sees its own straddler (200) below cut 20", + ); + assert_eq!( + ids(&remapped[2]), + vec![200u64, 300, 400], + "P2 (middle) sees siblings from BOTH halo bands — 200 via halo_lo, 400 via halo_hi — plus its own 300", + ); + assert_eq!( + ids(&remapped[3]), + vec![400u64], + "P3 sees its own straddler (400) above cut 30", + ); + assert_eq!( + ids(&remapped[4]), + vec![500u64], + "P4 sees only its own bucket — 400's halo band belongs to P2/P3, not here", + ); + } } diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index fcfd52514c..e673a51c8e 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -23,7 +23,9 @@ use ballista_core::client_pool::BallistaClientPool; use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec; -use ballista_core::execution_plans::{ShuffleReaderExec, ShuffleWriterExec}; +use ballista_core::execution_plans::{ + RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriterExec, +}; use ballista_core::serde::protobuf::ShuffleWritePartition; use ballista_core::serde::scheduler::PartitionStats; use ballista_core::{JobId, utils}; @@ -135,10 +137,6 @@ impl ExecutionEngine for DefaultExecutionEngine { ) -> Result> { let plan = plan .transform(|p| { - // TODO: RangeShuffleReaderExec needs the same late-bind - // (with_work_dir + with_client_pool) once a planner rule - // plants it; without it, the first task carrying one will - // fail with "work dir should have been set by executor". if let Some(reader) = p.downcast_ref::() { match &self.client_pool { Some(client_pool) => Ok(Transformed::yes(Arc::new( @@ -150,6 +148,17 @@ impl ExecutionEngine for DefaultExecutionEngine { reader.with_work_dir(work_dir.to_string()), ))), } + } else if let Some(reader) = p.downcast_ref::() { + match &self.client_pool { + Some(client_pool) => Ok(Transformed::yes(Arc::new( + reader + .with_work_dir(work_dir.to_string()) + .with_client_pool(client_pool.clone()), + ))), + None => Ok(Transformed::yes(Arc::new( + reader.with_work_dir(work_dir.to_string()), + ))), + } } else { // Scan restriction is scheduler-side (see // ballista/scheduler/src/state/task_builder.rs). The plan diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs index 43de627535..46a3c4bd4f 100644 --- a/ballista/scheduler/src/cluster/mod.rs +++ b/ballista/scheduler/src/cluster/mod.rs @@ -24,7 +24,7 @@ use crate::state::execution_graph::{ use crate::state::task_manager::JobInfoCache; use ballista_core::config::BallistaConfig; use ballista_core::error::Result; -use ballista_core::execution_plans::ShuffleReaderExec; +use ballista_core::execution_plans::{RangeShuffleReaderExec, ShuffleReaderExec}; use ballista_core::serde::protobuf::{ AvailableVcores, ExecutorHeartbeat, JobStatus, job_status, }; @@ -369,15 +369,17 @@ pub trait JobState: Send + Sync { /// (e.g. `UnorderedRangeRepartitionExec`). Stops at leaves, multi-child /// operators (fan-in / joins), and stage boundaries. /// -/// The stage-boundary stop is currently a `ShuffleReaderExec` downcast — the -/// only kind of stage-boundary leaf that appears in a resolved stage plan. -/// The *general* rule is "stop at any stage boundary"; if new stage-boundary -/// operators appear, add them here (or, better, get `ExecutionPlan` upstream -/// to expose an `is_stage_boundary()` property so we don't keep -/// enumerating). +/// The stage-boundary stop enumerates the leaf readers that terminate a +/// resolved stage plan: `ShuffleReaderExec` (regular / broadcast / coalesced) +/// and `RangeShuffleReaderExec` (ordering-preserving). The *general* rule is +/// "stop at any stage boundary"; if new stage-boundary operators appear, add +/// them here (or, better, get `ExecutionPlan` upstream to expose an +/// `is_stage_boundary()` property so we don't keep enumerating). fn stage_has_input_collapse(plan_root: &Arc) -> bool { fn walk(node: &Arc) -> bool { - if node.downcast_ref::().is_some() { + if node.downcast_ref::().is_some() + || node.downcast_ref::().is_some() + { return false; } if node.properties().output_partitioning().partition_count() == 1 { @@ -665,6 +667,7 @@ mod test { ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, }; + use crate::cluster::stage_has_input_collapse; use crate::cluster::{BoundTask, bind_task_bias, bind_task_round_robin}; use crate::state::execution_graph::{ExecutionGraph, StaticExecutionGraph}; use crate::state::task_manager::JobInfoCache; @@ -891,4 +894,37 @@ mod test { }, ] } + + /// Both shuffle reader kinds are stage boundaries — walking through them + /// to detect an input collapse would mis-classify the *next* stage's leaf + /// as this stage's collapse. `stage_has_input_collapse` must return false + /// as soon as a reader is seen. Guard the range variant explicitly since + /// `UnknownPartitioning(1)` would otherwise trigger the single-partition + /// arm. + #[test] + fn stage_has_input_collapse_stops_at_range_reader() { + use ballista_core::execution_plans::RangeShuffleReaderExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("v", 0))); + let merge_ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + // Single output partition — the case where the `partition_count == 1` + // arm would fire without the reader guard. + let reader = Arc::new( + RangeShuffleReaderExec::try_new(1, vec![vec![]], schema, merge_ordering) + .unwrap(), + ) as Arc; + let root: Arc = Arc::new(CoalescePartitionsExec::new(reader)); + + assert!( + !stage_has_input_collapse(&root), + "a range-shuffle reader is a stage boundary, not an input collapse", + ); + } } diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 0b3af8fc00..765b181fee 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -27,8 +27,8 @@ use ballista_core::execution_plans::ShuffleWriter; use ballista_core::execution_plans::sort_shuffle::SortShuffleConfig; use ballista_core::{ execution_plans::{ - ShuffleReaderExec, ShuffleWriterExec, SortShuffleWriterExec, - UnresolvedShuffleExec, + RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriterExec, + SortShuffleWriterExec, UnresolvedShuffleExec, }, serde::scheduler::PartitionLocation, }; @@ -785,6 +785,11 @@ pub fn remove_unresolved_shuffles( /// Rollback the ShuffleReaderExec to UnresolvedShuffleExec. /// Used when the input stages are finished but some partitions are missing due to executor lost. /// The entire stage need to be rolled back and rescheduled. +/// +/// `RangeShuffleReaderExec` rolls back to a plain `UnresolvedShuffleExec` — its +/// range-ness is a derived property of the child's declared ordering at plan +/// time, not intrinsic reader metadata. Re-planning walks the adapter, which +/// re-detects the ordering and plants a fresh `RangeShuffleReaderExec`. pub fn rollback_resolved_shuffles( stage: Arc, ) -> Result> { @@ -806,6 +811,14 @@ pub fn rollback_resolved_shuffles( )) }; new_children.push(unresolved); + } else if let Some(range_reader) = child.downcast_ref::() + { + let unresolved = Arc::new(UnresolvedShuffleExec::new( + range_reader.stage_id, + range_reader.schema(), + range_reader.properties().partitioning.clone(), + )); + new_children.push(unresolved); } else { new_children.push(rollback_resolved_shuffles(child.clone())?); } @@ -2041,6 +2054,51 @@ order by Ok(()) } + /// `RangeShuffleReaderExec` rolls back to a plain `UnresolvedShuffleExec` + /// (info-losing on the range-ness). Re-planning walks the adapter, which + /// re-detects the child's ordering and plants a fresh range reader. + #[tokio::test] + async fn rollback_resolved_shuffles_reduces_range_reader_to_plain_unresolved() + -> Result<(), BallistaError> { + use ballista_core::execution_plans::RangeShuffleReaderExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("v", 0))); + let merge_ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let reader = Arc::new( + RangeShuffleReaderExec::try_new( + 7, + vec![vec![]; 4], + schema.clone(), + merge_ordering, + ) + .unwrap(), + ) as Arc; + let parent: Arc = Arc::new( + datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec::new( + reader, + ), + ); + + let rolled_back = crate::planner::rollback_resolved_shuffles(parent)?; + let child = rolled_back.children()[0].clone(); + let unresolved = child + .downcast_ref::() + .expect("expected rolled-back UnresolvedShuffleExec"); + // The range-ness is derived at plan time; the rolled-back node carries + // no ordering, no broadcast, no coalesce. + assert!(!unresolved.broadcast); + assert!(unresolved.coalesce.is_none()); + assert_eq!(unresolved.stage_id, 7); + assert_eq!(unresolved.output_partition_count, 4); + + Ok(()) + } + #[tokio::test] async fn distributed_window_plan() -> Result<(), BallistaError> { let ctx = datafusion_test_context("testdata").await?; diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index b6ec54866b..29d137615d 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -16,22 +16,24 @@ // under the License. use crate::planner::create_shuffle_writer_with_config; -use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; +use crate::state::aqe::execution_plan::{ + AdaptiveDatafusionExec, ExchangeExec, RangeRepartitionRouting, +}; use crate::state::aqe::planner::AdaptiveStageInfo; use crate::state::execution_graph::StageOutput; use ballista_core::JobId; use ballista_core::execution_plans::{ - PerPartitionFilterExec, ShuffleReaderExec, range_partition_predicates, + RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, }; use datafusion::common::exec_err; use datafusion::config::ConfigOptions; use datafusion::error::DataFusionError; use datafusion::physical_plan::{ExecutionPlanProperties, Partitioning}; +use datafusion::scalar::ScalarValue; use datafusion::{ - common::tree_node::{Transformed, TreeNode}, + common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}, physical_plan::ExecutionPlan, }; -use log::debug; use std::collections::HashMap; use std::sync::Arc; @@ -71,67 +73,65 @@ impl BallistaAdapter { self.inputs.insert(stage_id, stage_output); let partitioning = exchange.properties().partitioning.clone(); - let reader = match (exchange.coalesce(), exchange.broadcast) { - (Some(cp), false) => { - // Concatenate M-shape locations into K-shape per CoalescePlan.groups. - let k_shape: Vec> = cp - .groups - .iter() - .map(|pg| { - let mut concat = Vec::new(); - for &idx in &pg.upstream_indices { - if let Some(inner) = partitions.get(idx as usize) { - concat.extend_from_slice(inner); + let reader: Arc = + match (exchange.coalesce(), exchange.broadcast) { + (Some(cp), false) => { + // Concatenate M-shape locations into K-shape per CoalescePlan.groups. + let k_shape: Vec> = cp + .groups + .iter() + .map(|pg| { + let mut concat = Vec::new(); + for &idx in &pg.upstream_indices { + if let Some(inner) = partitions.get(idx as usize) { + concat.extend_from_slice(inner); + } } + concat + }) + .collect(); + let new_partitioning = match &partitioning { + Partitioning::Hash(keys, _m) => { + Partitioning::Hash(keys.clone(), cp.groups.len()) } - concat - }) - .collect(); - let new_partitioning = match &partitioning { - Partitioning::Hash(keys, _m) => { - Partitioning::Hash(keys.clone(), cp.groups.len()) + _ => Partitioning::UnknownPartitioning(cp.groups.len()), + }; + Arc::new(ShuffleReaderExec::try_new_coalesced( + stage_id, + k_shape, + (*cp).clone(), + schema, + new_partitioning, + )?) + } + (None, false) => { + // Ordered-writer path: when the child declared an output + // ordering, preserve it across the shuffle boundary with a + // k-way merge instead of the arrival-order concat that the + // regular reader does. + if let Some(ordering) = exchange.input().output_ordering() { + Arc::new(RangeShuffleReaderExec::try_new( + stage_id, + partitions, + schema, + ordering.clone(), + )?) + } else { + Arc::new(ShuffleReaderExec::try_new( + stage_id, + partitions, + schema, + partitioning, + )?) } - _ => Partitioning::UnknownPartitioning(cp.groups.len()), - }; - ShuffleReaderExec::try_new_coalesced( + } + (_, true) => Arc::new(ShuffleReaderExec::try_new_broadcast( stage_id, - k_shape, - (*cp).clone(), + exchange.shuffle_partitions_flattened(), schema, - new_partitioning, - )? - } - (None, false) => ShuffleReaderExec::try_new( - stage_id, - partitions, - schema, - partitioning, - )?, - (_, true) => ShuffleReaderExec::try_new_broadcast( - stage_id, - exchange.shuffle_partitions_flattened(), - schema, - exchange.input().output_partitioning().partition_count(), - )?, - }; - - let reader: Arc = Arc::new(reader); - // Without a per-partition filter, straddling sub-parts from a - // range-repartitioned upstream would feed multiple downstream - // partitions and `FinalPartitioned` would split their partial sums. - if let Some(routing) = exchange.range_repartition_routing() { - let predicates = - range_partition_predicates(routing.routing_expr, &routing.cuts); - debug!( - "range-repartition: injecting PerPartitionFilterExec above \ - ShuffleReader for stage {} — {} predicates over {} cuts", - stage_id, - predicates.len(), - routing.cuts.len(), - ); - let filtered = PerPartitionFilterExec::try_new(reader, predicates)?; - return Ok(Transformed::yes(Arc::new(filtered))); - } + exchange.input().output_partitioning().partition_count(), + )?), + }; Ok(Transformed::yes(reader)) } else { Ok(Transformed::no(plan)) @@ -148,6 +148,7 @@ impl BallistaAdapter { ) -> datafusion::error::Result { if let Some(root) = plan.downcast_ref::() { let mut adapter = BallistaAdapter::default(); + resolve_range_filter_cuts(root.input())?; let plan = root .input() .clone() @@ -175,6 +176,7 @@ impl BallistaAdapter { }) } else if let Some(root) = plan.downcast_ref::() { let mut adapter = BallistaAdapter::default(); + resolve_range_filter_cuts(root.input())?; let plan = root .input() .clone() @@ -201,3 +203,185 @@ impl BallistaAdapter { } } } + +/// Walk `plan` and resolve every pending [`RangeFilterExec`]'s bounds from +/// its own descendant boundary `ExchangeExec`'s stored routing. Called at +/// `adapt_to_ballista` time, once the upstream stage's sketches have +/// merged into cuts and been parked on the boundary `ExchangeExec` via +/// `set_repartition_routing`. +/// +/// Pairing is by tree structure — each RFE descends its own subtree to +/// the first `ExchangeExec` and takes that exchange's routing. Using +/// `PhysicalExpr::eq` on routing exprs alone would collide in +/// multi-legged shapes (e.g. SMJ with range-repartition on both sides +/// sharing the same `Column(name, idx)` after independent projections); +/// descent is unique by construction because every RFE has exactly one +/// child. The RFE's own `routing_expr` is then cross-checked against +/// the descendant exchange's for a plant-time invariant assert — if they +/// disagree, the rule wired the wrong RFE to the boundary. +/// +/// RFE receives *unwidened* half-open ranges (`(cuts[k-1], cuts[k])` with ±∞ +/// sentinels at ends). RFE widens by its own halos internally at +/// `resolve_bounds` time — see the separation-of-concerns note on +/// [`RangeFilterExec`]. The scheduler stays halo-blind at this boundary. +/// +/// Errors if descent hits a fork (multi-child op) before reaching an +/// `ExchangeExec`, if the RFE has anything other than 1 child, if the +/// descendant `ExchangeExec` has no resolved routing yet, or if its +/// routing_expr disagrees with the RFE's — all four are plant-time or +/// stage-progress bugs. +fn resolve_range_filter_cuts( + plan: &Arc, +) -> Result<(), DataFusionError> { + plan.apply(|node| { + let Some(rf) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + if rf.raw_bounds().is_some() { + return Ok(TreeNodeRecursion::Continue); + } + let children = rf.children(); + let [child] = children.as_slice() else { + return datafusion::common::internal_err!( + "RangeFilterExec must have exactly 1 child, got {}", + children.len() + ); + }; + let routing = descend_to_boundary_routing(child)?; + if !rf.routing_expr().eq(&routing.routing_expr) { + return datafusion::common::internal_err!( + "RangeFilterExec routing_expr `{}` disagrees with its descendant \ + boundary ExchangeExec's routing_expr `{}` — plant-time invariant \ + broken", + rf.routing_expr(), + routing.routing_expr + ); + } + let raw_bounds = raw_bounds_from_cuts(&routing.cuts); + rf.resolve_bounds(raw_bounds)?; + Ok(TreeNodeRecursion::Continue) + })?; + Ok(()) +} + +/// Descend the single-child spine below a `RangeFilterExec` until we +/// hit an `ExchangeExec`, and return its `range_repartition_routing`. +/// Forks in the spine are shape violations because a range-filter only +/// makes sense above a single-input boundary. +fn descend_to_boundary_routing( + start: &Arc, +) -> Result { + let mut node = Arc::clone(start); + loop { + if let Some(exchange) = node.downcast_ref::() { + return exchange.range_repartition_routing().ok_or_else(|| { + DataFusionError::Internal( + "RangeFilterExec's descendant ExchangeExec has no resolved \ + range-repartition routing yet — stage progress skipped a step" + .into(), + ) + }); + } + let children = node.children(); + let [child] = children.as_slice() else { + return datafusion::common::internal_err!( + "RangeFilterExec descent hit a fork at `{}` ({} children) — cannot \ + pair with a single boundary", + node.name(), + children.len() + ); + }; + node = Arc::clone(*child); + } +} + +/// Project K-1 cuts to K half-open `(cuts[k-1], cuts[k])` ranges with `None` +/// sentinels at ±∞. This is the pure range-partitioning projection — no halo +/// arithmetic here (RFE widens internally at resolve time). +fn raw_bounds_from_cuts(cuts: &[f64]) -> Vec<(Option, Option)> { + let k = cuts.len() + 1; + (0..k) + .map(|i| { + let lo = i + .checked_sub(1) + .and_then(|j| cuts.get(j).copied()) + .map(|v| ScalarValue::Float64(Some(v))); + let hi = cuts.get(i).copied().map(|v| ScalarValue::Float64(Some(v))); + (lo, hi) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use ballista_core::execution_plans::RangeShuffleReaderExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::sorts::sort::SortExec; + + fn f64_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])) + } + + fn asc(schema: &Arc, col: &str) -> PhysicalSortExpr { + let column = Column::new_with_schema(col, schema).unwrap(); + PhysicalSortExpr::new_default(Arc::new(column)) + } + + /// When the exchange's child declares an output ordering, the adapter + /// must plant `RangeShuffleReaderExec` so the shuffle boundary preserves + /// sortedness via k-way merge. + #[test] + fn plants_range_reader_when_child_declares_ordering() { + let schema = f64_schema(); + let empty: Vec> = vec![vec![]]; + let source = + MemorySourceConfig::try_new_exec(&empty, schema.clone(), None).unwrap(); + let sort_lex = LexOrdering::new(vec![asc(&schema, "v")]).unwrap(); + let sorted = + Arc::new(SortExec::new(sort_lex, source).with_preserve_partitioning(true)) + as Arc; + + let exchange = ExchangeExec::new(sorted, None, 0); + exchange.set_stage_id(1); + exchange.resolve_shuffle_partitions(vec![vec![]]); + let plan: Arc = Arc::new(exchange); + + let mut adapter = BallistaAdapter::default(); + let out = adapter.transform_children(plan).unwrap().data; + + assert!( + out.downcast_ref::().is_some(), + "expected RangeShuffleReaderExec but got {}", + out.name() + ); + } + + /// The unordered path must still plant the regular `ShuffleReaderExec`. + #[test] + fn plants_regular_reader_when_no_ordering() { + let schema = f64_schema(); + let empty: Vec> = vec![vec![]]; + let source = + MemorySourceConfig::try_new_exec(&empty, schema.clone(), None).unwrap(); + + let exchange = ExchangeExec::new(source, None, 0); + exchange.set_stage_id(1); + exchange.resolve_shuffle_partitions(vec![vec![]]); + let plan: Arc = Arc::new(exchange); + + let mut adapter = BallistaAdapter::default(); + let out = adapter.transform_children(plan).unwrap().data; + + assert!( + out.downcast_ref::().is_some(), + "expected ShuffleReaderExec but got {}", + out.name() + ); + assert!(out.downcast_ref::().is_none()); + } +} diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index dac8482214..0b70cf612a 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -18,7 +18,7 @@ use crate::display::print_stage_metrics; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::scheduler_server::timestamp_millis; -use crate::state::aqe::execution_plan::RangeRepartitionRouting; +use crate::state::aqe::execution_plan::{ExchangeExec, RangeRepartitionRouting}; use crate::state::aqe::planner::{AdaptivePlanner, AdaptiveStageInfo}; use crate::state::execution_graph::{ ExecutionGraph, ExecutionGraphBox, ExecutionStage, ResolvedStage, RunningTaskInfo, @@ -29,7 +29,8 @@ use crate::state::task_manager::UpdatedStages; use ballista_core::JobId; use ballista_core::error::BallistaError; use ballista_core::execution_plans::{ - cut_partitions, merge_runtime_stats_reports, repartition_routing_expr, + RangeFilterExec, cut_partitions, merge_runtime_stats_reports, + repartition_routing_expr, }; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; @@ -38,10 +39,12 @@ use ballista_core::serde::protobuf::{ job_status, task_status, }; use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionLocation}; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; +use datafusion::scalar::ScalarValue; use log::{debug, error, info, warn}; use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -372,12 +375,31 @@ impl AdaptiveExecutionGraph { && let Some(routing) = Self::repartition_routing(stage, routing_expr)? { let reports = &stage.runtime_stats_reports; - let remapped = cut_partitions(partitions, reports, &routing.cuts) + // Reader-side halos widen each partition's effective + // range to `[cuts[k-1] - halo_lo, cuts[k] + halo_hi)`; + // files straddling that band must route to both sides. + // Range-repartition boundaries always have a consuming + // RFE (writer produces straddler duplicates that only + // the reader-side filter can trim), so the walker errors + // if none is found. + let (halo_lo, halo_hi) = downstream_halos(&self.planner.plan, stage_id) .map_err(|err| { - BallistaError::General(format!( - "range-repartition stage {stage_id}: overlap remap failed: {err}" - )) - })?; + BallistaError::General(format!( + "range-repartition stage {stage_id}: halo lookup failed: {err}" + )) + })?; + let remapped = cut_partitions( + partitions, + reports, + &routing.cuts, + halo_lo, + halo_hi, + ) + .map_err(|err| { + BallistaError::General(format!( + "range-repartition stage {stage_id}: overlap remap failed: {err}" + )) + })?; // Save boundaries to ExchangeExec so they are there for resolve_stage_partitions self.planner.set_repartition_routing(stage_id, routing)?; remapped @@ -1429,6 +1451,80 @@ impl ExecutionGraph for AdaptiveExecutionGraph { } } +/// Find the halos of the downstream `RangeFilterExec` planted directly on +/// the boundary `ExchangeExec` produced by `producer_stage_id`. A +/// range-repartition boundary always has a consuming RFE — the writer +/// produces straddler duplicates that only the reader-side filter can +/// trim — so absence is a shape bug, not a runtime default. +/// +/// Boundaries are disambiguated by `ExchangeExec::stage_id()` (the id +/// of the stage that produces to the exchange), not by routing +/// expression: two range-repartition stages could share the same +/// `Column(name, index)` shape after independent projections, so +/// `PhysicalExpr::eq` isn't unique across stages. The producer stage id +/// is. +/// +/// The rule may plant multiple RFEs (e.g. wide directly on the boundary +/// and narrow above the window operator). Only the one whose immediate +/// child is the boundary `ExchangeExec` describes the reader-visible +/// halo band and matters for straddler routing. +/// +/// # Arguments +/// +/// * `full_plan` — the AdaptivePlanner's current plan tree; the walker +/// descends the whole tree looking for the RFE on this boundary. +/// * `producer_stage_id` — the id of the completing upstream stage; +/// matches `ExchangeExec::stage_id()` on the boundary exchange. +fn downstream_halos( + full_plan: &Arc, + producer_stage_id: usize, +) -> datafusion::common::Result<(f64, f64)> { + let mut result: Option<(f64, f64)> = None; + full_plan.apply(|node| { + let Some(rf) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + let children = rf.children(); + let [child] = children.as_slice() else { + return datafusion::common::internal_err!( + "RangeFilterExec must have exactly 1 child, got {}", + children.len() + ); + }; + // RFE above the window op (narrow, halo=[0,0]) has a + // `PartitionedBoundedWindowAggExec` child, not the boundary + // exchange — keep looking for the wide RFE below. + let Some(exchange) = child.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + if exchange.stage_id() != Some(producer_stage_id) { + return Ok(TreeNodeRecursion::Continue); + } + result = Some((scalar_to_f64(rf.halo_lo())?, scalar_to_f64(rf.halo_hi())?)); + Ok(TreeNodeRecursion::Stop) + })?; + result.ok_or_else(|| { + datafusion::common::DataFusionError::Internal(format!( + "range-repartition boundary for producer stage {producer_stage_id} has \ + no consuming RangeFilterExec — straddler files would corrupt \ + downstream partial aggregates. Check the rule that planted this ORRE." + )) + }) +} + +/// Halos travel through the RFE public API as `ScalarValue` for future +/// type widening (Interval, timestamps under KLL). Today the internal +/// routing math is `f64`; any other variant is a shape violation upstream +/// and we fail loud rather than silently zero-widen. +fn scalar_to_f64(sv: &ScalarValue) -> datafusion::common::Result { + match sv { + ScalarValue::Float64(Some(v)) => Ok(*v), + other => datafusion::common::internal_err!( + "only f64 halos are implemented, got: {other:?}" + ), + } +} + /// Checks is the plan same as expected string representation #[cfg(test)] #[macro_export] diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs index 80e325566f..c3a7b9962b 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs @@ -17,7 +17,8 @@ use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use ballista_core::execution_plans::{ - OrderedRangeRepartitionExec, UnorderedRangeRepartitionExec, preserves_partitioning, + OrderedRangeRepartitionExec, RangeFilterExec, UnorderedRangeRepartitionExec, + preserves_partitioning, }; use datafusion::common::plan_err; use datafusion::common::tree_node::{Transformed, TreeNode}; @@ -104,7 +105,7 @@ impl DistributedExchangeRule { execution_plan.downcast_ref::() { let input = sort_preserving_merge.input(); - if input.downcast_ref::().is_none() + if !is_stage_boundary(input) && !matches!(nearest_exchange_status(input), ExchangeStatus::Unresolved) { let exchange_exec = ExchangeExec::new( @@ -261,6 +262,22 @@ impl PhysicalOptimizerRule for DistributedExchangeRule { } } +/// True when `node` is (or transparently sits on) a stage boundary. +/// `RangeFilterExec` counts because we chose not to fold range-filtering +/// into `ShuffleReader`/`ExchangeExec` — the operator is part of the +/// boundary shape by design. +fn is_stage_boundary(node: &Arc) -> bool { + if node.is::() { + return true; + } + if node.is::() + && let [child] = node.children().as_slice() + { + return child.is::(); + } + false +} + /// Scans the subtree for the nearest `ExchangeExec` in each path and returns the /// aggregate status. Stops recursing at `ExchangeExec` boundaries so that only the /// shallowest exchange in each branch is considered. @@ -534,6 +551,52 @@ mod tests { ); } + #[test] + fn spm_skips_when_range_filter_covers_exchange() { + // ParallelWindowRule plants a RangeFilterExec directly on the + // resolved range-repartition ExchangeExec with SPM above. The + // filter must count as part of the boundary — otherwise DE + // inserts another ExchangeExec between SPM and the filter, + // collapsing K partitions into a single outer-stage task. + use ballista_core::execution_plans::RangeFilterExec; + use datafusion::scalar::ScalarValue; + + let rule = DistributedExchangeRule::default(); + let exchange = resolved_exchange(float_leaf_exec()); + let filter: Arc = Arc::new( + RangeFilterExec::try_new_pending( + exchange, + Arc::new(Column::new("v", 0)), + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + ) + .unwrap(), + ); + let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("v", 0))); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let spm: Arc = + Arc::new(SortPreservingMergeExec::new(ordering, filter)); + + let result = rule.optimize(spm, &config()).unwrap(); + + let adaptive = result.downcast_ref::().unwrap(); + let spm_out = adaptive + .input() + .downcast_ref::() + .unwrap(); + let below_spm = spm_out.children()[0]; + assert!( + below_spm.downcast_ref::().is_some(), + "SPM's direct child should remain RangeFilterExec, not a new ExchangeExec" + ); + assert!( + below_spm.children()[0] + .downcast_ref::() + .is_some(), + "resolved ExchangeExec should remain under the RangeFilterExec" + ); + } + // --- RepartitionExec --- #[test] @@ -830,7 +893,7 @@ mod tests { /// range-repartition-inserting rule emits, with nothing above it — /// must still get an `ExchangeExec` wrapped above it. Without it, /// `set_repartition_routing` has no parking slot for the recovered - /// cuts and downstream never gets a `PerPartitionFilterExec` to + /// cuts and downstream never gets a `RangeFilterExec` to /// trim straddler duplication. #[test] fn range_repartition_at_plan_root_gets_exchange_inserted() { @@ -920,7 +983,7 @@ mod tests { /// A `ProjectionExec` between (O/U)RRE and the boundary could /// reindex, drop, or shadow the routing expression's referenced - /// columns — the read-side `PerPartitionFilterExec` would evaluate + /// columns — the read-side `RangeFilterExec` would evaluate /// against the wrong column and silently misroute. DER rejects the /// shape at plan time; the fix will be revisited when arbitrary /// routing expressions replace the current single-column form. diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs index b4ae1fbc83..392c4bc373 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs @@ -19,9 +19,11 @@ pub mod chaos_exec; pub mod coalesce_partitions; pub mod distributed_exchange; pub mod join_selection; +pub mod parallel_window; pub mod propagate_empty; pub use coalesce_partitions::*; pub use distributed_exchange::*; pub use join_selection::*; +pub use parallel_window::*; pub use propagate_empty::*; diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs new file mode 100644 index 0000000000..fb0038206e --- /dev/null +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs @@ -0,0 +1,567 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Rewrite bounded-RANGE-frame windows into a distributed range-shuffle so +//! `BoundedWindowAggExec`'s single-partition constraint isn't a serial +//! bottleneck. +//! +//! # Matched shape +//! +//! ```text +//! BoundedWindowAggExec [Sorted, RANGE frame, finite bounds] +//! SortPreservingMergeExec [ORDER BY] +//! SortExec (preserve_partitioning=true) [ORDER BY] +//! +//! ``` +//! +//! Restricted to: +//! - single window expression +//! - no PARTITION BY +//! - single-column ORDER BY on a physical `Column` (widening: multi-key, +//! computed exprs — separate rewrites) +//! - `RANGE` frame with finite `PRECEDING` / `FOLLOWING` / `CurrentRow` +//! bounds (UNBOUNDED frames go down a different path) +//! - ORDER BY column is `Float64` today (T-Digest restriction; lifts when +//! the sketch swaps to KLL) +//! +//! # Rewrite +//! +//! ```text +//! RangeFilterExec (narrow, halo_lo=0, halo_hi=0, cuts=pending) +//! PartitionedBoundedWindowAggExec (wraps BWAG; declares UnspecifiedDistribution) +//! RangeFilterExec (wide, halo_lo, halo_hi, cuts=pending) +//! RuntimeStatsExec #2 (per-ORRE-output-partition sketch → scheduler) +//! OrderedRangeRepartitionExec (K outputs, walks child for RSE #1) +//! SortExec (planted here, preserve_partitioning=true) +//! RuntimeStatsExec #1 (local sketch — feeds ORRE's cut walker) +//! +//! ``` +//! +//! RSE#1 sits *below* SortExec so the local sketch ingests the whole +//! partition while Sort buffers, giving the scheduler full-fidelity cuts +//! to hand ORRE before it starts routing. RSE#1 above Sort would force +//! ORRE to route against a still-being-built sketch → skewed shuffle files. +//! +//! The rule runs *after* DF's optimizer chain so `EnforceSorting` / +//! `RepartitionFileScans` have already materialized the SortExec placement +//! we peel here. Running earlier hits two failure modes: (a) sources with +//! `sort_order_for_reorder` set have no SortExec yet at all, and (b) DF's +//! sort-pushdown later moves any Sort we plant down through the +//! passthrough RSE#1, undoing the intended order. +//! +//! Any SPM the DF planner inserted above BWAG for its `SinglePartition` +//! requirement is dropped: the wrapper flips that declaration to +//! `UnspecifiedDistribution`, and `EnforceDistribution` doesn't re-add one. +//! +//! Both `RangeFilterExec` operators are planted with `cuts=None`. The +//! scheduler-side `resolve_range_filter_cuts` walker fills them in once +//! stage-0's `RuntimeStatsExec` reports have been merged into cuts. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use ballista_core::execution_plans::{ + OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec, RangeFilterExec, + RuntimeStatsExec, +}; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::logical_expr::{WindowFrameBound, WindowFrameUnits}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::windows::BoundedWindowAggExec; +use datafusion::scalar::ScalarValue; +use log::debug; + +/// Physical optimizer pass: match the parallel-window shape and rewrite each +/// hit to insert `RSE#1 → ORRE → RSE#2 → RangeFilterExec_wide` below the +/// existing SPM. +#[derive(Default, Debug)] +pub struct ParallelWindowRule; + +impl PhysicalOptimizerRule for ParallelWindowRule { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> datafusion::common::Result> { + let bc = config + .extensions + .get::() + .cloned() + .unwrap_or_default(); + if !bc.parallel_window_enabled() { + return Ok(plan); + } + // K = the number of range-disjoint output partitions the ORRE will + // produce. At rule-fire time DataFusion's initial physical plan is + // still "loose" — the DataSourceExec below BWAG has 1 file_group, + // not the eventual `target_partitions` split (RepartitionFileScans + // and friends run later in the AQE chain). So we can't ask the + // plan tree for the true source width; we use the config knob that + // those later rules also target. + let output_partitions = config.execution.target_partitions.max(2); + plan.transform_up(|node| match maybe_rewrite_bwag(&node, output_partitions)? { + Some(rewritten) => Ok(Transformed::yes(rewritten)), + None => Ok(Transformed::no(node)), + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "ParallelWindow" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// True if any descendant of `nodes` is an `OrderedRangeRepartitionExec` +/// or `RangeFilterExec`. Used as an idempotency guard: those ops are what +/// our own rewrite plants, so seeing them below a BWAG means we've already +/// rewritten this candidate on a previous optimizer pass. +fn subtree_contains_our_rewrite(nodes: &[&Arc]) -> bool { + for node in nodes { + if node.is::() || node.is::() { + return true; + } + if subtree_contains_our_rewrite(node.children().as_slice()) { + return true; + } + } + false +} + +/// True when the bound is `CurrentRow` or a non-null scalar offset. +/// `UNBOUNDED PRECEDING/FOLLOWING` is a typed-null scalar and returns `false`. +fn is_finite(bound: &WindowFrameBound) -> bool { + match bound { + WindowFrameBound::CurrentRow => true, + WindowFrameBound::Preceding(scalar) | WindowFrameBound::Following(scalar) => { + !scalar.is_null() + } + } +} + +fn fmt_bound(bound: &WindowFrameBound) -> String { + match bound { + WindowFrameBound::CurrentRow => "CURRENT ROW".to_string(), + WindowFrameBound::Preceding(scalar) => format!("{scalar} PRECEDING"), + WindowFrameBound::Following(scalar) => format!("{scalar} FOLLOWING"), + } +} + +/// Extract the halo width in `f64` from a bound. `CurrentRow` → `Some(0.0)`. +/// Non-numeric scalars (e.g. Interval bounds) return `None` — a shape gate, +/// widened alongside KLL. +fn halo_from_bound(bound: &WindowFrameBound) -> Option { + let scalar = match bound { + WindowFrameBound::CurrentRow => return Some(0.0), + WindowFrameBound::Preceding(s) | WindowFrameBound::Following(s) => s, + }; + match scalar { + ScalarValue::Int8(Some(v)) => Some(*v as f64), + ScalarValue::Int16(Some(v)) => Some(*v as f64), + ScalarValue::Int32(Some(v)) => Some(*v as f64), + ScalarValue::Int64(Some(v)) => Some(*v as f64), + ScalarValue::UInt8(Some(v)) => Some(*v as f64), + ScalarValue::UInt16(Some(v)) => Some(*v as f64), + ScalarValue::UInt32(Some(v)) => Some(*v as f64), + ScalarValue::UInt64(Some(v)) => Some(*v as f64), + ScalarValue::Float32(Some(v)) => Some(*v as f64), + ScalarValue::Float64(Some(v)) => Some(*v), + _ => None, + } +} + +/// Match the parallel-window shape rooted at `node` and, if it fits, splice +/// `RFE_narrow → PBWAG(BWAG) → RFE_wide → RSE#2 → ORRE → SortExec → RSE#1 → +/// ` in place of the DF-planted `BWAG → SPM → SortExec → ` +/// subtree. +/// +/// - `Ok(None)`: shape gate missed (not a BWAG, PARTITION BY present, ROWS +/// frame, UNBOUNDED bound, non-Float64 ORDER BY, non-numeric bound scalar, +/// or subtree already rewritten). No log noise on the hot path. +/// - `Ok(Some(_))`: rewrite happened. +/// - `Err(_)`: an invariant the shape gates should have upheld didn't — BWAG +/// with ≠1 child, SPM/SortExec with ≠1 child, schema-lookup failure on the +/// ORDER BY expression, or a constructor `try_new` error. +/// +/// The rule runs after DF's optimizer chain, so BWAG's descendants have the +/// fully-materialized `SPM → SortExec → source` shape by the time we peel +/// here (see module doc for why the SPM and SortExec get stripped). +fn maybe_rewrite_bwag( + node: &Arc, + output_partitions: usize, +) -> datafusion::common::Result>> { + let Some(window) = node.downcast_ref::() else { + return Ok(None); + }; + // Shape gates as slice patterns: 0 or 2+ elements simply don't match. + let [expr] = window.window_expr() else { + return Ok(None); + }; + let [] = expr.partition_by() else { + return Ok(None); + }; + let [order] = expr.order_by() else { + return Ok(None); + }; + let Some(column) = order.expr.downcast_ref::() else { + return Ok(None); + }; + // TODO: DESC support. SQL RANGE frame semantics invert with the sort + // direction — `k PRECEDING` refers to *larger* values under a DESC + // ORDER BY — so the halo must widen the upper side of each bucket + // rather than the lower. `RangeFilterExec::sorted_on_key` also refuses + // DESC input, so the fast path would need to grow a mirrored branch. + // Land both together; until then, gate DESC to the serial path. + if order.options.descending { + return Ok(None); + } + let frame = expr.get_window_frame(); + let WindowFrameUnits::Range = frame.units else { + return Ok(None); + }; + if !is_finite(&frame.start_bound) || !is_finite(&frame.end_bound) { + return Ok(None); + } + // Idempotency: re-plans (AQE fires the optimizer chain again for stage + // N+1) would otherwise wrap another ORRE around the previous rewrite's + // RangeFilterExec+ShuffleReader — and that ORRE's child doesn't claim + // ordering, blowing up at execute-time. + if subtree_contains_our_rewrite(window.children().as_slice()) { + return Ok(None); + } + let (Some(halo_lo), Some(halo_hi)) = ( + halo_from_bound(&frame.start_bound), + halo_from_bound(&frame.end_bound), + ) else { + return Ok(None); + }; + + let node_children = node.children(); + let [immediate] = node_children.as_slice() else { + return datafusion::common::internal_err!( + "ParallelWindowRule: BWAG must have exactly 1 child" + ); + }; + // Loop tolerates any order (SPM→Sort or Sort→SPM) or partial shapes + // (source that claims ordering natively via `sort_order_for_reorder` + // skips the Sort entirely). + let mut base_source: Arc = (*immediate).clone(); + while base_source.is::() || base_source.is::() { + let children = base_source.children(); + let [inner] = children.as_slice() else { + return datafusion::common::internal_err!( + "ParallelWindowRule: SPM/SortExec must have exactly 1 child" + ); + }; + base_source = (*inner).clone(); + } + let source_schema = base_source.schema(); + + // Route on the ORDER BY column. ORRE requires Float64 today (T-Digest + // restriction; lifts when the sketch swaps to KLL). + let routing_type = order.expr.data_type(&source_schema)?; + if !matches!(routing_type, DataType::Float64) { + return Ok(None); + } + + let sort_expr = normalize_sort_expr(order); + let rse1: Arc = Arc::new(RuntimeStatsExec::try_new( + base_source, + Some(vec![sort_expr.clone()]), + )?); + // Plant a fresh SortExec above RSE#1 as the pipeline break: Sort + // consumes all input before emitting the first row, so RSE#1's sketch + // fully ingests and reports while Sort buffers — ORRE then routes + // against final cuts instead of approximate ones (which would produce + // skewed shuffle files). + let sort_lex = LexOrdering::new(vec![sort_expr.clone()]).ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "ParallelWindowRule: could not build LexOrdering from ORDER BY".into(), + ) + })?; + let sorted_over_rse1: Arc = + Arc::new(SortExec::new(sort_lex, rse1).with_preserve_partitioning(true)); + let orre: Arc = Arc::new(OrderedRangeRepartitionExec::try_new( + sorted_over_rse1, + vec![sort_expr.clone()], + output_partitions, + )?); + let rse2: Arc = Arc::new(RuntimeStatsExec::try_new( + orre, + Some(vec![sort_expr.clone()]), + )?); + let wide_filter: Arc = Arc::new(RangeFilterExec::try_new_pending( + rse2, + sort_expr.expr.clone(), + ScalarValue::Float64(Some(halo_lo)), + ScalarValue::Float64(Some(halo_hi)), + )?); + + // Wrap BWAG in PartitionedBoundedWindowAggExec instead of collapsing + // K→1 with SPM. The wrapper declares `UnspecifiedDistribution` so + // EnforceDistribution won't reinsert an SPM below, and BWAG's own + // per-partition execute() runs each of the K sub-ranges independently. + // See execution_plans::partitioned_bounded_window_agg for what makes + // this safe (range-repartition upstream + halo). + let partitioned_bwag: Arc = + Arc::new(PartitionedBoundedWindowAggExec::try_new( + window.window_expr().to_vec(), + wide_filter, + )?); + // Narrow filter above BWAG drops the halo rows the wide filter let in + // for BWAG's frame-context. `halo_lo == halo_hi == 0.0` collapses the + // predicate to `cuts[k-1] <= v < cuts[k]` — task k's own range. + let narrow_filter: Arc = + Arc::new(RangeFilterExec::try_new_pending( + partitioned_bwag, + sort_expr.expr.clone(), + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + )?); + + debug!( + "ParallelWindowRule: rewrote BWAG on `{}` (RANGE {} - {})", + column.name(), + fmt_bound(&frame.start_bound), + fmt_bound(&frame.end_bound), + ); + Ok(Some(narrow_filter)) +} + +/// ORRE requires `nulls_first == false` today (T-Digest has no NULL slot). +/// The BWAG's `NULLS LAST` sort expressions arrive with `nulls_first: false` +/// already, but explicit sanitization keeps the invariant obvious to future +/// readers. +fn normalize_sort_expr(expr: &PhysicalSortExpr) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: expr.expr.clone(), + options: SortOptions { + descending: expr.options.descending, + nulls_first: false, + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::config::ExtensionOptions; + use datafusion::datasource::empty::EmptyTable; + use datafusion::physical_plan::displayable; + use datafusion::prelude::SessionContext; + + async fn plan(sql: &str) -> datafusion::common::Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id1", DataType::Int64, false), + Field::new("id2", DataType::Int64, false), + Field::new("id3", DataType::Int64, false), + Field::new("v2", DataType::Float64, false), + ])); + let ctx = SessionContext::new(); + ctx.register_table("large", Arc::new(EmptyTable::new(schema)))?; + ctx.sql(sql).await?.create_physical_plan().await + } + + fn optimize( + plan: Arc, + ) -> datafusion::common::Result> { + let mut config = ConfigOptions::default(); + config.execution.target_partitions = 8; + let mut bc = BallistaConfig::default(); + bc.set("planner.parallel_window.enabled", "true").unwrap(); + config.extensions.insert(bc); + ParallelWindowRule.optimize(plan, &config) + } + + #[tokio::test] + async fn disabled_by_default() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + // No BallistaConfig extension registered → default is `false`. + let mut config = ConfigOptions::default(); + config.execution.target_partitions = 8; + let out = ParallelWindowRule.optimize(plan.clone(), &config)?; + let rendered = format!("{}", displayable(out.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "flag off: rewrite must not fire:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn rewrites_q8_shape() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + // The rewrite must plant each of these ops. Cheap string contains — + // exhaustive plan-shape assertions in follow-up integration tests. + for expected in [ + "PartitionedBoundedWindowAggExec", + "BoundedWindowAggExec", + "RangeFilterExec", + "RuntimeStatsExec", + "OrderedRangeRepartitionExec", + "SortExec", + ] { + assert!( + rendered.contains(expected), + "expected `{expected}` in rewritten plan:\n{rendered}" + ); + } + // BWAG's SinglePartition collapse is what this whole rewrite + // avoids — any SPM in the output would defeat that. + assert!( + !rendered.contains("SortPreservingMergeExec"), + "SortPreservingMergeExec must NOT appear in the rewritten plan:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_rows_frame() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT avg(v2) OVER (ORDER BY id3 \ + ROWS BETWEEN 100 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "ROWS frames should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_partition_by() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (PARTITION BY id1 ORDER BY v2 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "PARTITION BY should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_unbounded_frame() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "UNBOUNDED PRECEDING should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_descending_order_by() -> datafusion::common::Result<()> { + // DESC + RANGE flips the SQL frame semantics: `k PRECEDING` refers to + // larger values, not smaller. The current halo widens the lower side + // of each bucket, so DESC would silently miss frame ancestors that + // land in the next-higher bucket. Gate out until the halo-swap + + // RFE fast-path DESC support land together. + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 DESC \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "DESC ORDER BY should not be rewritten (halo direction bug):\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_non_float64_order_key() -> datafusion::common::Result<()> { + // id3 is Int64; ORRE requires Float64 today (T-Digest restriction). + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY id3 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "non-Float64 order key should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[test] + fn halo_from_bound_reads_all_numeric_variants() { + assert_eq!(halo_from_bound(&WindowFrameBound::CurrentRow), Some(0.0)); + assert_eq!( + halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Int64(Some(3)))), + Some(3.0) + ); + assert_eq!( + halo_from_bound(&WindowFrameBound::Following(ScalarValue::Float64(Some( + 2.5 + )))), + Some(2.5) + ); + assert_eq!( + halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Utf8(Some( + "x".into() + )))), + None + ); + } +} diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 6dc08dab6a..99b44a0057 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -22,7 +22,7 @@ use crate::state::aqe::execution_plan::{ use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule; use crate::state::aqe::optimizer_rule::{ CoalescePartitionsRule, DelayJoinSelectionRule, DistributedExchangeRule, - PropagateEmptyExecRule, SelectJoinRule, + ParallelWindowRule, PropagateEmptyExecRule, SelectJoinRule, }; use crate::state::distributed_explain::handle_explain_plan; use crate::state::execution_stage::StageOutput; @@ -555,6 +555,16 @@ impl AdaptivePlanner { // physical_optimizers.extend(Self::datafusion_optimizers()); + // Rewrite bounded RANGE-frame windows into a range-shuffle so BWAG's + // single-partition constraint is not a serial bottleneck. Runs AFTER + // DataFusion's optimizer chain (EnforceSorting, RepartitionFileScans, + // …) so we see the fully-materialized SortExec placement — placement + // we peel and re-plant so RSE#1 sits *below* the pipeline-break Sort, + // letting the sketch fully report before ORRE routes. Must still run + // before DistributedExchangeRule — the rule emits an ORRE that DE + // picks up as the shuffle-boundary K-space source. + physical_optimizers.push(Arc::new(ParallelWindowRule)); + // `DistributedExchangeRule` should be the last plan mutator rule in the chain physical_optimizers .push(Arc::new(DistributedExchangeRule::new(plan_id_generator))); diff --git a/ballista/scheduler/src/state/aqe/test/range_repartition.rs b/ballista/scheduler/src/state/aqe/test/range_repartition.rs index 8d25a8b83e..129e90d2d3 100644 --- a/ballista/scheduler/src/state/aqe/test/range_repartition.rs +++ b/ballista/scheduler/src/state/aqe/test/range_repartition.rs @@ -149,7 +149,8 @@ async fn routing_parks_when_range_repartition_is_plan_root() }, }]; let cuts = vec![15.0]; - let remapped = cut_partitions(vec![vec![location(0, 7, 3)]], &reports, &cuts)?; + let remapped = + cut_partitions(vec![vec![location(0, 7, 3)]], &reports, &cuts, 0.0, 0.0)?; // `cut_partitions` must duplicate the straddler into both partitions — // the read-side filter is expected to trim on read. diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index 68811762e0..fc5367e579 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -36,7 +36,9 @@ //! flows from parent to descendants via function arguments, so sibling //! subtrees never share state and there's no traversal-order dependency. -use ballista_core::execution_plans::{PerPartitionFilterExec, ShuffleReaderExec}; +use ballista_core::execution_plans::{ + RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, +}; use datafusion::common::internal_err; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{ @@ -84,27 +86,40 @@ fn restrict( return Ok(rewritten); } - // PerPartitionFilterExec: its `predicates` vec is positionally aligned - // with the child's output partitions (predicates[k] filters - // input.execute(k)). When we restrict the child from K partitions to - // `partitions.len()`, the predicate vec must be sliced by the same - // indices in the same order - if !under_collect && let Some(ppf) = plan.downcast_ref::() { + // RangeFilterExec: raw_bounds is indexed by input partition; restriction + // slices bounds parallel to the input's partition subset. Halos + routing + // are carried over verbatim; RFE re-widens on the fresh operator. + if !under_collect && let Some(rf) = plan.downcast_ref::() { let children = plan.children(); let [child] = children.as_slice() else { return internal_err!( - "PerPartitionFilterExec must have exactly 1 child, got {}", + "RangeFilterExec must have exactly 1 child, got {}", children.len() ); }; let new_child = restrict((*child).clone(), partitions, false)?; - let new_predicates: Vec<_> = partitions + let raw_bounds = rf.raw_bounds().ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "RangeFilterExec: task-restriction before resolve_bounds()".into(), + ) + })?; + let sliced_bounds: Vec<_> = partitions .iter() - .map(|&part_idx| ppf.predicates()[part_idx].clone()) - .collect(); - return Ok(Arc::new(PerPartitionFilterExec::try_new( + .map(|&global| { + raw_bounds.get(global).cloned().ok_or_else(|| { + datafusion::common::DataFusionError::Internal(format!( + "RangeFilterExec: partition index {global} out of bounds ({} raw bounds)", + raw_bounds.len() + )) + }) + }) + .collect::>()?; + return Ok(Arc::new(RangeFilterExec::try_new_resolved( new_child, - new_predicates, + rf.routing_expr().clone(), + rf.halo_lo().clone(), + rf.halo_hi().clone(), + sliced_bounds, )?)); } @@ -288,6 +303,24 @@ fn select_output_partitions( return Some(Arc::new(restricted)); } + // RangeShuffleReaderExec: cross-stage inputs, ordering-preserving. Same + // partition-slice restriction as ShuffleReaderExec — carry the merge + // ordering unchanged. + if let Some(reader) = plan.downcast_ref::() { + let kept: Vec> = indices + .iter() + .filter_map(|&p| reader.partition.get(p).cloned()) + .collect(); + let restricted = RangeShuffleReaderExec::try_new( + reader.stage_id, + kept, + reader.schema(), + reader.merge_ordering().clone(), + ) + .ok()?; + return Some(Arc::new(restricted)); + } + // DataSourceExec: file-backed or in-memory scans. if let Some(exec) = plan.downcast_ref::() { let source: &dyn Any = exec.data_source().as_ref(); @@ -650,23 +683,19 @@ mod tests { ); } - /// A `PerPartitionFilterExec` restricted to a subset of partitions must - /// slice its `predicates` vector by the same indices, in the same order, - /// as its child. Otherwise the operator's construction invariant - /// (`predicates.len() == child.partition_count()`) breaks and - /// task-local partition `j` would filter through a global predicate - /// that no longer matches. + /// A `RangeFilterExec` restricted to a subset of partitions must slice + /// its `raw_bounds` by the same indices in the same order as its child. + /// Halos + routing_expr carry over unchanged. #[test] - fn per_partition_filter_predicates_are_sliced_with_partitions() { - use ballista_core::execution_plans::PerPartitionFilterExec; - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::physical_expr::{Partitioning, PhysicalExpr}; - use datafusion::scalar::ScalarValue; - - // 4 upstream partitions, each with its own bespoke predicate so we - // can assert the slice ordering survives. - let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + fn range_filter_bounds_are_sliced_with_partitions() { + use ballista_core::execution_plans::RangeFilterExec; + use datafusion::physical_expr::Partitioning; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::Column; + + // 4 upstream partitions with a global 4-way range partition. + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); let partitions_locs: Vec> = (0..4).map(|i| vec![create_partition(i)]).collect(); let reader = ShuffleReaderExec::try_new( @@ -676,41 +705,38 @@ mod tests { Partitioning::UnknownPartitioning(4), ) .unwrap(); - let make_pred = |lo: i64| -> Arc { - Arc::new(BinaryExpr::new( - Arc::new(Column::new("v", 0)), - Operator::GtEq, - Arc::new(Literal::new(ScalarValue::Int64(Some(lo)))), - )) - }; - let predicates: Vec> = - (0..4).map(|i| make_pred(i as i64 * 100)).collect(); + use datafusion::scalar::ScalarValue; + let routing_expr: Arc = Arc::new(Column::new("v", 0)); + // K=4 raw bounds derived from cuts [100, 200, 300]. + let sv = |v: f64| ScalarValue::Float64(Some(v)); + let raw_bounds: Vec<(Option, Option)> = vec![ + (None, Some(sv(100.0))), + (Some(sv(100.0)), Some(sv(200.0))), + (Some(sv(200.0)), Some(sv(300.0))), + (Some(sv(300.0)), None), + ]; let plan: Arc = Arc::new( - PerPartitionFilterExec::try_new( + RangeFilterExec::try_new_resolved( Arc::new(reader) as Arc, - predicates.clone(), + routing_expr, + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + raw_bounds.clone(), ) .unwrap(), ); let restricted = restrict_plan_to_partitions(plan, &[1, 3]).unwrap(); - let ppf = restricted - .downcast_ref::() - .expect("top must remain PerPartitionFilterExec"); - assert_eq!(ppf.predicates().len(), 2); - assert_eq!( - ppf.predicates()[0].to_string(), - predicates[1].to_string(), - "local partition 0 must carry the global-partition-1 predicate" - ); - assert_eq!( - ppf.predicates()[1].to_string(), - predicates[3].to_string(), - "local partition 1 must carry the global-partition-3 predicate" - ); + let rf = restricted + .downcast_ref::() + .expect("top must remain RangeFilterExec"); + let restricted_bounds = rf.raw_bounds().unwrap(); + assert_eq!(restricted_bounds.len(), 2); + assert_eq!(restricted_bounds[0], raw_bounds[1]); + assert_eq!(restricted_bounds[1], raw_bounds[3]); // Reader below must have been restricted in the same order. - let child = ppf.children()[0].clone(); + let child = rf.children()[0].clone(); let reader = child .downcast_ref::() .expect("child must be a ShuffleReaderExec"); diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c0c721cf09..215fb6fb11 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -129,6 +129,7 @@ standard DataFusion settings. | ballista.planner.coalesce.merged_partition_factor | Float64 | 1.2 | Two adjacent partitions are merged when their combined size is below target_partition_bytes times this factor. Mirrors Spark's legacy coalesce semantics. | | ballista.planner.coalesce.small_partition_factor | Float64 | 0.2 | A coalesced partition smaller than target_partition_bytes times this factor counts as small and is merged into its neighbour. Mirrors Spark's legacy coalesce semantics. | | ballista.planner.coalesce.target_partition_bytes | UInt64 | 67108864 | Target post-coalesce partition size in bytes. Mirrors Spark's advisoryPartitionSizeInBytes. | +| ballista.planner.parallel_window.enabled | Boolean | false | Enables the AQE parallel-window rule (ParallelWindowRule), which rewrites bounded-RANGE-frame windows into a distributed range-shuffle so BoundedWindowAggExec's single-partition constraint is not a serial bottleneck. Disabled by default — opt in when the workload contains matching window shapes. | | ballista.planner.propagate_empty.enabled | Boolean | true | Enables the AQE propagate-empty-relation rule. Injects EmptyExec into the plan where an input is known to be empty, such as one side of a join, allowing downstream work to be skipped. | | ballista.scheduler.max_partitions_per_task | UInt64 | 1 | Upper bound on the number of input partitions packed into a single task's `partition_slice`. `1` (default) means one task per input partition. Raise to enable multi-partition tasks (fewer tasks, parallel-sort / parallel-join wins); `0` means unbounded — the scheduler fills each task up to the executor's free vcore count. Does not apply to collapse stages, which must pack their full pending queue into a single task for correctness. | | ballista.standalone.parallelism | UInt16 | number of available CPU cores | Number of concurrent tasks a standalone in-process executor will run. |