feat(scheduler): range-repartition rule primitives (batteries-included)

ballista/core/src/execution_plans/mod.rs

@@ -24,6 +24,7 @@ mod distributed_explain_analyze;
 mod distributed_query;
 mod ordered_range_repartition;
 mod per_partition_filter;
+pub mod plan_algebra;
 mod range_repartition_common;
 mod runtime_stats;
 mod shuffle_reader;
@@ -41,11 +42,13 @@ 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;
+pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates};
+pub use plan_algebra::{preserves_distribution, preserves_partitioning};
 pub use runtime_stats::{
     MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats,
-    collect_reports as collect_runtime_stats_reports, log_merged_runtime_stats,
-    merge_reports as merge_runtime_stats_reports, sketch_from_proto, sketch_to_proto,
+    collect_reports as collect_runtime_stats_reports, cut_partitions,
+    log_merged_runtime_stats, merge_reports as merge_runtime_stats_reports,
+    repartition_routing_expr, sketch_from_proto, sketch_to_proto,
 };
 pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec};
 pub use shuffle_reader::{stats_for_partition, stats_for_partitions};

ballista/core/src/execution_plans/per_partition_filter.rs

@@ -18,10 +18,11 @@
 //! Filter with a distinct predicate per input partition.
 //!
 //! `FilterExec` in DataFusion carries a single predicate applied to every
-//! partition. That's wrong for the URRE-consuming reader in the adaptive
-//! range shuffle: each downstream partition `k` needs a range predicate
-//! `cuts[k-1] <= key < cuts[k]` unique to that partition, so straddling
-//! sub-parts from the producer are trimmed to just partition `k`'s slice.
+//! partition. That's wrong for the range-repartition-consuming reader in
+//! the adaptive range shuffle: each downstream partition `k` needs a range
+//! predicate `cuts[k-1] <= key < cuts[k]` unique to that partition, so
+//! straddling sub-parts from the producer are trimmed to just partition
+//! `k`'s slice.
 //!
 //! One-task-per-downstream-partition + plain `FilterExec` would work but
 //! defeats vcore packing (`K` tasks instead of `K / vcores`). This operator
@@ -270,6 +271,73 @@ impl RecordBatchStream for PerPartitionFilterStream {
     }
 }
 
+/// Build the `K = cuts.len() + 1` half-open range predicates a
+/// `PerPartitionFilterExec` needs to reproduce the range repartition's
+/// write-side routing on the read side.
+///
+/// Partition `i` receives the predicate
+///
+/// ```text
+///   i = 0        →  routing_expr < cuts[0]
+///   0 < i < K-1  →  cuts[i-1] <= routing_expr AND routing_expr < cuts[i]
+///   i = K-1      →  routing_expr >= cuts[K-2]
+///   K = 1        →  lit(true)   // empty cuts, single-bucket range repartition
+/// ```
+///
+/// Consistent with the private `range_repartition_common::split_batch_by_range`
+/// helper, which uses the same half-open convention on the write side.
+/// Callers pass the range repartition's routing expression verbatim
+/// (`CAST(order_by[0] AS Float64)` today).
+///
+/// Non-null routing expressions only. Both `UnorderedRangeRepartitionExec`
+/// and `OrderedRangeRepartitionExec` refuse nullable routing exprs at
+/// `try_new`, so any expression that reaches this helper via
+/// `RangeRepartitionRouting` is guaranteed non-null — no `IS NULL` branch
+/// needed.
+pub fn range_partition_predicates(
+    routing_expr: Arc<dyn PhysicalExpr>,
+    cuts: &[f64],
+) -> Vec<Arc<dyn PhysicalExpr>> {
+    use datafusion::logical_expr::Operator;
+    use datafusion::physical_expr::expressions::{BinaryExpr, Literal};
+    use datafusion::scalar::ScalarValue;
+
+    let partition_count = cuts.len() + 1;
+    let lit = |v: f64| -> Arc<dyn PhysicalExpr> {
+        Arc::new(Literal::new(ScalarValue::Float64(Some(v))))
+    };
+    let ge = |lo: f64| -> Arc<dyn PhysicalExpr> {
+        Arc::new(BinaryExpr::new(
+            routing_expr.clone(),
+            Operator::GtEq,
+            lit(lo),
+        ))
+    };
+    let lt = |hi: f64| -> Arc<dyn PhysicalExpr> {
+        Arc::new(BinaryExpr::new(routing_expr.clone(), Operator::Lt, lit(hi)))
+    };
+    (0..partition_count)
+        .map(|partition_idx| {
+            let lo = partition_idx
+                .checked_sub(1)
+                .and_then(|cut_idx| cuts.get(cut_idx).copied());
+            let hi = cuts.get(partition_idx).copied();
+            match (lo, hi) {
+                (None, None) => {
+                    // K == 1: single bucket covers everything.
+                    Arc::new(Literal::new(ScalarValue::Boolean(Some(true))))
+                        as Arc<dyn PhysicalExpr>
+                }
+                (None, Some(hi)) => lt(hi),
+                (Some(lo), None) => ge(lo),
+                (Some(lo), Some(hi)) => {
+                    Arc::new(BinaryExpr::new(ge(lo), Operator::And, lt(hi)))
+                }
+            }
+        })
+        .collect()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -400,6 +468,83 @@ mod tests {
         );
     }
 
