From da2c6b524147272649cda3c72bafaa7414c9ce22 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 21 Jul 2026 08:47:29 -0600 Subject: [PATCH] =?UTF-8?q?feat(core):=20add=20UnorderedRangeRepartitionEx?= =?UTF-8?q?ec=20=E2=80=94=20value-range=20router=20over=20unordered=20inpu?= =?UTF-8?q?ts?= 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); + } }