From 9291bc59d832aa09265f695daadf7752da9659d2 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 9 Jul 2026 11:28:55 -0600 Subject: [PATCH 1/3] =?UTF-8?q?feat(core):=20add=20OrderedRangeRepartition?= =?UTF-8?q?Exec=20=E2=80=94=20N=20sorted=20overlapping=20=E2=86=92=20K=20s?= =?UTF-8?q?orted=20disjoint?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sibling of `UnorderedRangeRepartitionExec` (#2123). Same discovery mechanism (walk child subtree for matching `RuntimeStatsExec`, snapshot T-Digest, compute quantile cuts via `range_repartition_common::discover_cuts`); different execution model. Semantics: takes N locally-sorted input partitions whose value ranges may overlap; produces K range-disjoint output partitions where each output is fully sorted on the routing expression. Concatenating the K outputs in order yields a globally-sorted stream — exactly what Stage 2 of the parallel-window path needs to feed into Stage 3's BWAG. Execution model — N × K channels + K k-way merges: * N scatter tasks. Each reads one sorted input, calls `split_batch_by_range` on every batch, and forwards sub-batches to K per-input senders (one per output partition). * K output-partition streams. Each merges N sub-streams via DataFusion's `StreamingMergeBuilder` — SPM internals, no custom heap machinery here. * Channels: bounded `mpsc::channel(2)` per input × output for symmetric backpressure with the Unordered variant. Total buffered at rest ≤ N × K × 2 batches. Constructor requires `input.output_ordering()` to lead with the routing expression — otherwise the merger would produce garbled output. `PlanProperties::eq_properties` advertises each output partition as sorted on `order_by`, so downstream operators (BWAG, HaloDrop) can rely on that claim without an intervening `SortExec`. Proto + serde follow the Unordered pattern (`uint32 output_partitions` only; boundaries are runtime-discovered, stateless w.r.t. serde). Eight end-to-end tests: conservation across two overlapping inputs, each output sorted, outputs range-disjoint, no-stats fallback, constructor validation (empty order_by, unsorted input, mismatched sort key), and the double-execute guard. One serde roundtrip test on top. Drift from source branch (`brent/h2o-window-benchmarks`): * `RuntimeStatsExec` now rejects nullable routing expressions at construction (T-Digest has no NULL slot; tightened during #2123's upstreaming). Original tests declared `v2` as nullable and typed the batch builder as `Vec>` but never populated a `None`. Tightened schema + helper types to non-nullable `f64`; will loosen again when we swap T-Digest for KLL and NULL becomes a first-class routing key. Co-Authored-By: Claude Opus 4.7 (1M context) --- ballista/core/proto/ballista.proto | 17 + ballista/core/src/execution_plans/mod.rs | 2 + .../ordered_range_repartition.rs | 762 ++++++++++++++++++ ballista/core/src/serde/generated/ballista.rs | 24 +- ballista/core/src/serde/mod.rs | 94 ++- 5 files changed, 895 insertions(+), 4 deletions(-) create mode 100644 ballista/core/src/execution_plans/ordered_range_repartition.rs diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 76c47fc07..c2471c764 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -54,9 +54,26 @@ message BallistaPhysicalPlanNode { RuntimeStatsExecNode runtime_stats = 6; BufferExecNode buffer = 7; UnorderedRangeRepartitionExecNode unordered_range_repartition = 8; + OrderedRangeRepartitionExecNode ordered_range_repartition = 9; } } +// Value-range router over N locally-sorted overlapping input partitions. +// Redistributes them into K range-disjoint output partitions where each +// output is fully sorted on the routing expression. Uses N × K scatter +// channels feeding K per-output k-way merges (via DataFusion's SPM +// internals). Discovery of cut boundaries is identical to the unordered +// variant — walk the child subtree for a matching `RuntimeStatsExec`, snap +// its T-Digest. +message OrderedRangeRepartitionExecNode { + // Lexicographic ORDER BY. Routing keys off the first entry; the full + // ordering is preserved for downstream operators. Must match the + // input's declared `output_ordering` at plan time. + repeated datafusion.PhysicalSortExprNode order_by = 1; + // K — number of output partitions. Must be ≥ 2. + uint32 output_partitions = 2; +} + // Runtime-stats tap. Passes batches through unmodified while accumulating: // * per-partition row counts (always) // * a quantile sketch over the first ORDER BY expression's Float64 values diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index ce41627a7..7aec3cf47 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -22,6 +22,7 @@ mod buffer; mod chaos_exec; mod distributed_explain_analyze; mod distributed_query; +mod ordered_range_repartition; mod range_repartition_common; mod runtime_stats; mod shuffle_reader; @@ -38,6 +39,7 @@ pub use chaos_exec::ChaosExec; use datafusion::common::exec_err; pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; +pub use ordered_range_repartition::OrderedRangeRepartitionExec; pub use runtime_stats::{RuntimeStatsExec, sketch_from_proto, sketch_to_proto}; pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec}; pub use shuffle_reader::{stats_for_partition, stats_for_partitions}; diff --git a/ballista/core/src/execution_plans/ordered_range_repartition.rs b/ballista/core/src/execution_plans/ordered_range_repartition.rs new file mode 100644 index 000000000..7396c1303 --- /dev/null +++ b/ballista/core/src/execution_plans/ordered_range_repartition.rs @@ -0,0 +1,762 @@ +// 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 N locally-sorted overlapping input partitions. +//! Redistributes them into K range-disjoint output partitions where each +//! output is fully sorted on the routing expression. +//! +//! # Contrast with `UnorderedRangeRepartitionExec` +//! +//! The two operators share the same runtime discovery mechanism (walk child +//! subtree → find matching `RuntimeStatsExec` → quantile cuts from T-Digest, +//! see [`crate::execution_plans::range_repartition_common`]). They diverge +//! on execution model: +//! +//! - **Unordered**: N scatter tasks → K channels → K raw receivers. Rows +//! from different sources arrive interleaved, order is lost. +//! - **Ordered**: N scatter tasks → N × K channels → K [`StreamingMerge`] +//! streams. Each output p performs a heap-based k-way merge across N +//! sorted sub-streams (one per input source, pre-filtered to output p's +//! value range). Battle-tested merger via DataFusion's `SortPreservingMerge` +//! internals; no custom heap code here. +//! +//! # Discovery + fallback +//! +//! Identical to Unordered: walk `self.input`'s subtree through single-child +//! chains and the distribution-preserving whitelist; snapshot the matching +//! T-Digest on first batch; empty cut set = single-bucket fallback where all +//! rows land in output 0. +//! +//! # Ordering claim +//! +//! Constructor requires `input.output_ordering()` to lead with the routing +//! expression — otherwise the merger would produce garbled output. +//! [`PlanProperties::eq_properties`] declares each output partition sorted +//! on `order_by`, letting downstream operators (BWAG, HaloDrop) rely on the +//! claim without inserting a redundant `SortExec`. +//! +//! [`StreamingMerge`]: datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder + +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, LexOrdering, Partitioning, PhysicalExpr, PhysicalSortExpr, +}; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; +use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder; +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, per input source. Matches the +/// unordered variant's default; see the discussion there. Total buffered +/// batches at rest ≤ `N × K × CHANNEL_CAPACITY`. +const CHANNEL_CAPACITY: usize = 2; + +/// Value-range router that preserves the input's sort order. See the +/// module-level docs. +pub struct OrderedRangeRepartitionExec { + input: Arc, + /// Lexicographic ORDER BY. `try_new` guarantees the first entry evaluates + /// to `Float64` and matches the input's declared output ordering. + order_by: Vec, + /// K — number of output partitions. + output_partitions: usize, + /// Lazy channel + merge-stream setup guarded by `Mutex`. First + /// `execute()` call spawns the N input-reader tasks, wires N × K + /// channels, and constructs K `StreamingMerge` streams — one per + /// output partition. Subsequent `execute(p)` calls take + /// `merged_streams[p]`. + state: Arc>, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +/// Per-exec-instance lazy state. +struct DispatchState { + /// K slots. Each holds `Some(stream)` after setup, `None` once its + /// output partition has been consumed. Each stream is a k-way merge + /// (via `StreamingMerge`) across N sorted sub-streams — one per input + /// source, filtered to the output's value range on the scatter side. + merged_streams: Vec>, + initialized: bool, +} + +impl OrderedRangeRepartitionExec { + /// Wrap `input`. `order_by` must be non-empty, the first entry must + /// evaluate to `Float64`, and `input.output_ordering()` must lead with + /// the same expression (otherwise the merger produces garbled output). + 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!( + "OrderedRangeRepartitionExec 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!( + "OrderedRangeRepartitionExec routing expression `{}` must be Float64, got {:?}", + routing.expr, + routing_type + ); + } + // Input MUST claim to be sorted on our routing expression — otherwise + // the k-way merge produces garbled output. Sortedness of individual + // input partitions is enforced by the operator upstream (`SortExec` + // with `preserve_partitioning=true`); this check verifies the plan + // node declares that property. + let input_first_sort = input.output_ordering().map(|ordering| ordering.first()); + let Some(input_first) = input_first_sort else { + return internal_err!( + "OrderedRangeRepartitionExec requires sorted input — child plan claims no ordering" + ); + }; + if input_first.expr.as_ref() != routing.expr.as_ref() { + return internal_err!( + "OrderedRangeRepartitionExec: input's first sort key `{}` does not match \ + routing expression `{}`", + input_first.expr, + routing.expr + ); + } + // Advertise each output partition as sorted on `order_by`. Downstream + // operators (BWAG, HaloDrop) rely on this claim to skip redundant + // Sort insertions. + let eq_properties = EquivalenceProperties::new_with_orderings( + schema, + vec![LexOrdering::new(order_by.clone()).ok_or_else(|| { + internal_datafusion_err!( + "order_by is non-empty but LexOrdering rejected it" + ) + })?], + ); + let properties = Arc::new(PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(output_partitions), + input.pipeline_behavior(), + input.boundedness(), + )); + let state = Arc::new(Mutex::new(DispatchState { + merged_streams: Vec::new(), + initialized: false, + })); + Ok(Self { + input, + order_by, + output_partitions, + state, + properties, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + /// 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 OrderedRangeRepartitionExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("OrderedRangeRepartitionExec") + .field("order_by", &self.order_by) + .field("output_partitions", &self.output_partitions) + .finish() + } +} + +impl DisplayAs for OrderedRangeRepartitionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + let routing = &self.order_by[0]; + write!( + f, + "OrderedRangeRepartitionExec: routing={} {} → {} sorted partitions", + routing.expr, + if routing.options.descending { + "desc" + } else { + "asc" + }, + self.output_partitions, + ) + } +} + +impl ExecutionPlan for OrderedRangeRepartitionExec { + fn name(&self) -> &str { + "OrderedRangeRepartitionExec" + } + + 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!( + "OrderedRangeRepartitionExec expects exactly one child, got {}", + children.len() + ); + }; + Ok(Arc::new(OrderedRangeRepartitionExec::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 { + self.initialize_state(&mut state, &ctx)?; + } + let Some(slot) = state.merged_streams.get_mut(partition) else { + return internal_err!( + "OrderedRangeRepartitionExec: partition {} out of range (have {})", + partition, + self.output_partitions + ); + }; + slot.take().ok_or_else(|| { + internal_datafusion_err!( + "OrderedRangeRepartitionExec: execute({partition}) called twice \ + on the same instance" + ) + }) + } +} + +impl OrderedRangeRepartitionExec { + /// First-execute() setup: create N × K channels, spawn N scatter tasks, + /// build K `StreamingMerge` streams (one per output partition). Called + /// once, under the state `Mutex`. + fn initialize_state( + &self, + state: &mut DispatchState, + ctx: &Arc, + ) -> Result<()> { + let input_partitions = self.input.output_partitioning().partition_count(); + let schema = self.schema(); + let lex_ordering = LexOrdering::new(self.order_by.clone()).ok_or_else(|| { + internal_datafusion_err!("order_by is non-empty but LexOrdering rejected it") + })?; + + // Wire N × K channels. Two layouts of the same channels: + // • senders_per_input[i][p]: scatter task i's handle to output p. + // • receivers_per_output[p]: N receivers feeding output p's merge. + let mut senders_per_input: Vec>>> = + Vec::with_capacity(input_partitions); + let mut receivers_per_output: Vec>>> = (0 + ..self.output_partitions) + .map(|_| Vec::with_capacity(input_partitions)) + .collect(); + for _ in 0..input_partitions { + let mut senders = Vec::with_capacity(self.output_partitions); + for per_output_receivers in receivers_per_output.iter_mut() { + let (tx, rx) = mpsc::channel(CHANNEL_CAPACITY); + senders.push(tx); + per_output_receivers.push(rx); + } + senders_per_input.push(senders); + } + + // 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 routing_expr = self.order_by[0].expr.clone(); + for (input_partition, senders) in senders_per_input.into_iter().enumerate() { + let child = self.input.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; + }); + } + + // Build one `StreamingMerge` per output partition. Each merges N + // sub-streams (from each input source's channel to this output) + // in sort order. `StreamingMergeBuilder` is DataFusion's SPM + // internals — battle-tested cursor + heap machinery. + let mut merged_streams = Vec::with_capacity(self.output_partitions); + for (output_partition, receivers) in receivers_per_output.into_iter().enumerate() + { + let sub_streams: Vec = receivers + .into_iter() + .map(|rx| { + Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + ReceiverStream::new(rx), + )) as SendableRecordBatchStream + }) + .collect(); + let metrics = BaselineMetrics::new(&self.metrics, output_partition); + let merged = StreamingMergeBuilder::new() + .with_streams(sub_streams) + .with_schema(schema.clone()) + .with_expressions(&lex_ordering) + .with_batch_size(ctx.session_config().batch_size()) + .with_metrics(metrics) + .with_reservation( + datafusion::execution::memory_pool::MemoryConsumer::new(format!( + "OrderedRangeRepartitionExec[out={output_partition}]" + )) + .register(ctx.memory_pool()), + ) + .build()?; + merged_streams.push(Some(merged)); + } + state.merged_streams = merged_streams; + state.initialized = true; + Ok(()) + } +} + +/// One background task per input partition. Reads the sorted input; on the +/// first batch, the shared `OnceLock` learns the value cuts. Each batch is +/// then dispatched via `split_batch_by_range` to K per-input senders (one +/// per output partition). +/// +/// TODO: this is the naive path. Two follow-ups worth measuring: +/// 1. **Binary search on the batch head/tail against the cuts** to find +/// slice boundaries in O(log K) per batch instead of per-row +/// evaluation. Legal only because input is sorted — this operator's +/// whole precondition. +/// 2. **Send zero-copy Arrow slices** (`batch.slice(start, len)`) instead +/// of materialising via `take_arrays`. Turns per-batch scatter into a +/// few Arc bumps rather than N × K allocations under skew. +async fn scatter_input_partition( + child: Arc, + input_partition: usize, + ctx: Arc, + routing_expr: Arc, + senders: Vec>>, + 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 merger dropped its receiver — no one's listening. + // Stop reading so the input stream drops and backpressure propagates. + 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): input is sorted — this per-row `split_batch_by_range` + // is legal but wasteful. Two follow-ups worth measuring: + // 1. Binary-search batch head/tail against cuts to find slice + // boundaries in O(log K) per batch instead of per-row. + // 2. Send zero-copy `batch.slice(start, len)` instead of + // `take_arrays`-materialised sub-batches. Arc bumps replace + // allocations 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` provides backpressure: full channel + // suspends this task → suspends input read → propagates + // upstream. Error on send means downstream dropped its + // receiver; keep forwarding to other outputs. + let _ = senders[output].send(Ok(sub)).await; + } + } + Err(err) => { + broadcast_error(&senders, err).await; + return; + } + } + } + // Senders drop with this task → each output's merger sees EOF on that + // sub-stream. When all N sub-streams for output p reach EOF, the merger + // emits its final rows and returns None. +} + +#[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::sorts::sort::SortExec; + 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(), + } + } + + fn session() -> Arc { + Arc::new(SessionContext::new_with_state( + SessionStateBuilder::new().with_default_features().build(), + )) + } + + /// Feed pre-sorted per-partition test data through a `SortExec` + /// (preserve-partitioning) so the DRR's input carries the ordering + /// claim, then wrap in `RuntimeStatsExec` for the sketch source, then + /// the Ordered operator on top. + fn build_pipeline( + schema: &Arc, + partitions: Vec>, + k: usize, + ) -> Arc { + let batches: Vec> = partitions + .into_iter() + .map(|rows| { + let (keys, ids): (Vec<_>, Vec<_>) = rows.into_iter().unzip(); + vec![batch(schema, keys, ids)] + }) + .collect(); + let source = + MemorySourceConfig::try_new_exec(&batches, schema.clone(), None).unwrap(); + let sort_expr = asc(schema, "v2"); + let sort_lex = LexOrdering::new(vec![sort_expr.clone()]).unwrap(); + let sorted = + Arc::new(SortExec::new(sort_lex, source).with_preserve_partitioning(true)) + as Arc; + let stats = Arc::new( + RuntimeStatsExec::try_new(sorted, Some(vec![sort_expr.clone()])).unwrap(), + ) as Arc; + Arc::new(OrderedRangeRepartitionExec::try_new(stats, vec![sort_expr], k).unwrap()) + } + + async fn drain_concurrent( + exec: Arc, + ctx: Arc, + ) -> 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| { + let keys = b + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let ids = b + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + (0..b.num_rows()) + .map(|i| (keys.value(i), ids.value(i))) + .collect::>() + }) + .collect::>() + }) + }) + .collect(); + let mut per_output = Vec::with_capacity(output_partitions); + for handle in handles { + per_output.push(handle.await.unwrap()); + } + per_output + } + + // ---------- 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 = OrderedRangeRepartitionExec::try_new(empty_input(&schema), vec![], 4) + .expect_err("empty order_by must be rejected"); + assert!( + err.to_string().contains("at least one ORDER BY expression"), + "got: {err}" + ); + } + + #[test] + fn try_new_rejects_unsorted_input() { + let schema = schema_v2_id(); + // MemorySourceConfig with no declared ordering — output_ordering() is None. + let err = OrderedRangeRepartitionExec::try_new( + empty_input(&schema), + vec![asc(&schema, "v2")], + 4, + ) + .expect_err("input without ordering claim must be rejected"); + assert!( + err.to_string().contains("child plan claims no ordering"), + "got: {err}" + ); + } + + #[test] + fn try_new_rejects_mismatched_sort_key() { + let schema = schema_v2_id(); + // Sort input on `id` (Int64); DRR tries to route on `v2`. + let source = empty_input(&schema); + let id_sort = LexOrdering::new(vec![asc(&schema, "id")]).unwrap(); + let sorted_on_id = + Arc::new(SortExec::new(id_sort, source).with_preserve_partitioning(true)) + as Arc; + let err = OrderedRangeRepartitionExec::try_new( + sorted_on_id, + vec![asc(&schema, "v2")], + 4, + ) + .expect_err("mismatched sort key must be rejected"); + assert!( + err.to_string().contains("does not match routing"), + "got: {err}" + ); + } + + // ---------- End-to-end ----------------------------------------------- + + #[tokio::test] + async fn end_to_end_conserves_rows_across_two_overlapping_partitions() { + let schema = schema_v2_id(); + // Two partitions that overlap in value space (both cover [0, 40)): + // partition A: 0.0, 2.0, 4.0, ..., 38.0 (20 rows, even ids 0..20) + // partition B: 1.0, 3.0, 5.0, ..., 39.0 (20 rows, ids 100..120) + let part_a: Vec<(f64, i64)> = + (0..20).map(|i| ((i * 2) as f64, i as i64)).collect(); + let part_b: Vec<(f64, i64)> = (0..20) + .map(|i| ((i * 2 + 1) as f64, 100 + i as i64)) + .collect(); + let exec = build_pipeline(&schema, vec![part_a, part_b], 4); + let per_output = drain_concurrent(exec, session()).await; + let total: usize = per_output.iter().map(|v| v.len()).sum(); + assert_eq!(total, 40, "no row lost or duplicated"); + } + + /// Each output partition is *fully sorted* on the routing key — + /// that's the load-bearing claim of the operator. + #[tokio::test] + async fn end_to_end_each_output_is_sorted() { + let schema = schema_v2_id(); + let part_a: Vec<(f64, i64)> = + (0..25).map(|i| ((i * 2) as f64, i as i64)).collect(); + let part_b: Vec<(f64, i64)> = (0..25) + .map(|i| ((i * 2 + 1) as f64, 100 + i as i64)) + .collect(); + let exec = build_pipeline(&schema, vec![part_a, part_b], 4); + let per_output = drain_concurrent(exec, session()).await; + for (p, rows) in per_output.iter().enumerate() { + for pair in rows.windows(2) { + let [(left, _), (right, _)] = pair else { + unreachable!() + }; + assert!( + left <= right, + "output partition {p} not sorted: {left} > {right}" + ); + } + } + } + + /// Outputs must be range-disjoint — for every adjacent (p, p+1), + /// max(p) <= min(p+1). Combined with per-output sortedness this means + /// concatenating the K outputs produces a globally-sorted stream. + #[tokio::test] + async fn end_to_end_outputs_are_range_disjoint() { + let schema = schema_v2_id(); + let part_a: Vec<(f64, i64)> = + (0..30).map(|i| ((i * 2) as f64, i as i64)).collect(); + let part_b: Vec<(f64, i64)> = (0..30) + .map(|i| ((i * 2 + 1) as f64, 100 + i as i64)) + .collect(); + let part_c: Vec<(f64, i64)> = (0..30) + .map(|i| (((i * 2) + 100) as f64, 200 + i as i64)) + .collect(); + let exec = build_pipeline(&schema, vec![part_a, part_b, part_c], 4); + let per_output = drain_concurrent(exec, session()).await; + for window in per_output.windows(2) { + let [left, right] = window else { + unreachable!() + }; + if left.is_empty() || right.is_empty() { + continue; + } + let left_max = left + .iter() + .map(|(k, _)| *k) + .fold(f64::NEG_INFINITY, f64::max); + let right_min = right.iter().map(|(k, _)| *k).fold(f64::INFINITY, f64::min); + assert!( + left_max <= right_min, + "range-disjoint invariant broken: max(left)={left_max} > min(right)={right_min}" + ); + } + } + + /// Without a matching `RuntimeStatsExec` in the chain, discovery + /// returns empty cuts → single-bucket fallback → all rows to output 0. + #[tokio::test] + async fn end_to_end_fallback_when_no_runtime_stats() { + let schema = schema_v2_id(); + // Sort chain (needed for ordering claim) but no RuntimeStatsExec. + let rows: Vec<(f64, i64)> = (0..15).map(|i| (i as f64, i as i64)).collect(); + let source = MemorySourceConfig::try_new_exec( + &[vec![batch( + &schema, + rows.iter().map(|(k, _)| *k).collect(), + rows.iter().map(|(_, id)| *id).collect(), + )]], + schema.clone(), + None, + ) + .unwrap(); + let sort_expr = asc(&schema, "v2"); + let sort_lex = LexOrdering::new(vec![sort_expr.clone()]).unwrap(); + let sorted: Arc = + Arc::new(SortExec::new(sort_lex, source).with_preserve_partitioning(true)); + let exec: Arc = Arc::new( + OrderedRangeRepartitionExec::try_new(sorted, vec![sort_expr], 4).unwrap(), + ); + let per_output = drain_concurrent(exec, session()).await; + 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); + } + + #[tokio::test] + async fn execute_same_partition_twice_errors() { + let schema = schema_v2_id(); + let exec = build_pipeline(&schema, vec![vec![(1.0, 0)]], 3); + let ctx = session(); + let _ = drain_concurrent(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"), + "got: {err}" + ); + } +} diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 4143163b5..ca5b589b7 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, 8" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -57,8 +57,30 @@ pub mod ballista_physical_plan_node { Buffer(super::BufferExecNode), #[prost(message, tag = "8")] UnorderedRangeRepartition(super::UnorderedRangeRepartitionExecNode), + #[prost(message, tag = "9")] + OrderedRangeRepartition(super::OrderedRangeRepartitionExecNode), } } +/// Value-range router over N locally-sorted overlapping input partitions. +/// Redistributes them into K range-disjoint output partitions where each +/// output is fully sorted on the routing expression. Uses N × K scatter +/// channels feeding K per-output k-way merges (via DataFusion's SPM +/// internals). Discovery of cut boundaries is identical to the unordered +/// variant — walk the child subtree for a matching `RuntimeStatsExec`, snap +/// its T-Digest. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct OrderedRangeRepartitionExecNode { + /// Lexicographic ORDER BY. Routing keys off the first entry; the full + /// ordering is preserved for downstream operators. Must match the + /// input's declared `output_ordering` at plan time. + #[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, +} /// Runtime-stats tap. Passes batches through unmodified while accumulating: /// /// * per-partition row counts (always) diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index f5ab711d7..6be2b2984 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -56,9 +56,9 @@ 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, - UnorderedRangeRepartitionExec, UnresolvedShuffleExec, + BufferExec, BufferMode, ChaosExec, CoalescePlan, OrderedRangeRepartitionExec, + PartitionGroup, RuntimeStatsExec, ShuffleReaderExec, ShuffleWriterExec, + SortShuffleWriterExec, UnorderedRangeRepartitionExec, UnresolvedShuffleExec, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, @@ -617,6 +617,25 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { node.output_partitions as usize, )?)) } + PhysicalPlanType::OrderedRangeRepartition(node) => { + let [input] = inputs else { + return Err(DataFusionError::Internal(format!( + "OrderedRangeRepartitionExec 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(OrderedRangeRepartitionExec::try_new( + input.clone(), + order_by, + node.output_partitions as usize, + )?)) + } } } @@ -834,6 +853,27 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; 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::OrderedRangeRepartition( + protobuf::OrderedRangeRepartitionExecNode { + order_by, + output_partitions: exec.output_partitions() as u32, + }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode OrderedRangeRepartitionExec: {e:?}" + )) + })?; + Ok(()) } else { Err(DataFusionError::Internal(format!( "Unsupported plan node, name: [{}] ", @@ -1836,4 +1876,52 @@ mod test { ); assert!(!decoded_order_by[1].options.nulls_first); } + + /// `OrderedRangeRepartitionExec` requires input to declare a matching + /// sort ordering; wrap `EmptyExec` in a `SortExec` so the ordering + /// claim is present at construction time. + #[tokio::test] + async fn test_ordered_range_repartition_exec_roundtrip() { + use datafusion::arrow::compute::SortOptions; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::sorts::sort::SortExec; + + let schema = Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, false), + Field::new("id", DataType::Int64, false), + ])); + let leaf: 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 sort_lex = LexOrdering::new(vec![sort_expr.clone()]).unwrap(); + let input: Arc = + Arc::new(SortExec::new(sort_lex, leaf).with_preserve_partitioning(true)); + let original = OrderedRangeRepartitionExec::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 OrderedRangeRepartitionExec"); + let order_by = decoded.order_by(); + assert_eq!(order_by.len(), 1, "ORDER BY must round-trip"); + assert_eq!(order_by[0].expr.to_string(), sort_expr.expr.to_string()); + assert_eq!(decoded.output_partitions(), 8, "K must round-trip"); + } } From a20fc3e9b47c452eddf35b17ebbad45eb2497a97 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 24 Jul 2026 09:11:52 -0600 Subject: [PATCH 2/3] review: address review-preference audit for OrderedRangeRepartitionExec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the audit Phillip asked for on #2123 (commit a281667a on the unordered sibling). Preemptive polish before opening the PR for review. - Implement all 6 algebraic ExecutionPlan methods explicitly (required_input_ distribution, required_input_ordering, maintains_input_order, benefits_from_ input_partitioning, partition_statistics, cardinality_effect), each with a docstring justifying the value. Two differ materially from the unordered sibling: * required_input_ordering returns Some(order_by): the child MUST be sorted on the routing expression — the per-output StreamingMerge assumes each of its N sub-streams is locally sorted. Declaring the requirement lets the optimizer see it and skip a redundant SortExec above us. * maintains_input_order returns true: each output is a k-way merge across N sorted sub-streams filtered to a value range, so the relative order within each output faithfully preserves the child's sort. That's the whole point (contrast Unordered, which returns false because scattering breaks order). - Reject nullable routing expressions in try_new. Defense-in-depth against the no-stats-pair fallback path where discovery would silently fall back to single-bucket. TODO marker points at the KLL swap that lifts the restriction. The sibling RuntimeStatsExec already enforces the same invariant at plan-assembly time. - Add try_new_rejects_nullable_routing_key covering the new invariant. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ordered_range_repartition.rs | 99 ++++++++++++++++++- 1 file changed, 96 insertions(+), 3 deletions(-) diff --git a/ballista/core/src/execution_plans/ordered_range_repartition.rs b/ballista/core/src/execution_plans/ordered_range_repartition.rs index 7396c1303..6e693c2b1 100644 --- a/ballista/core/src/execution_plans/ordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/ordered_range_repartition.rs @@ -17,7 +17,28 @@ //! Value-range router over N locally-sorted overlapping input partitions. //! Redistributes them into K range-disjoint output partitions where each -//! output is fully sorted on the routing expression. +//! output is fully sorted on the routing expression. Concatenating the K +//! outputs in order yields a globally-sorted stream. +//! +//! ```text +//! N sorted input outputs (K = 4, each sorted) +//! partitions +//! ─┐ ┌─▶ 0: sorted, key < c₁ +//! ─┼──▶ RuntimeStatsExec ──▶ SortExec/preserve ──▶ OrderedRRE ─┼─▶ 1: sorted, c₁ ≤ key < c₂ +//! ─┤ (T-Digest tap) (per-partition sort) (K = 4) ├─▶ 2: sorted, c₂ ≤ key < c₃ +//! ─┘ │ └─▶ 3: sorted, c₃ ≤ key +//! ▲ │ +//! └──────────── walker ◀─────────────────┘ +//! (descends past sort-preserving +//! nodes, matches on routing expr, +//! reads K−1 quantile cuts) +//! +//! Inside OrderedRRE: N × K channels + K k-way merges. +//! Each of N scatter tasks splits its sorted input by the K-1 cuts +//! (K sub-batches per batch) and forwards each sub-batch to its output's +//! sender. Each of K outputs runs a `StreamingMerge` across the N +//! sub-streams targeting it → the output stream is sorted. +//! ``` //! //! # Contrast with `UnorderedRangeRepartitionExec` //! @@ -56,11 +77,13 @@ use std::sync::{Arc, Mutex, OnceLock}; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::{DataType, SchemaRef}; -use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{ - EquivalenceProperties, LexOrdering, Partitioning, PhysicalExpr, PhysicalSortExpr, + Distribution, EquivalenceProperties, LexOrdering, OrderingRequirements, Partitioning, + PhysicalExpr, PhysicalSortExpr, }; +use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; @@ -135,6 +158,14 @@ impl OrderedRangeRepartitionExec { 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!( + "OrderedRangeRepartitionExec: routing expression `{}` must be non-nullable", + routing.expr + ); + } // Input MUST claim to be sorted on our routing expression — otherwise // the k-way merge produces garbled output. Sortedness of individual // input partitions is enforced by the operator upstream (`SortExec` @@ -239,6 +270,50 @@ impl ExecutionPlan for OrderedRangeRepartitionExec { 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] + } + + /// The child MUST be sorted on the routing expression (and the full + /// `order_by` lex): the per-output `StreamingMerge` assumes each of its N + /// sub-streams is locally sorted. `try_new` also enforces this, but + /// declaring it here lets the optimizer see the requirement and avoid + /// inserting a redundant `SortExec` above us. + fn required_input_ordering(&self) -> Vec> { + let Some(lex) = LexOrdering::new(self.order_by.clone()) else { + return vec![None]; + }; + vec![Some(OrderingRequirements::new(lex.into()))] + } + + /// Each output partition is a k-way merge across N sorted sub-streams + /// pre-filtered to that output's value range. The relative order within + /// each output faithfully preserves the child's sort — that's the whole + /// point (contrast the Unordered variant, which returns `false`). + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + /// Extra input repartitioning above us buys nothing — we already read + /// all N 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>, @@ -612,6 +687,24 @@ mod tests { ); } + #[test] + fn try_new_rejects_nullable_routing_key() { + let schema = Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, true), // nullable + Field::new("id", DataType::Int64, false), + ])); + let err = OrderedRangeRepartitionExec::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 try_new_rejects_mismatched_sort_key() { let schema = schema_v2_id(); From a016902c03bb003f46b9512e5c7493ab93092faa Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 24 Jul 2026 09:28:55 -0600 Subject: [PATCH 3/3] review: surface scatter-task panics + set Eager/Cooperative on ORRE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports #2123's review-fix commits (6b300df2 + 422ee6a2) to the ordered sibling. 1. `scatter_input_partition` was spawned as a bare `tokio::spawn` whose `JoinHandle` was discarded. A panic inside the body silently dropped the sender, and downstream consumers saw a clean EOF while data was lost. Wrap it in `SpawnedTask::spawn(guarded_scatter(...))` (helper already shared via `range_repartition_common` — added for #2123 and consumed here verbatim). Hold the wait-tasks in `DispatchState._drop_helper` so dropping the exec aborts all N scatter tasks. Refactored the scatter body to return `Result<()>` instead of internally calling `broadcast_error` — guarded_scatter handles the broadcast on error/panic. 2. `PlanProperties::new` defaults to `Lazy` / `NonCooperative`; ORRE eagerly drives its inputs from the moment `execute()` is called and yields at per-output channel sends, so set `EvaluationType::Eager` + `SchedulingType::Cooperative` — same combination DataFusion's `RepartitionExec` uses. 3. New `scatter_task_panic_surfaces_as_error` test asserts every output partition surfaces the panic as `Err(..)` with the original message preserved. `OrderedRangeRepartitionExec::try_new` rejects inputs without a declared ordering, so extend `PanickingSourceExec` with a `with_ordering(order_by)` constructor that lets the panic-source satisfy the ordering-claim precondition. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ordered_range_repartition.rs | 168 +++++++++++++----- .../range_repartition_common.rs | 27 ++- 2 files changed, 148 insertions(+), 47 deletions(-) diff --git a/ballista/core/src/execution_plans/ordered_range_repartition.rs b/ballista/core/src/execution_plans/ordered_range_repartition.rs index 6e693c2b1..c040674da 100644 --- a/ballista/core/src/execution_plans/ordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/ordered_range_repartition.rs @@ -77,13 +77,16 @@ use std::sync::{Arc, Mutex, OnceLock}; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::runtime::SpawnedTask; use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{ Distribution, EquivalenceProperties, LexOrdering, OrderingRequirements, Partitioning, PhysicalExpr, PhysicalSortExpr, }; -use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, +}; use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; @@ -96,7 +99,7 @@ use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use crate::execution_plans::range_repartition_common::{ - broadcast_error, discover_cuts, split_batch_by_range, + discover_cuts, guarded_scatter, split_batch_by_range, }; /// Per-output-partition channel capacity, per input source. Matches the @@ -130,6 +133,13 @@ struct DispatchState { /// (via `StreamingMerge`) across N sorted sub-streams — one per input /// source, filtered to the output's value range on the scatter side. merged_streams: Vec>, + /// N 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, } @@ -196,14 +206,24 @@ impl OrderedRangeRepartitionExec { ) })?], ); - let properties = Arc::new(PlanProperties::new( - eq_properties, - Partitioning::UnknownPartitioning(output_partitions), - input.pipeline_behavior(), - input.boundedness(), - )); + let properties = Arc::new( + PlanProperties::new( + eq_properties, + 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 { merged_streams: Vec::new(), + _drop_helper: Vec::new(), initialized: false, })); Ok(Self { @@ -398,25 +418,37 @@ impl OrderedRangeRepartitionExec { // wins the `OnceLock::get_or_init` race. let cuts_cell: Arc>> = Arc::new(OnceLock::new()); let routing_expr = self.order_by[0].expr.clone(); + let mut drop_helper = Vec::with_capacity(input_partitions); for (input_partition, senders) in senders_per_input.into_iter().enumerate() { let child = self.input.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 { + // Move senders into an `Arc<[_]>` so scatter and guard can share + // the same slice — scatter uses them to forward sub-batches; + // guard uses them to broadcast panics/errors. + let scatter_senders: Arc<[mpsc::Sender>]> = + senders.into(); + let guard_senders = scatter_senders.clone(); + // `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, - senders, + scatter_senders, cuts_cell, output_partitions, - ) - .await; - }); + ), + guard_senders, + ))); } + state._drop_helper = drop_helper; // Build one `StreamingMerge` per output partition. Each merges N // sub-streams (from each input source's channel to this output) @@ -474,30 +506,18 @@ async fn scatter_input_partition( input_partition: usize, ctx: Arc, routing_expr: Arc, - senders: Vec>>, + senders: Arc<[mpsc::Sender>]>, cuts_cell: Arc>>, output_partitions: usize, -) { - let mut stream = match child.execute(input_partition, ctx) { - Ok(s) => s, - Err(err) => { - broadcast_error(&senders, err).await; - return; - } - }; +) -> Result<()> { + let mut stream = child.execute(input_partition, ctx)?; while let Some(batch_result) = stream.next().await { // Every downstream merger dropped its receiver — no one's listening. // Stop reading so the input stream drops and backpressure propagates. if senders.iter().all(|s| s.is_closed()) { - return; + return Ok(()); } - let batch = match batch_result { - Ok(b) => b, - Err(err) => { - broadcast_error(&senders, err).await; - return; - } - }; + let batch = batch_result?; let cuts = cuts_cell.get_or_init(|| { discover_cuts(&child, routing_expr.as_ref(), output_partitions) }); @@ -508,28 +528,22 @@ async fn scatter_input_partition( // 2. Send zero-copy `batch.slice(start, len)` instead of // `take_arrays`-materialised sub-batches. Arc bumps replace // allocations 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` provides backpressure: full channel - // suspends this task → suspends input read → propagates - // upstream. Error on send means downstream dropped its - // receiver; keep forwarding to other outputs. - let _ = senders[output].send(Ok(sub)).await; - } - } - Err(err) => { - broadcast_error(&senders, err).await; - return; + let splits = split_batch_by_range(&batch, &routing_expr, cuts)?; + for (output, sub) in splits.into_iter().enumerate() { + if sub.num_rows() == 0 { + continue; } + // `send().await` provides backpressure: full channel suspends + // this task → suspends input read → propagates upstream. Error + // on send means downstream dropped its receiver; keep forwarding + // to other outputs. + let _ = senders[output].send(Ok(sub)).await; } } // Senders drop with this task → each output's merger sees EOF on that // sub-stream. When all N sub-streams for output p reach EOF, the merger // emits its final rows and returns None. + Ok(()) } #[cfg(test)] @@ -852,4 +866,66 @@ mod tests { "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(); + // `OrderedRangeRepartitionExec::try_new` requires the child to + // declare an ordering; give the panic source one on `v2`. + let source: Arc = Arc::new( + PanickingSourceExec::with_ordering(&schema, vec![asc(&schema, "v2")]), + ); + let exec: Arc = Arc::new( + OrderedRangeRepartitionExec::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/execution_plans/range_repartition_common.rs b/ballista/core/src/execution_plans/range_repartition_common.rs index 4433172d1..30bdcb723 100644 --- a/ballista/core/src/execution_plans/range_repartition_common.rs +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -341,7 +341,9 @@ pub(super) mod test_util { use datafusion::arrow::datatypes::{Schema, SchemaRef}; use datafusion::common::Result; use datafusion::execution::TaskContext; - use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; + use datafusion::physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, PhysicalSortExpr, + }; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ @@ -377,6 +379,29 @@ pub(super) mod test_util { properties, } } + + /// Same as `new` but declares `order_by` as the output ordering. + /// `OrderedRangeRepartitionExec::try_new` rejects inputs without a + /// declared ordering, so its panic test needs this variant. + pub(crate) fn with_ordering( + schema: &Arc, + order_by: Vec, + ) -> Self { + let mut eq = EquivalenceProperties::new(schema.clone()); + if let Some(lex) = LexOrdering::new(order_by) { + eq.add_ordering(lex); + } + let properties = Arc::new(PlanProperties::new( + eq, + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + schema: schema.clone(), + properties, + } + } } impl DisplayAs for PanickingSourceExec {