+    /// The K=4 range predicates cover every value under the half-open
+    /// convention, and each row lands in exactly one predicate. Random
+    /// probe values are routed through the predicates and expected to
+    /// match the same partition assignment as the range repartition's
+    /// write-side `split_batch_by_range` would produce.
+    #[test]
+    fn range_partition_predicates_partition_every_value_exactly_once() {
+        use datafusion::arrow::array::Float64Array;
+        use datafusion::arrow::datatypes::Field;
+        use datafusion::physical_expr::expressions::Column;
+
+        let cuts = vec![10.0, 20.0, 30.0];
+        let k = cuts.len() + 1;
+        let routing: Arc<dyn PhysicalExpr> = Arc::new(Column::new("v", 0));
+        let preds = range_partition_predicates(routing, &cuts);
+        assert_eq!(preds.len(), k);
+
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)]));
+        let values: Vec<f64> = vec![
+            -5.0, 0.0, 9.999, 10.0, 15.0, 19.999, 20.0, 25.0, 30.0, 100.0,
+        ];
+        let arr = Float64Array::from_iter_values(values.iter().copied());
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap();
+
+        // For each row, find the unique partition whose predicate accepts it.
+        for (row, &v) in values.iter().enumerate() {
+            let mut hits = 0;
+            for pred in &preds {
+                let mask = pred
+                    .evaluate(&batch)
+                    .and_then(|v| v.into_array(batch.num_rows()))
+                    .unwrap();
+                let mask = as_boolean_array(&mask).unwrap();
+                if mask.value(row) {
+                    hits += 1;
+                }
+            }
+            assert_eq!(
+                hits, 1,
+                "value {v} matched {hits} predicates, expected exactly 1"
+            );
+        }
+
+        // Expected assignment mirrors split_batch_by_range: `partition_point`
+        // returns the count of cuts `<= key`, which is the partition index
+        // under the half-open convention.
+        let expected: Vec<usize> = values
+            .iter()
+            .map(|v| cuts.partition_point(|&c| c <= *v))
+            .collect();
+        for (row, want) in expected.iter().enumerate() {
+            let mask = preds[*want]
+                .evaluate(&batch)
+                .and_then(|v| v.into_array(batch.num_rows()))
+                .unwrap();
+            let mask = as_boolean_array(&mask).unwrap();
+            assert!(
+                mask.value(row),
+                "value {} should have landed in partition {}",
+                values[row],
+                want
+            );
+        }
+    }
+
+    /// Degenerate K=1 (empty cuts) yields a single lit(true) predicate.
+    #[test]
+    fn range_partition_predicates_single_bucket_when_cuts_empty() {
+        use datafusion::physical_expr::expressions::Column;
+
+        let routing: Arc<dyn PhysicalExpr> = Arc::new(Column::new("v", 0));
+        let preds = range_partition_predicates(routing, &[]);
+        assert_eq!(preds.len(), 1);
+        assert_eq!(preds[0].to_string(), "true");
+    }
+
     /// `with_new_children` swaps the input while preserving the predicate
     /// vector. Wrapping the original source in a `RepartitionExec` that
     /// keeps the partition count (RoundRobin(3)) gives a valid child; the

ballista/core/src/execution_plans/plan_algebra.rs

@@ -0,0 +1,70 @@
+// 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.
+
+//! Algebraic properties of physical plan nodes — do they preserve
+//! partitioning, do they preserve the distribution of row values, etc.
+//!
+//! DataFusion doesn't expose these as trait methods on `ExecutionPlan`
+//! (nice-to-haves like `ExecutionPlan::affects_partitioning()` /
+//! `ExecutionPlan::affects_distribution()` that would hopefully land one
+//! day), so we downcast against a hand-maintained whitelist. Being
+//! conservative is the safety net: unrecognized node → property assumed
+//! false → caller falls back to the safer path.
+
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::filter::FilterExec;
+use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion::physical_plan::sorts::sort::SortExec;
+use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
+
+use crate::execution_plans::{
+    BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec,
+};
+
+/// Whitelisted ops preserve the routing key's row set, values, and
+/// partitioning — an upstream sketch remains valid after the operator.
+pub fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool {
+    // Buffered batches replayed verbatim.
+    plan.downcast_ref::<BufferExec>().is_some()
+        // Per-partition sort: rows reorder within a partition, row set
+        // and counts unchanged. `preserve_partitioning=false` collapses
+        // N→1 (like SortPreservingMergeExec), so gate on the flag.
+        || plan
+            .downcast_ref::<SortExec>()
+            .is_some_and(|sort| sort.preserve_partitioning())
+        // Stage-boundary writers: batches to disk unchanged.
+        || plan.downcast_ref::<ShuffleWriterExec>().is_some()
+        || plan.downcast_ref::<SortShuffleWriterExec>().is_some()
+        // Pure row-annotation: one input row → one output row with an
+        // added column (window fn result); values, partitioning, count preserved.
+        || plan.downcast_ref::<BoundedWindowAggExec>().is_some()
+        || plan.downcast_ref::<WindowAggExec>().is_some()
+}
+
+/// Looser sibling of [`preserves_distribution`]: partitioning survives,
+/// but rows and values within a partition are fair game.
+pub fn preserves_partitioning(plan: &dyn ExecutionPlan) -> bool {
+    // Distribution-preserving is strictly stronger; compose to keep the
+    // whitelist deduplicated.
+    preserves_distribution(plan)
+        // Drops rows, but per-partition — no rows migrate.
+        || plan.downcast_ref::<FilterExec>().is_some()
+        // Rewrites columns; partition boundaries untouched.
+        || plan.downcast_ref::<ProjectionExec>().is_some()
+        // Stats tap; no data mutation.
+        || plan.downcast_ref::<RuntimeStatsExec>().is_some()
+}

