From da2c6b524147272649cda3c72bafaa7414c9ce22 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 21 Jul 2026 08:47:29 -0600 Subject: [PATCH 1/5] =?UTF-8?q?feat(core):=20add=20UnorderedRangeRepartiti?= =?UTF-8?q?onExec=20=E2=80=94=20value-range=20router=20over=20unordered=20?= =?UTF-8?q?inputs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reads P input partitions with no sort assumption, evaluates the first ORDER BY expression per row, and routes each row to one of K output partitions under the half-open convention: partition p owns [cut[p-1], cut[p]) with virtual -∞/+∞ sentinels on the ends. Boundaries are discovered at runtime, not baked in at plan time. On the first batch to arrive from any input partition, the operator walks its own child subtree for a sibling RuntimeStatsExec (#2094), snapshots its merged T-Digest, and computes K-1 quantile cuts. Every failure path (no matching stats, no sketch, empty sketch, mutex poisoning) falls back to a single-bucket routing that never crashes. Includes: - New `range_repartition_common.rs` with the shared building blocks (`discover_cuts`, `find_runtime_stats`, `preserves_distribution`, `split_batch_by_range`, `broadcast_error`). Reused by the pending OrderedRangeRepartitionExec follow-up. - `preserves_distribution` whitelists BufferExec (#2095) as a descend-through node so the discovery walker sees past a Dam. - Proto: `UnorderedRangeRepartitionExecNode` at BallistaPhysicalPlanNode tag 8. - Codec encode/decode + two serde roundtrip tests (single-key + multi-key). Co-Authored-By: Claude Opus 4.7 (1M context) --- ballista/core/proto/ballista.proto | 16 + ballista/core/src/execution_plans/mod.rs | 3 + .../range_repartition_common.rs | 374 +++++++++ .../unordered_range_repartition.rs | 785 ++++++++++++++++++ ballista/core/src/serde/generated/ballista.rs | 23 +- ballista/core/src/serde/mod.rs | 144 +++- 6 files changed, 1343 insertions(+), 2 deletions(-) create mode 100644 ballista/core/src/execution_plans/range_repartition_common.rs create mode 100644 ballista/core/src/execution_plans/unordered_range_repartition.rs diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 846fd708d..76c47fc07 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -53,6 +53,7 @@ message BallistaPhysicalPlanNode { ChaosExecNode chaos_exec = 5; RuntimeStatsExecNode runtime_stats = 6; BufferExecNode buffer = 7; + UnorderedRangeRepartitionExecNode unordered_range_repartition = 8; } } @@ -91,6 +92,21 @@ enum BufferMode { // Add here as they land in Rust. } +// Value-range router over unordered inputs. Reads P input partitions with +// no sort assumption, evaluates the first ORDER BY expression per row, and +// routes each row to one of K output partitions using cuts discovered at +// runtime from a sibling RuntimeStatsExec's quantile sketch. The child +// plan is plumbed by the framework as `inputs[0]`. +message UnorderedRangeRepartitionExecNode { + // Lexicographic ORDER BY carried through from the wrapping window + // operator. Routing keys off the first entry today; the full ordering is + // preserved so downstream operators (SortExec, BWAG) that need it keep + // working. + repeated datafusion.PhysicalSortExprNode order_by = 1; + // K — number of output partitions. Must be ≥ 2. + uint32 output_partitions = 2; +} + message ChaosExecNode { double failure_probability = 1; string fault_type = 2; diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index 0062f9709..ce41627a7 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -22,11 +22,13 @@ mod buffer; mod chaos_exec; mod distributed_explain_analyze; mod distributed_query; +mod range_repartition_common; mod runtime_stats; mod shuffle_reader; mod shuffle_writer; mod shuffle_writer_trait; pub mod sort_shuffle; +mod unordered_range_repartition; mod unresolved_shuffle; use std::path::{Path, PathBuf}; @@ -44,6 +46,7 @@ pub use shuffle_writer::ShuffleWriterExec; pub use shuffle_writer::compute_global_output_partition_ids; pub use shuffle_writer_trait::ShuffleWriter; pub use sort_shuffle::SortShuffleWriterExec; +pub use unordered_range_repartition::UnorderedRangeRepartitionExec; pub use unresolved_shuffle::UnresolvedShuffleExec; use crate::JobId; diff --git a/ballista/core/src/execution_plans/range_repartition_common.rs b/ballista/core/src/execution_plans/range_repartition_common.rs new file mode 100644 index 000000000..8a196b763 --- /dev/null +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -0,0 +1,374 @@ +// 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. + +//! Shared building blocks for the two range-repartition operators — +//! [`UnorderedRangeRepartitionExec`] and (soon) `OrderedRangeRepartitionExec`. +//! +//! The two operators disagree substantially on execution model — the +//! unordered variant is pure scatter, the ordered one is scatter + per-output +//! k-way merge — but they agree on: +//! +//! 1. **How to find the cut boundaries at runtime** — walk the child subtree +//! for a matching sibling [`RuntimeStatsExec`], snapshot its T-Digest, +//! compute quantile cuts. Only descend through whitelisted +//! distribution-preserving operators; refuse otherwise. +//! 2. **How to split one batch across K value ranges** — [`split_batch_by_range`]. +//! 3. **How to broadcast a terminal error to every output channel** — +//! [`broadcast_error`]. +//! +//! Everything in this module is `pub(super)` — visible to sibling +//! `execution_plans::*` modules that own the operators, invisible outside. +//! +//! [`UnorderedRangeRepartitionExec`]: super::UnorderedRangeRepartitionExec +//! [`RuntimeStatsExec`]: super::RuntimeStatsExec + +use std::sync::Arc; + +use datafusion::arrow::array::{Array, Float64Array, RecordBatch, UInt32Array}; +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 log::warn; +use tokio::sync::mpsc; + +use crate::execution_plans::{ + BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec, +}; + +/// Walk `child`'s subtree for a [`RuntimeStatsExec`] that sketches on our +/// routing expression, snapshot its merged T-Digest, and compute `K - 1` +/// quantile cuts. Any failure to find a matching sketch returns an empty +/// `Vec` — the caller's `split_batch_by_range(&[])` produces a single +/// bucket and every row lands in output partition 0. Never crashes. +pub(super) fn discover_cuts( + child: &Arc, + routing_expr: &dyn PhysicalExpr, + output_partitions: usize, +) -> Vec { + let Some(stats) = find_runtime_stats(child, routing_expr) else { + warn!( + "range-repartition: no matching RuntimeStatsExec found in child subtree — \ + single-bucket fallback" + ); + return Vec::new(); + }; + // Walker returned Some → stats.order_by()'s first entry matches our + // routing expression → RuntimeStatsExec's construction contract + // guarantees sketch is present. Belt-and-braces arms in case that + // invariant ever drifts, plus mutex-poisoning is theoretically possible. + let sketch = match stats.merged_quantile_sketch() { + Ok(Some(sketch)) => sketch, + Ok(None) => { + warn!( + "range-repartition: matching RuntimeStatsExec has no sketch \ + (RuntimeStatsExec contract broken?) — single-bucket fallback" + ); + return Vec::new(); + } + Err(e) => { + warn!( + "range-repartition: sketch snapshot failed ({e}) — single-bucket fallback" + ); + return Vec::new(); + } + }; + // `count()` is the sum of centroid weights — total observed row count + // that fed the digest. Zero means no samples arrived before the + // snapshot; degenerate cuts would follow. + if sketch.count() == 0.0 { + warn!( + "range-repartition: matching sketch has no samples yet — single-bucket fallback" + ); + return Vec::new(); + } + // K-1 cuts at 1/K, 2/K, ..., (K-1)/K. `estimate_quantile` is monotone by + // construction, so cuts are non-decreasing (ties possible on hot-value + // distributions — `split_batch_by_range` handles those correctly, it + // just skews the resulting distribution). + let k = output_partitions as f64; + (1..output_partitions) + .map(|i| sketch.estimate_quantile(i as f64 / k)) + .collect() +} + +/// Walks `plan`'s subtree through single-child chains only, returning the +/// first [`RuntimeStatsExec`] that sketches on `routing_expr`. +/// +/// Two invariants have to hold for a sketch to be trustworthy: +/// 1. **Expression match.** A sketch of column `foo` says nothing about +/// routing on column `bar`. A `RuntimeStatsExec` sketching on a different +/// expression is treated as a plain passthrough — the walker keeps +/// descending past it looking for a matching one deeper in the chain. +/// 2. **Distribution preservation.** Any operator between us and the stats +/// that drops rows (`FilterExec`, `LimitExec`), transforms the routing +/// value (`ProjectionExec` with a computed column), or duplicates rows +/// (`JoinExec`) makes the sketch stale — the count still holds but the +/// distribution has drifted. The walker consults [`preserves_distribution`] +/// and refuses to descend past anything it doesn't know is safe. +/// +/// Also stops at any branch (> 1 child) or leaf (0 children) — descending +/// into a join's sides would risk picking up a sketch of the wrong subtree. +pub(super) fn find_runtime_stats<'a>( + plan: &'a Arc, + routing_expr: &dyn PhysicalExpr, +) -> Option<&'a RuntimeStatsExec> { + if let Some(stats) = plan.downcast_ref::() { + let matches = stats + .order_by() + .and_then(|order_by| order_by.first()) + .is_some_and(|first| first.expr.as_ref() == routing_expr); + if matches { + return Some(stats); + } + // Non-matching stats is still a passthrough for our purposes — fall + // through to the descent step. + } else if !preserves_distribution(plan.as_ref()) { + // Unrecognized node type — could change the row set or value + // distribution of the routing key. Refuse to descend. + return None; + } + let children = plan.children(); + let [only_child] = children.as_slice() else { + return None; + }; + 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 isn't going to land), 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::().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::() + .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::().is_some() + || plan.downcast_ref::().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::().is_some() + || plan.downcast_ref::().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 +/// and at `+∞` on partition K-1). Output vector is always length K; empty +/// buckets produce empty `RecordBatch`es rather than being omitted, so +/// callers can index by partition id. +/// +/// NULL routing keys land in partition 0 today; a follow-up will honor +/// `sort_options.nulls_first`/`nulls_last` to match SQL semantics. +/// +/// Boundaries are pre-extracted `f64`s; widening to other routing key +/// types generalizes this function and the discovery path together. +pub(super) fn split_batch_by_range( + batch: &RecordBatch, + routing_expr: &Arc, + boundaries: &[f64], +) -> Result> { + let output_partitions = boundaries.len() + 1; + let schema = batch.schema(); + if batch.num_rows() == 0 { + return Ok((0..output_partitions) + .map(|_| RecordBatch::new_empty(schema.clone())) + .collect()); + } + let evaluated = routing_expr.evaluate(batch)?; + let array = evaluated.into_array(batch.num_rows())?; + let keys = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "range-repartition: routing expr produced {:?}, expected Float64", + array.data_type() + ) + })?; + + let mut buckets: Vec> = (0..output_partitions).map(|_| Vec::new()).collect(); + for row in 0..batch.num_rows() { + let target = if keys.is_null(row) { + 0 + } else { + // partition_point returns the count of elements matching the + // predicate — here `<= key` — which is the target partition + // index under the half-open convention. + boundaries.partition_point(|&cut| cut <= keys.value(row)) + }; + buckets[target].push(row as u32); + } + + let mut result = Vec::with_capacity(output_partitions); + for indices in buckets { + if indices.is_empty() { + result.push(RecordBatch::new_empty(schema.clone())); + } else { + let idx_array = UInt32Array::from(indices); + let taken = take_arrays(batch.columns(), &idx_array, None)?; + result.push(RecordBatch::try_new(schema.clone(), taken)?); + } + } + Ok(result) +} + +/// Best-effort broadcast of a terminal error to every output channel. +/// `DataFusionError` isn't `Clone`; serialize via `to_string` and re-wrap as +/// `Internal` on replicas beyond the first. +pub(super) async fn broadcast_error( + senders: &[mpsc::Sender>], + err: datafusion::error::DataFusionError, +) { + let message = err.to_string(); + let mut first = Some(err); + for sender in senders.iter() { + let payload = match first.take() { + Some(original) => Err(original), + None => Err(internal_datafusion_err!("{}", message)), + }; + let _ = sender.send(payload).await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Float64Array, Int64Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::col; + use std::sync::Arc; + + fn f64_col(schema: &Schema, name: &str) -> Arc { + col(name, schema).unwrap() + } + + fn schema_v2_id() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, true), + Field::new("id", DataType::Int64, false), + ])) + } + + fn batch(schema: &Arc, keys: Vec>, ids: Vec) -> RecordBatch { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Float64Array::from(keys)), + Arc::new(Int64Array::from(ids)), + ], + ) + .unwrap() + } + + #[test] + fn split_conserves_rows_across_buckets() { + let schema = schema_v2_id(); + let batch = batch( + &schema, + vec![ + Some(-3.0), + Some(0.0), + Some(1.5), + Some(5.0), + Some(9.9), + Some(10.0), + Some(100.0), + ], + vec![0, 1, 2, 3, 4, 5, 6], + ); + let routing = f64_col(&schema, "v2"); + let splits = split_batch_by_range(&batch, &routing, &[0.0, 10.0]).unwrap(); + assert_eq!(splits.len(), 3, "K = boundaries.len() + 1"); + let total: usize = splits.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, batch.num_rows(), "no row lost or duplicated"); + } + + #[test] + fn split_half_open_boundary_lands_in_higher_partition() { + let schema = schema_v2_id(); + // 0.0 lands in partition 1 (not 0); 10.0 lands in partition 2 (not 1). + let batch = batch( + &schema, + vec![Some(-0.1), Some(0.0), Some(9.999), Some(10.0)], + vec![0, 1, 2, 3], + ); + let routing = f64_col(&schema, "v2"); + let splits = split_batch_by_range(&batch, &routing, &[0.0, 10.0]).unwrap(); + assert_eq!(splits[0].num_rows(), 1); + assert_eq!(splits[1].num_rows(), 2); + assert_eq!(splits[2].num_rows(), 1); + } + + #[test] + fn split_routes_nulls_to_partition_zero() { + let schema = schema_v2_id(); + let batch = batch( + &schema, + vec![None, Some(5.0), None, Some(50.0)], + vec![0, 1, 2, 3], + ); + let routing = f64_col(&schema, "v2"); + let splits = split_batch_by_range(&batch, &routing, &[10.0]).unwrap(); + assert_eq!(splits[0].num_rows(), 3); + assert_eq!(splits[1].num_rows(), 1); + } + + #[test] + fn split_empty_batch_produces_k_empty_batches() { + let schema = schema_v2_id(); + let batch = batch(&schema, vec![], vec![]); + let routing = f64_col(&schema, "v2"); + let splits = split_batch_by_range(&batch, &routing, &[0.0, 10.0]).unwrap(); + assert_eq!(splits.len(), 3); + assert!(splits.iter().all(|b| b.num_rows() == 0)); + } +} diff --git a/ballista/core/src/execution_plans/unordered_range_repartition.rs b/ballista/core/src/execution_plans/unordered_range_repartition.rs new file mode 100644 index 000000000..cbe1b96c4 --- /dev/null +++ b/ballista/core/src/execution_plans/unordered_range_repartition.rs @@ -0,0 +1,785 @@ +// 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. + +//! Value-range router over unordered inputs. Reads `P` input partitions with +//! no sort assumption, evaluates the first ORDER BY expression per row, and +//! routes each row to one of `K` output partitions under the half-open +//! convention: partition `p` owns `[cut[p-1], cut[p])` with virtual `-∞`/`+∞` +//! sentinels on the ends. +//! +//! # Dynamic discovery +//! +//! The whole point of this operator is that boundaries are *discovered at +//! runtime*, not baked in at plan time. On the first batch to arrive from +//! any input partition, the operator walks its own child subtree to find a +//! sibling [`RuntimeStatsExec`], snapshots its `merged_quantile_sketch()`, +//! and computes `K - 1` quantile cuts at `1/K, 2/K, ..., (K-1)/K`. All +//! batches (including the one that triggered the snapshot) then route +//! through those cuts. +//! +//! # Fallback +//! +//! If no sibling [`RuntimeStatsExec`] exists, or it's in row-count-only mode +//! (no sketch), or the sketch is empty, discovery returns an empty cut set +//! and every row lands in output partition 0 — the natural single-bucket +//! outcome of `boundaries.len() + 1 = 1`. Runtime routing must never crash; +//! degraded-but-alive beats the alternative, and downstream sees an empty +//! stream on the K-1 partitions that got no data. +//! +//! # Type generality +//! +//! The impl hardcodes Float64 downcast internally (that's what DataFusion's +//! T-Digest speaks). The public API and the sibling [`RuntimeStatsExec`] +//! stay type-agnostic; widening to other `Ord` `ScalarValue` types replaces +//! the downcast + boundary computation, no API break. +//! +//! Sibling `OrderedRangeRepartitionExec` (not yet built) handles the sorted +//! case (N sorted → M sorted range-disjoint via k-way merge). See +//! `docs/source/contributors-guide/parallel-window-kll-adaptive.md`. +//! +//! [`RuntimeStatsExec`]: crate::execution_plans::RuntimeStatsExec + +use std::fmt::{self, Debug, Formatter}; +use std::sync::{Arc, Mutex, OnceLock}; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{ + EquivalenceProperties, Partitioning, PhysicalExpr, PhysicalSortExpr, +}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::stream::StreamExt; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; + +use crate::execution_plans::range_repartition_common::{ + broadcast_error, discover_cuts, split_batch_by_range, +}; + +/// Per-output-partition channel capacity. Small = tight backpressure; the +/// classic double-buffering shape (one batch in-flight while consumer works +/// on another). When a downstream drain lags, `send().await` suspends the +/// scatter task, which suspends its input read, which propagates upstream. +/// Bump if profiling shows scatter tasks blocked on `send` more than the +/// downstream is meaningfully doing work. +const CHANNEL_CAPACITY: usize = 2; + +/// Value-range router over unordered inputs. See the module-level docs. +pub struct UnorderedRangeRepartitionExec { + input: Arc, + /// Lexicographic ORDER BY carried through from the wrapping window + /// operator. `try_new` guarantees at least one element; the first entry + /// (a `Float64` column, until we widen) drives routing. + order_by: Vec, + /// K — number of output partitions. K=1 collapses all P inputs to a + /// single bucket (the same shape discovery-failure fallback produces); + /// larger K spreads rows by value range. + output_partitions: usize, + /// Lazy channel setup guarded by `Mutex`. First `execute()` call spawns + /// the P input-reader tasks and creates the K channels; subsequent + /// `execute(p)` calls take `channels[p]`. + state: Arc>, + properties: Arc, +} + +/// Per-exec-instance lazy state. Owns the K receivers between setup and the +/// K per-partition takes. +struct DispatchState { + /// K slots. Each holds `Some(receiver)` after setup, `None` once its + /// output partition has been consumed. + receivers: Vec>>>, + initialized: bool, +} + +impl UnorderedRangeRepartitionExec { + /// Wrap `input`. `order_by` must be non-empty and its first entry must + /// evaluate to `Float64` against `input.schema()`. `output_partitions` + /// is K; any value works, K=1 gives a coalesce-shaped passthrough. + pub fn try_new( + input: Arc, + order_by: Vec, + output_partitions: usize, + // TODO: support RANGE & ROW halos + ) -> Result { + let [routing, ..] = order_by.as_slice() else { + return internal_err!( + "UnorderedRangeRepartitionExec requires at least one ORDER BY expression" + ); + }; + let schema = input.schema(); + let routing_type = routing.expr.data_type(&schema)?; + if !matches!(routing_type, DataType::Float64) { + // TODO: support all continuous primitives + return internal_err!( + "UnorderedRangeRepartitionExec routing expression `{}` must be Float64, got {:?}", + routing.expr, + routing_type + ); + } + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(output_partitions), + input.pipeline_behavior(), + input.boundedness(), + )); + let state = Arc::new(Mutex::new(DispatchState { + receivers: Vec::new(), + initialized: false, + })); + Ok(Self { + input, + order_by, + output_partitions, + state, + properties, + }) + } + + /// Full ORDER BY carried through from the wrapping window operator. + pub fn order_by(&self) -> &[PhysicalSortExpr] { + &self.order_by + } + + /// K — the fixed output partition count. + pub fn output_partitions(&self) -> usize { + self.output_partitions + } +} + +impl Debug for UnorderedRangeRepartitionExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("UnorderedRangeRepartitionExec") + .field("order_by", &self.order_by) + .field("output_partitions", &self.output_partitions) + .finish() + } +} + +impl DisplayAs for UnorderedRangeRepartitionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + let routing = &self.order_by[0]; + write!( + f, + "UnorderedRangeRepartitionExec: routing={} {} → {} partitions", + routing.expr, + if routing.options.descending { + "desc" + } else { + "asc" + }, + self.output_partitions, + ) + } +} + +impl ExecutionPlan for UnorderedRangeRepartitionExec { + fn name(&self) -> &str { + "UnorderedRangeRepartitionExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [input] = children.as_slice() else { + return internal_err!( + "UnorderedRangeRepartitionExec expects exactly one child, got {}", + children.len() + ); + }; + Ok(Arc::new(UnorderedRangeRepartitionExec::try_new( + input.clone(), + self.order_by.clone(), + self.output_partitions, + )?)) + } + + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + let mut state = self + .state + .lock() + .map_err(|_| internal_datafusion_err!("dispatch state mutex poisoned"))?; + if !state.initialized { + let mut senders = Vec::with_capacity(self.output_partitions); + let mut receivers = Vec::with_capacity(self.output_partitions); + for _ in 0..self.output_partitions { + let (tx, rx) = mpsc::channel(CHANNEL_CAPACITY); + senders.push(tx); + receivers.push(Some(rx)); + } + state.receivers = receivers; + state.initialized = true; + let senders: Arc<[mpsc::Sender>]> = senders.into(); + // Empty `Vec` = discovery failed = single-bucket fallback. + // Populated once, on the first batch, by whichever scatter task + // wins the `OnceLock::get_or_init` race. + let cuts_cell: Arc>> = Arc::new(OnceLock::new()); + let input_partitions = self.input.output_partitioning().partition_count(); + let routing_expr = self.order_by[0].expr.clone(); // TODO: KLL for multi-column? + for input_partition in 0..input_partitions { + let child = self.input.clone(); + let senders = senders.clone(); + let cuts_cell = cuts_cell.clone(); + let routing_expr = routing_expr.clone(); + let ctx = ctx.clone(); + let output_partitions = self.output_partitions; + tokio::spawn(async move { + scatter_input_partition( + child, + input_partition, + ctx, + routing_expr, + senders, + cuts_cell, + output_partitions, + ) + .await; + }); + } + } + let Some(slot) = state.receivers.get_mut(partition) else { + return internal_err!( + "UnorderedRangeRepartitionExec: partition {} out of range (have {})", + partition, + self.output_partitions + ); + }; + let receiver = slot.take().ok_or_else(|| { + internal_datafusion_err!( + "UnorderedRangeRepartitionExec: execute({partition}) called twice \ + on the same instance" + ) + })?; + drop(state); + let stream = ReceiverStream::new(receiver); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} + +/// One background task per input partition. Reads the child's batches; on +/// the first batch, the shared `OnceLock` learns the value cuts (by walking +/// the child tree for `RuntimeStatsExec` and snapshotting its sketch). +/// Every batch — including the first — is then dispatched via +/// `split_batch_by_range`, which handles the degenerate empty-cuts case +/// (all rows to partition 0) transparently. +async fn scatter_input_partition( + child: Arc, + input_partition: usize, + ctx: Arc, + routing_expr: Arc, + senders: Arc<[mpsc::Sender>]>, + cuts_cell: Arc>>, + output_partitions: usize, +) { + let mut stream = match child.execute(input_partition, ctx) { + Ok(s) => s, + Err(err) => { + broadcast_error(&senders, err).await; + return; + } + }; + while let Some(batch_result) = stream.next().await { + // Every downstream consumer dropped its receiver — DRR itself was + // dropped, or every output partition had a `LIMIT` above and hit it. + // Either way, no one's listening; stop reading input so this scatter + // task exits promptly (drops its `stream`, which propagates upstream). + if senders.iter().all(|s| s.is_closed()) { + return; + } + let batch = match batch_result { + Ok(b) => b, + Err(err) => { + broadcast_error(&senders, err).await; + return; + } + }; + let cuts = cuts_cell.get_or_init(|| { + discover_cuts(&child, routing_expr.as_ref(), output_partitions) + }); + // TODO(perf): `split_batch_by_range` materialises K sub-batches per + // input batch via `take_arrays` — one copy per row into a fresh + // allocation. Unlike the ordered variant we can't slice + // contiguous ranges (input isn't sorted), but an + // `Arc` broadcast + receiver-side filter would skip + // the scatter-side allocations at the cost of duplicating the + // filter work K times. Worth measuring under skew. + match split_batch_by_range(&batch, &routing_expr, cuts) { + Ok(splits) => { + for (output, sub) in splits.into_iter().enumerate() { + if sub.num_rows() == 0 { + continue; + } + // `send().await` is where the backpressure lives: + // suspends when the channel is at capacity, which + // suspends this scatter task, which suspends its + // input read → propagates upstream. Errors mean the + // downstream dropped its receiver; keep forwarding + // to the other outputs. + let _ = senders[output].send(Ok(sub)).await; + } + } + Err(err) => { + broadcast_error(&senders, err).await; + return; + } + } + } + // All senders drop with this task → receivers see EOF. +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution_plans::RuntimeStatsExec; + use datafusion::arrow::array::{Float64Array, Int64Array}; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_expr::expressions::col; + use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + use datafusion::prelude::SessionContext; + use std::sync::Arc; + + fn schema_v2_id() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, false), + Field::new("id", DataType::Int64, false), + ])) + } + + fn batch(schema: &Arc, keys: Vec, ids: Vec) -> RecordBatch { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Float64Array::from(keys)), + Arc::new(Int64Array::from(ids)), + ], + ) + .unwrap() + } + + fn asc(schema: &Schema, name: &str) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col(name, schema).unwrap(), + options: Default::default(), + } + } + + // ---------- try_new: constructor validation ------------------------- + + fn empty_input(schema: &Arc) -> Arc { + MemorySourceConfig::try_new_exec(&[vec![]], schema.clone(), None).unwrap() + } + + #[test] + fn try_new_rejects_empty_order_by() { + let schema = schema_v2_id(); + let err = UnorderedRangeRepartitionExec::try_new(empty_input(&schema), vec![], 3) + .expect_err("empty order_by must be rejected"); + assert!( + err.to_string().contains("at least one ORDER BY expression"), + "error should name the missing invariant, got: {err}" + ); + } + + #[test] + fn try_new_rejects_non_float64_routing_key() { + let schema = schema_v2_id(); + let err = UnorderedRangeRepartitionExec::try_new( + empty_input(&schema), + vec![asc(&schema, "id")], // Int64 + 3, + ) + .expect_err("Int64 routing key must be rejected"); + assert!( + err.to_string().contains("must be Float64"), + "error should name the type mismatch, got: {err}" + ); + } + + #[test] + fn output_partitioning_reflects_k() { + let schema = schema_v2_id(); + let exec = UnorderedRangeRepartitionExec::try_new( + empty_input(&schema), + vec![asc(&schema, "v2")], + 4, + ) + .unwrap(); + assert_eq!(exec.properties().output_partitioning().partition_count(), 4); + } + + // ---------- execute(): end-to-end routing over Scan → Stats → DRR -- + + fn mem_input( + schema: &Arc, + partitions: Vec>, + ) -> Arc { + let batches: Vec> = partitions + .into_iter() + .map(|rows| { + let (keys, ids): (Vec<_>, Vec<_>) = rows.into_iter().unzip(); + vec![batch(schema, keys, ids)] + }) + .collect(); + MemorySourceConfig::try_new_exec(&batches, schema.clone(), None).unwrap() + } + + /// Wrap `source` in a `RuntimeStatsExec` (sketch mode on `v2`), then in + /// a DRR — matches the Stage-1 layout minus the Dam. + fn source_stats_drr( + schema: &Arc, + source: Arc, + k: usize, + ) -> Arc { + let stats = Arc::new( + RuntimeStatsExec::try_new(source, Some(vec![asc(schema, "v2")])).unwrap(), + ); + Arc::new( + UnorderedRangeRepartitionExec::try_new(stats, vec![asc(schema, "v2")], k) + .unwrap(), + ) + } + + /// Drain all K output partitions **concurrently** — otherwise the + /// bounded scatter channels can deadlock: while a test sequentially + /// polls partition 0, scatter tasks fill channels 1..K-1 to + /// `CHANNEL_CAPACITY` and block on send, never dropping their partition-0 + /// senders. In production K downstream tasks run concurrently and this + /// pattern never arises. + async fn drain_concurrent( + exec: Arc, + ctx: Arc, + column: usize, + extract: fn(&RecordBatch, usize) -> Vec, + ) -> Vec> { + let output_partitions = exec.output_partitioning().partition_count(); + let streams: Vec<_> = (0..output_partitions) + .map(|p| exec.execute(p, ctx.task_ctx()).unwrap()) + .collect(); + let handles: Vec<_> = streams + .into_iter() + .map(|stream| { + tokio::spawn(async move { + let batches: Vec = + <_ as futures::stream::StreamExt>::collect::< + Vec>, + >(stream) + .await + .into_iter() + .map(|r| r.unwrap()) + .collect(); + batches.iter().flat_map(|b| extract(b, column)).collect() + }) + }) + .collect(); + let mut per_output = Vec::with_capacity(output_partitions); + for handle in handles { + per_output.push(handle.await.unwrap()); + } + per_output + } + + fn extract_ids(batch: &RecordBatch, column: usize) -> Vec { + batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .collect() + } + + fn extract_keys(batch: &RecordBatch, column: usize) -> Vec { + batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .collect() + } + + async fn ids_per_output( + exec: Arc, + ctx: Arc, + ) -> Vec> { + drain_concurrent(exec, ctx, 1, extract_ids).await + } + + async fn keys_per_output( + exec: Arc, + ctx: Arc, + ) -> Vec> { + drain_concurrent(exec, ctx, 0, extract_keys).await + } + + fn session() -> Arc { + Arc::new(SessionContext::new_with_state( + SessionStateBuilder::new().with_default_features().build(), + )) + } + + /// Random-uniform-ish input over Scan → RuntimeStatsExec → DRR: verify + /// every row lands in *some* output, and no row lands in more than one. + #[tokio::test] + async fn end_to_end_conserves_rows() { + let schema = schema_v2_id(); + // 20 rows, roughly evenly spread across [0, 100). + let rows: Vec<(f64, i64)> = + (0..20).map(|i| ((i as f64) * 5.0, i as i64)).collect(); + let source = mem_input(&schema, vec![rows]); + let exec = source_stats_drr(&schema, source, 4); + let per_output = ids_per_output(exec, session()).await; + let total: usize = per_output.iter().map(|v| v.len()).sum(); + assert_eq!(total, 20, "every row must land in exactly one output"); + } + + /// The load-bearing property: outputs are range-disjoint. For every + /// adjacent pair `(p, p+1)`, `max(partition[p]) <= min(partition[p+1])`. + #[tokio::test] + async fn end_to_end_partitions_are_range_disjoint() { + let schema = schema_v2_id(); + let rows: Vec<(f64, i64)> = + (0..40).map(|i| ((i as f64) * 2.5, i as i64)).collect(); + let source = mem_input(&schema, vec![rows]); + let exec = source_stats_drr(&schema, source, 4); + let keys = keys_per_output(exec, session()).await; + for window in keys.windows(2) { + let [left, right] = window else { + unreachable!() + }; + if left.is_empty() || right.is_empty() { + continue; + } + let left_max = left.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + let right_min = right.iter().cloned().fold(f64::INFINITY, f64::min); + assert!( + left_max <= right_min, + "range-disjoint invariant broken: max(left)={left_max} > min(right)={right_min}" + ); + } + } + + /// Multiple input partitions still route correctly. Distinct values in + /// each input partition; conservation must hold across all of them. + #[tokio::test] + async fn end_to_end_multiple_input_partitions() { + let schema = schema_v2_id(); + let partitions = vec![ + (0..10).map(|i| (i as f64, i as i64)).collect(), + (0..10).map(|i| (i as f64 + 20.0, 100 + i as i64)).collect(), + (0..10).map(|i| (i as f64 + 50.0, 200 + i as i64)).collect(), + ]; + let source = mem_input(&schema, partitions); + let exec = source_stats_drr(&schema, source, 4); + let per_output = ids_per_output(exec, session()).await; + let total: usize = per_output.iter().map(|v| v.len()).sum(); + assert_eq!(total, 30); + } + + /// Fallback: without a sibling `RuntimeStatsExec`, every batch funnels + /// to `FALLBACK_PARTITION`. No crash, no error, no data loss. + #[tokio::test] + async fn end_to_end_fallback_when_no_runtime_stats() { + let schema = schema_v2_id(); + let rows: Vec<(f64, i64)> = (0..15).map(|i| (i as f64, i as i64)).collect(); + let source = mem_input(&schema, vec![rows]); + // No RuntimeStatsExec in the tree — DRR built directly on the source. + let exec: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(source, vec![asc(&schema, "v2")], 4) + .unwrap(), + ); + let per_output = ids_per_output(exec, session()).await; + assert_eq!(per_output.len(), 4); + // Empty cuts → single-bucket fallback → everything on partition 0. + assert_eq!(per_output[0].len(), 15); + assert_eq!(per_output[1].len(), 0); + assert_eq!(per_output[2].len(), 0); + assert_eq!(per_output[3].len(), 0); + } + + /// Fallback also fires when the sibling `RuntimeStatsExec` is in + /// row-count-only mode (no sketch to read cuts from). + #[tokio::test] + async fn end_to_end_fallback_when_stats_has_no_sketch() { + let schema = schema_v2_id(); + let rows: Vec<(f64, i64)> = (0..12).map(|i| (i as f64, i as i64)).collect(); + let source = mem_input(&schema, vec![rows]); + let stats = Arc::new(RuntimeStatsExec::try_new(source, None).unwrap()); // row-count only + let exec: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(stats, vec![asc(&schema, "v2")], 3) + .unwrap(), + ); + let per_output = ids_per_output(exec, session()).await; + assert_eq!(per_output.len(), 3); + assert_eq!(per_output[0].len(), 12); + assert_eq!(per_output[1].len(), 0); + assert_eq!(per_output[2].len(), 0); + } + + /// A `RuntimeStatsExec` hiding behind a branching node (here `UnionExec` + /// with two children) is *not* what we want — its sketch reflects only + /// one side of the branch, not the data actually flowing through DRR. + /// Discovery must refuse to descend and fall back cleanly. + #[tokio::test] + async fn end_to_end_fallback_when_runtime_stats_hides_behind_branch() { + use datafusion::physical_plan::union::UnionExec; + let schema = schema_v2_id(); + // Branch A: has a RuntimeStatsExec sketching v2, values in [0, 5). + let branch_a = mem_input( + &schema, + vec![(0..5).map(|i| (i as f64, i as i64)).collect()], + ); + let stats_on_branch_a = Arc::new( + RuntimeStatsExec::try_new(branch_a, Some(vec![asc(&schema, "v2")])).unwrap(), + ); + // Branch B: no stats; values shifted to [100, 105) so a "wrong" + // sketch (only branch A) would send them all to the top partition. + let branch_b = mem_input( + &schema, + vec![(0..5).map(|i| (100.0 + i as f64, 100 + i as i64)).collect()], + ); + let union: Arc = + UnionExec::try_new(vec![stats_on_branch_a, branch_b]).unwrap(); + let exec: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(union, vec![asc(&schema, "v2")], 3) + .unwrap(), + ); + let per_output = ids_per_output(exec, session()).await; + // Empty cuts → single-bucket fallback → all on partition 0. If + // find_runtime_stats had descended into the union, branch A's tiny + // sketch would have spread rows across partitions instead. + assert_eq!(per_output.len(), 3); + assert_eq!(per_output[0].len(), 10); + assert_eq!(per_output[1].len(), 0); + assert_eq!(per_output[2].len(), 0); + } + + /// `BufferExec` is on the distribution-preserving whitelist: the walker + /// must traverse it to reach the matching `RuntimeStatsExec`, and cuts + /// come out non-empty (data lands across multiple output partitions, + /// not funneled to bucket 0). + #[tokio::test] + async fn end_to_end_walker_descends_through_buffer() { + use crate::execution_plans::{BufferExec, BufferMode}; + let schema = schema_v2_id(); + let rows: Vec<(f64, i64)> = (0..40).map(|i| (i as f64, i as i64)).collect(); + let source = mem_input(&schema, vec![rows]); + let stats = Arc::new( + RuntimeStatsExec::try_new(source, Some(vec![asc(&schema, "v2")])).unwrap(), + ); + let buffer = Arc::new(BufferExec::try_new(stats, BufferMode::Dam).unwrap()); + let exec: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(buffer, vec![asc(&schema, "v2")], 4) + .unwrap(), + ); + let per_output = ids_per_output(exec, session()).await; + let non_empty = per_output.iter().filter(|v| !v.is_empty()).count(); + assert!( + non_empty > 1, + "walker should have traversed BufferExec, populated cuts, and \ + spread rows across multiple outputs — got per_output={per_output:?}" + ); + } + + /// `FilterExec` is *not* on the whitelist: it drops rows, so any upstream + /// sketch has stale distribution info. The walker must refuse to descend + /// through it → single-bucket fallback (all rows to partition 0). + #[tokio::test] + async fn end_to_end_walker_refuses_filter() { + use datafusion::physical_expr::expressions::{Literal, binary}; + use datafusion::physical_plan::filter::FilterExec; + use datafusion::scalar::ScalarValue; + let schema = schema_v2_id(); + let rows: Vec<(f64, i64)> = (0..15).map(|i| (i as f64, i as i64)).collect(); + let source = mem_input(&schema, vec![rows]); + let stats = Arc::new( + RuntimeStatsExec::try_new(source, Some(vec![asc(&schema, "v2")])).unwrap(), + ); + // `v2 >= 0` — passes every row (we just want a FilterExec node in the + // chain so the walker has to decide whether to descend past it). + let predicate = binary( + col("v2", schema.as_ref()).unwrap(), + datafusion::logical_expr::Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), + schema.as_ref(), + ) + .unwrap(); + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, stats).unwrap()); + let exec: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(filter, vec![asc(&schema, "v2")], 4) + .unwrap(), + ); + let per_output = ids_per_output(exec, session()).await; + // Every row funnels to bucket 0 — walker refused to descend past + // FilterExec, so cuts came back empty. + assert_eq!(per_output.len(), 4); + assert_eq!(per_output[0].len(), 15); + assert_eq!(per_output[1].len(), 0); + assert_eq!(per_output[2].len(), 0); + assert_eq!(per_output[3].len(), 0); + } + + /// First execute(p) takes channels[p]; second call must error rather than + /// silently returning nothing. Uses the concurrent-drain helper because + /// bounded scatter channels would deadlock under sequential collection. + #[tokio::test] + async fn execute_same_partition_twice_errors() { + let schema = schema_v2_id(); + let source = mem_input(&schema, vec![vec![(5.0, 0)]]); + let exec = source_stats_drr(&schema, source, 3); + let ctx = session(); + let _ = ids_per_output(exec.clone(), ctx.clone()).await; + let err = match exec.execute(0, ctx.task_ctx()) { + Ok(_) => panic!("second execute on same partition must error"), + Err(err) => err, + }; + assert!( + err.to_string().contains("execute(0) called twice"), + "error should name the invariant, got: {err}" + ); + } +} diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 9fa2aedb0..4143163b5 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode { pub struct BallistaPhysicalPlanNode { #[prost( oneof = "ballista_physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 5, 6, 7" + tags = "1, 2, 3, 4, 5, 6, 7, 8" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -55,6 +55,8 @@ pub mod ballista_physical_plan_node { RuntimeStats(super::RuntimeStatsExecNode), #[prost(message, tag = "7")] Buffer(super::BufferExecNode), + #[prost(message, tag = "8")] + UnorderedRangeRepartition(super::UnorderedRangeRepartitionExecNode), } } /// Runtime-stats tap. Passes batches through unmodified while accumulating: @@ -90,6 +92,25 @@ pub struct BufferExecNode { #[prost(enumeration = "BufferMode", tag = "1")] pub mode: i32, } +/// Value-range router over unordered inputs. Reads P input partitions with +/// no sort assumption, evaluates the first ORDER BY expression per row, and +/// routes each row to one of K output partitions using cuts discovered at +/// runtime from a sibling RuntimeStatsExec's quantile sketch. The child +/// plan is plumbed by the framework as `inputs\[0\]`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct UnorderedRangeRepartitionExecNode { + /// Lexicographic ORDER BY carried through from the wrapping window + /// operator. Routing keys off the first entry today; the full ordering is + /// preserved so downstream operators (SortExec, BWAG) that need it keep + /// working. + #[prost(message, repeated, tag = "1")] + pub order_by: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalSortExprNode, + >, + /// K — number of output partitions. Must be ≥ 2. + #[prost(uint32, tag = "2")] + pub output_partitions: u32, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct ChaosExecNode { #[prost(double, tag = "1")] diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index ae2b85c93..e4cb677bf 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -57,7 +57,8 @@ use std::{convert::TryInto, io::Cursor}; use crate::execution_plans::sort_shuffle::SortShuffleConfig; use crate::execution_plans::{ BufferExec, BufferMode, ChaosExec, CoalescePlan, PartitionGroup, RuntimeStatsExec, - ShuffleReaderExec, ShuffleWriterExec, SortShuffleWriterExec, UnresolvedShuffleExec, + ShuffleReaderExec, ShuffleWriterExec, SortShuffleWriterExec, + UnorderedRangeRepartitionExec, UnresolvedShuffleExec, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, @@ -594,6 +595,25 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { }; Ok(Arc::new(BufferExec::try_new(input.clone(), mode)?)) } + PhysicalPlanType::UnorderedRangeRepartition(node) => { + let [input] = inputs else { + return Err(DataFusionError::Internal(format!( + "UnorderedRangeRepartitionExec expects exactly 1 input, got {}", + inputs.len() + ))); + }; + let order_by = parse_physical_sort_exprs( + &node.order_by, + &decode_ctx, + input.schema().as_ref(), + &converter, + )?; + Ok(Arc::new(UnorderedRangeRepartitionExec::try_new( + input.clone(), + order_by, + node.output_partitions as usize, + )?)) + } } } @@ -808,6 +828,27 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { DataFusionError::Internal(format!("failed to encode BufferExec: {e:?}")) })?; Ok(()) + } else if let Some(exec) = node.downcast_ref::() { + let converter = DefaultPhysicalProtoConverter {}; + let order_by = serialize_physical_sort_exprs( + exec.order_by().iter().cloned(), + self.default_codec.as_ref(), + &converter, + )?; + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::UnorderedRangeRepartition( + protobuf::UnorderedRangeRepartitionExecNode { + order_by, + output_partitions: exec.output_partitions() as u32, + }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode UnorderedRangeRepartitionExec: {e:?}" + )) + })?; + Ok(()) } else { Err(DataFusionError::Internal(format!( "Unsupported plan node, name: [{}] ", @@ -1614,4 +1655,105 @@ mod test { .expect("Expected BufferExec"); assert!(matches!(decoded.mode(), BufferMode::Dam)); } + + /// The wrapper carries its ORDER BY across the network so the executor's + /// sampler knows which column to route on. Round-trip a single-key + /// ordering to guard the codec's happy path. + #[tokio::test] + async fn test_unordered_range_repartition_exec_roundtrip_single_key() { + use crate::execution_plans::UnorderedRangeRepartitionExec; + use datafusion::arrow::compute::SortOptions; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_plan::empty::EmptyExec; + + let schema = Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, false), + Field::new("id", DataType::Int64, false), + ])); + let input: Arc = Arc::new(EmptyExec::new(schema.clone())); + let sort_expr = PhysicalSortExpr { + expr: col("v2", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }; + let original = UnorderedRangeRepartitionExec::try_new( + input.clone(), + vec![sort_expr.clone()], + 8, + ) + .unwrap(); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec.try_encode(Arc::new(original), &mut buf).unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded_plan = codec.try_decode(&buf, &[input], &ctx).unwrap(); + + let decoded = decoded_plan + .downcast_ref::() + .expect("Expected UnorderedRangeRepartitionExec"); + let order_by = decoded.order_by(); + assert_eq!(order_by.len(), 1, "single-key ORDER BY should round-trip"); + assert_eq!(order_by[0].expr.to_string(), sort_expr.expr.to_string()); + assert!(!order_by[0].options.descending); + assert!(order_by[0].options.nulls_first); + assert_eq!(decoded.output_partitions(), 8, "K must round-trip"); + } + + /// The API accepts multi-key lexicographic ordering (RANGE-frame windows + /// only route on the first key, but the wire format must not drop the + /// rest — that's the ROWS-frame follow-up's substrate). + #[tokio::test] + async fn test_unordered_range_repartition_exec_roundtrip_multi_key() { + use crate::execution_plans::UnorderedRangeRepartitionExec; + use datafusion::arrow::compute::SortOptions; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_plan::empty::EmptyExec; + + let schema = Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, false), + Field::new("id", DataType::Int64, false), + ])); + let input: Arc = Arc::new(EmptyExec::new(schema.clone())); + let order_by = vec![ + PhysicalSortExpr { + expr: col("v2", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + PhysicalSortExpr { + expr: col("id", schema.as_ref()).unwrap(), + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + ]; + let original = + UnorderedRangeRepartitionExec::try_new(input.clone(), order_by.clone(), 16) + .unwrap(); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec.try_encode(Arc::new(original), &mut buf).unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded_plan = codec.try_decode(&buf, &[input], &ctx).unwrap(); + + let decoded = decoded_plan + .downcast_ref::() + .expect("Expected UnorderedRangeRepartitionExec"); + let decoded_order_by = decoded.order_by(); + assert_eq!(decoded_order_by.len(), 2, "multi-key ORDER BY must survive"); + assert!( + decoded_order_by[1].options.descending, + "sort options must survive per-key" + ); + assert!(!decoded_order_by[1].options.nulls_first); + } } From a281667a56b038ec0291e87d03b09c46052f8240 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 21 Jul 2026 10:17:59 -0600 Subject: [PATCH 2/5] review: address review-preference audit for UnorderedRangeRepartitionExec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Preemptive polish before opening the PR for review — matches the pattern of reviewer asks Phillip made on #2094 and #2095: - Implement all 6 algebraic ExecutionPlan methods explicitly (required_input_ distribution, required_input_ordering, maintains_input_order, benefits_from_ input_partitioning, partition_statistics, cardinality_effect), each with a docstring justifying the value. maintains_input_order=false and partition_statistics=Unknown are the two that materially differ from the merged passthrough ops (RuntimeStatsExec, BufferExec). - Reject nullable routing expressions in try_new. The sibling RuntimeStatsExec already enforces this at plan-assembly time, but defense-in-depth catches the no-stats-pair path where discovery would silently fall back to single-bucket. TODO marker points at the KLL swap that lifts the restriction. - Add try_new_rejects_nullable_routing_key covering the new invariant. Includes two authorial touch-ups picked up along the way: a tone softening on the preserves_distribution docstring and a TODO on split_batch_by_range's row-by-row loop pointing at a vectorized Arrow rewrite. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../range_repartition_common.rs | 3 +- .../unordered_range_repartition.rs | 74 +++++++++++++++++-- 2 files changed, 69 insertions(+), 8 deletions(-) diff --git a/ballista/core/src/execution_plans/range_repartition_common.rs b/ballista/core/src/execution_plans/range_repartition_common.rs index 8a196b763..7bc8b8dd5 100644 --- a/ballista/core/src/execution_plans/range_repartition_common.rs +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -160,7 +160,7 @@ pub(super) fn find_runtime_stats<'a>( /// 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 isn't going to land), we maintain this by hand. +/// 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 @@ -235,6 +235,7 @@ pub(super) fn split_batch_by_range( ) })?; + // TODO: vectorized arrow computation let mut buckets: Vec> = (0..output_partitions).map(|_| Vec::new()).collect(); for row in 0..batch.num_rows() { let target = if keys.is_null(row) { diff --git a/ballista/core/src/execution_plans/unordered_range_repartition.rs b/ballista/core/src/execution_plans/unordered_range_repartition.rs index cbe1b96c4..aabb5b815 100644 --- a/ballista/core/src/execution_plans/unordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/unordered_range_repartition.rs @@ -23,8 +23,8 @@ //! //! # Dynamic discovery //! -//! The whole point of this operator is that boundaries are *discovered at -//! runtime*, not baked in at plan time. On the first batch to arrive from +//! The point of this operator is that boundaries are discovered at +//! runtime, not baked in at plan time. On the first batch to arrive from //! any input partition, the operator walks its own child subtree to find a //! sibling [`RuntimeStatsExec`], snapshots its `merged_quantile_sketch()`, //! and computes `K - 1` quantile cuts at `1/K, 2/K, ..., (K-1)/K`. All @@ -58,11 +58,13 @@ use std::sync::{Arc, Mutex, OnceLock}; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::{DataType, SchemaRef}; -use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{ - EquivalenceProperties, Partitioning, PhysicalExpr, PhysicalSortExpr, + Distribution, EquivalenceProperties, OrderingRequirements, Partitioning, + PhysicalExpr, PhysicalSortExpr, }; +use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, @@ -136,6 +138,14 @@ impl UnorderedRangeRepartitionExec { routing_type ); } + // TODO: fixed by KLL — a NULL-aware sketch lifts this restriction and + // lets `split_batch_by_range` honor SortOptions::nulls_first properly. + if routing.expr.nullable(&schema)? { + return internal_err!( + "UnorderedRangeRepartitionExec: routing expression `{}` must be non-nullable", + routing.expr + ); + } let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(schema), Partitioning::UnknownPartitioning(output_partitions), @@ -209,6 +219,41 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { vec![&self.input] } + /// Input distribution is irrelevant — the operator re-routes every row + /// by value range regardless of how the child organized them. + fn required_input_distribution(&self) -> Vec { + vec![Distribution::UnspecifiedDistribution] + } + + /// "Unordered" in the name: no ordering requirement on the child. + fn required_input_ordering(&self) -> Vec> { + vec![None] + } + + /// Scattering by value range across K outputs definitively breaks the + /// input order — that's the whole point. + fn maintains_input_order(&self) -> Vec { + vec![false] + } + + /// Extra input repartitioning above us buys nothing — we already read + /// all P input partitions concurrently and scatter them to K outputs. + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + /// Row count is preserved but redistributed across K new output + /// partitions; per-partition breakdown depends on the runtime sketch, + /// which isn't known at plan time. + fn partition_statistics(&self, _partition: Option) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema()))) + } + + /// Every input row is emitted exactly once. Overrides default `Unknown`. + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + fn with_new_children( self: Arc, children: Vec>, @@ -246,9 +291,6 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { state.receivers = receivers; state.initialized = true; let senders: Arc<[mpsc::Sender>]> = senders.into(); - // Empty `Vec` = discovery failed = single-bucket fallback. - // Populated once, on the first batch, by whichever scatter task - // wins the `OnceLock::get_or_init` race. let cuts_cell: Arc>> = Arc::new(OnceLock::new()); let input_partitions = self.input.output_partitioning().partition_count(); let routing_expr = self.order_by[0].expr.clone(); // TODO: KLL for multi-column? @@ -436,6 +478,24 @@ mod tests { ); } + #[test] + fn try_new_rejects_nullable_routing_key() { + let schema = Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, true), // nullable + Field::new("id", DataType::Int64, false), + ])); + let err = UnorderedRangeRepartitionExec::try_new( + empty_input(&schema), + vec![asc(&schema, "v2")], + 3, + ) + .expect_err("nullable routing key must be rejected"); + assert!( + err.to_string().contains("must be non-nullable"), + "error should name the nullability constraint, got: {err}" + ); + } + #[test] fn output_partitioning_reflects_k() { let schema = schema_v2_id(); From 312d6ee0d42c817a0637afd924be06c07819c71a Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 21 Jul 2026 11:21:47 -0600 Subject: [PATCH 3/5] docs(core): add data-flow ASCII diagram to UnorderedRangeRepartitionExec module docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors DataFusion's RepartitionExec style. Shows P input partitions → the RuntimeStatsExec/BufferExec/URRE chain → K half-open value-range outputs, with the discovery walker's descent traced explicitly against the actual child operators rather than a side channel. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../execution_plans/unordered_range_repartition.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/ballista/core/src/execution_plans/unordered_range_repartition.rs b/ballista/core/src/execution_plans/unordered_range_repartition.rs index aabb5b815..e4fa9dccc 100644 --- a/ballista/core/src/execution_plans/unordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/unordered_range_repartition.rs @@ -21,6 +21,20 @@ //! convention: partition `p` owns `[cut[p-1], cut[p])` with virtual `-∞`/`+∞` //! sentinels on the ends. //! +//! ```text +//! P input outputs (K = 4) +//! partitions +//! ─┐ ┌─▶ 0: key < c₁ +//! ─┼──▶ RuntimeStatsExec ──▶ BufferExec ──▶ UnorderedRRE ─┼─▶ 1: c₁ ≤ key < c₂ +//! ─┤ (T-Digest tap) (Dam, (K = 4) ├─▶ 2: c₂ ≤ key < c₃ +//! ─┘ whitelisted) │ └─▶ 3: c₃ ≤ key +//! ▲ │ +//! └────────── walker ◀────────────┘ +//! (descends past BufferExec, +//! matches on routing expr, +//! reads K−1 quantile cuts) +//! ``` +//! //! # Dynamic discovery //! //! The point of this operator is that boundaries are discovered at From 6b300df2fc3bb310e5f128623422c15c7857dea5 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Wed, 22 Jul 2026 04:53:34 -0600 Subject: [PATCH 4/5] review: surface scatter-task panics + set Eager/Cooperative on DRR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two review comments on PR #2123: 1. `scatter_input_partition` was spawned as a bare `tokio::spawn` whose `JoinHandle` was discarded. A panic inside the body silently dropped the sender, and downstream consumers saw a clean EOF while data was lost. Wrap it in `SpawnedTask` and add a `wait_for_task` helper (in `range_repartition_common.rs` — the ordered variant will share it) that awaits the scatter task and broadcasts any panic or DFError to every output channel. Hold the wait-tasks in `DispatchState._drop_helper` so dropping the exec aborts all P scatter tasks. 2. `PlanProperties::new` defaults to `Lazy` / `NonCooperative`; DRR eagerly drives its inputs from the moment `execute()` is called and yields at per-output channel sends, so set `EvaluationType::Eager` + `SchedulingType::Cooperative` — same combination DataFusion's `RepartitionExec` uses. New test `scatter_task_panic_surfaces_as_error` uses an inline `PanickingSourceExec` whose stream panics on first poll, and asserts at least one output partition surfaces the panic as `Err(..)` rather than a spurious clean EOF. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../range_repartition_common.rs | 38 ++++ .../unordered_range_repartition.rs | 215 ++++++++++++++---- 2 files changed, 207 insertions(+), 46 deletions(-) diff --git a/ballista/core/src/execution_plans/range_repartition_common.rs b/ballista/core/src/execution_plans/range_repartition_common.rs index 7bc8b8dd5..155ad90e3 100644 --- a/ballista/core/src/execution_plans/range_repartition_common.rs +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -40,7 +40,9 @@ use std::sync::Arc; use datafusion::arrow::array::{Array, Float64Array, RecordBatch, UInt32Array}; use datafusion::arrow::compute::take_arrays; +use datafusion::common::runtime::SpawnedTask; use datafusion::common::{Result, internal_datafusion_err}; +use datafusion::error::DataFusionError; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::sorts::sort::SortExec; @@ -280,6 +282,42 @@ pub(super) async fn broadcast_error( } } +/// Await a scatter task and turn every terminal state into something the K +/// downstream consumers can observe. Modelled on DataFusion's +/// `RepartitionExec::wait_for_task`. +/// +/// Without this, a panic in a scatter task closes its senders with no data +/// sent — every output partition sees a clean EOF and the overall task +/// reports success while silently dropping all remaining rows. Here every +/// terminal state (panic, DFError, clean end) becomes an explicit signal +/// on every output channel. +pub(super) async fn wait_for_task( + input_task: SpawnedTask>, + senders: Arc<[mpsc::Sender>]>, +) { + match input_task.join().await { + Err(join_err) => { + let ctx = if join_err.is_panic() { + "scatter task panicked" + } else { + "scatter task cancelled" + }; + broadcast_error( + &senders, + DataFusionError::Context( + ctx.to_string(), + Box::new(DataFusionError::External(Box::new(join_err))), + ), + ) + .await; + } + Ok(Err(err)) => broadcast_error(&senders, err).await, + Ok(Ok(())) => { + // Drop the senders — receivers see clean EOF. + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/ballista/core/src/execution_plans/unordered_range_repartition.rs b/ballista/core/src/execution_plans/unordered_range_repartition.rs index e4fa9dccc..b4f0b06bf 100644 --- a/ballista/core/src/execution_plans/unordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/unordered_range_repartition.rs @@ -72,13 +72,16 @@ use std::sync::{Arc, Mutex, OnceLock}; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::runtime::SpawnedTask; use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{ Distribution, EquivalenceProperties, OrderingRequirements, Partitioning, PhysicalExpr, PhysicalSortExpr, }; -use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, +}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, @@ -89,7 +92,7 @@ use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use crate::execution_plans::range_repartition_common::{ - broadcast_error, discover_cuts, split_batch_by_range, + discover_cuts, split_batch_by_range, wait_for_task, }; /// Per-output-partition channel capacity. Small = tight backpressure; the @@ -124,6 +127,12 @@ struct DispatchState { /// K slots. Each holds `Some(receiver)` after setup, `None` once its /// output partition has been consumed. receivers: Vec>>>, + /// P wait-tasks (one per input partition). Each awaits its scatter task + /// and turns panic/error/clean-end into an explicit signal on every + /// output channel. `SpawnedTask` aborts its inner tokio task on drop, so + /// holding these here ties the background work's lifetime to this exec's + /// — dropping the exec cancels all scatter work. + _drop_helper: Vec>, initialized: bool, } @@ -160,14 +169,24 @@ impl UnorderedRangeRepartitionExec { routing.expr ); } - let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(schema), - Partitioning::UnknownPartitioning(output_partitions), - input.pipeline_behavior(), - input.boundedness(), - )); + let properties = Arc::new( + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(output_partitions), + input.pipeline_behavior(), + input.boundedness(), + ) + // Scatter tasks eagerly drive their inputs the moment `execute()` + // is first called (not on-demand from downstream polls), and their + // per-output channels are the yield points the scheduler needs to + // trade CPU across tasks — same combination DataFusion's + // RepartitionExec uses. + .with_evaluation_type(EvaluationType::Eager) + .with_scheduling_type(SchedulingType::Cooperative), + ); let state = Arc::new(Mutex::new(DispatchState { receivers: Vec::new(), + _drop_helper: Vec::new(), initialized: false, })); Ok(Self { @@ -308,26 +327,36 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { let cuts_cell: Arc>> = Arc::new(OnceLock::new()); let input_partitions = self.input.output_partitioning().partition_count(); let routing_expr = self.order_by[0].expr.clone(); // TODO: KLL for multi-column? + let mut drop_helper = Vec::with_capacity(input_partitions); for input_partition in 0..input_partitions { let child = self.input.clone(); - let senders = senders.clone(); + let scatter_senders = senders.clone(); let cuts_cell = cuts_cell.clone(); let routing_expr = routing_expr.clone(); let ctx = ctx.clone(); let output_partitions = self.output_partitions; - tokio::spawn(async move { + let input_task = SpawnedTask::spawn(async move { scatter_input_partition( child, input_partition, ctx, routing_expr, - senders, + scatter_senders, cuts_cell, output_partitions, ) - .await; + .await }); + // A separate task awaits the scatter and surfaces panic / + // error / clean-end onto every output channel. Without this + // layer, a scatter panic drops its senders silently and + // downstream sees a spurious EOF. + let wait_senders = senders.clone(); + drop_helper.push(SpawnedTask::spawn(async move { + wait_for_task(input_task, wait_senders).await + })); } + state._drop_helper = drop_helper; } let Some(slot) = state.receivers.get_mut(partition) else { return internal_err!( @@ -357,6 +386,12 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { /// Every batch — including the first — is then dispatched via /// `split_batch_by_range`, which handles the degenerate empty-cuts case /// (all rows to partition 0) transparently. +/// +/// Errors are surfaced by returning `Err(..)`; the wrapping `wait_for_task` +/// broadcasts them to every output channel so downstream consumers observe +/// a failure rather than a silent EOF. Panics inside the body — including +/// panics from `split_batch_by_range` or the discovery walker — do the same +/// via the JoinError path. async fn scatter_input_partition( child: Arc, input_partition: usize, @@ -365,29 +400,17 @@ async fn scatter_input_partition( senders: Arc<[mpsc::Sender>]>, cuts_cell: Arc>>, output_partitions: usize, -) { - let mut stream = match child.execute(input_partition, ctx) { - Ok(s) => s, - Err(err) => { - broadcast_error(&senders, err).await; - return; - } - }; +) -> Result<()> { + let mut stream = child.execute(input_partition, ctx)?; while let Some(batch_result) = stream.next().await { // Every downstream consumer dropped its receiver — DRR itself was // dropped, or every output partition had a `LIMIT` above and hit it. // Either way, no one's listening; stop reading input so this scatter // task exits promptly (drops its `stream`, which propagates upstream). if senders.iter().all(|s| s.is_closed()) { - return; + return Ok(()); } - let batch = match batch_result { - Ok(b) => b, - Err(err) => { - broadcast_error(&senders, err).await; - return; - } - }; + let batch = batch_result?; let cuts = cuts_cell.get_or_init(|| { discover_cuts(&child, routing_expr.as_ref(), output_partitions) }); @@ -398,28 +421,21 @@ async fn scatter_input_partition( // `Arc` broadcast + receiver-side filter would skip // the scatter-side allocations at the cost of duplicating the // filter work K times. Worth measuring under skew. - match split_batch_by_range(&batch, &routing_expr, cuts) { - Ok(splits) => { - for (output, sub) in splits.into_iter().enumerate() { - if sub.num_rows() == 0 { - continue; - } - // `send().await` is where the backpressure lives: - // suspends when the channel is at capacity, which - // suspends this scatter task, which suspends its - // input read → propagates upstream. Errors mean the - // downstream dropped its receiver; keep forwarding - // to the other outputs. - let _ = senders[output].send(Ok(sub)).await; - } - } - Err(err) => { - broadcast_error(&senders, err).await; - return; + let splits = split_batch_by_range(&batch, &routing_expr, cuts)?; + for (output, sub) in splits.into_iter().enumerate() { + if sub.num_rows() == 0 { + continue; } + // `send().await` is where the backpressure lives: suspends + // when the channel is at capacity, which suspends this scatter + // task, which suspends its input read → propagates upstream. + // Errors mean the downstream dropped its receiver; keep + // forwarding to the other outputs. + let _ = senders[output].send(Ok(sub)).await; } } // All senders drop with this task → receivers see EOF. + Ok(()) } #[cfg(test)] @@ -856,4 +872,111 @@ mod tests { "error should name the invariant, got: {err}" ); } + + /// A panic inside the scatter task's polling of the child stream must + /// surface to every output partition as an `Err(..)` — not a silent EOF. + /// The wrapping wait-task converts `JoinError` on panic into a + /// `DataFusionError::Context("scatter task panicked", ..)` broadcast + /// across all K senders. + #[tokio::test] + async fn scatter_task_panic_surfaces_as_error() { + let schema = schema_v2_id(); + let source: Arc = Arc::new(PanickingSourceExec::new(&schema)); + let exec: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(source, vec![asc(&schema, "v2")], 3) + .unwrap(), + ); + let ctx = session(); + let output_partitions = exec.output_partitioning().partition_count(); + let streams: Vec<_> = (0..output_partitions) + .map(|p| exec.execute(p, ctx.task_ctx()).unwrap()) + .collect(); + let mut saw_panic_err = false; + for stream in streams { + let results: Vec> = + <_ as futures::stream::StreamExt>::collect(stream).await; + for r in results { + if let Err(err) = r + && err.to_string().contains("panicked") + { + saw_panic_err = true; + } + } + } + assert!( + saw_panic_err, + "at least one output partition must surface the scatter panic \ + as an Err — otherwise consumers see a spurious clean EOF" + ); + } + + /// Test-only source: its stream panics on the first poll. Used to + /// exercise the wait-task's JoinError → broadcast-error path. + #[derive(Debug)] + struct PanickingSourceExec { + schema: SchemaRef, + properties: Arc, + } + + impl PanickingSourceExec { + fn new(schema: &Arc) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(1), + datafusion::physical_plan::execution_plan::EmissionType::Incremental, + datafusion::physical_plan::execution_plan::Boundedness::Bounded, + )); + Self { + schema: schema.clone(), + properties, + } + } + } + + impl DisplayAs for PanickingSourceExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "PanickingSourceExec") + } + } + + impl ExecutionPlan for PanickingSourceExec { + fn name(&self) -> &str { + "PanickingSourceExec" + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _ctx: Arc, + ) -> Result { + let stream = futures::stream::once(async { + panic!("PanickingSourceExec stream boom"); + #[allow(unreachable_code)] + Ok(RecordBatch::new_empty(Arc::new(Schema::empty()))) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) + } + } } From 422ee6a2489e74f4ac5f6140c453e7e3085fa5c5 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Wed, 22 Jul 2026 05:04:03 -0600 Subject: [PATCH 5/5] refactor(core): collapse scatter guard to one task; share PanickingSourceExec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follows 6b300df2. Two changes: 1. Collapse two SpawnedTasks to one. DataFusion's model spawned an inner task for the scatter body and an outer task that awaited it — the outer task only existed to observe panics via `JoinError::is_panic()`. A single task using `AssertUnwindSafe(fut).catch_unwind()` does the same job in one spawn. Replaces `wait_for_task` with `guarded_scatter(fut, senders)` and adds `panic_payload_message` for `&str` / `String` payload downcasts so the original panic message survives. 2. Move `PanickingSourceExec` into `range_repartition_common::test_util`. The ordered variant will exercise the same `guarded_scatter` panic path; sharing the source keeps the two operator test files from re-declaring it. Also renamed the synthetic message from a joke to `PanickingSourceExec: synthetic test panic` (surfaced via a `SYNTHETIC_PANIC_MESSAGE` const so tests can assert-round-trip through the payload downcast). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../range_repartition_common.rs | 163 +++++++++++++++--- .../unordered_range_repartition.rs | 149 ++++++---------- 2 files changed, 185 insertions(+), 127 deletions(-) diff --git a/ballista/core/src/execution_plans/range_repartition_common.rs b/ballista/core/src/execution_plans/range_repartition_common.rs index 155ad90e3..4433172d1 100644 --- a/ballista/core/src/execution_plans/range_repartition_common.rs +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -36,17 +36,19 @@ //! [`UnorderedRangeRepartitionExec`]: super::UnorderedRangeRepartitionExec //! [`RuntimeStatsExec`]: super::RuntimeStatsExec +use std::any::Any; +use std::future::Future; +use std::panic::AssertUnwindSafe; use std::sync::Arc; use datafusion::arrow::array::{Array, Float64Array, RecordBatch, UInt32Array}; use datafusion::arrow::compute::take_arrays; -use datafusion::common::runtime::SpawnedTask; use datafusion::common::{Result, internal_datafusion_err}; -use datafusion::error::DataFusionError; 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; @@ -282,38 +284,145 @@ pub(super) async fn broadcast_error( } } -/// Await a scatter task and turn every terminal state into something the K -/// downstream consumers can observe. Modelled on DataFusion's -/// `RepartitionExec::wait_for_task`. +/// Run a scatter body and convert every terminal state — clean EOF, DFError, +/// or panic — into an explicit signal on the K output channels. Wraps the +/// future in `catch_unwind` so a panic inside `fut` becomes a broadcast +/// error rather than a silent sender drop (which downstream would misread +/// as a clean EOF). /// -/// Without this, a panic in a scatter task closes its senders with no data -/// sent — every output partition sees a clean EOF and the overall task -/// reports success while silently dropping all remaining rows. Here every -/// terminal state (panic, DFError, clean end) becomes an explicit signal -/// on every output channel. -pub(super) async fn wait_for_task( - input_task: SpawnedTask>, +/// `AssertUnwindSafe` is required because async futures aren't `UnwindSafe` +/// by default. Safe here because the scatter futures' captured state is +/// `Arc`s and owned locals — nothing observes post-panic state after we +/// broadcast and return. +pub(super) async fn guarded_scatter( + fut: F, senders: Arc<[mpsc::Sender>]>, -) { - match input_task.join().await { - Err(join_err) => { - let ctx = if join_err.is_panic() { - "scatter task panicked" - } else { - "scatter task cancelled" - }; +) where + F: Future>, +{ + match AssertUnwindSafe(fut).catch_unwind().await { + Ok(Ok(())) => { + // Drop the senders — receivers see clean EOF. + } + Ok(Err(err)) => broadcast_error(&senders, err).await, + Err(panic_payload) => { + let msg = panic_payload_message(panic_payload); broadcast_error( &senders, - DataFusionError::Context( - ctx.to_string(), - Box::new(DataFusionError::External(Box::new(join_err))), - ), + internal_datafusion_err!("scatter task panicked: {msg}"), ) .await; } - Ok(Err(err)) => broadcast_error(&senders, err).await, - Ok(Ok(())) => { - // Drop the senders — receivers see clean EOF. + } +} + +/// Extract a human-readable message from a `catch_unwind` panic payload. +/// `panic!("literal")` yields `&'static str`; `panic!("{x}")` yields +/// `String`. Anything else falls back to a placeholder. +fn panic_payload_message(payload: Box) -> String { + if let Some(s) = payload.downcast_ref::<&'static str>() { + (*s).to_string() + } else if let Some(s) = payload.downcast_ref::() { + s.clone() + } else { + "unknown panic payload".to_string() + } +} + +/// Shared test-only utilities for the two range-repartition operators. +/// The ordered variant will exercise the same `guarded_scatter` panic-path, +/// so the panic-source lives here rather than being duplicated per operator. +#[cfg(test)] +pub(super) mod test_util { + use std::fmt::{self, Formatter}; + use std::sync::Arc; + + use datafusion::arrow::array::RecordBatch; + use datafusion::arrow::datatypes::{Schema, SchemaRef}; + use datafusion::common::Result; + use datafusion::execution::TaskContext; + use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; + use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, + }; + + /// Message the synthetic panic will surface. Exposed so tests can + /// assert-round-trip: after `guarded_scatter`'s payload downcast, this + /// exact string must appear in the broadcast error. + pub(crate) const SYNTHETIC_PANIC_MESSAGE: &str = + "PanickingSourceExec: synthetic test panic"; + + /// Test-only source whose stream panics on the first poll. Used to + /// exercise `guarded_scatter`'s `catch_unwind` → broadcast-error path + /// end-to-end through a real `ExecutionPlan::execute()` call. + #[derive(Debug)] + pub(crate) struct PanickingSourceExec { + schema: SchemaRef, + properties: Arc, + } + + impl PanickingSourceExec { + pub(crate) fn new(schema: &Arc) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + schema: schema.clone(), + properties, + } + } + } + + impl DisplayAs for PanickingSourceExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "PanickingSourceExec") + } + } + + impl ExecutionPlan for PanickingSourceExec { + fn name(&self) -> &str { + "PanickingSourceExec" + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _ctx: Arc, + ) -> Result { + let stream = futures::stream::once(async { + panic!("{SYNTHETIC_PANIC_MESSAGE}"); + #[allow(unreachable_code)] + Ok(RecordBatch::new_empty(Arc::new(Schema::empty()))) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) } } } diff --git a/ballista/core/src/execution_plans/unordered_range_repartition.rs b/ballista/core/src/execution_plans/unordered_range_repartition.rs index b4f0b06bf..5d7719e80 100644 --- a/ballista/core/src/execution_plans/unordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/unordered_range_repartition.rs @@ -92,7 +92,7 @@ use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use crate::execution_plans::range_repartition_common::{ - discover_cuts, split_batch_by_range, wait_for_task, + discover_cuts, guarded_scatter, split_batch_by_range, }; /// Per-output-partition channel capacity. Small = tight backpressure; the @@ -127,11 +127,12 @@ struct DispatchState { /// K slots. Each holds `Some(receiver)` after setup, `None` once its /// output partition has been consumed. receivers: Vec>>>, - /// P wait-tasks (one per input partition). Each awaits its scatter task - /// and turns panic/error/clean-end into an explicit signal on every - /// output channel. `SpawnedTask` aborts its inner tokio task on drop, so - /// holding these here ties the background work's lifetime to this exec's - /// — dropping the exec cancels all scatter work. + /// P scatter tasks (one per input partition). Each runs its scatter + /// body inside `guarded_scatter`, which converts panic/error/clean-end + /// into an explicit signal on every output channel. `SpawnedTask` + /// aborts its inner tokio task on drop, so holding these here ties the + /// background work's lifetime to this exec's — dropping the exec + /// cancels all scatter work. _drop_helper: Vec>, initialized: bool, } @@ -331,11 +332,16 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { for input_partition in 0..input_partitions { let child = self.input.clone(); let scatter_senders = senders.clone(); + let guard_senders = senders.clone(); let cuts_cell = cuts_cell.clone(); let routing_expr = routing_expr.clone(); let ctx = ctx.clone(); let output_partitions = self.output_partitions; - let input_task = SpawnedTask::spawn(async move { + // `guarded_scatter` wraps the body in `catch_unwind`: a + // panic inside `scatter_input_partition` (or anything it + // calls) becomes a broadcast error rather than a silent + // sender-drop that downstream would misread as clean EOF. + drop_helper.push(SpawnedTask::spawn(guarded_scatter( scatter_input_partition( child, input_partition, @@ -344,17 +350,9 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { scatter_senders, cuts_cell, output_partitions, - ) - .await - }); - // A separate task awaits the scatter and surfaces panic / - // error / clean-end onto every output channel. Without this - // layer, a scatter panic drops its senders silently and - // downstream sees a spurious EOF. - let wait_senders = senders.clone(); - drop_helper.push(SpawnedTask::spawn(async move { - wait_for_task(input_task, wait_senders).await - })); + ), + guard_senders, + ))); } state._drop_helper = drop_helper; } @@ -875,11 +873,14 @@ mod tests { /// A panic inside the scatter task's polling of the child stream must /// surface to every output partition as an `Err(..)` — not a silent EOF. - /// The wrapping wait-task converts `JoinError` on panic into a - /// `DataFusionError::Context("scatter task panicked", ..)` broadcast - /// across all K senders. + /// `guarded_scatter`'s `catch_unwind` wraps the panic into a + /// `DataFusionError` broadcast across all K senders, and the payload + /// downcast preserves the original panic message. #[tokio::test] async fn scatter_task_panic_surfaces_as_error() { + use crate::execution_plans::range_repartition_common::test_util::{ + PanickingSourceExec, SYNTHETIC_PANIC_MESSAGE, + }; let schema = schema_v2_id(); let source: Arc = Arc::new(PanickingSourceExec::new(&schema)); let exec: Arc = Arc::new( @@ -891,92 +892,40 @@ mod tests { let streams: Vec<_> = (0..output_partitions) .map(|p| exec.execute(p, ctx.task_ctx()).unwrap()) .collect(); - let mut saw_panic_err = false; + let mut outputs_with_err = 0; + let mut outputs_carrying_original_message = 0; for stream in streams { let results: Vec> = <_ as futures::stream::StreamExt>::collect(stream).await; + let mut err_seen = false; + let mut msg_seen = false; for r in results { - if let Err(err) = r - && err.to_string().contains("panicked") - { - saw_panic_err = true; + if let Err(err) = r { + let text = err.to_string(); + if text.contains("panicked") { + err_seen = true; + } + if text.contains(SYNTHETIC_PANIC_MESSAGE) { + msg_seen = true; + } } } - } - assert!( - saw_panic_err, - "at least one output partition must surface the scatter panic \ - as an Err — otherwise consumers see a spurious clean EOF" - ); - } - - /// Test-only source: its stream panics on the first poll. Used to - /// exercise the wait-task's JoinError → broadcast-error path. - #[derive(Debug)] - struct PanickingSourceExec { - schema: SchemaRef, - properties: Arc, - } - - impl PanickingSourceExec { - fn new(schema: &Arc) -> Self { - let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(schema.clone()), - Partitioning::UnknownPartitioning(1), - datafusion::physical_plan::execution_plan::EmissionType::Incremental, - datafusion::physical_plan::execution_plan::Boundedness::Bounded, - )); - Self { - schema: schema.clone(), - properties, + if err_seen { + outputs_with_err += 1; + } + if msg_seen { + outputs_carrying_original_message += 1; } } - } - - impl DisplayAs for PanickingSourceExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "PanickingSourceExec") - } - } - - impl ExecutionPlan for PanickingSourceExec { - fn name(&self) -> &str { - "PanickingSourceExec" - } - - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - - fn properties(&self) -> &Arc { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Ok(self) - } - - fn execute( - &self, - _partition: usize, - _ctx: Arc, - ) -> Result { - let stream = futures::stream::once(async { - panic!("PanickingSourceExec stream boom"); - #[allow(unreachable_code)] - Ok(RecordBatch::new_empty(Arc::new(Schema::empty()))) - }); - Ok(Box::pin(RecordBatchStreamAdapter::new( - self.schema.clone(), - stream, - ))) - } + // Every output partition must see the panic — not a spurious EOF — + // and the original panic message must survive the payload downcast. + assert_eq!( + outputs_with_err, output_partitions, + "all output partitions must surface the scatter panic as Err" + ); + assert_eq!( + outputs_carrying_original_message, output_partitions, + "original panic message must survive catch_unwind payload extraction" + ); } }