apache
diff --git a/‎ballista/core/proto/ballista.proto‎
Lines changed: 4 additions & 4 deletions b/‎ballista/core/proto/ballista.proto‎
Lines changed: 4 additions & 4 deletions
diff --git a/‎ballista/core/src/config.rs‎
Lines changed: 18 additions & 0 deletions b/‎ballista/core/src/config.rs‎
Lines changed: 18 additions & 0 deletions
diff --git a/‎ballista/core/src/execution_plans/mod.rs‎
Lines changed: 1 addition & 1 deletion b/‎ballista/core/src/execution_plans/mod.rs‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎ballista/core/src/execution_plans/ordered_range_repartition.rs‎
Lines changed: 51 additions & 28 deletions b/‎ballista/core/src/execution_plans/ordered_range_repartition.rs‎
Lines changed: 51 additions & 28 deletions
diff --git a/‎ballista/core/src/execution_plans/range_filter.rs‎
Lines changed: 8 additions & 3 deletions b/‎ballista/core/src/execution_plans/range_filter.rs‎
Lines changed: 8 additions & 3 deletions
diff --git a/‎ballista/core/src/execution_plans/runtime_stats.rs‎
Lines changed: 11 additions & 1 deletion b/‎ballista/core/src/execution_plans/runtime_stats.rs‎
Lines changed: 11 additions & 1 deletion
diff --git a/‎ballista/core/src/serde/generated/ballista.rs‎
Lines changed: 4 additions & 4 deletions b/‎ballista/core/src/serde/generated/ballista.rs‎
Lines changed: 4 additions & 4 deletions
diff --git a/‎ballista/executor/src/execution_engine.rs‎
Lines changed: 14 additions & 1 deletion b/‎ballista/executor/src/execution_engine.rs‎
Lines changed: 14 additions & 1 deletion
diff --git a/‎ballista/scheduler/src/cluster/mod.rs‎
Lines changed: 44 additions & 8 deletions b/‎ballista/scheduler/src/cluster/mod.rs‎
Lines changed: 44 additions & 8 deletions
@@ -136,10 +136,10 @@ message PerPartitionFilterExecNode {
   repeated datafusion.PhysicalExprNode predicates = 1;
 }
 
-// Filter inputs with a per-input-partition half-open range predicate
-// widened by `halo_lo` / `halo_hi`. `raw_bounds[k]` is the cut range for
-// input partition `k` before halo widening; RFE widens internally at
-// resolve time. Zero halo recovers the exact range-repartition trim used
+// Filter inputs with a per-input-partition half-open range
+// predicate widened by `halo_lo` / `halo_hi`. `raw_bounds[k]` is the cut
+// range for input partition `k` before halo widening; RFE widens internally
+// at resolve time. Zero halo recovers the exact range-repartition trim used
 // above `ShuffleReaderExec`; non-zero halo widens each partition's read
 // range to include a boundary "context" band (bounded RANGE-frame windows).
 // The child plan is plumbed by the framework as `inputs[0]` during decode.
 
@@ -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,
 
@@ -48,7 +48,7 @@ pub use ordered_range_repartition::OrderedRangeRepartitionExec;
 pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec;
 pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates};
 pub use plan_algebra::{preserves_distribution, preserves_partitioning};
-pub use range_filter::{RangeBound, RangeFilterExec, WidenedBound};
+pub use range_filter::RangeFilterExec;
 pub use range_shuffle_reader::RangeShuffleReaderExec;
 pub use runtime_stats::{
     MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats,
 
@@ -178,25 +178,13 @@ impl OrderedRangeRepartitionExec {
                 routing.expr
             );
         }
