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..4433172d1 --- /dev/null +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -0,0 +1,522 @@ +// 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::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::{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, +}; + +/// 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 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::().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() + ) + })?; + + // 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) { + 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; + } +} + +/// 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). +/// +/// `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>]>, +) 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, + internal_datafusion_err!("scatter task panicked: {msg}"), + ) + .await; + } + } +} + +/// 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, + ))) + } + } +} + +#[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..5d7719e80 --- /dev/null +++ b/ballista/core/src/execution_plans/unordered_range_repartition.rs @@ -0,0 +1,931 @@ +// 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. +//! +//! ```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 +//! 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::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, EvaluationType, SchedulingType, +}; +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::{ + discover_cuts, guarded_scatter, 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>>>, + /// 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, +} + +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 + ); + } + // 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), + 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 { + 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] + } + + /// 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>, + ) -> 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(); + 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 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; + // `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, + ctx, + routing_expr, + scatter_senders, + cuts_cell, + output_partitions, + ), + guard_senders, + ))); + } + state._drop_helper = drop_helper; + } + 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. +/// +/// 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, + ctx: Arc, + routing_expr: Arc, + senders: Arc<[mpsc::Sender>]>, + cuts_cell: Arc>>, + output_partitions: usize, +) -> 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 Ok(()); + } + let batch = batch_result?; + 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. + 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)] +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 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(); + 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}" + ); + } + + /// 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. + /// `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( + 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 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 { + let text = err.to_string(); + if text.contains("panicked") { + err_seen = true; + } + if text.contains(SYNTHETIC_PANIC_MESSAGE) { + msg_seen = true; + } + } + } + if err_seen { + outputs_with_err += 1; + } + if msg_seen { + outputs_carrying_original_message += 1; + } + } + // 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" + ); + } +} 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); + } }