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

ballista/core/proto/ballista.proto

@@ -55,7 +55,9 @@ message BallistaPhysicalPlanNode {
     BufferExecNode buffer = 7;
     UnorderedRangeRepartitionExecNode unordered_range_repartition = 8;
     OrderedRangeRepartitionExecNode ordered_range_repartition = 9;
-    PerPartitionFilterExecNode per_partition_filter = 10;
+    RangeFilterExecNode range_filter = 10;
+    PartitionedBoundedWindowAggExecNode partitioned_bounded_window_agg = 11;
+    RangeShuffleReaderExecNode range_shuffle_reader = 12;
   }
 }
 
@@ -125,12 +127,33 @@ message UnorderedRangeRepartitionExecNode {
   uint32 output_partitions = 2;
 }
 
-// Filter with per-input-partition predicates. `predicates[k]` is the
-// boolean expression applied to input partition `k`. Requires
-// `predicates.len() == input_partition_count`. The child plan is
-// plumbed by the framework as `inputs[0]` during decode.
-message PerPartitionFilterExecNode {
-  repeated datafusion.PhysicalExprNode predicates = 1;
+// Filter over ordered inputs with per-partition half-open range predicates
+// derived from `cuts` + `halo_lo` / `halo_hi`. 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. Serialization requires cuts to be resolved.
+//
+// `partition_indices` maps this operator's local partition index to the
+// global partition index in the original K-shape defined by `cuts`. Under
+// task-level restriction the input is sliced to a subset of the K global
+// partitions; each entry stays < `cuts.len() + 1`.
+message RangeFilterExecNode {
+  datafusion.PhysicalExprNode routing_expr = 1;
+  repeated datafusion_common.ScalarValue cuts = 2;
+  datafusion_common.ScalarValue halo_lo = 3;
+  datafusion_common.ScalarValue halo_hi = 4;
+  repeated uint32 partition_indices = 5;
+}
+
+// Wrapper for `BoundedWindowAggExec` that overrides
+// `required_input_distribution` to `Unspecified` — see the module doc on
+// `execution_plans::partitioned_bounded_window_agg` for what makes that safe.
+// The child plan is plumbed by the framework as `inputs[0]` during decode.
+// `input_order_mode` and `can_repartition` are hardcoded on the decode side
+// per the rule's shape gates; only `window_expr` needs to cross the wire.
+message PartitionedBoundedWindowAggExecNode {
+  repeated datafusion.PhysicalWindowExprNode window_expr = 1;
 }
 
 message ChaosExecNode {
@@ -190,6 +213,18 @@ message ShuffleReaderPartition {
   repeated PartitionLocation location = 1;
 }
 
+// Ordering-preserving shuffle reader. Reuses `ShuffleReaderPartition` for the
+// M-shape source layout. Partitioning is derived from `partition.len()`
+// (always `UnknownPartitioning`, range-partitioned by `merge_ordering`).
+message RangeShuffleReaderExecNode {
+  repeated ShuffleReaderPartition partition = 1;
+  datafusion_common.Schema schema = 2;
+  uint32 stage_id = 3;
+  // Sort key the reader's k-way merge preserves. Advertised on the reader's
+  // `PlanProperties.eq_properties` for downstream consumers.
+  repeated datafusion.PhysicalSortExprNode merge_ordering = 4;
+}
+
 // CoalescePartitionsRule output: groups upstream partitions into coalesced output partitions.
 // Empty when no coalesce is applied (the optional field on the parent message is absent).
 message CoalescePlan {

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/mod.rs

@@ -23,9 +23,11 @@ mod chaos_exec;
 mod distributed_explain_analyze;
 mod distributed_query;
 mod ordered_range_repartition;
-mod per_partition_filter;
+mod partitioned_bounded_window_agg;
 pub mod plan_algebra;
+mod range_filter;
 mod range_repartition_common;
+mod range_shuffle_reader;
 mod runtime_stats;
 mod shuffle_reader;
 mod shuffle_writer;
@@ -42,8 +44,10 @@ use datafusion::common::exec_err;
 pub use distributed_explain_analyze::DistributedExplainAnalyzeExec;
 pub use distributed_query::{DistributedQueryExec, execute_physical_plan};
 pub use ordered_range_repartition::OrderedRangeRepartitionExec;
-pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates};
+pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec;
 pub use plan_algebra::{preserves_distribution, preserves_partitioning};
+pub use range_filter::RangeFilterExec;
+pub use range_shuffle_reader::RangeShuffleReaderExec;
 pub use runtime_stats::{
     MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats,
     collect_reports as collect_runtime_stats_reports, cut_partitions,

ballista/core/src/execution_plans/ordered_range_repartition.rs

@@ -176,25 +176,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.
@@ -356,6 +344,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()
@@ -686,17 +698,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}"
         );
     }
@@ -720,20 +739,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}"