-        // Input MUST claim to be sorted on our routing expression — otherwise
-        // the k-way merge produces garbled output. Sortedness of individual
-        // input partitions is enforced by the operator upstream (`SortExec`
-        // with `preserve_partitioning=true`); this check verifies the plan
-        // node declares that property.
-        let input_first_sort = input.output_ordering().map(|ordering| ordering.first());
-        let Some(input_first) = input_first_sort else {
-            return internal_err!(
-                "OrderedRangeRepartitionExec requires sorted input — child plan claims no ordering"
-            );
-        };
-        if input_first.expr.as_ref() != routing.expr.as_ref() {
-            return internal_err!(
-                "OrderedRangeRepartitionExec: input's first sort key `{}` does not match \
-                 routing expression `{}`",
-                input_first.expr,
-                routing.expr
-            );
-        }
+        // NB: input sortedness is NOT checked at try_new. The k-way merge in
+        // execute() needs each input partition sorted on the routing key, but
+        // that guarantee comes from `required_input_ordering()` below +
+        // `EnforceSorting`: DataFusion inserts a `SortExec` above the source
+        // when the declared requirement isn't satisfied. Checking at try_new
+        // races with rule-time construction (rules run before EnforceSorting),
+        // so the runtime check has moved to `execute()`.
         // Advertise each output partition as sorted on `order_by`. Downstream
         // operators (BWAG, HaloDrop) rely on this claim to skip redundant
         // Sort insertions.