ballista/core/src/execution_plans/range_repartition_common.rs

@@ -46,15 +46,12 @@ use datafusion::arrow::compute::take_arrays;
 use datafusion::common::{Result, internal_datafusion_err};
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_plan::ExecutionPlan;
-use datafusion::physical_plan::sorts::sort::SortExec;
-use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
 use futures::FutureExt;
 use log::warn;
 use tokio::sync::mpsc;
 
-use crate::execution_plans::{
-    BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec,
-};
+use crate::execution_plans::RuntimeStatsExec;
+use crate::execution_plans::plan_algebra::preserves_distribution;
 
 /// Walk `child`'s subtree for a [`RuntimeStatsExec`] that sketches on our
 /// routing expression, snapshot its merged T-Digest, and compute `K - 1`
@@ -155,54 +152,6 @@ pub(super) fn find_runtime_stats<'a>(
     find_runtime_stats(only_child, routing_expr)
 }
 
-/// Whitelist of pass-through operator types the walker will descend through
-/// on its way to a matching [`RuntimeStatsExec`]. Unlisted operators might
-/// drop rows, duplicate rows, or transform the routing key's value — any of
-/// which would make an upstream sketch stale by the time data reaches us.
-///
-/// Being conservative is the safety net: unrecognized node → walker gives
-/// up → single-bucket fallback. Extending this list requires positive
-/// verification that the operator is a distribution-preserving passthrough
-/// for the routing key. Absent an upstream `ExecutionPlan::affects_distribution()`
-/// method (nice-to-have that hopefully lands one day), we maintain this by hand.
-pub(super) fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool {
-    // Every entry here is a *claim* that the operator (1) doesn't drop
-    // rows, (2) doesn't duplicate rows, (3) doesn't transform the routing
-    // key's value, AND (4) doesn't change partitioning (per-partition
-    // slots downstream still map to the same partitions upstream). Losing
-    // any of those invalidates the sketch/count on the other side.
-    //
-    // Notable *exclusions*:
-    //   • `SortPreservingMergeExec` — collapses N partitions to 1, so
-    //     per-partition slots below it don't align with the single
-    //     partition above. Values are preserved, but the partitioning
-    //     invariant fails.
-    //   • `ProjectionExec` — might compute a new column that shadows or
-    //     replaces the routing key, transforming values invisibly.
-    //   • `FilterExec`, `LimitExec`, joins — drop or duplicate rows.
-    //   • Our own DRRs — repartition by value; that's the whole point.
-    plan.downcast_ref::<BufferExec>().is_some()
-        // `SortExec` reorders rows within each partition; row set and per-
-        // partition counts unchanged — but ONLY when `preserve_partitioning`
-        // is true. The `preserve_partitioning=false` variant collapses N
-        // partitions to 1 (like `SortPreservingMergeExec`), which would
-        // invalidate per-partition slot alignment.
-        || plan
-            .downcast_ref::<SortExec>()
-            .is_some_and(|sort| sort.preserve_partitioning())
-        // `ShuffleWriterExec` / `SortShuffleWriterExec` sit at the top of
-        // every stage's plan on the executor side. They write batches to
-        // disk unchanged — no transformation, no filtering.
-        || plan.downcast_ref::<ShuffleWriterExec>().is_some()
-        || plan.downcast_ref::<SortShuffleWriterExec>().is_some()
-        // `BoundedWindowAggExec` / `WindowAggExec` are pure row-annotation:
-        // each input row emits exactly one output row with a new column
-        // (the window function's result). The routing key's values,
-        // partitioning, and row count are all preserved verbatim.
-        || plan.downcast_ref::<BoundedWindowAggExec>().is_some()
-        || plan.downcast_ref::<WindowAggExec>().is_some()
-}
-
 /// Split `batch` into `K = boundaries.len() + 1` sub-batches under the
 /// half-open convention: partition `p` receives rows where
 /// `boundaries[p-1] <= key < boundaries[p]` (open at `-∞` on partition 0