feat: parallel BWAG for the range-window shape (h2o Q8)

.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

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<HashMap<String, ConfigEntry>> = 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,

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;
 

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<Option<Arc<dyn PhysicalExpr>>> {
@@ -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<Vec<PartitionLocation>>` 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<Vec<PartitionLocation>>,
     reports: &[TaskRuntimeStats],
     global_cuts: &[f64],
+    halo_lo: f64,
+    halo_hi: f64,
 ) -> Result<Vec<Vec<PartitionLocation>>> {
     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,15 +1798,16 @@ 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.
         let reports = vec![sketch_report(300, vec![vec![5.0, 15.0, 25.0]])];
         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<u64> = b.iter().map(|l| l.file_id.unwrap()).collect();
@@ -1900,12 +1947,90 @@ 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!(
             msg.contains("num_rows=None") && msg.contains("no usable sketch"),
             "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<u64> = 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",
+        );
+    }
 }

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<Arc<dyn QueryStageExecutor>> {
         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::<ShuffleReaderExec>() {
                     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::<RangeShuffleReaderExec>() {
+                    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