@@ -362,6 +350,30 @@ impl ExecutionPlan for OrderedRangeRepartitionExec {
         partition: usize,
         ctx: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
+        // Invariant: EnforceSorting must have satisfied our
+        // `required_input_ordering` — the k-way merge assumes each input
+        // partition is sorted on the routing key. Rules that emit ORRE run
+        // before EnforceSorting, so this check lives at execute-time rather
+        // than construction-time.
+        let input_first = self
+            .input
+            .output_ordering()
+            .map(|ordering| ordering.first())
+            .ok_or_else(|| {
+                internal_datafusion_err!(
+                    "OrderedRangeRepartitionExec: input claims no ordering at execute — \
+                     EnforceSorting should have planted a SortExec"
+                )
+            })?;
+        let routing = &self.order_by[0];
+        if input_first.expr.as_ref() != routing.expr.as_ref() {
+            return internal_err!(
+                "OrderedRangeRepartitionExec: input's first sort key `{}` does not \
+                 match routing expression `{}`",
+                input_first.expr,
+                routing.expr
+            );
+        }
         let mut state = self
             .state
             .lock()
@@ -756,17 +768,24 @@ mod tests {
     }
 
     #[test]
-    fn try_new_rejects_unsorted_input() {
+    fn execute_rejects_unsorted_input() {
+        // try_new no longer checks input ordering — EnforceSorting is
+        // trusted to plant a SortExec after rule-time construction. If
+        // the invariant is broken by the time we get to execute(), the
+        // runtime check fires.
         let schema = schema_v2_id();
-        // MemorySourceConfig with no declared ordering — output_ordering() is None.
-        let err = OrderedRangeRepartitionExec::try_new(
+        let orre = OrderedRangeRepartitionExec::try_new(
             empty_input(&schema),
             vec![asc(&schema, "v2")],
             4,
         )
-        .expect_err("input without ordering claim must be rejected");
+        .expect("construction succeeds; check moved to execute()");
+        let ctx = datafusion::prelude::SessionContext::new().task_ctx();
+        let Err(err) = orre.execute(0, ctx) else {
+            panic!("execute() should reject input without ordering claim");
+        };
         assert!(
-            err.to_string().contains("child plan claims no ordering"),
+            err.to_string().contains("input claims no ordering"),
             "got: {err}"
         );
     }
@@ -790,20 +809,24 @@ mod tests {
     }
 
     #[test]
-    fn try_new_rejects_mismatched_sort_key() {
+    fn execute_rejects_mismatched_sort_key() {
         let schema = schema_v2_id();
-        // Sort input on `id` (Int64); DRR tries to route on `v2`.
+        // Sort input on `id` (Int64); ORRE tries to route on `v2`.
         let source = empty_input(&schema);
         let id_sort = LexOrdering::new(vec![asc(&schema, "id")]).unwrap();
         let sorted_on_id =
             Arc::new(SortExec::new(id_sort, source).with_preserve_partitioning(true))
                 as Arc<dyn ExecutionPlan>;
-        let err = OrderedRangeRepartitionExec::try_new(
+        let orre = OrderedRangeRepartitionExec::try_new(
             sorted_on_id,
             vec![asc(&schema, "v2")],
             4,
         )
-        .expect_err("mismatched sort key must be rejected");
+        .expect("construction succeeds; check moved to execute()");
+        let ctx = datafusion::prelude::SessionContext::new().task_ctx();
+        let Err(err) = orre.execute(0, ctx) else {
+            panic!("execute() should reject mismatched sort key");
+        };
         assert!(
             err.to_string().contains("does not match routing"),
             "got: {err}"
 
@@ -62,7 +62,7 @@ use std::task::{Context, Poll};
 
 use datafusion::arrow::array::{Array, RecordBatch};
 use datafusion::arrow::compute::filter_record_batch;
-use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::datatypes::{DataType, SchemaRef};
 use datafusion::common::cast::{as_boolean_array, as_float64_array};
 use datafusion::common::{Result, Statistics, internal_err};
 use datafusion::execution::TaskContext;
@@ -90,7 +90,7 @@ pub type RangeBound = (Option<ScalarValue>, Option<ScalarValue>);
 
 /// Bounds after halo widening. Float64-only internally today — see the
 /// "Type generality" section in the module doc.
-pub type WidenedBound = (Option<f64>, Option<f64>);
+type WidenedBound = (Option<f64>, Option<f64>);
 
 /// Both raw and widened bounds. `raw` is preserved for serialization; the
 /// executor consumes `widened`.
@@ -654,12 +654,17 @@ impl RecordBatchStream for RangeFilterStream {
     }
 }
 
+// Silence the unused import warning when the file is compiled without
+// arrow — DataType is only exercised via `data_type(...)` return checks.
+#[allow(dead_code)]
+fn _touch_datatype(_: DataType) {}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use datafusion::arrow::array::Float64Array;
     use datafusion::arrow::compute::SortOptions;
-    use datafusion::arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::arrow::datatypes::{Field, Schema};
     use datafusion::datasource::memory::MemorySourceConfig;
     use datafusion::datasource::source::DataSourceExec;
     use datafusion::physical_expr::LexOrdering;
 
@@ -874,6 +874,16 @@ pub fn repartition_routing_expr(
 /// 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.
+///
+/// TODO(halo-aware routing): when the downstream stage has a
+/// `RangeFilterExec` with non-zero halo (bounded RANGE-frame windows), each
+/// partition's *effective* read range is `[cuts[k-1] - halo_lo, cuts[k] +
+/// halo_hi)`. This function currently uses the raw cut range, so files
+/// straddling the halo boundary aren't routed to their halo-widened
+/// consumer. Boundary rows near cuts can be missing from downstream
+/// window sums — a correctness gap for RANGE frames that this refactor
+/// does not resolve. Fix requires reaching across stages to read the
+/// consumer RFE's halos and widening the overlap check here.
 pub fn cut_partitions(
     original_partitions: Vec<Vec<PartitionLocation>>,
     reports: &[TaskRuntimeStats],
@@ -1755,7 +1765,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.
 
@@ -152,10 +152,10 @@ pub struct PerPartitionFilterExecNode {
         ::datafusion_proto::protobuf::PhysicalExprNode,
     >,
 }
-/// Filter inputs with a per-input-partition half-open range predicate
-/// widened by `halo_lo` / `halo_hi`. `raw_bounds\[k\]` is the cut range for
-/// input partition `k` before halo widening; RFE widens internally at
-/// resolve time. Zero halo recovers the exact range-repartition trim used
+/// Filter inputs with a per-input-partition half-open range
+/// predicate widened by `halo_lo` / `halo_hi`. `raw_bounds\[k\]` is the cut
+/// range for input partition `k` before halo widening; RFE widens internally
+/// at resolve time. Zero halo recovers the exact range-repartition trim used
 /// above `ShuffleReaderExec`; non-zero halo widens each partition's read
 /// range to include a boundary "context" band (bounded RANGE-frame windows).
 /// The child plan is plumbed by the framework as `inputs\[0\]` during decode.
 
@@ -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};
@@ -150,6 +152,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
 
@@ -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<dyn ExecutionPlan>) -> bool {
     fn walk(node: &Arc<dyn ExecutionPlan>) -> bool {
-        if node.downcast_ref::<ShuffleReaderExec>().is_some() {
+        if node.downcast_ref::<ShuffleReaderExec>().is_some()
+            || node.downcast_ref::<RangeShuffleReaderExec>().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<dyn ExecutionPlan>;
+        let root: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(reader));
+
+        assert!(
+            !stage_has_input_collapse(&root),
+            "a range-shuffle reader is a stage boundary, not an input collapse",
+        );
+    }
 }