From 8d858e2f689be16323dd53046e1ba308343feba0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 1 Jul 2026 11:16:11 -0600 Subject: [PATCH 1/4] feat(scheduler): add opt-in datafusion-scheduler crate modeling in-process distributed execution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `datafusion-scheduler`, an opt-in (workspace-excluded) crate that runs a physical `ExecutionPlan` the way a distributed engine would — split into stages at shuffle boundaries, each stage serialized via datafusion-proto, each task executed from a freshly deserialized plan with its own TaskContext, data crossing stage boundaries as Arrow-IPC frames over a streaming in-memory exchange — entirely in one process. It models the execution contract common to Ballista and datafusion-distributed (serialize per stage, instantiate a fresh plan per task, execute one partition in isolation) without any network/gRPC, so it can serve as a small reference example and, once wired into CI, a regression guard for physical-plan changes that assume shared in-process state. - Scheduler splits the plan at RepartitionExec(Hash) / CoalescePartitionsExec / SortPreservingMergeExec into a StageGraph (producers before consumers). - ExchangeSinkExec / ExchangeSourceExec move data over an InMemoryExchange of bounded channels carrying Arrow-IPC frames; ExchangeCodec (de)serializes them. - Spawn-all, no-barrier streaming executor; each task decodes its own plan with its own SessionContext / TaskContext. - run_distributed + assert_distributed_eq differential harness vs single-node collect; unit + equivalence + streaming tests. The crate is added to [workspace] exclude and only builds when explicitly targeted, so it does not affect normal builds or CI. --- Cargo.toml | 2 +- datafusion-scheduler/.gitignore | 1 + datafusion-scheduler/Cargo.toml | 40 ++ datafusion-scheduler/src/config.rs | 50 ++ datafusion-scheduler/src/exchange/codec.rs | 328 +++++++++++++ datafusion-scheduler/src/exchange/mod.rs | 415 ++++++++++++++++ datafusion-scheduler/src/exchange/sink.rs | 244 ++++++++++ datafusion-scheduler/src/exchange/source.rs | 205 ++++++++ datafusion-scheduler/src/executor.rs | 176 +++++++ datafusion-scheduler/src/lib.rs | 70 +++ datafusion-scheduler/src/scheduler.rs | 207 ++++++++ datafusion-scheduler/src/serde.rs | 43 ++ datafusion-scheduler/src/stage.rs | 59 +++ datafusion-scheduler/src/test_util.rs | 45 ++ datafusion-scheduler/tests/equivalence.rs | 299 ++++++++++++ .../tests/error_propagation.rs | 111 +++++ datafusion-scheduler/tests/single_stage.rs | 68 +++ datafusion-scheduler/tests/stage_splitting.rs | 458 ++++++++++++++++++ datafusion-scheduler/tests/streaming.rs | 125 +++++ datafusion-scheduler/tests/two_stage.rs | 118 +++++ 20 files changed, 3063 insertions(+), 1 deletion(-) create mode 100644 datafusion-scheduler/.gitignore create mode 100644 datafusion-scheduler/Cargo.toml create mode 100644 datafusion-scheduler/src/config.rs create mode 100644 datafusion-scheduler/src/exchange/codec.rs create mode 100644 datafusion-scheduler/src/exchange/mod.rs create mode 100644 datafusion-scheduler/src/exchange/sink.rs create mode 100644 datafusion-scheduler/src/exchange/source.rs create mode 100644 datafusion-scheduler/src/executor.rs create mode 100644 datafusion-scheduler/src/lib.rs create mode 100644 datafusion-scheduler/src/scheduler.rs create mode 100644 datafusion-scheduler/src/serde.rs create mode 100644 datafusion-scheduler/src/stage.rs create mode 100644 datafusion-scheduler/src/test_util.rs create mode 100644 datafusion-scheduler/tests/equivalence.rs create mode 100644 datafusion-scheduler/tests/error_propagation.rs create mode 100644 datafusion-scheduler/tests/single_stage.rs create mode 100644 datafusion-scheduler/tests/stage_splitting.rs create mode 100644 datafusion-scheduler/tests/streaming.rs create mode 100644 datafusion-scheduler/tests/two_stage.rs diff --git a/Cargo.toml b/Cargo.toml index ff5d3afcf48f5..389b736056fe3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,7 +67,7 @@ members = [ "datafusion/macros", "datafusion/doc", ] -exclude = ["dev/depcheck"] +exclude = ["dev/depcheck", "datafusion-scheduler"] resolver = "2" [workspace.package] diff --git a/datafusion-scheduler/.gitignore b/datafusion-scheduler/.gitignore new file mode 100644 index 0000000000000..03314f77b5aa4 --- /dev/null +++ b/datafusion-scheduler/.gitignore @@ -0,0 +1 @@ +Cargo.lock diff --git a/datafusion-scheduler/Cargo.toml b/datafusion-scheduler/Cargo.toml new file mode 100644 index 0000000000000..2a3cb3519f0b6 --- /dev/null +++ b/datafusion-scheduler/Cargo.toml @@ -0,0 +1,40 @@ +# 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. + +[package] +name = "datafusion-scheduler" +description = "In-process model of stage-based distributed execution for Apache DataFusion (test/regression harness)" +version = "0.1.0" +edition = "2024" +rust-version = "1.88.0" +license = "Apache-2.0" +publish = false + +[dependencies] +datafusion = { path = "../datafusion/core" } +datafusion-common = { path = "../datafusion/common" } +datafusion-common-runtime = { path = "../datafusion/common-runtime" } +datafusion-proto = { path = "../datafusion/proto" } +futures = "0.3" +tokio = { version = "1", features = ["rt-multi-thread", "macros"] } +prost = "0.14.1" # pinned to match datafusion-proto's resolved prost (verified via `cargo tree -i prost`) + +[dev-dependencies] +tempfile = "3" # only tests need a temp dir now (the #1907 CSV-scan regression guard) + +[lints.clippy] +all = "warn" diff --git a/datafusion-scheduler/src/config.rs b/datafusion-scheduler/src/config.rs new file mode 100644 index 0000000000000..063b2276b02f5 --- /dev/null +++ b/datafusion-scheduler/src/config.rs @@ -0,0 +1,50 @@ +// 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. + +use std::sync::Arc; + +use datafusion::prelude::SessionContext; + +/// Default per-channel frame capacity (backpressure bound). A small constant: +/// large enough to pipeline without pointless blocking, small enough to keep +/// the streaming model honest. +const DEFAULT_CHANNEL_CAPACITY: usize = 16; + +#[derive(Clone)] +pub struct SchedulerConfig { + /// Rebuilds an isolated SessionContext per task (UDFs, object stores, config). + pub session_builder: Arc SessionContext + Send + Sync>, + /// Per-channel frame capacity for the in-memory exchange (backpressure + /// bound). Bounded channels give the streaming analogue of a memory budget: + /// a fast producer blocks until the consumer drains. + pub channel_capacity: usize, +} + +impl SchedulerConfig { + /// Convenience for tests: a `session_builder` that clones the driver + /// context's state (so registered tables/UDFs survive rebuild) plus a + /// sensible default `channel_capacity`. + pub fn in_memory(ctx: &SessionContext) -> Self { + let state = ctx.state(); + Self { + session_builder: Arc::new(move || { + SessionContext::new_with_state(state.clone()) + }), + channel_capacity: DEFAULT_CHANNEL_CAPACITY, + } + } +} diff --git a/datafusion-scheduler/src/exchange/codec.rs b/datafusion-scheduler/src/exchange/codec.rs new file mode 100644 index 0000000000000..072b8c9ace0be --- /dev/null +++ b/datafusion-scheduler/src/exchange/codec.rs @@ -0,0 +1,328 @@ +// 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. + +//! `ExchangeCodec`: a `PhysicalExtensionCodec` that serializes/deserializes +//! [`ExchangeSinkExec`] and [`ExchangeSourceExec`] so a stage plan containing +//! them survives a `datafusion-proto` round-trip. +//! +//! The codec holds the shared [`InMemoryExchange`] and injects it on decode — +//! the exchange itself is not carried in the serialized bytes. This mirrors how +//! an executor knows its own transport endpoint independently of the plan it +//! receives (exactly as the old `ShuffleCodec` injected the shuffle dir). + +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; +use datafusion_common::not_impl_err; +use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; +use datafusion_proto::physical_plan::to_proto::serialize_partitioning; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalPlanDecodeContext, + PhysicalProtoConverterExtension, +}; +use datafusion_proto::protobuf; +use prost::Message; + +use super::InMemoryExchange; +use super::sink::ExchangeSinkExec; +use super::source::ExchangeSourceExec; + +const SINK_TAG: u8 = 0; +const SOURCE_TAG: u8 = 1; + +/// `PhysicalExtensionCodec` for [`ExchangeSinkExec`] and [`ExchangeSourceExec`]. +/// +/// Holds the shared [`InMemoryExchange`] that gets injected into decoded +/// operators; the exchange itself is not part of the serialized bytes — it +/// mirrors an executor knowing its own transport endpoint independently of the +/// plan it receives. +pub struct ExchangeCodec { + pub exchange: Arc, +} + +impl std::fmt::Debug for ExchangeCodec { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + // The exchange holds live channels and is intentionally opaque here. + f.debug_struct("ExchangeCodec").finish_non_exhaustive() + } +} + +fn write_u64(buf: &mut Vec, v: u64) { + buf.extend_from_slice(&v.to_le_bytes()); +} + +fn read_u64(buf: &[u8], offset: &mut usize) -> Result { + let end = offset.checked_add(8).ok_or_else(|| { + DataFusionError::Internal( + "ExchangeCodec: buffer too short reading u64".to_string(), + ) + })?; + let bytes: [u8; 8] = buf + .get(*offset..end) + .ok_or_else(|| { + DataFusionError::Internal( + "ExchangeCodec: buffer too short reading u64".to_string(), + ) + })? + .try_into() + .unwrap(); + *offset = end; + Ok(u64::from_le_bytes(bytes)) +} + +fn write_len_prefixed(buf: &mut Vec, bytes: &[u8]) { + write_u64(buf, bytes.len() as u64); + buf.extend_from_slice(bytes); +} + +fn read_len_prefixed<'a>(buf: &'a [u8], offset: &mut usize) -> Result<&'a [u8]> { + let len = read_u64(buf, offset)? as usize; + let end = offset.checked_add(len).ok_or_else(|| { + DataFusionError::Internal( + "ExchangeCodec: buffer too short reading length-prefixed bytes".to_string(), + ) + })?; + let bytes = buf.get(*offset..end).ok_or_else(|| { + DataFusionError::Internal( + "ExchangeCodec: buffer too short reading length-prefixed bytes".to_string(), + ) + })?; + *offset = end; + Ok(bytes) +} + +impl PhysicalExtensionCodec for ExchangeCodec { + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + if let Some(sink) = node.downcast_ref::() { + buf.push(SINK_TAG); + write_u64(buf, sink.stage_id() as u64); + let part_proto = serialize_partitioning( + sink.output_partitioning_spec(), + &DefaultPhysicalExtensionCodec {}, + proto_converter, + )?; + write_len_prefixed(buf, &part_proto.encode_to_vec()); + return Ok(()); + } + + if let Some(source) = node.downcast_ref::() { + buf.push(SOURCE_TAG); + write_u64(buf, source.from_stage_id() as u64); + write_u64(buf, source.num_producer_tasks() as u64); + write_u64(buf, source.output_partition_count() as u64); + let schema_proto: protobuf::Schema = source.schema().as_ref().try_into()?; + write_len_prefixed(buf, &schema_proto.encode_to_vec()); + return Ok(()); + } + + not_impl_err!("ExchangeCodec: unsupported plan node {}", node.name()) + } + + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let (tag, rest) = buf.split_first().ok_or_else(|| { + DataFusionError::Internal("ExchangeCodec: empty buffer".to_string()) + })?; + + match *tag { + SINK_TAG => { + let mut offset = 0usize; + let stage_id = read_u64(rest, &mut offset)? as usize; + let part_bytes = read_len_prefixed(rest, &mut offset)?; + let part_proto = + protobuf::Partitioning::decode(part_bytes).map_err(|e| { + DataFusionError::Internal(format!( + "ExchangeCodec: failed to decode Partitioning: {e}" + )) + })?; + + let input = inputs.first().cloned().ok_or_else(|| { + DataFusionError::Internal( + "ExchangeCodec: ExchangeSinkExec requires exactly one input" + .to_string(), + ) + })?; + + let default_codec = DefaultPhysicalExtensionCodec {}; + let decode_ctx = PhysicalPlanDecodeContext::new(ctx, &default_codec); + let partitioning: Partitioning = parse_protobuf_partitioning( + Some(&part_proto), + &decode_ctx, + input.schema().as_ref(), + proto_converter, + )? + .ok_or_else(|| { + DataFusionError::Internal( + "ExchangeCodec: missing Partitioning".to_string(), + ) + })?; + + let sink = ExchangeSinkExec::try_new( + stage_id, + input, + partitioning, + self.exchange.clone(), + )?; + Ok(Arc::new(sink)) + } + SOURCE_TAG => { + let mut offset = 0usize; + let from_stage_id = read_u64(rest, &mut offset)? as usize; + let num_producer_tasks = read_u64(rest, &mut offset)? as usize; + let output_partition_count = read_u64(rest, &mut offset)? as usize; + let schema_bytes = read_len_prefixed(rest, &mut offset)?; + let schema_proto = + protobuf::Schema::decode(schema_bytes).map_err(|e| { + DataFusionError::Internal(format!( + "ExchangeCodec: failed to decode Schema: {e}" + )) + })?; + let schema: datafusion::arrow::datatypes::Schema = + (&schema_proto).try_into()?; + + let source = ExchangeSourceExec::try_new( + from_stage_id, + Arc::new(schema), + num_producer_tasks, + output_partition_count, + self.exchange.clone(), + )?; + Ok(Arc::new(source)) + } + other => not_impl_err!("ExchangeCodec: unknown node tag {other}"), + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::{Column, col}; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::{ExecutionPlan, Partitioning}; + use datafusion::prelude::SessionContext; + + use super::{ExchangeCodec, read_len_prefixed, write_u64}; + use crate::exchange::{ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange}; + use crate::serde::{decode_plan, encode_plan}; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + #[tokio::test] + async fn sink_round_trips_through_exchange_codec() { + let schema = schema(); + let empty: Arc = Arc::new(EmptyExec::new(schema.clone())); + let partitioning = Partitioning::Hash(vec![col("a", &schema).unwrap()], 3); + + let exchange = InMemoryExchange::new(); + let plan: Arc = Arc::new( + ExchangeSinkExec::try_new(1, empty, partitioning, exchange.clone()).unwrap(), + ); + + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + + let codec = ExchangeCodec { + exchange: exchange.clone(), + }; + let bytes = encode_plan(&plan, &codec).unwrap(); + let decoded = decode_plan(&bytes, &task_ctx, &codec).unwrap(); + + let sink = decoded + .downcast_ref::() + .expect("decoded plan should be an ExchangeSinkExec"); + assert_eq!(sink.stage_id(), 1); + match sink.output_partitioning_spec() { + Partitioning::Hash(exprs, n) => { + assert_eq!(*n, 3); + assert_eq!(exprs.len(), 1); + let col = exprs[0] + .downcast_ref::() + .expect("hash partitioning expr should be a Column"); + assert_eq!(col.name(), "a"); + } + other => panic!( + "expected Partitioning::Hash([a], 3) after round-trip, got {other:?}" + ), + } + assert!( + sink.input().downcast_ref::().is_some(), + "child should round-trip as EmptyExec" + ); + } + + #[tokio::test] + async fn source_round_trips_through_exchange_codec() { + let schema = schema(); + let exchange = InMemoryExchange::new(); + + let plan: Arc = Arc::new( + ExchangeSourceExec::try_new(1, schema.clone(), 2, 3, exchange.clone()) + .unwrap(), + ); + + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + + let codec = ExchangeCodec { + exchange: exchange.clone(), + }; + let bytes = encode_plan(&plan, &codec).unwrap(); + let decoded = decode_plan(&bytes, &task_ctx, &codec).unwrap(); + + let source = decoded + .downcast_ref::() + .expect("decoded plan should be an ExchangeSourceExec"); + assert_eq!(source.from_stage_id(), 1); + assert_eq!(source.num_producer_tasks(), 2); + assert_eq!(source.output_partition_count(), 3); + assert_eq!(source.schema(), schema); + } + + #[test] + fn read_len_prefixed_rejects_malformed_length_without_panicking() { + // Length prefix claims far more bytes than remain in the buffer (and + // would overflow `offset + len` if computed unchecked). This must + // return an `Err`, not panic, even in debug builds. + let mut buf = Vec::new(); + write_u64(&mut buf, u64::MAX - 1); + buf.extend_from_slice(&[1, 2, 3]); + + let mut offset = 0usize; + let result = read_len_prefixed(&buf, &mut offset); + assert!( + result.is_err(), + "expected Err for malformed length prefix, got {result:?}" + ); + } +} diff --git a/datafusion-scheduler/src/exchange/mod.rs b/datafusion-scheduler/src/exchange/mod.rs new file mode 100644 index 0000000000000..87c30e7425003 --- /dev/null +++ b/datafusion-scheduler/src/exchange/mod.rs @@ -0,0 +1,415 @@ +// 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. + +//! Streaming in-memory exchange: an `Arc`-shared registry of bounded +//! single-producer / single-consumer channels that carry Arrow-IPC-encoded +//! batch frames between producer and consumer stages, with no disk and no +//! barrier. This is the in-process analogue of a network shuffle transport. +//! +//! A shuffle edge from a producer stage with `T` tasks to a consumer stage +//! that reads `P` output partitions is a grid of `T × P` channels, each keyed +//! `(producer_stage_id, producer_task, consumer_partition)`. The producer +//! task `t` owns the `Sender` of channel `(s, t, d)` for every output bucket +//! `d`; the consumer partition `d` owns the `Receiver` of channel `(s, t, d)` +//! for every producer task `t`. + +pub mod codec; +pub mod sink; +pub mod source; + +use std::collections::{HashMap, HashSet}; +use std::sync::Mutex; + +use datafusion::error::{DataFusionError, Result}; +use tokio::sync::mpsc::{Receiver, Sender, channel}; + +pub use codec::ExchangeCodec; +pub use sink::ExchangeSinkExec; +pub use source::ExchangeSourceExec; + +/// Key identifying a single channel: `(producer_stage_id, producer_task, +/// consumer_partition)`. +type ChannelKey = (usize, usize, usize); + +/// One IPC frame per message (`Vec` = one IPC-encoded [`RecordBatch`]). +/// +/// [`RecordBatch`]: datafusion::arrow::record_batch::RecordBatch +pub type Frame = Vec; + +/// A single SPSC channel's endpoints, either of which can be moved out exactly +/// once via [`InMemoryExchange::take_sender`] / [`InMemoryExchange::take_receiver`]. +struct ChannelSlot { + sender: Option>, + receiver: Option>, +} + +struct Inner { + channels: HashMap, + registered_stages: HashSet, +} + +/// A registry of bounded SPSC channels — "the wire". Shared as an +/// `Arc` by every task of every stage; it is the one piece +/// of state tasks share (like a Flight endpoint), and it is never serialized +/// into a plan. +pub struct InMemoryExchange { + inner: Mutex, +} + +impl InMemoryExchange { + pub fn new() -> std::sync::Arc { + std::sync::Arc::new(Self { + inner: Mutex::new(Inner { + channels: HashMap::new(), + registered_stages: HashSet::new(), + }), + }) + } + + /// Pre-create the `num_tasks × num_partitions` channels for one shuffle + /// edge, each bounded to `capacity` frames. Idempotent per `stage_id`: a + /// re-registration of an already-registered stage is a no-op (the existing + /// channels are left untouched). + pub fn register_stage( + &self, + stage_id: usize, + num_tasks: usize, + num_partitions: usize, + capacity: usize, + ) { + let mut inner = self.inner.lock().unwrap(); + if !inner.registered_stages.insert(stage_id) { + // Already registered — no-op. + return; + } + for task in 0..num_tasks { + for partition in 0..num_partitions { + let (tx, rx) = channel(capacity); + inner.channels.insert( + (stage_id, task, partition), + ChannelSlot { + sender: Some(tx), + receiver: Some(rx), + }, + ); + } + } + } + + /// Producer side: move out the `Sender` for `(stage_id, task, partition)`. + /// Errors if the channel was never registered or its sender was already + /// taken. + pub fn take_sender( + &self, + stage_id: usize, + task: usize, + partition: usize, + ) -> Result> { + let mut inner = self.inner.lock().unwrap(); + let slot = inner.channels.get_mut(&(stage_id, task, partition)).ok_or_else(|| { + DataFusionError::Internal(format!( + "exchange: no channel registered for sender (stage={stage_id}, task={task}, partition={partition})" + )) + })?; + slot.sender.take().ok_or_else(|| { + DataFusionError::Internal(format!( + "exchange: sender already taken for (stage={stage_id}, task={task}, partition={partition})" + )) + }) + } + + /// Consumer side: move out the `Receiver` for `(stage_id, task, + /// partition)`. Errors if the channel was never registered or its receiver + /// was already taken. + pub fn take_receiver( + &self, + stage_id: usize, + task: usize, + partition: usize, + ) -> Result> { + let mut inner = self.inner.lock().unwrap(); + let slot = inner.channels.get_mut(&(stage_id, task, partition)).ok_or_else(|| { + DataFusionError::Internal(format!( + "exchange: no channel registered for receiver (stage={stage_id}, task={task}, partition={partition})" + )) + })?; + slot.receiver.take().ok_or_else(|| { + DataFusionError::Internal(format!( + "exchange: receiver already taken for (stage={stage_id}, task={task}, partition={partition})" + )) + }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::sync::Arc; + use std::time::Duration; + + use datafusion::arrow::array::{Int32Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion::physical_expr::expressions::col; + use datafusion::physical_plan::test::TestMemoryExec; + use datafusion::physical_plan::{ExecutionPlan, Partitioning}; + use datafusion::prelude::SessionContext; + use datafusion_common_runtime::SpawnedTask; + use futures::StreamExt; + + use super::InMemoryExchange; + use super::sink::ExchangeSinkExec; + use super::source::ExchangeSourceExec; + + fn int32_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + /// Drives a sink stream to completion, returning the number of data + /// batches it yielded (must be 0). + async fn drain_sink( + mut stream: datafusion::physical_plan::SendableRecordBatchStream, + ) -> usize { + let mut n = 0; + while let Some(res) = stream.next().await { + res.unwrap(); + n += 1; + } + n + } + + /// Drives a source stream to completion, returning all Int32 values read. + async fn collect_source( + mut stream: datafusion::physical_plan::SendableRecordBatchStream, + ) -> Vec { + let mut values = Vec::new(); + while let Some(res) = stream.next().await { + let batch = res.unwrap(); + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend(col.iter().flatten()); + } + values + } + + /// Real streaming round-trip: 2 producer tasks hash-partition into 3 + /// buckets; 3 consumer partitions merge across both producers. Sink and + /// source tasks run CONCURRENTLY — the source blocks on its channels until + /// the sink feeds them, so both must be spawned together. + #[tokio::test] + async fn streaming_round_trip_across_concurrent_sink_and_source() { + let exchange = InMemoryExchange::new(); + let stage_id = 9; + let num_tasks = 2; + let num_partitions = 3; + exchange.register_stage(stage_id, num_tasks, num_partitions, 4); + + let schema = int32_schema(); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![0, 1, 2]))], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![3, 4, 5]))], + ) + .unwrap(); + let input = TestMemoryExec::try_new_exec( + &[vec![batch0], vec![batch1]], + schema.clone(), + None, + ) + .unwrap(); + + let partitioning = + Partitioning::Hash(vec![col("a", &schema).unwrap()], num_partitions); + let sink = Arc::new( + ExchangeSinkExec::try_new(stage_id, input, partitioning, exchange.clone()) + .unwrap(), + ); + let source = Arc::new( + ExchangeSourceExec::try_new( + stage_id, + schema.clone(), + num_tasks, + num_partitions, + exchange.clone(), + ) + .unwrap(), + ); + + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + + // Spawn ALL tasks (sinks + sources) before joining any: a sink blocks + // on `send().await` until a source drains, so joining sinks first + // would deadlock. + let mut sink_handles = Vec::new(); + for t in 0..num_tasks { + let stream = sink.execute(t, task_ctx.clone()).unwrap(); + sink_handles.push(SpawnedTask::spawn(drain_sink(stream))); + } + let mut source_handles = Vec::new(); + for d in 0..num_partitions { + let stream = source.execute(d, task_ctx.clone()).unwrap(); + source_handles.push(SpawnedTask::spawn(collect_source(stream))); + } + + // Sinks must yield zero data batches. + for h in sink_handles { + let n = h.join().await.unwrap(); + assert_eq!(n, 0, "ExchangeSinkExec must yield zero batches"); + } + + // Collect per-partition outputs and check hash-partition consistency: + // every value must land in the bucket its hash selects. + let mut partitioner = + datafusion::physical_plan::repartition::BatchPartitioner::try_new( + Partitioning::Hash(vec![col("a", &schema).unwrap()], num_partitions), + datafusion::physical_plan::metrics::MetricBuilder::new( + &datafusion::physical_plan::metrics::ExecutionPlanMetricsSet::new(), + ) + .subset_time("t", 0), + 0, + 1, + ) + .unwrap(); + + let mut all_values: Vec = Vec::new(); + for (d, h) in source_handles.into_iter().enumerate() { + let values = h.join().await.unwrap(); + for v in &values { + let one = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![*v]))], + ) + .unwrap(); + let mut bucket = None; + for res in partitioner.partition_iter(one).unwrap() { + let (b, batch) = res.unwrap(); + if batch.num_rows() > 0 { + bucket = Some(b); + } + } + assert_eq!( + bucket, + Some(d), + "value {v} arrived at partition {d} but hashes to bucket {bucket:?}" + ); + } + all_values.extend(values); + } + + assert_eq!(all_values.len(), 6, "expected all 6 input rows"); + let got: HashSet = all_values.into_iter().collect(); + let expected: HashSet = [0, 1, 2, 3, 4, 5].into_iter().collect(); + assert_eq!(got, expected); + } + + /// Tight-backpressure test: capacity=1 with many rows, so producers block + /// on full channels while consumers drain. Must complete (not deadlock); + /// a `timeout` turns a hang into a failure instead of hanging CI. + #[tokio::test] + async fn tight_backpressure_does_not_deadlock() { + let result = tokio::time::timeout(Duration::from_secs(30), async { + let exchange = InMemoryExchange::new(); + let stage_id = 3; + let num_tasks = 2; + let num_partitions = 4; + // Capacity 1: the tightest bound that still makes progress. + exchange.register_stage(stage_id, num_tasks, num_partitions, 1); + + let schema = int32_schema(); + // 200 rows per task, spread across many small batches so the sink + // blocks on `send().await` repeatedly. + let mut task0 = Vec::new(); + let mut task1 = Vec::new(); + for i in 0..100i32 { + task0.push( + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![i, i + 1000]))], + ) + .unwrap(), + ); + task1.push( + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![i + 100_000, i + 200_000]))], + ) + .unwrap(), + ); + } + let input = + TestMemoryExec::try_new_exec(&[task0, task1], schema.clone(), None) + .unwrap(); + + let partitioning = + Partitioning::Hash(vec![col("a", &schema).unwrap()], num_partitions); + let sink = Arc::new( + ExchangeSinkExec::try_new( + stage_id, + input, + partitioning, + exchange.clone(), + ) + .unwrap(), + ); + let source = Arc::new( + ExchangeSourceExec::try_new( + stage_id, + schema.clone(), + num_tasks, + num_partitions, + exchange.clone(), + ) + .unwrap(), + ); + + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + + let mut sink_handles = Vec::new(); + for t in 0..num_tasks { + let stream = sink.execute(t, task_ctx.clone()).unwrap(); + sink_handles.push(SpawnedTask::spawn(drain_sink(stream))); + } + let mut source_handles = Vec::new(); + for d in 0..num_partitions { + let stream = source.execute(d, task_ctx.clone()).unwrap(); + source_handles.push(SpawnedTask::spawn(collect_source(stream))); + } + + for h in sink_handles { + assert_eq!(h.join().await.unwrap(), 0); + } + let mut total = 0usize; + for h in source_handles { + total += h.join().await.unwrap().len(); + } + total + }) + .await; + + let total = result.expect("backpressure round-trip deadlocked / timed out"); + // 2 tasks × 200 rows each. + assert_eq!(total, 400, "expected all 400 rows to round-trip"); + } +} diff --git a/datafusion-scheduler/src/exchange/sink.rs b/datafusion-scheduler/src/exchange/sink.rs new file mode 100644 index 0000000000000..eb9d7bab003cf --- /dev/null +++ b/datafusion-scheduler/src/exchange/sink.rs @@ -0,0 +1,244 @@ +// 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. + +//! `ExchangeSinkExec`: the producer side of a streaming in-memory exchange. +//! Runs one input partition (task), hash-partitions its rows into `P` output +//! buckets, IPC-encodes each bucket batch into a frame, and pushes it into the +//! corresponding channel of the shared [`InMemoryExchange`]. Yields zero data +//! batches — it is a pure side-effecting sink, the streaming replacement for +//! `ShuffleWriterExec`. + +use std::fmt; +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::ipc::writer::StreamWriter; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::error::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder}; +use datafusion::physical_plan::repartition::BatchPartitioner; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, SendableRecordBatchStream, +}; +use futures::StreamExt; +use tokio::sync::mpsc::Sender; + +use super::{Frame, InMemoryExchange}; + +/// Producer operator for a shuffle edge. For task `t` it takes the `P` senders +/// `(stage_id, t, d)` for `d in 0..P`, runs input partition `t`, routes each +/// batch through a [`BatchPartitioner`] into `P` buckets, IPC-encodes each +/// bucket batch to a frame, and awaits `send` on the bucket's channel (the +/// await is the backpressure point). When the input ends it drops all senders, +/// signalling end-of-stream to the consumers. +/// +/// This operator's own output partitioning mirrors its *input's* partitioning +/// (task count follows the input); `output_partitioning` only controls how many +/// buckets each task fans out into. +pub struct ExchangeSinkExec { + stage_id: usize, + input: Arc, + output_partitioning: Partitioning, + exchange: Arc, + properties: Arc, +} + +impl ExchangeSinkExec { + pub fn try_new( + stage_id: usize, + input: Arc, + output_partitioning: Partitioning, + exchange: Arc, + ) -> Result { + let properties = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + Ok(Self { + stage_id, + input, + output_partitioning, + exchange, + properties, + }) + } + + pub fn stage_id(&self) -> usize { + self.stage_id + } + + pub fn input(&self) -> &Arc { + &self.input + } + + pub fn output_partitioning_spec(&self) -> &Partitioning { + &self.output_partitioning + } +} + +impl fmt::Debug for ExchangeSinkExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("ExchangeSinkExec") + .field("stage_id", &self.stage_id) + .field("output_partitioning", &self.output_partitioning) + .finish() + } +} + +impl DisplayAs for ExchangeSinkExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "ExchangeSinkExec: stage={}, output_partitioning={:?}", + self.stage_id, self.output_partitioning + ) + } + DisplayFormatType::TreeRender => { + writeln!(f, "stage={}", self.stage_id) + } + } + } +} + +impl ExecutionPlan for ExchangeSinkExec { + fn name(&self) -> &'static str { + "ExchangeSinkExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(ExchangeSinkExec::try_new( + self.stage_id, + children.swap_remove(0), + self.output_partitioning.clone(), + Arc::clone(&self.exchange), + )?)) + } + + /// Streams input partition `partition` into the exchange, then completes + /// with zero data batches. + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let num_output_partitions = self.output_partitioning.partition_count(); + + // Take this task's `P` senders up front; a missing/taken sender is a + // wiring bug and fails here, before any input is polled. + let mut senders = Vec::with_capacity(num_output_partitions); + for d in 0..num_output_partitions { + senders.push(self.exchange.take_sender(self.stage_id, partition, d)?); + } + + let input_stream = self.input.execute(partition, context)?; + let schema = self.input.schema(); + let num_input_partitions = self.input.output_partitioning().partition_count(); + + let fut = run_sink( + input_stream, + schema.clone(), + self.output_partitioning.clone(), + senders, + partition, + num_input_partitions, + ); + + let stream = futures::stream::once(fut).filter_map(|res| async move { + match res { + Ok(()) => None, + Err(e) => Some(Err(e)), + } + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} + +/// Drains `input`, hash-partitioning each batch into buckets, IPC-encoding +/// each bucket batch to a frame and awaiting its send on the bucket's channel. +/// Returns after the input ends; dropping `senders` on return closes the +/// channels, signalling end-of-stream to the consumers. +async fn run_sink( + mut input: SendableRecordBatchStream, + schema: SchemaRef, + partitioning: Partitioning, + senders: Vec>, + task_id: usize, + num_input_partitions: usize, +) -> Result<()> { + let metrics_set = ExecutionPlanMetricsSet::new(); + let timer = MetricBuilder::new(&metrics_set).subset_time("repartition_time", task_id); + let mut partitioner = + BatchPartitioner::try_new(partitioning, timer, task_id, num_input_partitions)?; + + while let Some(batch) = input.next().await { + // A genuine error from this sink's own input stream still propagates + // as a query error. + let batch = batch?; + for res in partitioner.partition_iter(batch)? { + let (bucket, part_batch) = res?; + let frame = encode_frame(&part_batch, &schema)?; + // `Sender::send` errs ONLY when the receiver is gone: the consumer + // hung up (e.g. a LIMIT upstream stopped pulling, or the query is + // tearing down). That is not a query error — it is a graceful stop + // of this sink. Stop sending and end cleanly; dropping `senders` + // closes the remaining channels for the other consumers. + if senders[bucket].send(frame).await.is_err() { + return Ok(()); + } + } + } + + // Explicit for clarity: closing every sender signals end-of-stream. + drop(senders); + Ok(()) +} + +/// IPC-encodes a single [`RecordBatch`] into one self-contained frame using an +/// Arrow `StreamWriter` (schema + one batch + end marker) written into a +/// `Vec`. +fn encode_frame(batch: &RecordBatch, schema: &SchemaRef) -> Result { + let mut buf: Vec = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, schema)?; + writer.write(batch)?; + writer.finish()?; + } + Ok(buf) +} diff --git a/datafusion-scheduler/src/exchange/source.rs b/datafusion-scheduler/src/exchange/source.rs new file mode 100644 index 0000000000000..252d96316c1b6 --- /dev/null +++ b/datafusion-scheduler/src/exchange/source.rs @@ -0,0 +1,205 @@ +// 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. + +//! `ExchangeSourceExec`: the consumer leaf of a streaming in-memory exchange. +//! For output partition `d` it takes the receivers `(from_stage_id, t, d)` for +//! every producer task `t`, merges them, IPC-decodes each frame back into +//! `RecordBatch`es, and streams them out. The streaming replacement for +//! `ShuffleReaderExec`. + +use std::fmt; +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::ipc::reader::StreamReader; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::error::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + SendableRecordBatchStream, +}; +use futures::stream::{self, StreamExt}; +use tokio::sync::mpsc::Receiver; + +use super::{Frame, InMemoryExchange}; + +/// Consumer leaf for a shuffle edge. For output partition `d` it merges the +/// `T` channels `(from_stage_id, t, d)` (one per producer task) into a single +/// stream, decoding each IPC frame back into `RecordBatch`es. The merged +/// stream ends only when all `T` producers have closed their senders. +pub struct ExchangeSourceExec { + from_stage_id: usize, + schema: SchemaRef, + num_producer_tasks: usize, + output_partition_count: usize, + exchange: Arc, + properties: Arc, +} + +impl ExchangeSourceExec { + pub fn try_new( + from_stage_id: usize, + schema: SchemaRef, + num_producer_tasks: usize, + output_partition_count: usize, + exchange: Arc, + ) -> Result { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(output_partition_count), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Self { + from_stage_id, + schema, + num_producer_tasks, + output_partition_count, + exchange, + properties, + }) + } + + pub fn from_stage_id(&self) -> usize { + self.from_stage_id + } + + pub fn num_producer_tasks(&self) -> usize { + self.num_producer_tasks + } + + pub fn output_partition_count(&self) -> usize { + self.output_partition_count + } +} + +impl fmt::Debug for ExchangeSourceExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("ExchangeSourceExec") + .field("from_stage_id", &self.from_stage_id) + .field("num_producer_tasks", &self.num_producer_tasks) + .field("output_partition_count", &self.output_partition_count) + .finish() + } +} + +impl DisplayAs for ExchangeSourceExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "ExchangeSourceExec: from_stage={}, num_producer_tasks={}, output_partition_count={}", + self.from_stage_id, + self.num_producer_tasks, + self.output_partition_count + ) + } + DisplayFormatType::TreeRender => { + writeln!(f, "from_stage={}", self.from_stage_id) + } + } + } +} + +impl ExecutionPlan for ExchangeSourceExec { + fn name(&self) -> &'static str { + "ExchangeSourceExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + // Leaf node: nothing to swap. + Ok(self) + } + + /// Merges every producer task's channel for output partition `partition`, + /// decoding IPC frames back into record batches as they arrive. + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + // Take one receiver per producer task, up front — a missing/taken + // receiver is a wiring bug and fails here. + let mut receivers = Vec::with_capacity(self.num_producer_tasks); + for t in 0..self.num_producer_tasks { + receivers.push(self.exchange.take_receiver( + self.from_stage_id, + t, + partition, + )?); + } + + // Turn each receiver into a stream of frames via `unfold` (no + // `tokio-stream` dep), then `select_all`-merge them: the merged stream + // is exhausted only once every producer has closed its sender. + let frame_streams = receivers + .into_iter() + .map(|rx| receiver_to_stream(rx).boxed()); + let merged = stream::select_all(frame_streams); + + // Decode each frame (one frame -> one or more batches) and flatten. + let batch_stream = merged.flat_map(|frame| { + let decoded: Vec> = match decode_frame(&frame) { + Ok(batches) => batches.into_iter().map(Ok).collect(), + Err(e) => vec![Err(e)], + }; + stream::iter(decoded) + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + batch_stream, + ))) + } +} + +/// Turns a bounded mpsc `Receiver` into a `Stream` that ends +/// when the channel is closed. +fn receiver_to_stream(rx: Receiver) -> impl futures::Stream { + stream::unfold(rx, |mut rx| async move { rx.recv().await.map(|f| (f, rx)) }) +} + +/// IPC-decodes one frame (produced by the sink's `StreamWriter`) back into its +/// record batches. +fn decode_frame(frame: &[u8]) -> Result> { + let reader = StreamReader::try_new(std::io::Cursor::new(frame), None)?; + let mut batches = Vec::new(); + for batch in reader { + batches.push(batch?); + } + Ok(batches) +} diff --git a/datafusion-scheduler/src/executor.rs b/datafusion-scheduler/src/executor.rs new file mode 100644 index 0000000000000..64b82e8cbe5a4 --- /dev/null +++ b/datafusion-scheduler/src/executor.rs @@ -0,0 +1,176 @@ +// 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. + +//! Streaming, no-barrier executor for a [`StageGraph`]. Registers every +//! shuffle edge's channels on the shared [`InMemoryExchange`], then spawns +//! EVERY task of EVERY stage concurrently and only then awaits them. +//! +//! There is no per-stage barrier. A consumer task blocked on a channel is +//! unblocked by a producer task that is already running; the stage DAG is +//! acyclic (producers have lower ids) and the channels are bounded SPSC, so +//! spawn-all-then-await over that dataflow cannot deadlock. Only the final +//! stage's task outputs are collected (in ascending task order); producer +//! stages drive the exchange as a side effect. + +use std::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::error::{DataFusionError, Result}; +use datafusion_common_runtime::SpawnedTask; +use futures::StreamExt; + +use crate::config::SchedulerConfig; +use crate::exchange::{ExchangeCodec, InMemoryExchange}; +use crate::serde::{decode_plan, encode_plan}; +use crate::stage::StageGraph; + +/// Number of output partitions (== number of tasks) a stage's plan produces. +fn stage_task_count(plan: &Arc) -> usize { + plan.properties().output_partitioning().partition_count() +} + +/// Executes `graph` over the shared streaming `exchange`. +/// +/// 1. Register the channels for every producer stage (every stage except the +/// final one) on `exchange` BEFORE spawning any task — a consumer must be +/// able to take its receiver the moment it starts. +/// 2. Encode each stage's plan once, then spawn ALL tasks of ALL stages +/// concurrently (`SpawnedTask`). Each task rebuilds its own +/// `SessionContext`, decodes its own fresh plan instance with a fresh +/// `ExchangeCodec`, executes its one partition, and drains fully. Final +/// tasks keep their batches; producer tasks discard theirs (their work is +/// the frames they push into the exchange). +/// 3. Await ALL handles — no barrier. Propagate the first genuine +/// execution/join error; otherwise return the final stage's batches in +/// ascending task order. +pub async fn execute_stage_graph( + graph: &StageGraph, + exchange: &Arc, + config: &SchedulerConfig, +) -> Result> { + // Register every producer stage's channels up front, before any task is + // spawned. The final stage produces the query result and has no sink, so + // it registers nothing. + for stage in &graph.stages { + if stage.id != graph.final_stage_id { + let num_tasks = stage_task_count(&stage.plan); + exchange.register_stage( + stage.id, + num_tasks, + stage.output_partition_count, + config.channel_capacity, + ); + } + } + + // Spawn every task of every stage. Producer tasks go in one bucket (awaited + // only for errors); the final stage's tasks go in another (awaited for + // their batches, in ascending task order). + let mut producer_handles: Vec>> = Vec::new(); + let mut final_handles: Option>>>> = None; + + for stage in &graph.stages { + // Encode this stage's plan ONCE; every task decodes its own instance. + let codec = ExchangeCodec { + exchange: exchange.clone(), + }; + let bytes = encode_plan(&stage.plan, &codec)?; + let ntasks = stage_task_count(&stage.plan); + let is_final = stage.id == graph.final_stage_id; + + if is_final { + let mut handles = Vec::with_capacity(ntasks); + for t in 0..ntasks { + let bytes = bytes.clone(); + let session_builder = config.session_builder.clone(); + let exchange = exchange.clone(); + handles.push(SpawnedTask::spawn(async move { + let task_ctx = (session_builder)().task_ctx(); + let codec = ExchangeCodec { exchange }; + let plan = decode_plan(&bytes, &task_ctx, &codec)?; + let mut stream = plan.execute(t, task_ctx)?; + let mut batches = Vec::new(); + while let Some(b) = stream.next().await { + batches.push(b?); + } + Ok::<_, DataFusionError>(batches) + })); + } + final_handles = Some(handles); + } else { + for t in 0..ntasks { + let bytes = bytes.clone(); + let session_builder = config.session_builder.clone(); + let exchange = exchange.clone(); + producer_handles.push(SpawnedTask::spawn(async move { + let task_ctx = (session_builder)().task_ctx(); + let codec = ExchangeCodec { exchange }; + let plan = decode_plan(&bytes, &task_ctx, &codec)?; + let mut stream = plan.execute(t, task_ctx)?; + // Drive the sink to completion; it yields no data rows. + while let Some(b) = stream.next().await { + b?; + } + Ok::<_, DataFusionError>(()) + })); + } + } + } + + let final_handles = final_handles.ok_or_else(|| { + DataFusionError::Internal("no stage matched final_stage_id".to_string()) + })?; + + // Await ALL handles (no barrier — everything is already running). Collect + // the final batches in ascending task order and the first genuine error + // from any task, whichever bucket it came from. + let mut first_error: Option = None; + + for h in producer_handles { + match h.await { + Ok(Ok(())) => {} + Ok(Err(e)) => { + first_error.get_or_insert(e); + } + Err(e) => { + first_error.get_or_insert(DataFusionError::Execution(format!( + "producer stage task panicked: {e}" + ))); + } + } + } + + let mut out = Vec::new(); + for h in final_handles { + match h.await { + Ok(Ok(batches)) => out.extend(batches), + Ok(Err(e)) => { + first_error.get_or_insert(e); + } + Err(e) => { + first_error.get_or_insert(DataFusionError::Execution(format!( + "final stage task panicked: {e}" + ))); + } + } + } + + if let Some(e) = first_error { + return Err(e); + } + Ok(out) +} diff --git a/datafusion-scheduler/src/lib.rs b/datafusion-scheduler/src/lib.rs new file mode 100644 index 0000000000000..4d03c6100c0de --- /dev/null +++ b/datafusion-scheduler/src/lib.rs @@ -0,0 +1,70 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! In-process model of stage-based distributed execution for Apache DataFusion. +//! +//! Splits a physical plan into stages at shuffle boundaries, serializes each +//! stage, and runs each task from a freshly deserialized plan with its own +//! `TaskContext` — surfacing physical-plan changes that assume shared +//! in-process state. Data crosses stage boundaries as IPC-serialized frames +//! over a streaming in-memory exchange (no disk, no barrier). See +//! `specs/datafusion-scheduler.md` in ballista-ai-dev. + +mod config; +pub mod exchange; +mod executor; +mod scheduler; +mod serde; +mod stage; +pub mod test_util; + +pub use config::SchedulerConfig; +pub use exchange::{ + ExchangeCodec, ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange, +}; +pub use executor::execute_stage_graph; +pub use scheduler::create_stages; +pub use stage::{QueryStage, StageGraph, StageId}; +pub use test_util::assert_distributed_eq; + +use std::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::error::Result; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; + +/// Splits `plan` into a [`StageGraph`] at shuffle boundaries and runs it +/// through the streaming, no-barrier stage executor +/// ([`execute_stage_graph`]): all tasks of all stages are spawned +/// concurrently, data crosses each boundary as IPC frames over a single shared +/// [`InMemoryExchange`], and the final stage's collected output is returned. +/// +/// `ctx` is the driver session; each task rebuilds an isolated session via +/// `config.session_builder` and never touches `ctx`. +pub async fn run_distributed( + _ctx: &SessionContext, + plan: Arc, + config: SchedulerConfig, +) -> Result> { + // One exchange instance is threaded through both stage splitting (so + // sinks/sources capture it) and the executor (so it registers the channels + // and drives the tasks) — producers and consumers must share ONE wire. + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange)?; + execute_stage_graph(&graph, &exchange, &config).await +} diff --git a/datafusion-scheduler/src/scheduler.rs b/datafusion-scheduler/src/scheduler.rs new file mode 100644 index 0000000000000..f4d596ff77d21 --- /dev/null +++ b/datafusion-scheduler/src/scheduler.rs @@ -0,0 +1,207 @@ +// 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. + +//! Splits a physical plan into a [`StageGraph`] at shuffle boundaries. +//! +//! Mirrors the shape of Ballista's +//! `DefaultDistributedPlanner::plan_query_stages_internal` +//! (`ballista/scheduler/src/planner.rs`), but simplified for this in-process +//! model: no broadcast-join promotion, no subquery special-casing, no AQE. + +use std::sync::Arc; + +use datafusion::error::Result; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::{ + ExecutionPlan, Partitioning, with_new_children_if_necessary, +}; + +use crate::exchange::{ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange}; +use crate::stage::{QueryStage, StageGraph, StageId}; + +/// Splits `plan` into a [`StageGraph`] at shuffle boundaries (hash +/// repartitioning, coalescing to a single partition, and sort-preserving +/// merges), wiring each boundary through the shared streaming `exchange`. +/// +/// Stage ids are allocated in the order stages are produced during the +/// recursion, so producer stages always have a lower id than the consumer +/// stages that read from them. The final stage (the rewritten top-level +/// plan, representing the "collect final result" step) is pushed last and +/// is not wrapped in an `ExchangeSinkExec`. +/// +/// The `exchange` is cloned into every `ExchangeSinkExec`/`ExchangeSourceExec` +/// built here; the executor registers the boundary channels on that same +/// instance before spawning tasks, so producers and consumers share one wire. +pub fn create_stages( + plan: Arc, + exchange: &Arc, +) -> Result { + let mut next_id: StageId = 0; + let (final_plan, mut stages) = + split_at_shuffle_boundaries(plan, exchange, &mut next_id)?; + + let final_stage_id = next_id; + let input_stage_ids = collect_source_stage_ids(&final_plan); + let output_partition_count = final_plan + .properties() + .output_partitioning() + .partition_count(); + stages.push(QueryStage { + id: final_stage_id, + plan: final_plan, + output_partition_count, + input_stage_ids, + }); + + Ok(StageGraph { + stages, + final_stage_id, + }) +} + +/// Recursively rewrites `plan`, replacing every shuffle boundary with an +/// `ExchangeSourceExec` leaf fed by a new producer `QueryStage`. Returns the +/// rewritten plan along with every producer stage created while recursing +/// through it (in the order they were created). +fn split_at_shuffle_boundaries( + plan: Arc, + exchange: &Arc, + next_id: &mut StageId, +) -> Result<(Arc, Vec)> { + if plan.children().is_empty() { + return Ok((plan, Vec::new())); + } + + let mut stages = Vec::new(); + let mut children = Vec::new(); + for child in plan.children() { + let (new_child, mut child_stages) = + split_at_shuffle_boundaries(Arc::clone(child), exchange, next_id)?; + children.push(new_child); + stages.append(&mut child_stages); + } + + // CoalescePartitionsExec and SortPreservingMergeExec both collapse their + // (possibly many) input partitions down to a single output partition; + // they get identical shuffle-boundary treatment here. + if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + { + let child = children[0].clone(); + let source = shuffle_boundary( + child, + Partitioning::RoundRobinBatch(1), + 1, + exchange, + next_id, + &mut stages, + )?; + return Ok((source, stages)); + } + + if let Some(repart) = plan.downcast_ref::() { + return match repart.partitioning().clone() { + Partitioning::Hash(exprs, p) => { + let child = children[0].clone(); + let source = shuffle_boundary( + child, + Partitioning::Hash(exprs, p), + p, + exchange, + next_id, + &mut stages, + )?; + Ok((source, stages)) + } + // v1 simplification: non-hash repartitions (e.g. round-robin) + // exist purely to change intra-stage parallelism, not to cross a + // shuffle boundary. This simplified scheduler has no notion of + // "reshape partition count without a shuffle", so we simply drop + // the node and keep the child's own partitioning/task count. + _ => Ok((children[0].clone(), stages)), + }; + } + + let new_plan = with_new_children_if_necessary(plan, children)?; + Ok((new_plan, stages)) +} + +/// Wraps `child` in a producer `QueryStage` rooted at an `ExchangeSinkExec` +/// (pushed into `stages`), and returns an `ExchangeSourceExec` that stands in +/// for `child` at the shuffle boundary. +fn shuffle_boundary( + child: Arc, + sink_partitioning: Partitioning, + output_partition_count: usize, + exchange: &Arc, + next_id: &mut StageId, + stages: &mut Vec, +) -> Result> { + // Number of producer tasks = the child's own partition count, computed + // before it gets wrapped in the sink (the sink's output partitioning + // controls how many *buckets* each task fans out to, not how many tasks + // the sink itself has -- see `ExchangeSinkExec`). + let num_producer_tasks = child.properties().output_partitioning().partition_count(); + let child_schema = child.schema(); + let input_stage_ids = collect_source_stage_ids(&child); + + let stage_id = *next_id; + *next_id += 1; + + let sink: Arc = Arc::new(ExchangeSinkExec::try_new( + stage_id, + child, + sink_partitioning, + exchange.clone(), + )?); + + stages.push(QueryStage { + id: stage_id, + plan: sink, + output_partition_count, + input_stage_ids, + }); + + let source: Arc = Arc::new(ExchangeSourceExec::try_new( + stage_id, + child_schema, + num_producer_tasks, + output_partition_count, + exchange.clone(), + )?); + + Ok(source) +} + +/// Walks `plan`'s tree looking for `ExchangeSourceExec` leaves and collects +/// their producer stage ids -- these are the upstream stages `plan` depends on. +fn collect_source_stage_ids(plan: &Arc) -> Vec { + let mut ids = Vec::new(); + walk_collect_source_stage_ids(plan, &mut ids); + ids +} + +fn walk_collect_source_stage_ids(plan: &Arc, ids: &mut Vec) { + if let Some(source) = plan.downcast_ref::() { + ids.push(source.from_stage_id()); + } + for child in plan.children() { + walk_collect_source_stage_ids(child, ids); + } +} diff --git a/datafusion-scheduler/src/serde.rs b/datafusion-scheduler/src/serde.rs new file mode 100644 index 0000000000000..81d6ad849574e --- /dev/null +++ b/datafusion-scheduler/src/serde.rs @@ -0,0 +1,43 @@ +// 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. + +use std::sync::Arc; + +use datafusion::error::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::physical_plan::{AsExecutionPlan, PhysicalExtensionCodec}; +use datafusion_proto::protobuf::PhysicalPlanNode; + +pub fn encode_plan( + plan: &Arc, + codec: &dyn PhysicalExtensionCodec, +) -> Result> { + let node = PhysicalPlanNode::try_from_physical_plan(plan.clone(), codec)?; + let mut buf = Vec::new(); + node.try_encode(&mut buf)?; // Vec: BufMut + Ok(buf) +} + +pub fn decode_plan( + bytes: &[u8], + ctx: &TaskContext, + codec: &dyn PhysicalExtensionCodec, +) -> Result> { + let node = PhysicalPlanNode::try_decode(bytes)?; + node.try_into_physical_plan(ctx, codec) +} diff --git a/datafusion-scheduler/src/stage.rs b/datafusion-scheduler/src/stage.rs new file mode 100644 index 0000000000000..c8e90a0b64f1c --- /dev/null +++ b/datafusion-scheduler/src/stage.rs @@ -0,0 +1,59 @@ +// 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. + +//! A [`StageGraph`] is a physical plan split into [`QueryStage`]s at shuffle +//! boundaries. Each stage other than the final one is rooted at an +//! `ExchangeSinkExec` (its "producer" tasks push IPC frames into the exchange); +//! the final stage is the plan that produces the query's overall result and is +//! not wrapped in a sink. A stage's `input_stage_ids` name the upstream stages +//! whose exchange output it reads via `ExchangeSourceExec` leaves. + +use std::sync::Arc; + +use datafusion::physical_plan::ExecutionPlan; + +/// Identifies a [`QueryStage`] within a [`StageGraph`]. Stable only within a +/// single `StageGraph` — not a global identifier. +pub type StageId = usize; + +/// One stage of a [`StageGraph`]: a subtree of the original physical plan, +/// plus the bookkeeping needed to schedule and connect it to its inputs. +pub struct QueryStage { + /// This stage's id, unique within its `StageGraph`. + pub id: StageId, + /// The stage's root plan. For every stage but the final one, this is an + /// `ExchangeSinkExec`; the final stage's plan is the rewritten original + /// plan with no sink wrapping. + pub plan: Arc, + /// Number of output partitions this stage produces (i.e. how a + /// downstream `ExchangeSourceExec` should be sized). For the final stage + /// this is simply `plan`'s own output partition count. + pub output_partition_count: usize, + /// Ids of upstream stages this stage reads from via `ExchangeSourceExec` + /// leaves in `plan`. + pub input_stage_ids: Vec, +} + +/// A physical plan split into [`QueryStage`]s at shuffle boundaries. +/// +/// Stages are ordered so that a producer's id is always lower than the ids +/// of the consumer stages that read from it. +pub struct StageGraph { + pub stages: Vec, + /// Id of the stage that produces the overall query result. + pub final_stage_id: StageId, +} diff --git a/datafusion-scheduler/src/test_util.rs b/datafusion-scheduler/src/test_util.rs new file mode 100644 index 0000000000000..762d99897a657 --- /dev/null +++ b/datafusion-scheduler/src/test_util.rs @@ -0,0 +1,45 @@ +// 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. + +use std::sync::Arc; + +use datafusion::physical_plan::{ExecutionPlan, collect}; +use datafusion::prelude::SessionContext; + +use crate::{SchedulerConfig, run_distributed}; + +/// Assert distributed execution yields the same multiset of rows as `collect`. +pub async fn assert_distributed_eq(ctx: &SessionContext, plan: Arc) { + let expected = collect(plan.clone(), ctx.task_ctx()) + .await + .expect("collect"); + let actual = run_distributed(ctx, plan, SchedulerConfig::in_memory(ctx)) + .await + .expect("run_distributed"); + // Compare as sorted string rows via pretty-format for schema-agnostic equality. + let e = datafusion::arrow::util::pretty::pretty_format_batches(&expected) + .unwrap() + .to_string(); + let a = datafusion::arrow::util::pretty::pretty_format_batches(&actual) + .unwrap() + .to_string(); + let mut el: Vec<&str> = e.lines().collect(); + el.sort(); + let mut al: Vec<&str> = a.lines().collect(); + al.sort(); + assert_eq!(el, al, "distributed output != collect output"); +} diff --git a/datafusion-scheduler/tests/equivalence.rs b/datafusion-scheduler/tests/equivalence.rs new file mode 100644 index 0000000000000..d0498d5db31de --- /dev/null +++ b/datafusion-scheduler/tests/equivalence.rs @@ -0,0 +1,299 @@ +// 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. + +//! End-to-end equivalence tests: run representative plans (joins, sorts, +//! multi-stage chains) through `run_distributed` and assert the result is the +//! same multiset of rows as single-node `collect`, via `assert_distributed_eq`. +//! +//! Plus the #1907-class regression guard: prove that when a file scan is split +//! into isolated per-task plans, each task reads ONLY its own partition rather +//! than draining a plan-instance-shared pool of file work and scanning the +//! whole table (the failure class of Ballista issue #1907). +//! +//! Plan shapes quoted in each test's doc comment were captured with +//! `displayable(plan.as_ref()).indent(false)` against this DataFusion checkout +//! at `target_partitions = 4`. + +use std::sync::Arc; + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::MemTable; +use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; +use datafusion_scheduler::test_util::assert_distributed_eq; + +/// Registers two 2-partition `MemTable`s, `l(k INT, v INT)` and +/// `r(k INT, v INT)`, against a `SessionContext` configured with +/// `target_partitions = 4`. The hash-join single-partition thresholds are +/// dropped to 0 so a join on `k` plans as a `Partitioned` join with both sides +/// hash-repartitioned (two shuffle inputs into the join stage) rather than a +/// `CollectLeft`/broadcast that these tiny tables would otherwise qualify for. +/// +/// `prefer_hash_join` selects a `HashJoinExec` (true) vs a `SortMergeJoinExec` +/// (false); both still hash-repartition each side. +async fn context_with_two_partitioned_tables(prefer_hash_join: bool) -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + + let lbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let lbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![40, 50, 60])), + ], + ) + .unwrap(); + let rbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![100, 200, 300])), + ], + ) + .unwrap(); + let rbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![400, 500, 600])), + ], + ) + .unwrap(); + + let mut config = SessionConfig::new().with_target_partitions(4); + config.options_mut().optimizer.prefer_hash_join = prefer_hash_join; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold = 0; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold_rows = 0; + let ctx = SessionContext::new_with_config(config); + + let ltable = + MemTable::try_new(schema.clone(), vec![vec![lbatch0], vec![lbatch1]]).unwrap(); + let rtable = MemTable::try_new(schema, vec![vec![rbatch0], vec![rbatch1]]).unwrap(); + ctx.register_table("l", Arc::new(ltable)).unwrap(); + ctx.register_table("r", Arc::new(rtable)).unwrap(); + ctx +} + +/// Registers table `t(a INT)` backed by a 2-partition `MemTable`, against a +/// `SessionContext` configured with `target_partitions = 4`. +async fn context_with_partitioned_table() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 2, 5]))], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![2, 3, 4, 4, 1]))], + ) + .unwrap(); + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx +} + +/// **Test 1 — hash join.** `SELECT ... FROM l JOIN r ON l.k = r.k` with +/// `prefer_hash_join = true`. Both sides hash-repartition on `k`, so the join +/// stage reads two shuffle inputs. +/// +/// Plan (`displayable(..).indent(false)`, target_partitions = 4): +/// +/// ```text +/// HashJoinExec: mode=Partitioned, join_type=Inner, on=[(k@0, k@0)], projection=[k@0, v@1, v@3] +/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// ``` +#[tokio::test] +async fn hash_join_matches_collect() { + let ctx = context_with_two_partitioned_tables(true).await; + let df = ctx + .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; +} + +/// **Test 2 — sort-merge join.** Same query with `prefer_hash_join = false`, +/// which plans a `SortMergeJoinExec`; each side is still hash-repartitioned on +/// `k` and then sorted, so the join stage again reads two shuffle inputs. +#[tokio::test] +async fn sort_merge_join_matches_collect() { + let ctx = context_with_two_partitioned_tables(false).await; + let df = ctx + .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; +} + +/// **Test 3 — sort (multiset equivalence).** `SELECT a FROM t ORDER BY a` +/// WITHOUT a LIMIT. The top `SortPreservingMergeExec` is a shuffle boundary +/// (gather to 1 partition). `assert_distributed_eq` compares order-insensitively +/// and v1 models `SortPreservingMergeExec` as a plain gather (no global-order +/// preservation, a deliberate spec deferral — see the spec's "Open questions"), +/// so a full ORDER BY yields the same MULTISET and must pass. +/// +/// Plan (`displayable(..).indent(false)`, target_partitions = 4): +/// +/// ```text +/// SortPreservingMergeExec: [a@0 ASC NULLS LAST] +/// SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// ``` +#[tokio::test] +async fn sort_without_limit_matches_collect_as_multiset() { + let ctx = context_with_partitioned_table().await; + let df = ctx.sql("SELECT a FROM t ORDER BY a").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; +} + +/// Documents a known v1 limitation, kept as an `#[ignore]`d test so it never +/// gates CI. `SELECT a FROM t ORDER BY a LIMIT 5` (top-N) selects a specific +/// SET of rows that depends on GLOBAL order. Because v1 models +/// `SortPreservingMergeExec` as a plain gather that does NOT preserve global +/// order (see the spec's "Open questions": SortPreservingMergeExec global-order +/// fidelity), the distributed path can legitimately return a different set of 5 +/// rows than single-node `collect`. This is expected until order-preserving +/// shuffle lands, so the test is ignored rather than asserted. +#[tokio::test] +#[ignore = "top-N after sort is a known v1 limitation: SortPreservingMergeExec \ + is modeled as an unordered gather (see spec Open questions)"] +async fn top_n_after_sort_is_a_known_v1_limitation() { + let ctx = context_with_partitioned_table().await; + let df = ctx.sql("SELECT a FROM t ORDER BY a LIMIT 5").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; +} + +/// **Test 4 — three-stage chain.** `SELECT a, count(*) AS c FROM t GROUP BY a +/// ORDER BY a`: a hash-repartition boundary for the final aggregate plus a +/// `SortPreservingMergeExec` gather for the sort → two shuffle boundaries, three +/// stages. Multiset equivalence (order-insensitive). +/// +/// Plan (`displayable(..).indent(false)`, target_partitions = 4): +/// +/// ```text +/// SortPreservingMergeExec: [a@0 ASC NULLS LAST] +/// SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] +/// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as c] +/// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] +/// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 +/// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// ``` +#[tokio::test] +async fn group_by_then_order_by_three_stage_chain_matches_collect() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; +} + +/// **#1907 regression guard — CAUGHT REGRESSION, currently `#[ignore]`d.** +/// +/// Writes 100 rows (`x = 0..24, 100..124, 200..224, 300..324`) across 4 CSV +/// files, registers them as one table, and runs `SELECT count(*), sum(x)`. When +/// the scan splits into isolated per-task plans, each task must read ONLY its +/// own file group. This is the exact failure class of Ballista issue +/// [#1907](https://github.com/apache/datafusion-ballista/issues/1907). In this +/// checkout the mechanism is the **morsel-driven file scan** +/// (`datafusion/datasource/src/morsel/`): a `DataSourceExec` hands file work out +/// from a pool **shared across the plan instance's output partitions**, drained +/// cooperatively only when *all* partitions are polled concurrently. A +/// distributed executor polls **one partition per isolated task**, so that task +/// drains the whole pool and scans the entire table. (Reproduces with a bare +/// DataFusion plan and a single `execute(0, ..)` call — no scheduler involved.) +/// +/// This harness surfaces exactly that. Observed at `target_partitions = 4` +/// (see the captured plan/stage dump in the Task 7 report): +/// +/// - single-node `collect` (all 4 partitions polled together, queue drained +/// cooperatively): `n = 100`, `s = 16200` — correct. +/// - `run_distributed` (4 isolated tasks, each polling its 1 partition alone): +/// `n = 400`, `s = 64800` — **exactly 4× (= num_producer_tasks) inflation.** +/// +/// The stage split is correct (stage 0 = `Partial` aggregate over the +/// `DataSourceExec: file_groups={4 groups}`, `num_producer_tasks = 4`; stage 1 = +/// `Final` aggregate reading the shuffle) and each task's file paths are +/// correct — the inflation is DataFusion's morsel-driven scan behavior, NOT a +/// defect in this harness. Per the crate's purpose we do NOT weaken the +/// assertion or force it green; the test is `#[ignore]`d as a caught regression +/// to be triaged. Run with `--ignored` to reproduce; it should pass once the +/// underlying DataFusion scan-isolation behavior is fixed. +#[tokio::test] +#[ignore = "CAUGHT #1907-class regression: DataFusion's morsel-driven file \ + scan makes each isolated per-task scan read the whole table \ + (observed 4x inflation: n=400 vs 100, s=64800 vs 16200). Genuine \ + DataFusion behavior surfaced by the harness, not a test defect."] +async fn regression_isolated_scan_reads_only_its_partition() { + let dir = tempfile::TempDir::new().unwrap(); + // 4 files, 25 rows each = 100 rows total; distinct x per file so sum is + // exact and any per-task double-read is visible as an integer multiple. + for f in 0..4 { + let mut body = String::from("x\n"); + for i in 0..25 { + let x = f * 100 + i; + body.push_str(&x.to_string()); + body.push('\n'); + } + std::fs::write(dir.path().join(format!("part_{f}.csv")), body).unwrap(); + } + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + ctx.register_csv("t", dir.path().to_str().unwrap(), CsvReadOptions::new()) + .await + .unwrap(); + + let df = ctx + .sql("SELECT count(*) AS n, sum(x) AS s FROM t") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + // Distributed sum/count must equal single-node collect. A whole-table-per + // -task scan inflates both by the task count; this assertion catches it. + assert_distributed_eq(&ctx, plan).await; +} diff --git a/datafusion-scheduler/tests/error_propagation.rs b/datafusion-scheduler/tests/error_propagation.rs new file mode 100644 index 0000000000000..17c61feff4ae7 --- /dev/null +++ b/datafusion-scheduler/tests/error_propagation.rs @@ -0,0 +1,111 @@ +// 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. + +//! Proves a task execution error surfaces as a `DataFusionError` out of +//! `run_distributed` -- not swallowed, not hung, not a panic escaping the +//! spawned task -- for the simplest case: a single (no-shuffle) final stage +//! whose plan errors while executing. +//! +//! The failing node here is a `ScalarUDF` that always returns `Err`, wired +//! into a plan of otherwise-ordinary, `datafusion-proto`-native nodes (table +//! scan + projection). That's deliberate: the executor always round-trips +//! every stage's plan through `encode_plan`/`decode_plan` with +//! [`ExchangeCodec`](datafusion_scheduler::ExchangeCodec), which only knows how +//! to serialize `ExchangeSinkExec`/`ExchangeSourceExec`. A hand-rolled leaf +//! `ExecutionPlan` (as opposed to a UDF referenced by name and resolved from +//! the rebuilt session's function registry on decode) would fail at the +//! encode step instead, which would only prove plan serialization errors +//! propagate -- not task execution errors. + +use std::sync::Arc; + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; +use datafusion::prelude::SessionContext; +use datafusion_scheduler::{SchedulerConfig, run_distributed}; + +/// A scalar UDF that unconditionally fails when invoked, with a message +/// ("boom") distinctive enough to assert on without false-matching some +/// unrelated error path. +#[derive(Debug, PartialEq, Eq, Hash)] +struct BoomUdf { + signature: Signature, +} + +impl BoomUdf { + fn new() -> Self { + Self { + signature: Signature::uniform(1, vec![DataType::Int32], Volatility::Volatile), + } + } +} + +impl ScalarUDFImpl for BoomUdf { + fn name(&self) -> &str { + "boom" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type( + &self, + _arg_types: &[DataType], + ) -> datafusion::error::Result { + Ok(DataType::Int32) + } + + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> datafusion::error::Result { + Err(DataFusionError::Execution("boom".to_string())) + } +} + +#[tokio::test] +async fn task_execution_error_surfaces_from_run_distributed() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_batch("t", batch).unwrap(); + ctx.register_udf(ScalarUDF::from(BoomUdf::new())); + + // A single-partition, no-shuffle plan: `create_stages` produces exactly + // one (final) stage, so this exercises `execute_stage_graph`'s + // single-stage path end to end, including the encode/decode round-trip. + let df = ctx.sql("SELECT boom(a) FROM t").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let config = SchedulerConfig::in_memory(&ctx); + let result = run_distributed(&ctx, plan, config).await; + + let err = result + .expect_err("task execution error must surface as an Err, not be swallowed"); + let msg = err.to_string(); + assert!( + msg.contains("boom"), + "expected error message to contain \"boom\", got: {msg}" + ); +} diff --git a/datafusion-scheduler/tests/single_stage.rs b/datafusion-scheduler/tests/single_stage.rs new file mode 100644 index 0000000000000..b8b7ff75342ce --- /dev/null +++ b/datafusion-scheduler/tests/single_stage.rs @@ -0,0 +1,68 @@ +// 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. + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::prelude::*; +use datafusion_scheduler::{SchedulerConfig, run_distributed}; +use std::sync::Arc; + +#[tokio::test] +async fn single_stage_scan_filter_matches_collect() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + ) + .unwrap(); + + let ctx = SessionContext::new(); + ctx.register_batch("t", batch).unwrap(); + + let df = ctx.sql("SELECT a FROM t WHERE a > 2").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let expected = datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()) + .await + .unwrap(); + + let config = SchedulerConfig::in_memory(&ctx); + let actual = run_distributed(&ctx, plan, config).await.unwrap(); + + // order-insensitive compare + assert_eq!(total_rows(&expected), total_rows(&actual)); + assert_eq!(sorted_values(&expected), sorted_values(&actual)); +} + +fn total_rows(b: &[RecordBatch]) -> usize { + b.iter().map(|b| b.num_rows()).sum() +} +fn sorted_values(b: &[RecordBatch]) -> Vec { + let mut v: Vec = b + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + v.sort(); + v +} diff --git a/datafusion-scheduler/tests/stage_splitting.rs b/datafusion-scheduler/tests/stage_splitting.rs new file mode 100644 index 0000000000000..2e38ec56ffff4 --- /dev/null +++ b/datafusion-scheduler/tests/stage_splitting.rs @@ -0,0 +1,458 @@ +// 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. + +//! Tests for `datafusion_scheduler::create_stages`, which splits a physical +//! plan into a `StageGraph` at shuffle boundaries. +//! +//! Expected stage shape below was derived by printing +//! `displayable(plan.as_ref()).indent(false)` for each query against this +//! DataFusion checkout before writing assertions -- see the plans quoted in +//! each test's doc comment. + +use std::sync::Arc; + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::MemTable; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_scheduler::{ + ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange, create_stages, +}; + +/// Registers table `t(a INT)` backed by a 2-partition `MemTable`, against a +/// `SessionContext` configured with `target_partitions = 4`. +async fn context_with_partitioned_table() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![2, 3, 4]))], + ) + .unwrap(); + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx +} + +fn contains_exchange_sink(plan: &Arc) -> bool { + plan.downcast_ref::().is_some() + || plan.children().iter().any(|c| contains_exchange_sink(c)) +} + +fn contains_exchange_source(plan: &Arc) -> bool { + plan.downcast_ref::().is_some() + || plan.children().iter().any(|c| contains_exchange_source(c)) +} + +/// Registers two 2-partition `MemTable`s, `l(k INT, v INT)` and +/// `r(k INT, v INT)`, against a `SessionContext` configured with +/// `target_partitions = 4` and the hash-join collect-left thresholds +/// dropped to 0 so a join on `k` plans as a `Partitioned` hash join +/// (`RepartitionExec: partitioning=Hash(...)` on both sides) rather than +/// `CollectLeft`, which these tiny in-memory tables would otherwise qualify +/// for. +async fn context_with_two_partitioned_tables() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + + let lbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let lbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![40, 50, 60])), + ], + ) + .unwrap(); + let rbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![100, 200, 300])), + ], + ) + .unwrap(); + let rbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![400, 500, 600])), + ], + ) + .unwrap(); + + let mut config = SessionConfig::new().with_target_partitions(4); + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold = 0; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold_rows = 0; + let ctx = SessionContext::new_with_config(config); + + let ltable = + MemTable::try_new(schema.clone(), vec![vec![lbatch0], vec![lbatch1]]).unwrap(); + let rtable = MemTable::try_new(schema, vec![vec![rbatch0], vec![rbatch1]]).unwrap(); + ctx.register_table("l", Arc::new(ltable)).unwrap(); + ctx.register_table("r", Arc::new(rtable)).unwrap(); + ctx +} + +/// Asserts the topological-order invariant `StageGraph` promises: every +/// producer stage id referenced via `input_stage_ids` is strictly lower than +/// the id of the stage that reads it, and `final_stage_id` is both the +/// highest id in the graph and a stage no other stage depends on. +fn assert_topological_order(graph: &datafusion_scheduler::StageGraph) { + for stage in &graph.stages { + for &producer_id in &stage.input_stage_ids { + assert!( + producer_id < stage.id, + "stage {} depends on stage {} which does not have a lower id", + stage.id, + producer_id + ); + } + } + + let max_id = graph.stages.iter().map(|s| s.id).max().unwrap(); + assert_eq!( + graph.final_stage_id, max_id, + "final_stage_id should be the highest stage id in the graph" + ); + assert!( + !graph + .stages + .iter() + .any(|s| s.input_stage_ids.contains(&graph.final_stage_id)), + "no stage should read from the final stage" + ); +} + +/// `SELECT a, count(*) FROM t GROUP BY a` with `target_partitions = 4` over a +/// 2-partition input produces (verified via `displayable(..).indent(false)` +/// against this checkout): +/// +/// ```text +/// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as count(*)] +/// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] +/// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 +/// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// ``` +/// +/// The single `RepartitionExec: partitioning=Hash(...)` is the only shuffle +/// boundary, so this should split into exactly 2 stages: a producer stage +/// rooted at `ExchangeSinkExec` (wrapping the partial-aggregate subtree), +/// and a final stage (the projection/final-aggregate subtree) reading the +/// producer's output via a `ExchangeSourceExec` leaf. +#[tokio::test] +async fn group_by_splits_into_producer_and_final_stage_at_hash_repartition() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) FROM t GROUP BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!(graph.stages.len(), 2, "expected exactly 2 stages"); + assert_eq!( + graph.final_stage_id, + graph.stages.last().unwrap().id, + "final stage should be the last one pushed" + ); + + // Producer stage (id 0): the partial-aggregate subtree wrapped in a + // ExchangeSinkExec, hash-partitioned into 4 buckets, with no upstream + // shuffle dependency of its own. + let producer = &graph.stages[0]; + assert_eq!(producer.id, 0); + assert!( + producer.plan.downcast_ref::().is_some(), + "producer stage's plan root should be a ExchangeSinkExec, got: {}", + producer.plan.name() + ); + assert_eq!(producer.output_partition_count, 4); + assert!( + producer.input_stage_ids.is_empty(), + "producer stage should have no upstream shuffle dependencies" + ); + + // Final stage: the projection/final-aggregate subtree, containing a + // ExchangeSourceExec leaf that reads the producer stage's output, and + // NOT itself wrapped in a ExchangeSinkExec. + let final_stage = &graph.stages[1]; + assert_eq!(final_stage.id, graph.final_stage_id); + assert!( + final_stage + .plan + .downcast_ref::() + .is_none(), + "final stage must not be wrapped in a ExchangeSinkExec" + ); + assert!( + contains_exchange_source(&final_stage.plan), + "final stage should contain a ExchangeSourceExec leaf" + ); + assert_eq!( + final_stage.input_stage_ids, + vec![producer.id], + "final stage should depend on the producer stage" + ); + assert_eq!(final_stage.output_partition_count, 4); +} + +/// `SELECT a FROM t WHERE a > 2` has no repartitioning at all (verified via +/// `displayable(..).indent(false)`): +/// +/// ```text +/// FilterExec: a@0 > 2 +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// ``` +/// +/// so it should produce exactly 1 stage containing neither a +/// `ExchangeSinkExec` nor a `ExchangeSourceExec` anywhere in its plan. +#[tokio::test] +async fn filter_query_with_no_repartition_stays_a_single_stage() { + let ctx = context_with_partitioned_table().await; + let df = ctx.sql("SELECT a FROM t WHERE a > 2").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!(graph.stages.len(), 1, "expected exactly 1 stage"); + let only_stage = &graph.stages[0]; + assert_eq!(only_stage.id, graph.final_stage_id); + assert!(only_stage.input_stage_ids.is_empty()); + assert!( + !contains_exchange_sink(&only_stage.plan), + "no-shuffle plan must not contain a ExchangeSinkExec" + ); + assert!( + !contains_exchange_source(&only_stage.plan), + "no-shuffle plan must not contain a ExchangeSourceExec" + ); +} + +/// `SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY c` with +/// `target_partitions = 4` over a 2-partition input produces (verified via +/// `displayable(..).indent(false)` against this checkout): +/// +/// ```text +/// SortPreservingMergeExec: [c@1 ASC NULLS LAST] +/// SortExec: expr=[c@1 ASC NULLS LAST], preserve_partitioning=[true] +/// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as c] +/// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] +/// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 +/// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// ``` +/// +/// There are TWO shuffle boundaries here: the `RepartitionExec: Hash(...)` +/// feeding the final aggregate, and the top-level `SortPreservingMergeExec` +/// collapsing the sorted partitions down to 1. That means the split walks a +/// chain of dependent producer stages (stage 1's subtree embeds stage 0's +/// `ExchangeSourceExec`) rather than just a single boundary, which is what +/// Task 6 needs the ascending-id topological-order invariant to hold for. +/// +/// This should split into exactly 3 stages: +/// - stage 0: producer wrapping the partial-aggregate subtree (hash +/// repartition boundary), no upstream dependency. +/// - stage 1: producer wrapping the final-aggregate/projection/sort subtree +/// (sort-preserving-merge boundary), depending on stage 0. +/// - stage 2 (final): a bare `ExchangeSourceExec` reading stage 1's output. +#[tokio::test] +async fn multi_boundary_chain_has_ascending_topological_order() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY c") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!( + graph.stages.len(), + 3, + "expected 2 producer stages plus a final stage" + ); + + // General invariant Task 6 depends on: every producer id referenced by + // a stage is lower than that stage's own id, and final_stage_id is the + // max id and is read by nobody. + assert_topological_order(&graph); + + let stage0 = &graph.stages[0]; + assert_eq!(stage0.id, 0); + assert!( + stage0.plan.downcast_ref::().is_some(), + "stage 0 should be rooted at a ExchangeSinkExec" + ); + assert!( + stage0.input_stage_ids.is_empty(), + "stage 0 (partial-aggregate producer) has no upstream shuffle dependency" + ); + + let stage1 = &graph.stages[1]; + assert_eq!(stage1.id, 1); + assert!( + stage1.plan.downcast_ref::().is_some(), + "stage 1 should be rooted at a ExchangeSinkExec" + ); + assert_eq!( + stage1.input_stage_ids, + vec![stage0.id], + "stage 1 (sort-preserving-merge producer) depends on stage 0" + ); + assert!( + contains_exchange_source(&stage1.plan), + "stage 1's subtree should embed a ExchangeSourceExec reading stage 0" + ); + + let final_stage = &graph.stages[2]; + assert_eq!(final_stage.id, graph.final_stage_id); + assert!( + final_stage + .plan + .downcast_ref::() + .is_none(), + "final stage must not be wrapped in a ExchangeSinkExec" + ); + assert!( + final_stage + .plan + .downcast_ref::() + .is_some(), + "final stage's plan should itself be the ExchangeSourceExec reading stage 1" + ); + assert_eq!( + final_stage.input_stage_ids, + vec![stage1.id], + "final stage depends on stage 1" + ); +} + +/// `SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k` with +/// `target_partitions = 4` and the hash-join collect-left thresholds forced +/// to 0 (see `context_with_two_partitioned_tables`) produces (verified via +/// `displayable(..).indent(false)` against this checkout): +/// +/// ```text +/// HashJoinExec: mode=Partitioned, join_type=Inner, on=[(k@0, k@0)], projection=[k@0, v@1, v@3] +/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 +/// DataSourceExec: partitions=2, partition_sizes=[1, 1] +/// ``` +/// +/// Both join inputs are independently hash-repartitioned, so this should +/// split into exactly 3 stages: two independent producer stages (one per +/// join side, neither depending on the other) and a final stage rooted at +/// the (rewritten) `HashJoinExec` itself, whose plan contains two +/// `ExchangeSourceExec` leaves -- one per producer. This is the multi-input +/// wiring shape (`input_stage_ids.len() == 2`) Task 6 needs to schedule a +/// stage that depends on more than one upstream stage. +#[tokio::test] +async fn join_with_two_shuffled_sides_has_two_input_stage_ids() { + let ctx = context_with_two_partitioned_tables().await; + let df = ctx + .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!( + graph.stages.len(), + 3, + "expected 2 independent producer stages plus a final (join) stage" + ); + + assert_topological_order(&graph); + + let final_stage = &graph.stages[2]; + assert_eq!(final_stage.id, graph.final_stage_id); + assert!( + final_stage + .plan + .downcast_ref::() + .is_none(), + "final (join) stage must not be wrapped in a ExchangeSinkExec" + ); + assert_eq!( + final_stage.input_stage_ids.len(), + 2, + "join stage should depend on exactly 2 upstream producer stages" + ); + + // Both referenced producer ids must be valid stages in the graph, with + // ids lower than the join stage's own id (already checked generically + // by assert_topological_order, reasserted here for clarity), and each + // must itself have no upstream dependency (leaf hash-repartition + // producers reading straight from a DataSourceExec). + let stage_ids: std::collections::HashSet<_> = + graph.stages.iter().map(|s| s.id).collect(); + for &producer_id in &final_stage.input_stage_ids { + assert!( + stage_ids.contains(&producer_id), + "referenced producer stage {producer_id} must exist in the graph" + ); + assert!( + producer_id < final_stage.id, + "producer stage {producer_id} must have a lower id than the join stage" + ); + let producer = graph + .stages + .iter() + .find(|s| s.id == producer_id) + .expect("producer stage must exist"); + assert!( + producer.plan.downcast_ref::().is_some(), + "producer stage {producer_id} should be rooted at a ExchangeSinkExec" + ); + assert!( + producer.input_stage_ids.is_empty(), + "producer stage {producer_id} (leaf hash-repartition) has no upstream dependency" + ); + } +} diff --git a/datafusion-scheduler/tests/streaming.rs b/datafusion-scheduler/tests/streaming.rs new file mode 100644 index 0000000000000..e384795009565 --- /dev/null +++ b/datafusion-scheduler/tests/streaming.rs @@ -0,0 +1,125 @@ +// 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. + +//! Streaming / backpressure integration test. +//! +//! Runs a real multi-stage query end-to-end through `run_distributed` with +//! `channel_capacity = 1` — the tightest possible backpressure through the FULL +//! executor: producers block on full channels while downstream consumers drain. +//! The whole run is wrapped in a `tokio::time::timeout` so a deadlock FAILS the +//! test rather than hanging CI. +//! +//! This is the end-to-end proof that the spawn-all, no-barrier executor +//! pipelines correctly under backpressure: a consumer blocked on a channel is +//! unblocked by a producer that is already running (all tasks of all stages are +//! spawned before any is awaited). If the executor used a per-stage barrier +//! instead, capacity-1 channels would deadlock and this test would time out. + +use std::sync::Arc; +use std::time::Duration; + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::MemTable; +use datafusion::physical_plan::collect; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_scheduler::{SchedulerConfig, run_distributed}; + +/// Registers table `t(a INT)` backed by a 2-partition `MemTable` with enough +/// rows/batches that capacity-1 channels actually force repeated blocking, +/// against a `SessionContext` at `target_partitions = 4`. +async fn context_with_many_rows() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Many small batches per partition so the sink blocks on `send().await` + // over and over against the capacity-1 channels. + let mut p0 = Vec::new(); + let mut p1 = Vec::new(); + for i in 0..200i32 { + p0.push( + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![i % 17]))], + ) + .unwrap(), + ); + p1.push( + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![(i + 5) % 17]))], + ) + .unwrap(), + ); + } + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![p0, p1]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx +} + +/// Three-stage chain (`GROUP BY` hash boundary + `ORDER BY` gather boundary) +/// run under `channel_capacity = 1`, inside a timeout. Asserts multiset +/// equivalence with single-node `collect`. A hang (deadlock) fails via timeout; +/// a wrong result fails via the equivalence check. +#[tokio::test] +async fn tight_backpressure_multi_stage_pipeline_matches_collect() { + let result = tokio::time::timeout(Duration::from_secs(30), async { + let ctx = context_with_many_rows().await; + let df = ctx + .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let expected = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + + // capacity = 1: the tightest bound that still makes progress, exercised + // through the full run_distributed executor (not just sink/source). + let state = ctx.state(); + let config = SchedulerConfig { + session_builder: Arc::new(move || { + SessionContext::new_with_state(state.clone()) + }), + channel_capacity: 1, + }; + let actual = run_distributed(&ctx, plan, config).await.unwrap(); + + (expected, actual) + }) + .await; + + let (expected, actual) = + result.expect("multi-stage pipeline deadlocked / timed out under capacity=1"); + + // Order-insensitive multiset equivalence via sorted pretty-printed rows. + let e = datafusion::arrow::util::pretty::pretty_format_batches(&expected) + .unwrap() + .to_string(); + let a = datafusion::arrow::util::pretty::pretty_format_batches(&actual) + .unwrap() + .to_string(); + let mut el: Vec<&str> = e.lines().collect(); + el.sort(); + let mut al: Vec<&str> = a.lines().collect(); + al.sort(); + assert_eq!( + el, al, + "streaming distributed output != collect output under capacity=1" + ); +} diff --git a/datafusion-scheduler/tests/two_stage.rs b/datafusion-scheduler/tests/two_stage.rs new file mode 100644 index 0000000000000..ebffc7606476d --- /dev/null +++ b/datafusion-scheduler/tests/two_stage.rs @@ -0,0 +1,118 @@ +// 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. + +//! End-to-end test that a real 2-stage plan (partial aggregate + streaming +//! in-memory exchange + final aggregate) executes correctly through +//! `run_distributed`/`execute_stage_graph`, matching single-process `collect` +//! output. Data crosses the stage boundary as IPC frames over the +//! `InMemoryExchange` — there is no disk shuffle to inspect, so correctness of +//! the 2-stage GROUP BY across the exchange is the assertion. + +use std::sync::Arc; + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::MemTable; +use datafusion::physical_plan::collect; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_scheduler::test_util::assert_distributed_eq; +use datafusion_scheduler::{SchedulerConfig, create_stages, run_distributed}; + +/// Registers table `t(a INT)` backed by a 2-partition `MemTable`, against a +/// `SessionContext` configured with `target_partitions = 4` -- matching the +/// setup in `tests/stage_splitting.rs` that is known to split `GROUP BY` +/// into a producer stage (partial aggregate, hash-repartitioned) and a +/// final stage (final aggregate) at a real shuffle boundary. +async fn context_with_partitioned_table() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![2, 3, 4]))], + ) + .unwrap(); + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx +} + +#[tokio::test] +async fn group_by_across_a_real_shuffle_matches_collect() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) FROM t GROUP BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + assert_distributed_eq(&ctx, plan).await; +} + +/// Proves `run_distributed` is genuinely routed through `create_stages` + the +/// stage-graph executor (data really crossing the exchange), not coincidentally +/// producing correct output by re-running the original in-process plan. +/// +/// The old disk model let a test assert shuffle files landed under `stage_0/`. +/// The streaming exchange leaves nothing on disk, so instead we assert on the +/// stage split itself — this GROUP BY splits into a producer stage (id 0, +/// hash-repartitioned partial aggregate, whose output must cross the exchange) +/// and a final stage reading it — and that the end-to-end result over that real +/// two-stage exchange execution matches single-node `collect`. +#[tokio::test] +async fn group_by_data_crosses_the_exchange_between_two_stages() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) FROM t GROUP BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let expected = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + + // The split must produce a producer stage feeding a final stage over the + // exchange — i.e. a genuine 2-stage plan, not a single re-executed tree. + let probe_exchange = datafusion_scheduler::InMemoryExchange::new(); + let graph = create_stages(plan.clone(), &probe_exchange).unwrap(); + assert_eq!( + graph.stages.len(), + 2, + "GROUP BY should split into a producer stage and a final stage" + ); + assert_ne!( + graph.final_stage_id, 0, + "there must be a producer stage (id 0) feeding the final stage" + ); + + // Execute end-to-end: the producer stage's partial-aggregate output only + // reaches the final stage by crossing the in-memory exchange as IPC frames. + let config = SchedulerConfig::in_memory(&ctx); + let actual = run_distributed(&ctx, plan, config).await.unwrap(); + + let expected_rows: usize = expected.iter().map(|b| b.num_rows()).sum(); + let actual_rows: usize = actual.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + expected_rows, actual_rows, + "row count across the exchange must match single-node collect" + ); +} From 5f1e61318dd7350a0c3b53678aec92691924b351 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 14:32:02 -0600 Subject: [PATCH 2/4] refactor(scheduler): convert opt-in crate into a datafusion-examples example Address PR #23282 review feedback: fold the datafusion-scheduler crate into datafusion-examples/examples/scheduler/ as a runnable example, with all integration tests preserved as `#[cfg(test)]` modules so `cargo test --example scheduler` still exercises the isolated / serialized execution path. - New example at `datafusion-examples/examples/scheduler/` mirrors the original crate layout (config, exchange/{sink,source,codec}, executor, scheduler, serde, stage, test_util) plus a `main.rs` runner exposing a `distributed_pipeline` subcommand. - All tests folded inline: stage-splitting tests in `scheduler.rs`, exchange round-trip / backpressure / codec tests in `exchange/{mod,codec}.rs`, and the equivalence / error / single-stage / two-stage / streaming suites as separate test modules in `main.rs`. Same 18 passing + 2 ignored tests as before, plus a smoke test for the runnable example. - Remove the `datafusion-scheduler/` crate and its workspace `exclude`. - Add a "Scheduler Examples" entry to `datafusion-examples/README.md`. --- Cargo.toml | 2 +- datafusion-examples/README.md | 10 + .../examples/scheduler}/config.rs | 0 .../examples/scheduler}/exchange/codec.rs | 3 +- .../examples/scheduler}/exchange/mod.rs | 2 +- .../examples/scheduler}/exchange/sink.rs | 1 + .../examples/scheduler}/exchange/source.rs | 1 + .../examples/scheduler}/executor.rs | 2 +- .../examples/scheduler/main.rs | 833 ++++++++++++++++++ .../examples/scheduler/scheduler.rs | 647 ++++++++++++++ .../examples/scheduler}/serde.rs | 0 .../examples/scheduler}/stage.rs | 1 + .../examples/scheduler}/test_util.rs | 7 +- datafusion-scheduler/.gitignore | 1 - datafusion-scheduler/Cargo.toml | 40 - datafusion-scheduler/src/lib.rs | 70 -- datafusion-scheduler/src/scheduler.rs | 207 ----- datafusion-scheduler/tests/equivalence.rs | 299 ------- .../tests/error_propagation.rs | 111 --- datafusion-scheduler/tests/single_stage.rs | 68 -- datafusion-scheduler/tests/stage_splitting.rs | 458 ---------- datafusion-scheduler/tests/streaming.rs | 125 --- datafusion-scheduler/tests/two_stage.rs | 118 --- 23 files changed, 1501 insertions(+), 1505 deletions(-) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/config.rs (100%) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/exchange/codec.rs (99%) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/exchange/mod.rs (99%) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/exchange/sink.rs (99%) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/exchange/source.rs (99%) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/executor.rs (99%) create mode 100644 datafusion-examples/examples/scheduler/main.rs create mode 100644 datafusion-examples/examples/scheduler/scheduler.rs rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/serde.rs (100%) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/stage.rs (98%) rename {datafusion-scheduler/src => datafusion-examples/examples/scheduler}/test_util.rs (89%) delete mode 100644 datafusion-scheduler/.gitignore delete mode 100644 datafusion-scheduler/Cargo.toml delete mode 100644 datafusion-scheduler/src/lib.rs delete mode 100644 datafusion-scheduler/src/scheduler.rs delete mode 100644 datafusion-scheduler/tests/equivalence.rs delete mode 100644 datafusion-scheduler/tests/error_propagation.rs delete mode 100644 datafusion-scheduler/tests/single_stage.rs delete mode 100644 datafusion-scheduler/tests/stage_splitting.rs delete mode 100644 datafusion-scheduler/tests/streaming.rs delete mode 100644 datafusion-scheduler/tests/two_stage.rs diff --git a/Cargo.toml b/Cargo.toml index 389b736056fe3..ff5d3afcf48f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,7 +67,7 @@ members = [ "datafusion/macros", "datafusion/doc", ] -exclude = ["dev/depcheck", "datafusion-scheduler"] +exclude = ["dev/depcheck"] resolver = "2" [workspace.package] diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md index 073f269d4a35d..3a99734cc44f0 100644 --- a/datafusion-examples/README.md +++ b/datafusion-examples/README.md @@ -199,6 +199,16 @@ cargo run --example dataframe -- dataframe | pivot_unpivot | [`relation_planner/pivot_unpivot.rs`](examples/relation_planner/pivot_unpivot.rs) | Implement PIVOT / UNPIVOT | | table_sample | [`relation_planner/table_sample.rs`](examples/relation_planner/table_sample.rs) | Implement TABLESAMPLE | +## Scheduler Examples + +### Group: `scheduler` + +#### Category: Distributed + +| Subcommand | File Path | Description | +| --------------------- | ----------------------------------------------------- | ---------------------------------------------------------------------------------------------------- | +| distributed_pipeline | [`scheduler/main.rs`](examples/scheduler/main.rs) | In-process model of stage-based distributed execution (stage splitting, per-task plan serialization) | + ## SQL Ops Examples ### Group: `sql_ops` diff --git a/datafusion-scheduler/src/config.rs b/datafusion-examples/examples/scheduler/config.rs similarity index 100% rename from datafusion-scheduler/src/config.rs rename to datafusion-examples/examples/scheduler/config.rs diff --git a/datafusion-scheduler/src/exchange/codec.rs b/datafusion-examples/examples/scheduler/exchange/codec.rs similarity index 99% rename from datafusion-scheduler/src/exchange/codec.rs rename to datafusion-examples/examples/scheduler/exchange/codec.rs index 072b8c9ace0be..2ced6f3d3b143 100644 --- a/datafusion-scheduler/src/exchange/codec.rs +++ b/datafusion-examples/examples/scheduler/exchange/codec.rs @@ -203,8 +203,7 @@ impl PhysicalExtensionCodec for ExchangeCodec { "ExchangeCodec: failed to decode Schema: {e}" )) })?; - let schema: datafusion::arrow::datatypes::Schema = - (&schema_proto).try_into()?; + let schema: arrow::datatypes::Schema = (&schema_proto).try_into()?; let source = ExchangeSourceExec::try_new( from_stage_id, diff --git a/datafusion-scheduler/src/exchange/mod.rs b/datafusion-examples/examples/scheduler/exchange/mod.rs similarity index 99% rename from datafusion-scheduler/src/exchange/mod.rs rename to datafusion-examples/examples/scheduler/exchange/mod.rs index 87c30e7425003..b2fd75aafe63b 100644 --- a/datafusion-scheduler/src/exchange/mod.rs +++ b/datafusion-examples/examples/scheduler/exchange/mod.rs @@ -163,11 +163,11 @@ mod tests { use datafusion::arrow::array::{Int32Array, RecordBatch}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion::common::runtime::SpawnedTask; use datafusion::physical_expr::expressions::col; use datafusion::physical_plan::test::TestMemoryExec; use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion::prelude::SessionContext; - use datafusion_common_runtime::SpawnedTask; use futures::StreamExt; use super::InMemoryExchange; diff --git a/datafusion-scheduler/src/exchange/sink.rs b/datafusion-examples/examples/scheduler/exchange/sink.rs similarity index 99% rename from datafusion-scheduler/src/exchange/sink.rs rename to datafusion-examples/examples/scheduler/exchange/sink.rs index eb9d7bab003cf..970a6666d643c 100644 --- a/datafusion-scheduler/src/exchange/sink.rs +++ b/datafusion-examples/examples/scheduler/exchange/sink.rs @@ -86,6 +86,7 @@ impl ExchangeSinkExec { self.stage_id } + #[cfg(test)] pub fn input(&self) -> &Arc { &self.input } diff --git a/datafusion-scheduler/src/exchange/source.rs b/datafusion-examples/examples/scheduler/exchange/source.rs similarity index 99% rename from datafusion-scheduler/src/exchange/source.rs rename to datafusion-examples/examples/scheduler/exchange/source.rs index 252d96316c1b6..453fa836a0e38 100644 --- a/datafusion-scheduler/src/exchange/source.rs +++ b/datafusion-examples/examples/scheduler/exchange/source.rs @@ -78,6 +78,7 @@ impl ExchangeSourceExec { }) } + #[expect(clippy::wrong_self_convention)] pub fn from_stage_id(&self) -> usize { self.from_stage_id } diff --git a/datafusion-scheduler/src/executor.rs b/datafusion-examples/examples/scheduler/executor.rs similarity index 99% rename from datafusion-scheduler/src/executor.rs rename to datafusion-examples/examples/scheduler/executor.rs index 64b82e8cbe5a4..490a42aa74ba9 100644 --- a/datafusion-scheduler/src/executor.rs +++ b/datafusion-examples/examples/scheduler/executor.rs @@ -29,8 +29,8 @@ use std::sync::Arc; use datafusion::arrow::array::RecordBatch; +use datafusion::common::runtime::SpawnedTask; use datafusion::error::{DataFusionError, Result}; -use datafusion_common_runtime::SpawnedTask; use futures::StreamExt; use crate::config::SchedulerConfig; diff --git a/datafusion-examples/examples/scheduler/main.rs b/datafusion-examples/examples/scheduler/main.rs new file mode 100644 index 0000000000000..8ebbedd1e6330 --- /dev/null +++ b/datafusion-examples/examples/scheduler/main.rs @@ -0,0 +1,833 @@ +// 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. + +//! # Scheduler Examples — in-process model of stage-based distributed execution +//! +//! This example models how a distributed engine built on DataFusion (such as +//! Ballista or datafusion-distributed) actually executes a query: it splits +//! the physical plan into stages at shuffle boundaries, serializes each stage +//! via `datafusion-proto`, and runs each task from a freshly deserialized +//! plan with its own `TaskContext`. Data crosses stage boundaries as +//! Arrow-IPC frames over a streaming in-memory exchange — no network, no +//! disk, no barrier — so it runs entirely inside one process while still +//! exercising the isolate-and-serialize contract that a real distributed +//! executor depends on. +//! +//! The motivating problem: DataFusion's physical plans are written and +//! tested assuming everything runs in one process, with all partitions of an +//! operator polled together by cooperative threads on a shared runtime. A +//! change that quietly relies on that in-process model still passes +//! DataFusion's own tests but breaks downstream distributed engines. The +//! tests under `#[cfg(test)]` in this example exercise the isolated / +//! serialized / spawned-in-isolation execution path, giving a place to catch +//! that class of regression in-tree. +//! +//! ## Usage +//! ```bash +//! cargo run --example scheduler -- [all|distributed_pipeline] +//! ``` +//! +//! Each subcommand runs a corresponding example: +//! - `all` — run all examples included in this module +//! +//! - `distributed_pipeline` +//! (desc: Run a multi-stage `GROUP BY ... ORDER BY` query through the +//! stage-graph executor and print the result alongside single-node +//! `collect` for comparison) +//! +//! ## Tests +//! +//! Run the full harness (unit + integration tests folded into this example) +//! with: +//! ```bash +//! cargo test --example scheduler +//! ``` + +mod config; +mod exchange; +mod executor; +mod scheduler; +mod serde; +mod stage; +#[cfg(test)] +mod test_util; + +use std::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::error::{DataFusionError, Result}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use strum::{IntoEnumIterator, VariantNames}; +use strum_macros::{Display, EnumIter, EnumString, VariantNames}; + +use crate::config::SchedulerConfig; +use crate::exchange::InMemoryExchange; +use crate::executor::execute_stage_graph; +use crate::scheduler::create_stages; + +/// Splits `plan` into a stage graph at shuffle boundaries and runs it through +/// the streaming, no-barrier stage executor: all tasks of all stages are +/// spawned concurrently, data crosses each boundary as IPC frames over a +/// single shared [`InMemoryExchange`], and the final stage's collected output +/// is returned. +/// +/// `ctx` is the driver session; each task rebuilds an isolated session via +/// `config.session_builder` and never touches `ctx`. +pub async fn run_distributed( + _ctx: &SessionContext, + plan: Arc, + config: SchedulerConfig, +) -> Result> { + // One exchange instance is threaded through both stage splitting (so + // sinks/sources capture it) and the executor (so it registers the + // channels and drives the tasks) — producers and consumers must share + // ONE wire. + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange)?; + execute_stage_graph(&graph, &exchange, &config).await +} + +#[derive(EnumIter, EnumString, Display, VariantNames)] +#[strum(serialize_all = "snake_case")] +enum ExampleKind { + All, + DistributedPipeline, +} + +impl ExampleKind { + const EXAMPLE_NAME: &'static str = "scheduler"; + + fn runnable() -> impl Iterator { + ExampleKind::iter().filter(|v| !matches!(v, ExampleKind::All)) + } + + async fn run(&self) -> Result<()> { + match self { + ExampleKind::All => { + for example in ExampleKind::runnable() { + println!("Running example: {example}"); + Box::pin(example.run()).await?; + } + } + ExampleKind::DistributedPipeline => distributed_pipeline().await?, + } + Ok(()) + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let usage = format!( + "Usage: cargo run --example {} -- [{}]", + ExampleKind::EXAMPLE_NAME, + ExampleKind::VARIANTS.join("|") + ); + + let example: ExampleKind = std::env::args() + .nth(1) + .unwrap_or_else(|| ExampleKind::All.to_string()) + .parse() + .map_err(|_| DataFusionError::Execution(format!("Unknown example. {usage}")))?; + + example.run().await +} + +/// Runs a 3-stage `GROUP BY ... ORDER BY` query through the scheduler and +/// compares it to single-node `collect`, printing both. Illustrates the full +/// path: stage splitting, per-stage plan serialization, isolated per-task +/// execution, and Arrow-IPC frames flowing over the in-memory exchange +/// between three stages (partial-aggregate producer → final-aggregate +/// producer → gather-sort final stage). +async fn distributed_pipeline() -> Result<()> { + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::collect; + use datafusion::prelude::SessionConfig; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 2, 5]))], + )?; + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![2, 3, 4, 4, 1]))], + )?; + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]])?; + ctx.register_table("t", Arc::new(table))?; + + let df = ctx + .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a") + .await?; + let plan = df.create_physical_plan().await?; + + println!("--- single-node collect ---"); + let expected = collect(plan.clone(), ctx.task_ctx()).await?; + println!("{}", pretty_format_batches(&expected)?); + + println!("--- run_distributed (3-stage streaming exchange) ---"); + let actual = run_distributed(&ctx, plan, SchedulerConfig::in_memory(&ctx)).await?; + println!("{}", pretty_format_batches(&actual)?); + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_distributed_pipeline() { + distributed_pipeline().await.unwrap(); + } +} + +/// End-to-end equivalence tests: run representative plans (joins, sorts, +/// multi-stage chains) through `run_distributed` and assert the result is +/// the same multiset of rows as single-node `collect`, via +/// `assert_distributed_eq`. +/// +/// Plus the #1907-class regression guard: prove that when a file scan is +/// split into isolated per-task plans, each task reads ONLY its own +/// partition rather than draining a plan-instance-shared pool of file work +/// and scanning the whole table (the failure class of Ballista issue +/// #1907). +/// +/// Plan shapes quoted in each test's doc comment were captured with +/// `displayable(plan.as_ref()).indent(false)` against this DataFusion +/// checkout at `target_partitions = 4`. +#[cfg(test)] +mod equivalence_tests { + use std::sync::Arc; + + use datafusion::arrow::array::{Int32Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::MemTable; + use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + + use crate::test_util::assert_distributed_eq; + + /// Registers two 2-partition `MemTable`s, `l(k INT, v INT)` and + /// `r(k INT, v INT)`, against a `SessionContext` configured with + /// `target_partitions = 4`. The hash-join single-partition thresholds are + /// dropped to 0 so a join on `k` plans as a `Partitioned` join with both + /// sides hash-repartitioned (two shuffle inputs into the join stage) + /// rather than a `CollectLeft`/broadcast that these tiny tables would + /// otherwise qualify for. + /// + /// `prefer_hash_join` selects a `HashJoinExec` (true) vs a + /// `SortMergeJoinExec` (false); both still hash-repartition each side. + async fn context_with_two_partitioned_tables( + prefer_hash_join: bool, + ) -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + + let lbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let lbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![40, 50, 60])), + ], + ) + .unwrap(); + let rbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![100, 200, 300])), + ], + ) + .unwrap(); + let rbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![400, 500, 600])), + ], + ) + .unwrap(); + + let mut config = SessionConfig::new().with_target_partitions(4); + config.options_mut().optimizer.prefer_hash_join = prefer_hash_join; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold = 0; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold_rows = 0; + let ctx = SessionContext::new_with_config(config); + + let ltable = + MemTable::try_new(schema.clone(), vec![vec![lbatch0], vec![lbatch1]]) + .unwrap(); + let rtable = + MemTable::try_new(schema, vec![vec![rbatch0], vec![rbatch1]]).unwrap(); + ctx.register_table("l", Arc::new(ltable)).unwrap(); + ctx.register_table("r", Arc::new(rtable)).unwrap(); + ctx + } + + /// Registers table `t(a INT)` backed by a 2-partition `MemTable`, + /// against a `SessionContext` configured with `target_partitions = 4`. + async fn context_with_partitioned_table() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 2, 5]))], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![2, 3, 4, 4, 1]))], + ) + .unwrap(); + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx + } + + /// **Test 1 — hash join.** `SELECT ... FROM l JOIN r ON l.k = r.k` with + /// `prefer_hash_join = true`. Both sides hash-repartition on `k`, so the + /// join stage reads two shuffle inputs. + #[tokio::test] + async fn hash_join_matches_collect() { + let ctx = context_with_two_partitioned_tables(true).await; + let df = ctx + .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; + } + + /// **Test 2 — sort-merge join.** Same query with + /// `prefer_hash_join = false`, which plans a `SortMergeJoinExec`; each + /// side is still hash-repartitioned on `k` and then sorted. + #[tokio::test] + async fn sort_merge_join_matches_collect() { + let ctx = context_with_two_partitioned_tables(false).await; + let df = ctx + .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; + } + + /// **Test 3 — sort (multiset equivalence).** + /// `SELECT a FROM t ORDER BY a` WITHOUT a LIMIT. + /// `assert_distributed_eq` compares order-insensitively, so this must + /// yield the same MULTISET as single-node `collect`. + #[tokio::test] + async fn sort_without_limit_matches_collect_as_multiset() { + let ctx = context_with_partitioned_table().await; + let df = ctx.sql("SELECT a FROM t ORDER BY a").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; + } + + /// Documents a known v1 limitation, kept as an `#[ignore]`d test so it + /// never gates CI. `SELECT a FROM t ORDER BY a LIMIT 5` (top-N) selects + /// a specific SET of rows that depends on GLOBAL order. Because v1 + /// models `SortPreservingMergeExec` as a plain gather that does NOT + /// preserve global order, the distributed path can legitimately return + /// a different set of 5 rows than single-node `collect`. + #[tokio::test] + #[ignore = "top-N after sort is a known v1 limitation: SortPreservingMergeExec \ + is modeled as an unordered gather"] + async fn top_n_after_sort_is_a_known_v1_limitation() { + let ctx = context_with_partitioned_table().await; + let df = ctx.sql("SELECT a FROM t ORDER BY a LIMIT 5").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; + } + + /// **Test 4 — three-stage chain.** + /// `SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a`: a + /// hash-repartition boundary for the final aggregate plus a + /// `SortPreservingMergeExec` gather for the sort → two shuffle + /// boundaries, three stages. Multiset equivalence (order-insensitive). + #[tokio::test] + async fn group_by_then_order_by_three_stage_chain_matches_collect() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + assert_distributed_eq(&ctx, plan).await; + } + + /// **#1907 regression guard — CAUGHT REGRESSION, currently `#[ignore]`d.** + /// + /// Writes 100 rows across 4 CSV files, registers them as one table, and + /// runs `SELECT count(*), sum(x)`. When the scan splits into isolated + /// per-task plans, each task must read ONLY its own file group. This is + /// the exact failure class of Ballista issue + /// [#1907](https://github.com/apache/datafusion-ballista/issues/1907). + /// In this checkout the mechanism is the **morsel-driven file scan** + /// (`datafusion/datasource/src/morsel/`): a `DataSourceExec` hands file + /// work out from a pool **shared across the plan instance's output + /// partitions**, drained cooperatively only when *all* partitions are + /// polled concurrently. A distributed executor polls **one partition per + /// isolated task**, so that task drains the whole pool and scans the + /// entire table. + /// + /// Observed at `target_partitions = 4`: + /// - single-node `collect`: `n = 100`, `s = 16200` — correct. + /// - `run_distributed`: `n = 400`, `s = 64800` — 4× inflation. + #[tokio::test] + #[ignore = "CAUGHT #1907-class regression: DataFusion's morsel-driven file \ + scan makes each isolated per-task scan read the whole table"] + async fn regression_isolated_scan_reads_only_its_partition() { + let dir = tempfile::TempDir::new().unwrap(); + // 4 files, 25 rows each = 100 rows total; distinct x per file so sum + // is exact and any per-task double-read is visible as an integer + // multiple. + for f in 0..4 { + let mut body = String::from("x\n"); + for i in 0..25 { + let x = f * 100 + i; + body.push_str(&x.to_string()); + body.push('\n'); + } + std::fs::write(dir.path().join(format!("part_{f}.csv")), body).unwrap(); + } + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + ctx.register_csv("t", dir.path().to_str().unwrap(), CsvReadOptions::new()) + .await + .unwrap(); + + let df = ctx + .sql("SELECT count(*) AS n, sum(x) AS s FROM t") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + assert_distributed_eq(&ctx, plan).await; + } +} + +/// Proves a task execution error surfaces as a `DataFusionError` out of +/// `run_distributed` -- not swallowed, not hung, not a panic escaping the +/// spawned task -- for the simplest case: a single (no-shuffle) final stage +/// whose plan errors while executing. +/// +/// The failing node here is a `ScalarUDF` that always returns `Err`, wired +/// into a plan of otherwise-ordinary, `datafusion-proto`-native nodes +/// (table scan + projection). That's deliberate: the executor always +/// round-trips every stage's plan through `encode_plan`/`decode_plan` with +/// `ExchangeCodec`, which only knows how to serialize +/// `ExchangeSinkExec`/`ExchangeSourceExec`. A hand-rolled leaf +/// `ExecutionPlan` (as opposed to a UDF referenced by name and resolved +/// from the rebuilt session's function registry on decode) would fail at +/// the encode step instead, which would only prove plan serialization +/// errors propagate -- not task execution errors. +#[cfg(test)] +mod error_propagation_tests { + use std::sync::Arc; + + use datafusion::arrow::array::{Int32Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::error::DataFusionError; + use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, + }; + use datafusion::prelude::SessionContext; + + use crate::config::SchedulerConfig; + use crate::run_distributed; + + /// A scalar UDF that unconditionally fails when invoked, with a message + /// ("boom") distinctive enough to assert on without false-matching some + /// unrelated error path. + #[derive(Debug, PartialEq, Eq, Hash)] + struct BoomUdf { + signature: Signature, + } + + impl BoomUdf { + fn new() -> Self { + Self { + signature: Signature::uniform( + 1, + vec![DataType::Int32], + Volatility::Volatile, + ), + } + } + } + + impl ScalarUDFImpl for BoomUdf { + fn name(&self) -> &str { + "boom" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type( + &self, + _arg_types: &[DataType], + ) -> datafusion::error::Result { + Ok(DataType::Int32) + } + + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> datafusion::error::Result { + Err(DataFusionError::Execution("boom".to_string())) + } + } + + #[tokio::test] + async fn task_execution_error_surfaces_from_run_distributed() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + + let ctx = SessionContext::new(); + ctx.register_batch("t", batch).unwrap(); + ctx.register_udf(ScalarUDF::from(BoomUdf::new())); + + // A single-partition, no-shuffle plan: `create_stages` produces + // exactly one (final) stage, so this exercises `execute_stage_graph`'s + // single-stage path end to end, including the encode/decode + // round-trip. + let df = ctx.sql("SELECT boom(a) FROM t").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let config = SchedulerConfig::in_memory(&ctx); + let result = run_distributed(&ctx, plan, config).await; + + let err = result + .expect_err("task execution error must surface as an Err, not be swallowed"); + let msg = err.to_string(); + assert!( + msg.contains("boom"), + "expected error message to contain \"boom\", got: {msg}" + ); + } +} + +#[cfg(test)] +mod single_stage_tests { + use std::sync::Arc; + + use datafusion::arrow::array::{Int32Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::prelude::*; + + use crate::config::SchedulerConfig; + use crate::run_distributed; + + #[tokio::test] + async fn single_stage_scan_filter_matches_collect() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + ) + .unwrap(); + + let ctx = SessionContext::new(); + ctx.register_batch("t", batch).unwrap(); + + let df = ctx.sql("SELECT a FROM t WHERE a > 2").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let expected = datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()) + .await + .unwrap(); + + let config = SchedulerConfig::in_memory(&ctx); + let actual = run_distributed(&ctx, plan, config).await.unwrap(); + + // order-insensitive compare + assert_eq!(total_rows(&expected), total_rows(&actual)); + assert_eq!(sorted_values(&expected), sorted_values(&actual)); + } + + fn total_rows(b: &[RecordBatch]) -> usize { + b.iter().map(|b| b.num_rows()).sum() + } + fn sorted_values(b: &[RecordBatch]) -> Vec { + let mut v: Vec = b + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + v.sort(); + v + } +} + +/// End-to-end test that a real 2-stage plan (partial aggregate + streaming +/// in-memory exchange + final aggregate) executes correctly through +/// `run_distributed`/`execute_stage_graph`, matching single-process +/// `collect` output. Data crosses the stage boundary as IPC frames over the +/// `InMemoryExchange` — there is no disk shuffle to inspect, so correctness +/// of the 2-stage GROUP BY across the exchange is the assertion. +#[cfg(test)] +mod two_stage_tests { + use std::sync::Arc; + + use datafusion::arrow::array::{Int32Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::collect; + use datafusion::prelude::{SessionConfig, SessionContext}; + + use crate::config::SchedulerConfig; + use crate::exchange::InMemoryExchange; + use crate::run_distributed; + use crate::scheduler::create_stages; + use crate::test_util::assert_distributed_eq; + + /// Registers table `t(a INT)` backed by a 2-partition `MemTable`, + /// against a `SessionContext` configured with `target_partitions = 4`. + async fn context_with_partitioned_table() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![2, 3, 4]))], + ) + .unwrap(); + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx + } + + #[tokio::test] + async fn group_by_across_a_real_shuffle_matches_collect() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) FROM t GROUP BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + assert_distributed_eq(&ctx, plan).await; + } + + /// Proves `run_distributed` is genuinely routed through `create_stages` + /// + the stage-graph executor (data really crossing the exchange), not + /// coincidentally producing correct output by re-running the original + /// in-process plan. + /// + /// We assert on the stage split itself — this GROUP BY splits into a + /// producer stage (id 0, hash-repartitioned partial aggregate, whose + /// output must cross the exchange) and a final stage reading it — and + /// that the end-to-end result over that real two-stage exchange + /// execution matches single-node `collect`. + #[tokio::test] + async fn group_by_data_crosses_the_exchange_between_two_stages() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) FROM t GROUP BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let expected = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + + // The split must produce a producer stage feeding a final stage over + // the exchange — i.e. a genuine 2-stage plan, not a single re-executed + // tree. + let probe_exchange = InMemoryExchange::new(); + let graph = create_stages(plan.clone(), &probe_exchange).unwrap(); + assert_eq!( + graph.stages.len(), + 2, + "GROUP BY should split into a producer stage and a final stage" + ); + assert_ne!( + graph.final_stage_id, 0, + "there must be a producer stage (id 0) feeding the final stage" + ); + + // Execute end-to-end: the producer stage's partial-aggregate output + // only reaches the final stage by crossing the in-memory exchange as + // IPC frames. + let config = SchedulerConfig::in_memory(&ctx); + let actual = run_distributed(&ctx, plan, config).await.unwrap(); + + let expected_rows: usize = expected.iter().map(|b| b.num_rows()).sum(); + let actual_rows: usize = actual.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + expected_rows, actual_rows, + "row count across the exchange must match single-node collect" + ); + } +} + +/// Streaming / backpressure integration test. +/// +/// Runs a real multi-stage query end-to-end through `run_distributed` with +/// `channel_capacity = 1` — the tightest possible backpressure through the +/// FULL executor: producers block on full channels while downstream +/// consumers drain. The whole run is wrapped in a `tokio::time::timeout` so +/// a deadlock FAILS the test rather than hanging CI. +/// +/// This is the end-to-end proof that the spawn-all, no-barrier executor +/// pipelines correctly under backpressure: a consumer blocked on a channel +/// is unblocked by a producer that is already running (all tasks of all +/// stages are spawned before any is awaited). +#[cfg(test)] +mod streaming_tests { + use std::sync::Arc; + use std::time::Duration; + + use datafusion::arrow::array::{Int32Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::collect; + use datafusion::prelude::{SessionConfig, SessionContext}; + + use crate::config::SchedulerConfig; + use crate::run_distributed; + + /// Registers table `t(a INT)` backed by a 2-partition `MemTable` with + /// enough rows/batches that capacity-1 channels actually force repeated + /// blocking, against a `SessionContext` at `target_partitions = 4`. + async fn context_with_many_rows() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Many small batches per partition so the sink blocks on + // `send().await` over and over against the capacity-1 channels. + let mut p0 = Vec::new(); + let mut p1 = Vec::new(); + for i in 0..200i32 { + p0.push( + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![i % 17]))], + ) + .unwrap(), + ); + p1.push( + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![(i + 5) % 17]))], + ) + .unwrap(), + ); + } + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![p0, p1]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx + } + + /// Three-stage chain (`GROUP BY` hash boundary + `ORDER BY` gather + /// boundary) run under `channel_capacity = 1`, inside a timeout. + /// Asserts multiset equivalence with single-node `collect`. A hang + /// (deadlock) fails via timeout; a wrong result fails via the + /// equivalence check. + #[tokio::test] + async fn tight_backpressure_multi_stage_pipeline_matches_collect() { + let result = tokio::time::timeout(Duration::from_secs(30), async { + let ctx = context_with_many_rows().await; + let df = ctx + .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let expected = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + + // capacity = 1: the tightest bound that still makes progress, + // exercised through the full run_distributed executor (not just + // sink/source). + let state = ctx.state(); + let config = SchedulerConfig { + session_builder: Arc::new(move || { + SessionContext::new_with_state(state.clone()) + }), + channel_capacity: 1, + }; + let actual = run_distributed(&ctx, plan, config).await.unwrap(); + + (expected, actual) + }) + .await; + + let (expected, actual) = + result.expect("multi-stage pipeline deadlocked / timed out under capacity=1"); + + // Order-insensitive multiset equivalence via sorted pretty-printed + // rows. + let e = arrow::util::pretty::pretty_format_batches(&expected) + .unwrap() + .to_string(); + let a = arrow::util::pretty::pretty_format_batches(&actual) + .unwrap() + .to_string(); + let mut el: Vec<&str> = e.lines().collect(); + el.sort(); + let mut al: Vec<&str> = a.lines().collect(); + al.sort(); + assert_eq!( + el, al, + "streaming distributed output != collect output under capacity=1" + ); + } +} diff --git a/datafusion-examples/examples/scheduler/scheduler.rs b/datafusion-examples/examples/scheduler/scheduler.rs new file mode 100644 index 0000000000000..7389b10656dc3 --- /dev/null +++ b/datafusion-examples/examples/scheduler/scheduler.rs @@ -0,0 +1,647 @@ +// 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. + +//! Splits a physical plan into a [`StageGraph`] at shuffle boundaries. +//! +//! Mirrors the shape of Ballista's +//! `DefaultDistributedPlanner::plan_query_stages_internal` +//! (`ballista/scheduler/src/planner.rs`), but simplified for this in-process +//! model: no broadcast-join promotion, no subquery special-casing, no AQE. + +use std::sync::Arc; + +use datafusion::error::Result; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::{ + ExecutionPlan, Partitioning, with_new_children_if_necessary, +}; + +use crate::exchange::{ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange}; +use crate::stage::{QueryStage, StageGraph, StageId}; + +/// Splits `plan` into a [`StageGraph`] at shuffle boundaries (hash +/// repartitioning, coalescing to a single partition, and sort-preserving +/// merges), wiring each boundary through the shared streaming `exchange`. +/// +/// Stage ids are allocated in the order stages are produced during the +/// recursion, so producer stages always have a lower id than the consumer +/// stages that read from them. The final stage (the rewritten top-level +/// plan, representing the "collect final result" step) is pushed last and +/// is not wrapped in an `ExchangeSinkExec`. +/// +/// The `exchange` is cloned into every `ExchangeSinkExec`/`ExchangeSourceExec` +/// built here; the executor registers the boundary channels on that same +/// instance before spawning tasks, so producers and consumers share one wire. +pub fn create_stages( + plan: Arc, + exchange: &Arc, +) -> Result { + let mut next_id: StageId = 0; + let (final_plan, mut stages) = + split_at_shuffle_boundaries(plan, exchange, &mut next_id)?; + + let final_stage_id = next_id; + let input_stage_ids = collect_source_stage_ids(&final_plan); + let output_partition_count = final_plan + .properties() + .output_partitioning() + .partition_count(); + stages.push(QueryStage { + id: final_stage_id, + plan: final_plan, + output_partition_count, + input_stage_ids, + }); + + Ok(StageGraph { + stages, + final_stage_id, + }) +} + +/// Recursively rewrites `plan`, replacing every shuffle boundary with an +/// `ExchangeSourceExec` leaf fed by a new producer `QueryStage`. Returns the +/// rewritten plan along with every producer stage created while recursing +/// through it (in the order they were created). +fn split_at_shuffle_boundaries( + plan: Arc, + exchange: &Arc, + next_id: &mut StageId, +) -> Result<(Arc, Vec)> { + if plan.children().is_empty() { + return Ok((plan, Vec::new())); + } + + let mut stages = Vec::new(); + let mut children = Vec::new(); + for child in plan.children() { + let (new_child, mut child_stages) = + split_at_shuffle_boundaries(Arc::clone(child), exchange, next_id)?; + children.push(new_child); + stages.append(&mut child_stages); + } + + // CoalescePartitionsExec and SortPreservingMergeExec both collapse their + // (possibly many) input partitions down to a single output partition; + // they get identical shuffle-boundary treatment here. + if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + { + let child = children[0].clone(); + let source = shuffle_boundary( + child, + Partitioning::RoundRobinBatch(1), + 1, + exchange, + next_id, + &mut stages, + )?; + return Ok((source, stages)); + } + + if let Some(repart) = plan.downcast_ref::() { + return match repart.partitioning().clone() { + Partitioning::Hash(exprs, p) => { + let child = children[0].clone(); + let source = shuffle_boundary( + child, + Partitioning::Hash(exprs, p), + p, + exchange, + next_id, + &mut stages, + )?; + Ok((source, stages)) + } + // v1 simplification: non-hash repartitions (e.g. round-robin) + // exist purely to change intra-stage parallelism, not to cross a + // shuffle boundary. This simplified scheduler has no notion of + // "reshape partition count without a shuffle", so we simply drop + // the node and keep the child's own partitioning/task count. + _ => Ok((children[0].clone(), stages)), + }; + } + + let new_plan = with_new_children_if_necessary(plan, children)?; + Ok((new_plan, stages)) +} + +/// Wraps `child` in a producer `QueryStage` rooted at an `ExchangeSinkExec` +/// (pushed into `stages`), and returns an `ExchangeSourceExec` that stands in +/// for `child` at the shuffle boundary. +fn shuffle_boundary( + child: Arc, + sink_partitioning: Partitioning, + output_partition_count: usize, + exchange: &Arc, + next_id: &mut StageId, + stages: &mut Vec, +) -> Result> { + // Number of producer tasks = the child's own partition count, computed + // before it gets wrapped in the sink (the sink's output partitioning + // controls how many *buckets* each task fans out to, not how many tasks + // the sink itself has -- see `ExchangeSinkExec`). + let num_producer_tasks = child.properties().output_partitioning().partition_count(); + let child_schema = child.schema(); + let input_stage_ids = collect_source_stage_ids(&child); + + let stage_id = *next_id; + *next_id += 1; + + let sink: Arc = Arc::new(ExchangeSinkExec::try_new( + stage_id, + child, + sink_partitioning, + exchange.clone(), + )?); + + stages.push(QueryStage { + id: stage_id, + plan: sink, + output_partition_count, + input_stage_ids, + }); + + let source: Arc = Arc::new(ExchangeSourceExec::try_new( + stage_id, + child_schema, + num_producer_tasks, + output_partition_count, + exchange.clone(), + )?); + + Ok(source) +} + +/// Walks `plan`'s tree looking for `ExchangeSourceExec` leaves and collects +/// their producer stage ids -- these are the upstream stages `plan` depends on. +fn collect_source_stage_ids(plan: &Arc) -> Vec { + let mut ids = Vec::new(); + walk_collect_source_stage_ids(plan, &mut ids); + ids +} + +fn walk_collect_source_stage_ids(plan: &Arc, ids: &mut Vec) { + if let Some(source) = plan.downcast_ref::() { + ids.push(source.from_stage_id()); + } + for child in plan.children() { + walk_collect_source_stage_ids(child, ids); + } +} + +#[cfg(test)] +mod tests { + //! Tests for `create_stages`, which splits a physical plan into a + //! `StageGraph` at shuffle boundaries. + //! + //! Expected stage shape below was derived by printing + //! `displayable(plan.as_ref()).indent(false)` for each query against this + //! DataFusion checkout before writing assertions -- see the plans quoted + //! in each test's doc comment. + + use std::sync::Arc; + + use datafusion::arrow::array::{Int32Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::prelude::{SessionConfig, SessionContext}; + + use crate::exchange::{ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange}; + use crate::scheduler::create_stages; + use crate::stage::StageGraph; + + /// Registers table `t(a INT)` backed by a 2-partition `MemTable`, against + /// a `SessionContext` configured with `target_partitions = 4`. + async fn context_with_partitioned_table() -> SessionContext { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![2, 3, 4]))], + ) + .unwrap(); + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx + } + + fn contains_exchange_sink(plan: &Arc) -> bool { + plan.downcast_ref::().is_some() + || plan.children().iter().any(|c| contains_exchange_sink(c)) + } + + fn contains_exchange_source(plan: &Arc) -> bool { + plan.downcast_ref::().is_some() + || plan.children().iter().any(|c| contains_exchange_source(c)) + } + + /// Registers two 2-partition `MemTable`s, `l(k INT, v INT)` and + /// `r(k INT, v INT)`, against a `SessionContext` configured with + /// `target_partitions = 4` and the hash-join collect-left thresholds + /// dropped to 0 so a join on `k` plans as a `Partitioned` hash join + /// (`RepartitionExec: partitioning=Hash(...)` on both sides) rather than + /// `CollectLeft`, which these tiny in-memory tables would otherwise + /// qualify for. + async fn context_with_two_partitioned_tables() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + + let lbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let lbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![40, 50, 60])), + ], + ) + .unwrap(); + let rbatch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![100, 200, 300])), + ], + ) + .unwrap(); + let rbatch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(Int32Array::from(vec![400, 500, 600])), + ], + ) + .unwrap(); + + let mut config = SessionConfig::new().with_target_partitions(4); + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold = 0; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold_rows = 0; + let ctx = SessionContext::new_with_config(config); + + let ltable = + MemTable::try_new(schema.clone(), vec![vec![lbatch0], vec![lbatch1]]) + .unwrap(); + let rtable = + MemTable::try_new(schema, vec![vec![rbatch0], vec![rbatch1]]).unwrap(); + ctx.register_table("l", Arc::new(ltable)).unwrap(); + ctx.register_table("r", Arc::new(rtable)).unwrap(); + ctx + } + + /// Asserts the topological-order invariant `StageGraph` promises: every + /// producer stage id referenced via `input_stage_ids` is strictly lower + /// than the id of the stage that reads it, and `final_stage_id` is both + /// the highest id in the graph and a stage no other stage depends on. + fn assert_topological_order(graph: &StageGraph) { + for stage in &graph.stages { + for &producer_id in &stage.input_stage_ids { + assert!( + producer_id < stage.id, + "stage {} depends on stage {} which does not have a lower id", + stage.id, + producer_id + ); + } + } + + let max_id = graph.stages.iter().map(|s| s.id).max().unwrap(); + assert_eq!( + graph.final_stage_id, max_id, + "final_stage_id should be the highest stage id in the graph" + ); + assert!( + !graph + .stages + .iter() + .any(|s| s.input_stage_ids.contains(&graph.final_stage_id)), + "no stage should read from the final stage" + ); + } + + /// `SELECT a, count(*) FROM t GROUP BY a` with `target_partitions = 4` + /// over a 2-partition input produces (verified via + /// `displayable(..).indent(false)` against this checkout): + /// + /// ```text + /// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as count(*)] + /// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] + /// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + /// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] + /// DataSourceExec: partitions=2, partition_sizes=[1, 1] + /// ``` + /// + /// The single `RepartitionExec: partitioning=Hash(...)` is the only + /// shuffle boundary, so this should split into exactly 2 stages: a + /// producer stage rooted at `ExchangeSinkExec` (wrapping the + /// partial-aggregate subtree), and a final stage (the + /// projection/final-aggregate subtree) reading the producer's output via + /// a `ExchangeSourceExec` leaf. + #[tokio::test] + async fn group_by_splits_into_producer_and_final_stage_at_hash_repartition() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) FROM t GROUP BY a") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!(graph.stages.len(), 2, "expected exactly 2 stages"); + assert_eq!( + graph.final_stage_id, + graph.stages.last().unwrap().id, + "final stage should be the last one pushed" + ); + + // Producer stage (id 0): the partial-aggregate subtree wrapped in a + // ExchangeSinkExec, hash-partitioned into 4 buckets, with no upstream + // shuffle dependency of its own. + let producer = &graph.stages[0]; + assert_eq!(producer.id, 0); + assert!( + producer.plan.downcast_ref::().is_some(), + "producer stage's plan root should be a ExchangeSinkExec, got: {}", + producer.plan.name() + ); + assert_eq!(producer.output_partition_count, 4); + assert!( + producer.input_stage_ids.is_empty(), + "producer stage should have no upstream shuffle dependencies" + ); + + // Final stage: the projection/final-aggregate subtree, containing a + // ExchangeSourceExec leaf that reads the producer stage's output, and + // NOT itself wrapped in a ExchangeSinkExec. + let final_stage = &graph.stages[1]; + assert_eq!(final_stage.id, graph.final_stage_id); + assert!( + final_stage + .plan + .downcast_ref::() + .is_none(), + "final stage must not be wrapped in a ExchangeSinkExec" + ); + assert!( + contains_exchange_source(&final_stage.plan), + "final stage should contain a ExchangeSourceExec leaf" + ); + assert_eq!( + final_stage.input_stage_ids, + vec![producer.id], + "final stage should depend on the producer stage" + ); + assert_eq!(final_stage.output_partition_count, 4); + } + + /// `SELECT a FROM t WHERE a > 2` has no repartitioning at all (verified + /// via `displayable(..).indent(false)`): + /// + /// ```text + /// FilterExec: a@0 > 2 + /// DataSourceExec: partitions=2, partition_sizes=[1, 1] + /// ``` + /// + /// so it should produce exactly 1 stage containing neither a + /// `ExchangeSinkExec` nor a `ExchangeSourceExec` anywhere in its plan. + #[tokio::test] + async fn filter_query_with_no_repartition_stays_a_single_stage() { + let ctx = context_with_partitioned_table().await; + let df = ctx.sql("SELECT a FROM t WHERE a > 2").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!(graph.stages.len(), 1, "expected exactly 1 stage"); + let only_stage = &graph.stages[0]; + assert_eq!(only_stage.id, graph.final_stage_id); + assert!(only_stage.input_stage_ids.is_empty()); + assert!( + !contains_exchange_sink(&only_stage.plan), + "no-shuffle plan must not contain a ExchangeSinkExec" + ); + assert!( + !contains_exchange_source(&only_stage.plan), + "no-shuffle plan must not contain a ExchangeSourceExec" + ); + } + + /// `SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY c` with + /// `target_partitions = 4` over a 2-partition input produces (verified + /// via `displayable(..).indent(false)` against this checkout): + /// + /// ```text + /// SortPreservingMergeExec: [c@1 ASC NULLS LAST] + /// SortExec: expr=[c@1 ASC NULLS LAST], preserve_partitioning=[true] + /// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as c] + /// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] + /// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + /// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] + /// DataSourceExec: partitions=2, partition_sizes=[1, 1] + /// ``` + /// + /// There are TWO shuffle boundaries here: the `RepartitionExec: Hash(...)` + /// feeding the final aggregate, and the top-level + /// `SortPreservingMergeExec` collapsing the sorted partitions down to 1. + /// That means the split walks a chain of dependent producer stages + /// (stage 1's subtree embeds stage 0's `ExchangeSourceExec`) rather than + /// just a single boundary. + /// + /// This should split into exactly 3 stages: + /// - stage 0: producer wrapping the partial-aggregate subtree (hash + /// repartition boundary), no upstream dependency. + /// - stage 1: producer wrapping the final-aggregate/projection/sort + /// subtree (sort-preserving-merge boundary), depending on stage 0. + /// - stage 2 (final): a bare `ExchangeSourceExec` reading stage 1's + /// output. + #[tokio::test] + async fn multi_boundary_chain_has_ascending_topological_order() { + let ctx = context_with_partitioned_table().await; + let df = ctx + .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY c") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!( + graph.stages.len(), + 3, + "expected 2 producer stages plus a final stage" + ); + + assert_topological_order(&graph); + + let stage0 = &graph.stages[0]; + assert_eq!(stage0.id, 0); + assert!( + stage0.plan.downcast_ref::().is_some(), + "stage 0 should be rooted at a ExchangeSinkExec" + ); + assert!( + stage0.input_stage_ids.is_empty(), + "stage 0 (partial-aggregate producer) has no upstream shuffle dependency" + ); + + let stage1 = &graph.stages[1]; + assert_eq!(stage1.id, 1); + assert!( + stage1.plan.downcast_ref::().is_some(), + "stage 1 should be rooted at a ExchangeSinkExec" + ); + assert_eq!( + stage1.input_stage_ids, + vec![stage0.id], + "stage 1 (sort-preserving-merge producer) depends on stage 0" + ); + assert!( + contains_exchange_source(&stage1.plan), + "stage 1's subtree should embed a ExchangeSourceExec reading stage 0" + ); + + let final_stage = &graph.stages[2]; + assert_eq!(final_stage.id, graph.final_stage_id); + assert!( + final_stage + .plan + .downcast_ref::() + .is_none(), + "final stage must not be wrapped in a ExchangeSinkExec" + ); + assert!( + final_stage + .plan + .downcast_ref::() + .is_some(), + "final stage's plan should itself be the ExchangeSourceExec reading stage 1" + ); + assert_eq!( + final_stage.input_stage_ids, + vec![stage1.id], + "final stage depends on stage 1" + ); + } + + /// `SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k` with + /// `target_partitions = 4` and the hash-join collect-left thresholds + /// forced to 0 (see `context_with_two_partitioned_tables`) produces + /// (verified via `displayable(..).indent(false)` against this checkout): + /// + /// ```text + /// HashJoinExec: mode=Partitioned, join_type=Inner, on=[(k@0, k@0)], projection=[k@0, v@1, v@3] + /// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 + /// DataSourceExec: partitions=2, partition_sizes=[1, 1] + /// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 + /// DataSourceExec: partitions=2, partition_sizes=[1, 1] + /// ``` + /// + /// Both join inputs are independently hash-repartitioned, so this should + /// split into exactly 3 stages: two independent producer stages (one per + /// join side, neither depending on the other) and a final stage rooted + /// at the (rewritten) `HashJoinExec` itself, whose plan contains two + /// `ExchangeSourceExec` leaves -- one per producer. + #[tokio::test] + async fn join_with_two_shuffled_sides_has_two_input_stage_ids() { + let ctx = context_with_two_partitioned_tables().await; + let df = ctx + .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let exchange = InMemoryExchange::new(); + let graph = create_stages(plan, &exchange).unwrap(); + + assert_eq!( + graph.stages.len(), + 3, + "expected 2 independent producer stages plus a final (join) stage" + ); + + assert_topological_order(&graph); + + let final_stage = &graph.stages[2]; + assert_eq!(final_stage.id, graph.final_stage_id); + assert!( + final_stage + .plan + .downcast_ref::() + .is_none(), + "final (join) stage must not be wrapped in a ExchangeSinkExec" + ); + assert_eq!( + final_stage.input_stage_ids.len(), + 2, + "join stage should depend on exactly 2 upstream producer stages" + ); + + let stage_ids: std::collections::HashSet<_> = + graph.stages.iter().map(|s| s.id).collect(); + for &producer_id in &final_stage.input_stage_ids { + assert!( + stage_ids.contains(&producer_id), + "referenced producer stage {producer_id} must exist in the graph" + ); + assert!( + producer_id < final_stage.id, + "producer stage {producer_id} must have a lower id than the join stage" + ); + let producer = graph + .stages + .iter() + .find(|s| s.id == producer_id) + .expect("producer stage must exist"); + assert!( + producer.plan.downcast_ref::().is_some(), + "producer stage {producer_id} should be rooted at a ExchangeSinkExec" + ); + assert!( + producer.input_stage_ids.is_empty(), + "producer stage {producer_id} (leaf hash-repartition) has no upstream dependency" + ); + } + } +} diff --git a/datafusion-scheduler/src/serde.rs b/datafusion-examples/examples/scheduler/serde.rs similarity index 100% rename from datafusion-scheduler/src/serde.rs rename to datafusion-examples/examples/scheduler/serde.rs diff --git a/datafusion-scheduler/src/stage.rs b/datafusion-examples/examples/scheduler/stage.rs similarity index 98% rename from datafusion-scheduler/src/stage.rs rename to datafusion-examples/examples/scheduler/stage.rs index c8e90a0b64f1c..6f8db8f261cb3 100644 --- a/datafusion-scheduler/src/stage.rs +++ b/datafusion-examples/examples/scheduler/stage.rs @@ -45,6 +45,7 @@ pub struct QueryStage { pub output_partition_count: usize, /// Ids of upstream stages this stage reads from via `ExchangeSourceExec` /// leaves in `plan`. + #[cfg_attr(not(test), expect(dead_code))] pub input_stage_ids: Vec, } diff --git a/datafusion-scheduler/src/test_util.rs b/datafusion-examples/examples/scheduler/test_util.rs similarity index 89% rename from datafusion-scheduler/src/test_util.rs rename to datafusion-examples/examples/scheduler/test_util.rs index 762d99897a657..078c213680e03 100644 --- a/datafusion-scheduler/src/test_util.rs +++ b/datafusion-examples/examples/scheduler/test_util.rs @@ -20,7 +20,8 @@ use std::sync::Arc; use datafusion::physical_plan::{ExecutionPlan, collect}; use datafusion::prelude::SessionContext; -use crate::{SchedulerConfig, run_distributed}; +use crate::config::SchedulerConfig; +use crate::run_distributed; /// Assert distributed execution yields the same multiset of rows as `collect`. pub async fn assert_distributed_eq(ctx: &SessionContext, plan: Arc) { @@ -31,10 +32,10 @@ pub async fn assert_distributed_eq(ctx: &SessionContext, plan: Arc = e.lines().collect(); diff --git a/datafusion-scheduler/.gitignore b/datafusion-scheduler/.gitignore deleted file mode 100644 index 03314f77b5aa4..0000000000000 --- a/datafusion-scheduler/.gitignore +++ /dev/null @@ -1 +0,0 @@ -Cargo.lock diff --git a/datafusion-scheduler/Cargo.toml b/datafusion-scheduler/Cargo.toml deleted file mode 100644 index 2a3cb3519f0b6..0000000000000 --- a/datafusion-scheduler/Cargo.toml +++ /dev/null @@ -1,40 +0,0 @@ -# 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. - -[package] -name = "datafusion-scheduler" -description = "In-process model of stage-based distributed execution for Apache DataFusion (test/regression harness)" -version = "0.1.0" -edition = "2024" -rust-version = "1.88.0" -license = "Apache-2.0" -publish = false - -[dependencies] -datafusion = { path = "../datafusion/core" } -datafusion-common = { path = "../datafusion/common" } -datafusion-common-runtime = { path = "../datafusion/common-runtime" } -datafusion-proto = { path = "../datafusion/proto" } -futures = "0.3" -tokio = { version = "1", features = ["rt-multi-thread", "macros"] } -prost = "0.14.1" # pinned to match datafusion-proto's resolved prost (verified via `cargo tree -i prost`) - -[dev-dependencies] -tempfile = "3" # only tests need a temp dir now (the #1907 CSV-scan regression guard) - -[lints.clippy] -all = "warn" diff --git a/datafusion-scheduler/src/lib.rs b/datafusion-scheduler/src/lib.rs deleted file mode 100644 index 4d03c6100c0de..0000000000000 --- a/datafusion-scheduler/src/lib.rs +++ /dev/null @@ -1,70 +0,0 @@ -// 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. - -//! In-process model of stage-based distributed execution for Apache DataFusion. -//! -//! Splits a physical plan into stages at shuffle boundaries, serializes each -//! stage, and runs each task from a freshly deserialized plan with its own -//! `TaskContext` — surfacing physical-plan changes that assume shared -//! in-process state. Data crosses stage boundaries as IPC-serialized frames -//! over a streaming in-memory exchange (no disk, no barrier). See -//! `specs/datafusion-scheduler.md` in ballista-ai-dev. - -mod config; -pub mod exchange; -mod executor; -mod scheduler; -mod serde; -mod stage; -pub mod test_util; - -pub use config::SchedulerConfig; -pub use exchange::{ - ExchangeCodec, ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange, -}; -pub use executor::execute_stage_graph; -pub use scheduler::create_stages; -pub use stage::{QueryStage, StageGraph, StageId}; -pub use test_util::assert_distributed_eq; - -use std::sync::Arc; - -use datafusion::arrow::array::RecordBatch; -use datafusion::error::Result; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::prelude::SessionContext; - -/// Splits `plan` into a [`StageGraph`] at shuffle boundaries and runs it -/// through the streaming, no-barrier stage executor -/// ([`execute_stage_graph`]): all tasks of all stages are spawned -/// concurrently, data crosses each boundary as IPC frames over a single shared -/// [`InMemoryExchange`], and the final stage's collected output is returned. -/// -/// `ctx` is the driver session; each task rebuilds an isolated session via -/// `config.session_builder` and never touches `ctx`. -pub async fn run_distributed( - _ctx: &SessionContext, - plan: Arc, - config: SchedulerConfig, -) -> Result> { - // One exchange instance is threaded through both stage splitting (so - // sinks/sources capture it) and the executor (so it registers the channels - // and drives the tasks) — producers and consumers must share ONE wire. - let exchange = InMemoryExchange::new(); - let graph = create_stages(plan, &exchange)?; - execute_stage_graph(&graph, &exchange, &config).await -} diff --git a/datafusion-scheduler/src/scheduler.rs b/datafusion-scheduler/src/scheduler.rs deleted file mode 100644 index f4d596ff77d21..0000000000000 --- a/datafusion-scheduler/src/scheduler.rs +++ /dev/null @@ -1,207 +0,0 @@ -// 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. - -//! Splits a physical plan into a [`StageGraph`] at shuffle boundaries. -//! -//! Mirrors the shape of Ballista's -//! `DefaultDistributedPlanner::plan_query_stages_internal` -//! (`ballista/scheduler/src/planner.rs`), but simplified for this in-process -//! model: no broadcast-join promotion, no subquery special-casing, no AQE. - -use std::sync::Arc; - -use datafusion::error::Result; -use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::repartition::RepartitionExec; -use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; -use datafusion::physical_plan::{ - ExecutionPlan, Partitioning, with_new_children_if_necessary, -}; - -use crate::exchange::{ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange}; -use crate::stage::{QueryStage, StageGraph, StageId}; - -/// Splits `plan` into a [`StageGraph`] at shuffle boundaries (hash -/// repartitioning, coalescing to a single partition, and sort-preserving -/// merges), wiring each boundary through the shared streaming `exchange`. -/// -/// Stage ids are allocated in the order stages are produced during the -/// recursion, so producer stages always have a lower id than the consumer -/// stages that read from them. The final stage (the rewritten top-level -/// plan, representing the "collect final result" step) is pushed last and -/// is not wrapped in an `ExchangeSinkExec`. -/// -/// The `exchange` is cloned into every `ExchangeSinkExec`/`ExchangeSourceExec` -/// built here; the executor registers the boundary channels on that same -/// instance before spawning tasks, so producers and consumers share one wire. -pub fn create_stages( - plan: Arc, - exchange: &Arc, -) -> Result { - let mut next_id: StageId = 0; - let (final_plan, mut stages) = - split_at_shuffle_boundaries(plan, exchange, &mut next_id)?; - - let final_stage_id = next_id; - let input_stage_ids = collect_source_stage_ids(&final_plan); - let output_partition_count = final_plan - .properties() - .output_partitioning() - .partition_count(); - stages.push(QueryStage { - id: final_stage_id, - plan: final_plan, - output_partition_count, - input_stage_ids, - }); - - Ok(StageGraph { - stages, - final_stage_id, - }) -} - -/// Recursively rewrites `plan`, replacing every shuffle boundary with an -/// `ExchangeSourceExec` leaf fed by a new producer `QueryStage`. Returns the -/// rewritten plan along with every producer stage created while recursing -/// through it (in the order they were created). -fn split_at_shuffle_boundaries( - plan: Arc, - exchange: &Arc, - next_id: &mut StageId, -) -> Result<(Arc, Vec)> { - if plan.children().is_empty() { - return Ok((plan, Vec::new())); - } - - let mut stages = Vec::new(); - let mut children = Vec::new(); - for child in plan.children() { - let (new_child, mut child_stages) = - split_at_shuffle_boundaries(Arc::clone(child), exchange, next_id)?; - children.push(new_child); - stages.append(&mut child_stages); - } - - // CoalescePartitionsExec and SortPreservingMergeExec both collapse their - // (possibly many) input partitions down to a single output partition; - // they get identical shuffle-boundary treatment here. - if plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() - { - let child = children[0].clone(); - let source = shuffle_boundary( - child, - Partitioning::RoundRobinBatch(1), - 1, - exchange, - next_id, - &mut stages, - )?; - return Ok((source, stages)); - } - - if let Some(repart) = plan.downcast_ref::() { - return match repart.partitioning().clone() { - Partitioning::Hash(exprs, p) => { - let child = children[0].clone(); - let source = shuffle_boundary( - child, - Partitioning::Hash(exprs, p), - p, - exchange, - next_id, - &mut stages, - )?; - Ok((source, stages)) - } - // v1 simplification: non-hash repartitions (e.g. round-robin) - // exist purely to change intra-stage parallelism, not to cross a - // shuffle boundary. This simplified scheduler has no notion of - // "reshape partition count without a shuffle", so we simply drop - // the node and keep the child's own partitioning/task count. - _ => Ok((children[0].clone(), stages)), - }; - } - - let new_plan = with_new_children_if_necessary(plan, children)?; - Ok((new_plan, stages)) -} - -/// Wraps `child` in a producer `QueryStage` rooted at an `ExchangeSinkExec` -/// (pushed into `stages`), and returns an `ExchangeSourceExec` that stands in -/// for `child` at the shuffle boundary. -fn shuffle_boundary( - child: Arc, - sink_partitioning: Partitioning, - output_partition_count: usize, - exchange: &Arc, - next_id: &mut StageId, - stages: &mut Vec, -) -> Result> { - // Number of producer tasks = the child's own partition count, computed - // before it gets wrapped in the sink (the sink's output partitioning - // controls how many *buckets* each task fans out to, not how many tasks - // the sink itself has -- see `ExchangeSinkExec`). - let num_producer_tasks = child.properties().output_partitioning().partition_count(); - let child_schema = child.schema(); - let input_stage_ids = collect_source_stage_ids(&child); - - let stage_id = *next_id; - *next_id += 1; - - let sink: Arc = Arc::new(ExchangeSinkExec::try_new( - stage_id, - child, - sink_partitioning, - exchange.clone(), - )?); - - stages.push(QueryStage { - id: stage_id, - plan: sink, - output_partition_count, - input_stage_ids, - }); - - let source: Arc = Arc::new(ExchangeSourceExec::try_new( - stage_id, - child_schema, - num_producer_tasks, - output_partition_count, - exchange.clone(), - )?); - - Ok(source) -} - -/// Walks `plan`'s tree looking for `ExchangeSourceExec` leaves and collects -/// their producer stage ids -- these are the upstream stages `plan` depends on. -fn collect_source_stage_ids(plan: &Arc) -> Vec { - let mut ids = Vec::new(); - walk_collect_source_stage_ids(plan, &mut ids); - ids -} - -fn walk_collect_source_stage_ids(plan: &Arc, ids: &mut Vec) { - if let Some(source) = plan.downcast_ref::() { - ids.push(source.from_stage_id()); - } - for child in plan.children() { - walk_collect_source_stage_ids(child, ids); - } -} diff --git a/datafusion-scheduler/tests/equivalence.rs b/datafusion-scheduler/tests/equivalence.rs deleted file mode 100644 index d0498d5db31de..0000000000000 --- a/datafusion-scheduler/tests/equivalence.rs +++ /dev/null @@ -1,299 +0,0 @@ -// 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. - -//! End-to-end equivalence tests: run representative plans (joins, sorts, -//! multi-stage chains) through `run_distributed` and assert the result is the -//! same multiset of rows as single-node `collect`, via `assert_distributed_eq`. -//! -//! Plus the #1907-class regression guard: prove that when a file scan is split -//! into isolated per-task plans, each task reads ONLY its own partition rather -//! than draining a plan-instance-shared pool of file work and scanning the -//! whole table (the failure class of Ballista issue #1907). -//! -//! Plan shapes quoted in each test's doc comment were captured with -//! `displayable(plan.as_ref()).indent(false)` against this DataFusion checkout -//! at `target_partitions = 4`. - -use std::sync::Arc; - -use datafusion::arrow::array::{Int32Array, RecordBatch}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::MemTable; -use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; -use datafusion_scheduler::test_util::assert_distributed_eq; - -/// Registers two 2-partition `MemTable`s, `l(k INT, v INT)` and -/// `r(k INT, v INT)`, against a `SessionContext` configured with -/// `target_partitions = 4`. The hash-join single-partition thresholds are -/// dropped to 0 so a join on `k` plans as a `Partitioned` join with both sides -/// hash-repartitioned (two shuffle inputs into the join stage) rather than a -/// `CollectLeft`/broadcast that these tiny tables would otherwise qualify for. -/// -/// `prefer_hash_join` selects a `HashJoinExec` (true) vs a `SortMergeJoinExec` -/// (false); both still hash-repartition each side. -async fn context_with_two_partitioned_tables(prefer_hash_join: bool) -> SessionContext { - let schema = Arc::new(Schema::new(vec![ - Field::new("k", DataType::Int32, false), - Field::new("v", DataType::Int32, false), - ])); - - let lbatch0 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(Int32Array::from(vec![10, 20, 30])), - ], - ) - .unwrap(); - let lbatch1 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![4, 5, 6])), - Arc::new(Int32Array::from(vec![40, 50, 60])), - ], - ) - .unwrap(); - let rbatch0 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(Int32Array::from(vec![100, 200, 300])), - ], - ) - .unwrap(); - let rbatch1 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![4, 5, 6])), - Arc::new(Int32Array::from(vec![400, 500, 600])), - ], - ) - .unwrap(); - - let mut config = SessionConfig::new().with_target_partitions(4); - config.options_mut().optimizer.prefer_hash_join = prefer_hash_join; - config - .options_mut() - .optimizer - .hash_join_single_partition_threshold = 0; - config - .options_mut() - .optimizer - .hash_join_single_partition_threshold_rows = 0; - let ctx = SessionContext::new_with_config(config); - - let ltable = - MemTable::try_new(schema.clone(), vec![vec![lbatch0], vec![lbatch1]]).unwrap(); - let rtable = MemTable::try_new(schema, vec![vec![rbatch0], vec![rbatch1]]).unwrap(); - ctx.register_table("l", Arc::new(ltable)).unwrap(); - ctx.register_table("r", Arc::new(rtable)).unwrap(); - ctx -} - -/// Registers table `t(a INT)` backed by a 2-partition `MemTable`, against a -/// `SessionContext` configured with `target_partitions = 4`. -async fn context_with_partitioned_table() -> SessionContext { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let batch0 = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![1, 2, 3, 2, 5]))], - ) - .unwrap(); - let batch1 = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![2, 3, 4, 4, 1]))], - ) - .unwrap(); - - let config = SessionConfig::new().with_target_partitions(4); - let ctx = SessionContext::new_with_config(config); - let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); - ctx.register_table("t", Arc::new(table)).unwrap(); - ctx -} - -/// **Test 1 — hash join.** `SELECT ... FROM l JOIN r ON l.k = r.k` with -/// `prefer_hash_join = true`. Both sides hash-repartition on `k`, so the join -/// stage reads two shuffle inputs. -/// -/// Plan (`displayable(..).indent(false)`, target_partitions = 4): -/// -/// ```text -/// HashJoinExec: mode=Partitioned, join_type=Inner, on=[(k@0, k@0)], projection=[k@0, v@1, v@3] -/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// ``` -#[tokio::test] -async fn hash_join_matches_collect() { - let ctx = context_with_two_partitioned_tables(true).await; - let df = ctx - .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - assert_distributed_eq(&ctx, plan).await; -} - -/// **Test 2 — sort-merge join.** Same query with `prefer_hash_join = false`, -/// which plans a `SortMergeJoinExec`; each side is still hash-repartitioned on -/// `k` and then sorted, so the join stage again reads two shuffle inputs. -#[tokio::test] -async fn sort_merge_join_matches_collect() { - let ctx = context_with_two_partitioned_tables(false).await; - let df = ctx - .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - assert_distributed_eq(&ctx, plan).await; -} - -/// **Test 3 — sort (multiset equivalence).** `SELECT a FROM t ORDER BY a` -/// WITHOUT a LIMIT. The top `SortPreservingMergeExec` is a shuffle boundary -/// (gather to 1 partition). `assert_distributed_eq` compares order-insensitively -/// and v1 models `SortPreservingMergeExec` as a plain gather (no global-order -/// preservation, a deliberate spec deferral — see the spec's "Open questions"), -/// so a full ORDER BY yields the same MULTISET and must pass. -/// -/// Plan (`displayable(..).indent(false)`, target_partitions = 4): -/// -/// ```text -/// SortPreservingMergeExec: [a@0 ASC NULLS LAST] -/// SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// ``` -#[tokio::test] -async fn sort_without_limit_matches_collect_as_multiset() { - let ctx = context_with_partitioned_table().await; - let df = ctx.sql("SELECT a FROM t ORDER BY a").await.unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - assert_distributed_eq(&ctx, plan).await; -} - -/// Documents a known v1 limitation, kept as an `#[ignore]`d test so it never -/// gates CI. `SELECT a FROM t ORDER BY a LIMIT 5` (top-N) selects a specific -/// SET of rows that depends on GLOBAL order. Because v1 models -/// `SortPreservingMergeExec` as a plain gather that does NOT preserve global -/// order (see the spec's "Open questions": SortPreservingMergeExec global-order -/// fidelity), the distributed path can legitimately return a different set of 5 -/// rows than single-node `collect`. This is expected until order-preserving -/// shuffle lands, so the test is ignored rather than asserted. -#[tokio::test] -#[ignore = "top-N after sort is a known v1 limitation: SortPreservingMergeExec \ - is modeled as an unordered gather (see spec Open questions)"] -async fn top_n_after_sort_is_a_known_v1_limitation() { - let ctx = context_with_partitioned_table().await; - let df = ctx.sql("SELECT a FROM t ORDER BY a LIMIT 5").await.unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - assert_distributed_eq(&ctx, plan).await; -} - -/// **Test 4 — three-stage chain.** `SELECT a, count(*) AS c FROM t GROUP BY a -/// ORDER BY a`: a hash-repartition boundary for the final aggregate plus a -/// `SortPreservingMergeExec` gather for the sort → two shuffle boundaries, three -/// stages. Multiset equivalence (order-insensitive). -/// -/// Plan (`displayable(..).indent(false)`, target_partitions = 4): -/// -/// ```text -/// SortPreservingMergeExec: [a@0 ASC NULLS LAST] -/// SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] -/// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as c] -/// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] -/// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 -/// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// ``` -#[tokio::test] -async fn group_by_then_order_by_three_stage_chain_matches_collect() { - let ctx = context_with_partitioned_table().await; - let df = ctx - .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - assert_distributed_eq(&ctx, plan).await; -} - -/// **#1907 regression guard — CAUGHT REGRESSION, currently `#[ignore]`d.** -/// -/// Writes 100 rows (`x = 0..24, 100..124, 200..224, 300..324`) across 4 CSV -/// files, registers them as one table, and runs `SELECT count(*), sum(x)`. When -/// the scan splits into isolated per-task plans, each task must read ONLY its -/// own file group. This is the exact failure class of Ballista issue -/// [#1907](https://github.com/apache/datafusion-ballista/issues/1907). In this -/// checkout the mechanism is the **morsel-driven file scan** -/// (`datafusion/datasource/src/morsel/`): a `DataSourceExec` hands file work out -/// from a pool **shared across the plan instance's output partitions**, drained -/// cooperatively only when *all* partitions are polled concurrently. A -/// distributed executor polls **one partition per isolated task**, so that task -/// drains the whole pool and scans the entire table. (Reproduces with a bare -/// DataFusion plan and a single `execute(0, ..)` call — no scheduler involved.) -/// -/// This harness surfaces exactly that. Observed at `target_partitions = 4` -/// (see the captured plan/stage dump in the Task 7 report): -/// -/// - single-node `collect` (all 4 partitions polled together, queue drained -/// cooperatively): `n = 100`, `s = 16200` — correct. -/// - `run_distributed` (4 isolated tasks, each polling its 1 partition alone): -/// `n = 400`, `s = 64800` — **exactly 4× (= num_producer_tasks) inflation.** -/// -/// The stage split is correct (stage 0 = `Partial` aggregate over the -/// `DataSourceExec: file_groups={4 groups}`, `num_producer_tasks = 4`; stage 1 = -/// `Final` aggregate reading the shuffle) and each task's file paths are -/// correct — the inflation is DataFusion's morsel-driven scan behavior, NOT a -/// defect in this harness. Per the crate's purpose we do NOT weaken the -/// assertion or force it green; the test is `#[ignore]`d as a caught regression -/// to be triaged. Run with `--ignored` to reproduce; it should pass once the -/// underlying DataFusion scan-isolation behavior is fixed. -#[tokio::test] -#[ignore = "CAUGHT #1907-class regression: DataFusion's morsel-driven file \ - scan makes each isolated per-task scan read the whole table \ - (observed 4x inflation: n=400 vs 100, s=64800 vs 16200). Genuine \ - DataFusion behavior surfaced by the harness, not a test defect."] -async fn regression_isolated_scan_reads_only_its_partition() { - let dir = tempfile::TempDir::new().unwrap(); - // 4 files, 25 rows each = 100 rows total; distinct x per file so sum is - // exact and any per-task double-read is visible as an integer multiple. - for f in 0..4 { - let mut body = String::from("x\n"); - for i in 0..25 { - let x = f * 100 + i; - body.push_str(&x.to_string()); - body.push('\n'); - } - std::fs::write(dir.path().join(format!("part_{f}.csv")), body).unwrap(); - } - - let config = SessionConfig::new().with_target_partitions(4); - let ctx = SessionContext::new_with_config(config); - ctx.register_csv("t", dir.path().to_str().unwrap(), CsvReadOptions::new()) - .await - .unwrap(); - - let df = ctx - .sql("SELECT count(*) AS n, sum(x) AS s FROM t") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - // Distributed sum/count must equal single-node collect. A whole-table-per - // -task scan inflates both by the task count; this assertion catches it. - assert_distributed_eq(&ctx, plan).await; -} diff --git a/datafusion-scheduler/tests/error_propagation.rs b/datafusion-scheduler/tests/error_propagation.rs deleted file mode 100644 index 17c61feff4ae7..0000000000000 --- a/datafusion-scheduler/tests/error_propagation.rs +++ /dev/null @@ -1,111 +0,0 @@ -// 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. - -//! Proves a task execution error surfaces as a `DataFusionError` out of -//! `run_distributed` -- not swallowed, not hung, not a panic escaping the -//! spawned task -- for the simplest case: a single (no-shuffle) final stage -//! whose plan errors while executing. -//! -//! The failing node here is a `ScalarUDF` that always returns `Err`, wired -//! into a plan of otherwise-ordinary, `datafusion-proto`-native nodes (table -//! scan + projection). That's deliberate: the executor always round-trips -//! every stage's plan through `encode_plan`/`decode_plan` with -//! [`ExchangeCodec`](datafusion_scheduler::ExchangeCodec), which only knows how -//! to serialize `ExchangeSinkExec`/`ExchangeSourceExec`. A hand-rolled leaf -//! `ExecutionPlan` (as opposed to a UDF referenced by name and resolved from -//! the rebuilt session's function registry on decode) would fail at the -//! encode step instead, which would only prove plan serialization errors -//! propagate -- not task execution errors. - -use std::sync::Arc; - -use datafusion::arrow::array::{Int32Array, RecordBatch}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::error::DataFusionError; -use datafusion::logical_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, -}; -use datafusion::prelude::SessionContext; -use datafusion_scheduler::{SchedulerConfig, run_distributed}; - -/// A scalar UDF that unconditionally fails when invoked, with a message -/// ("boom") distinctive enough to assert on without false-matching some -/// unrelated error path. -#[derive(Debug, PartialEq, Eq, Hash)] -struct BoomUdf { - signature: Signature, -} - -impl BoomUdf { - fn new() -> Self { - Self { - signature: Signature::uniform(1, vec![DataType::Int32], Volatility::Volatile), - } - } -} - -impl ScalarUDFImpl for BoomUdf { - fn name(&self) -> &str { - "boom" - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type( - &self, - _arg_types: &[DataType], - ) -> datafusion::error::Result { - Ok(DataType::Int32) - } - - fn invoke_with_args( - &self, - _args: ScalarFunctionArgs, - ) -> datafusion::error::Result { - Err(DataFusionError::Execution("boom".to_string())) - } -} - -#[tokio::test] -async fn task_execution_error_surfaces_from_run_distributed() { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let batch = - RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); - - let ctx = SessionContext::new(); - ctx.register_batch("t", batch).unwrap(); - ctx.register_udf(ScalarUDF::from(BoomUdf::new())); - - // A single-partition, no-shuffle plan: `create_stages` produces exactly - // one (final) stage, so this exercises `execute_stage_graph`'s - // single-stage path end to end, including the encode/decode round-trip. - let df = ctx.sql("SELECT boom(a) FROM t").await.unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let config = SchedulerConfig::in_memory(&ctx); - let result = run_distributed(&ctx, plan, config).await; - - let err = result - .expect_err("task execution error must surface as an Err, not be swallowed"); - let msg = err.to_string(); - assert!( - msg.contains("boom"), - "expected error message to contain \"boom\", got: {msg}" - ); -} diff --git a/datafusion-scheduler/tests/single_stage.rs b/datafusion-scheduler/tests/single_stage.rs deleted file mode 100644 index b8b7ff75342ce..0000000000000 --- a/datafusion-scheduler/tests/single_stage.rs +++ /dev/null @@ -1,68 +0,0 @@ -// 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. - -use datafusion::arrow::array::{Int32Array, RecordBatch}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::prelude::*; -use datafusion_scheduler::{SchedulerConfig, run_distributed}; -use std::sync::Arc; - -#[tokio::test] -async fn single_stage_scan_filter_matches_collect() { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], - ) - .unwrap(); - - let ctx = SessionContext::new(); - ctx.register_batch("t", batch).unwrap(); - - let df = ctx.sql("SELECT a FROM t WHERE a > 2").await.unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let expected = datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()) - .await - .unwrap(); - - let config = SchedulerConfig::in_memory(&ctx); - let actual = run_distributed(&ctx, plan, config).await.unwrap(); - - // order-insensitive compare - assert_eq!(total_rows(&expected), total_rows(&actual)); - assert_eq!(sorted_values(&expected), sorted_values(&actual)); -} - -fn total_rows(b: &[RecordBatch]) -> usize { - b.iter().map(|b| b.num_rows()).sum() -} -fn sorted_values(b: &[RecordBatch]) -> Vec { - let mut v: Vec = b - .iter() - .flat_map(|b| { - b.column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec() - }) - .collect(); - v.sort(); - v -} diff --git a/datafusion-scheduler/tests/stage_splitting.rs b/datafusion-scheduler/tests/stage_splitting.rs deleted file mode 100644 index 2e38ec56ffff4..0000000000000 --- a/datafusion-scheduler/tests/stage_splitting.rs +++ /dev/null @@ -1,458 +0,0 @@ -// 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. - -//! Tests for `datafusion_scheduler::create_stages`, which splits a physical -//! plan into a `StageGraph` at shuffle boundaries. -//! -//! Expected stage shape below was derived by printing -//! `displayable(plan.as_ref()).indent(false)` for each query against this -//! DataFusion checkout before writing assertions -- see the plans quoted in -//! each test's doc comment. - -use std::sync::Arc; - -use datafusion::arrow::array::{Int32Array, RecordBatch}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::MemTable; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_scheduler::{ - ExchangeSinkExec, ExchangeSourceExec, InMemoryExchange, create_stages, -}; - -/// Registers table `t(a INT)` backed by a 2-partition `MemTable`, against a -/// `SessionContext` configured with `target_partitions = 4`. -async fn context_with_partitioned_table() -> SessionContext { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let batch0 = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - ) - .unwrap(); - let batch1 = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![2, 3, 4]))], - ) - .unwrap(); - - let config = SessionConfig::new().with_target_partitions(4); - let ctx = SessionContext::new_with_config(config); - let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); - ctx.register_table("t", Arc::new(table)).unwrap(); - ctx -} - -fn contains_exchange_sink(plan: &Arc) -> bool { - plan.downcast_ref::().is_some() - || plan.children().iter().any(|c| contains_exchange_sink(c)) -} - -fn contains_exchange_source(plan: &Arc) -> bool { - plan.downcast_ref::().is_some() - || plan.children().iter().any(|c| contains_exchange_source(c)) -} - -/// Registers two 2-partition `MemTable`s, `l(k INT, v INT)` and -/// `r(k INT, v INT)`, against a `SessionContext` configured with -/// `target_partitions = 4` and the hash-join collect-left thresholds -/// dropped to 0 so a join on `k` plans as a `Partitioned` hash join -/// (`RepartitionExec: partitioning=Hash(...)` on both sides) rather than -/// `CollectLeft`, which these tiny in-memory tables would otherwise qualify -/// for. -async fn context_with_two_partitioned_tables() -> SessionContext { - let schema = Arc::new(Schema::new(vec![ - Field::new("k", DataType::Int32, false), - Field::new("v", DataType::Int32, false), - ])); - - let lbatch0 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(Int32Array::from(vec![10, 20, 30])), - ], - ) - .unwrap(); - let lbatch1 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![4, 5, 6])), - Arc::new(Int32Array::from(vec![40, 50, 60])), - ], - ) - .unwrap(); - let rbatch0 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(Int32Array::from(vec![100, 200, 300])), - ], - ) - .unwrap(); - let rbatch1 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![4, 5, 6])), - Arc::new(Int32Array::from(vec![400, 500, 600])), - ], - ) - .unwrap(); - - let mut config = SessionConfig::new().with_target_partitions(4); - config - .options_mut() - .optimizer - .hash_join_single_partition_threshold = 0; - config - .options_mut() - .optimizer - .hash_join_single_partition_threshold_rows = 0; - let ctx = SessionContext::new_with_config(config); - - let ltable = - MemTable::try_new(schema.clone(), vec![vec![lbatch0], vec![lbatch1]]).unwrap(); - let rtable = MemTable::try_new(schema, vec![vec![rbatch0], vec![rbatch1]]).unwrap(); - ctx.register_table("l", Arc::new(ltable)).unwrap(); - ctx.register_table("r", Arc::new(rtable)).unwrap(); - ctx -} - -/// Asserts the topological-order invariant `StageGraph` promises: every -/// producer stage id referenced via `input_stage_ids` is strictly lower than -/// the id of the stage that reads it, and `final_stage_id` is both the -/// highest id in the graph and a stage no other stage depends on. -fn assert_topological_order(graph: &datafusion_scheduler::StageGraph) { - for stage in &graph.stages { - for &producer_id in &stage.input_stage_ids { - assert!( - producer_id < stage.id, - "stage {} depends on stage {} which does not have a lower id", - stage.id, - producer_id - ); - } - } - - let max_id = graph.stages.iter().map(|s| s.id).max().unwrap(); - assert_eq!( - graph.final_stage_id, max_id, - "final_stage_id should be the highest stage id in the graph" - ); - assert!( - !graph - .stages - .iter() - .any(|s| s.input_stage_ids.contains(&graph.final_stage_id)), - "no stage should read from the final stage" - ); -} - -/// `SELECT a, count(*) FROM t GROUP BY a` with `target_partitions = 4` over a -/// 2-partition input produces (verified via `displayable(..).indent(false)` -/// against this checkout): -/// -/// ```text -/// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as count(*)] -/// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] -/// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 -/// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// ``` -/// -/// The single `RepartitionExec: partitioning=Hash(...)` is the only shuffle -/// boundary, so this should split into exactly 2 stages: a producer stage -/// rooted at `ExchangeSinkExec` (wrapping the partial-aggregate subtree), -/// and a final stage (the projection/final-aggregate subtree) reading the -/// producer's output via a `ExchangeSourceExec` leaf. -#[tokio::test] -async fn group_by_splits_into_producer_and_final_stage_at_hash_repartition() { - let ctx = context_with_partitioned_table().await; - let df = ctx - .sql("SELECT a, count(*) FROM t GROUP BY a") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let exchange = InMemoryExchange::new(); - let graph = create_stages(plan, &exchange).unwrap(); - - assert_eq!(graph.stages.len(), 2, "expected exactly 2 stages"); - assert_eq!( - graph.final_stage_id, - graph.stages.last().unwrap().id, - "final stage should be the last one pushed" - ); - - // Producer stage (id 0): the partial-aggregate subtree wrapped in a - // ExchangeSinkExec, hash-partitioned into 4 buckets, with no upstream - // shuffle dependency of its own. - let producer = &graph.stages[0]; - assert_eq!(producer.id, 0); - assert!( - producer.plan.downcast_ref::().is_some(), - "producer stage's plan root should be a ExchangeSinkExec, got: {}", - producer.plan.name() - ); - assert_eq!(producer.output_partition_count, 4); - assert!( - producer.input_stage_ids.is_empty(), - "producer stage should have no upstream shuffle dependencies" - ); - - // Final stage: the projection/final-aggregate subtree, containing a - // ExchangeSourceExec leaf that reads the producer stage's output, and - // NOT itself wrapped in a ExchangeSinkExec. - let final_stage = &graph.stages[1]; - assert_eq!(final_stage.id, graph.final_stage_id); - assert!( - final_stage - .plan - .downcast_ref::() - .is_none(), - "final stage must not be wrapped in a ExchangeSinkExec" - ); - assert!( - contains_exchange_source(&final_stage.plan), - "final stage should contain a ExchangeSourceExec leaf" - ); - assert_eq!( - final_stage.input_stage_ids, - vec![producer.id], - "final stage should depend on the producer stage" - ); - assert_eq!(final_stage.output_partition_count, 4); -} - -/// `SELECT a FROM t WHERE a > 2` has no repartitioning at all (verified via -/// `displayable(..).indent(false)`): -/// -/// ```text -/// FilterExec: a@0 > 2 -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// ``` -/// -/// so it should produce exactly 1 stage containing neither a -/// `ExchangeSinkExec` nor a `ExchangeSourceExec` anywhere in its plan. -#[tokio::test] -async fn filter_query_with_no_repartition_stays_a_single_stage() { - let ctx = context_with_partitioned_table().await; - let df = ctx.sql("SELECT a FROM t WHERE a > 2").await.unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let exchange = InMemoryExchange::new(); - let graph = create_stages(plan, &exchange).unwrap(); - - assert_eq!(graph.stages.len(), 1, "expected exactly 1 stage"); - let only_stage = &graph.stages[0]; - assert_eq!(only_stage.id, graph.final_stage_id); - assert!(only_stage.input_stage_ids.is_empty()); - assert!( - !contains_exchange_sink(&only_stage.plan), - "no-shuffle plan must not contain a ExchangeSinkExec" - ); - assert!( - !contains_exchange_source(&only_stage.plan), - "no-shuffle plan must not contain a ExchangeSourceExec" - ); -} - -/// `SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY c` with -/// `target_partitions = 4` over a 2-partition input produces (verified via -/// `displayable(..).indent(false)` against this checkout): -/// -/// ```text -/// SortPreservingMergeExec: [c@1 ASC NULLS LAST] -/// SortExec: expr=[c@1 ASC NULLS LAST], preserve_partitioning=[true] -/// ProjectionExec: expr=[a@0 as a, count(Int64(1))@1 as c] -/// AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] -/// RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 -/// AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// ``` -/// -/// There are TWO shuffle boundaries here: the `RepartitionExec: Hash(...)` -/// feeding the final aggregate, and the top-level `SortPreservingMergeExec` -/// collapsing the sorted partitions down to 1. That means the split walks a -/// chain of dependent producer stages (stage 1's subtree embeds stage 0's -/// `ExchangeSourceExec`) rather than just a single boundary, which is what -/// Task 6 needs the ascending-id topological-order invariant to hold for. -/// -/// This should split into exactly 3 stages: -/// - stage 0: producer wrapping the partial-aggregate subtree (hash -/// repartition boundary), no upstream dependency. -/// - stage 1: producer wrapping the final-aggregate/projection/sort subtree -/// (sort-preserving-merge boundary), depending on stage 0. -/// - stage 2 (final): a bare `ExchangeSourceExec` reading stage 1's output. -#[tokio::test] -async fn multi_boundary_chain_has_ascending_topological_order() { - let ctx = context_with_partitioned_table().await; - let df = ctx - .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY c") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let exchange = InMemoryExchange::new(); - let graph = create_stages(plan, &exchange).unwrap(); - - assert_eq!( - graph.stages.len(), - 3, - "expected 2 producer stages plus a final stage" - ); - - // General invariant Task 6 depends on: every producer id referenced by - // a stage is lower than that stage's own id, and final_stage_id is the - // max id and is read by nobody. - assert_topological_order(&graph); - - let stage0 = &graph.stages[0]; - assert_eq!(stage0.id, 0); - assert!( - stage0.plan.downcast_ref::().is_some(), - "stage 0 should be rooted at a ExchangeSinkExec" - ); - assert!( - stage0.input_stage_ids.is_empty(), - "stage 0 (partial-aggregate producer) has no upstream shuffle dependency" - ); - - let stage1 = &graph.stages[1]; - assert_eq!(stage1.id, 1); - assert!( - stage1.plan.downcast_ref::().is_some(), - "stage 1 should be rooted at a ExchangeSinkExec" - ); - assert_eq!( - stage1.input_stage_ids, - vec![stage0.id], - "stage 1 (sort-preserving-merge producer) depends on stage 0" - ); - assert!( - contains_exchange_source(&stage1.plan), - "stage 1's subtree should embed a ExchangeSourceExec reading stage 0" - ); - - let final_stage = &graph.stages[2]; - assert_eq!(final_stage.id, graph.final_stage_id); - assert!( - final_stage - .plan - .downcast_ref::() - .is_none(), - "final stage must not be wrapped in a ExchangeSinkExec" - ); - assert!( - final_stage - .plan - .downcast_ref::() - .is_some(), - "final stage's plan should itself be the ExchangeSourceExec reading stage 1" - ); - assert_eq!( - final_stage.input_stage_ids, - vec![stage1.id], - "final stage depends on stage 1" - ); -} - -/// `SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k` with -/// `target_partitions = 4` and the hash-join collect-left thresholds forced -/// to 0 (see `context_with_two_partitioned_tables`) produces (verified via -/// `displayable(..).indent(false)` against this checkout): -/// -/// ```text -/// HashJoinExec: mode=Partitioned, join_type=Inner, on=[(k@0, k@0)], projection=[k@0, v@1, v@3] -/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=2 -/// DataSourceExec: partitions=2, partition_sizes=[1, 1] -/// ``` -/// -/// Both join inputs are independently hash-repartitioned, so this should -/// split into exactly 3 stages: two independent producer stages (one per -/// join side, neither depending on the other) and a final stage rooted at -/// the (rewritten) `HashJoinExec` itself, whose plan contains two -/// `ExchangeSourceExec` leaves -- one per producer. This is the multi-input -/// wiring shape (`input_stage_ids.len() == 2`) Task 6 needs to schedule a -/// stage that depends on more than one upstream stage. -#[tokio::test] -async fn join_with_two_shuffled_sides_has_two_input_stage_ids() { - let ctx = context_with_two_partitioned_tables().await; - let df = ctx - .sql("SELECT l.k, l.v, r.v FROM l JOIN r ON l.k = r.k") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let exchange = InMemoryExchange::new(); - let graph = create_stages(plan, &exchange).unwrap(); - - assert_eq!( - graph.stages.len(), - 3, - "expected 2 independent producer stages plus a final (join) stage" - ); - - assert_topological_order(&graph); - - let final_stage = &graph.stages[2]; - assert_eq!(final_stage.id, graph.final_stage_id); - assert!( - final_stage - .plan - .downcast_ref::() - .is_none(), - "final (join) stage must not be wrapped in a ExchangeSinkExec" - ); - assert_eq!( - final_stage.input_stage_ids.len(), - 2, - "join stage should depend on exactly 2 upstream producer stages" - ); - - // Both referenced producer ids must be valid stages in the graph, with - // ids lower than the join stage's own id (already checked generically - // by assert_topological_order, reasserted here for clarity), and each - // must itself have no upstream dependency (leaf hash-repartition - // producers reading straight from a DataSourceExec). - let stage_ids: std::collections::HashSet<_> = - graph.stages.iter().map(|s| s.id).collect(); - for &producer_id in &final_stage.input_stage_ids { - assert!( - stage_ids.contains(&producer_id), - "referenced producer stage {producer_id} must exist in the graph" - ); - assert!( - producer_id < final_stage.id, - "producer stage {producer_id} must have a lower id than the join stage" - ); - let producer = graph - .stages - .iter() - .find(|s| s.id == producer_id) - .expect("producer stage must exist"); - assert!( - producer.plan.downcast_ref::().is_some(), - "producer stage {producer_id} should be rooted at a ExchangeSinkExec" - ); - assert!( - producer.input_stage_ids.is_empty(), - "producer stage {producer_id} (leaf hash-repartition) has no upstream dependency" - ); - } -} diff --git a/datafusion-scheduler/tests/streaming.rs b/datafusion-scheduler/tests/streaming.rs deleted file mode 100644 index e384795009565..0000000000000 --- a/datafusion-scheduler/tests/streaming.rs +++ /dev/null @@ -1,125 +0,0 @@ -// 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. - -//! Streaming / backpressure integration test. -//! -//! Runs a real multi-stage query end-to-end through `run_distributed` with -//! `channel_capacity = 1` — the tightest possible backpressure through the FULL -//! executor: producers block on full channels while downstream consumers drain. -//! The whole run is wrapped in a `tokio::time::timeout` so a deadlock FAILS the -//! test rather than hanging CI. -//! -//! This is the end-to-end proof that the spawn-all, no-barrier executor -//! pipelines correctly under backpressure: a consumer blocked on a channel is -//! unblocked by a producer that is already running (all tasks of all stages are -//! spawned before any is awaited). If the executor used a per-stage barrier -//! instead, capacity-1 channels would deadlock and this test would time out. - -use std::sync::Arc; -use std::time::Duration; - -use datafusion::arrow::array::{Int32Array, RecordBatch}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::MemTable; -use datafusion::physical_plan::collect; -use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_scheduler::{SchedulerConfig, run_distributed}; - -/// Registers table `t(a INT)` backed by a 2-partition `MemTable` with enough -/// rows/batches that capacity-1 channels actually force repeated blocking, -/// against a `SessionContext` at `target_partitions = 4`. -async fn context_with_many_rows() -> SessionContext { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - - // Many small batches per partition so the sink blocks on `send().await` - // over and over against the capacity-1 channels. - let mut p0 = Vec::new(); - let mut p1 = Vec::new(); - for i in 0..200i32 { - p0.push( - RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![i % 17]))], - ) - .unwrap(), - ); - p1.push( - RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![(i + 5) % 17]))], - ) - .unwrap(), - ); - } - - let config = SessionConfig::new().with_target_partitions(4); - let ctx = SessionContext::new_with_config(config); - let table = MemTable::try_new(schema, vec![p0, p1]).unwrap(); - ctx.register_table("t", Arc::new(table)).unwrap(); - ctx -} - -/// Three-stage chain (`GROUP BY` hash boundary + `ORDER BY` gather boundary) -/// run under `channel_capacity = 1`, inside a timeout. Asserts multiset -/// equivalence with single-node `collect`. A hang (deadlock) fails via timeout; -/// a wrong result fails via the equivalence check. -#[tokio::test] -async fn tight_backpressure_multi_stage_pipeline_matches_collect() { - let result = tokio::time::timeout(Duration::from_secs(30), async { - let ctx = context_with_many_rows().await; - let df = ctx - .sql("SELECT a, count(*) AS c FROM t GROUP BY a ORDER BY a") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let expected = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); - - // capacity = 1: the tightest bound that still makes progress, exercised - // through the full run_distributed executor (not just sink/source). - let state = ctx.state(); - let config = SchedulerConfig { - session_builder: Arc::new(move || { - SessionContext::new_with_state(state.clone()) - }), - channel_capacity: 1, - }; - let actual = run_distributed(&ctx, plan, config).await.unwrap(); - - (expected, actual) - }) - .await; - - let (expected, actual) = - result.expect("multi-stage pipeline deadlocked / timed out under capacity=1"); - - // Order-insensitive multiset equivalence via sorted pretty-printed rows. - let e = datafusion::arrow::util::pretty::pretty_format_batches(&expected) - .unwrap() - .to_string(); - let a = datafusion::arrow::util::pretty::pretty_format_batches(&actual) - .unwrap() - .to_string(); - let mut el: Vec<&str> = e.lines().collect(); - el.sort(); - let mut al: Vec<&str> = a.lines().collect(); - al.sort(); - assert_eq!( - el, al, - "streaming distributed output != collect output under capacity=1" - ); -} diff --git a/datafusion-scheduler/tests/two_stage.rs b/datafusion-scheduler/tests/two_stage.rs deleted file mode 100644 index ebffc7606476d..0000000000000 --- a/datafusion-scheduler/tests/two_stage.rs +++ /dev/null @@ -1,118 +0,0 @@ -// 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. - -//! End-to-end test that a real 2-stage plan (partial aggregate + streaming -//! in-memory exchange + final aggregate) executes correctly through -//! `run_distributed`/`execute_stage_graph`, matching single-process `collect` -//! output. Data crosses the stage boundary as IPC frames over the -//! `InMemoryExchange` — there is no disk shuffle to inspect, so correctness of -//! the 2-stage GROUP BY across the exchange is the assertion. - -use std::sync::Arc; - -use datafusion::arrow::array::{Int32Array, RecordBatch}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::MemTable; -use datafusion::physical_plan::collect; -use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_scheduler::test_util::assert_distributed_eq; -use datafusion_scheduler::{SchedulerConfig, create_stages, run_distributed}; - -/// Registers table `t(a INT)` backed by a 2-partition `MemTable`, against a -/// `SessionContext` configured with `target_partitions = 4` -- matching the -/// setup in `tests/stage_splitting.rs` that is known to split `GROUP BY` -/// into a producer stage (partial aggregate, hash-repartitioned) and a -/// final stage (final aggregate) at a real shuffle boundary. -async fn context_with_partitioned_table() -> SessionContext { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let batch0 = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - ) - .unwrap(); - let batch1 = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![2, 3, 4]))], - ) - .unwrap(); - - let config = SessionConfig::new().with_target_partitions(4); - let ctx = SessionContext::new_with_config(config); - let table = MemTable::try_new(schema, vec![vec![batch0], vec![batch1]]).unwrap(); - ctx.register_table("t", Arc::new(table)).unwrap(); - ctx -} - -#[tokio::test] -async fn group_by_across_a_real_shuffle_matches_collect() { - let ctx = context_with_partitioned_table().await; - let df = ctx - .sql("SELECT a, count(*) FROM t GROUP BY a") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - assert_distributed_eq(&ctx, plan).await; -} - -/// Proves `run_distributed` is genuinely routed through `create_stages` + the -/// stage-graph executor (data really crossing the exchange), not coincidentally -/// producing correct output by re-running the original in-process plan. -/// -/// The old disk model let a test assert shuffle files landed under `stage_0/`. -/// The streaming exchange leaves nothing on disk, so instead we assert on the -/// stage split itself — this GROUP BY splits into a producer stage (id 0, -/// hash-repartitioned partial aggregate, whose output must cross the exchange) -/// and a final stage reading it — and that the end-to-end result over that real -/// two-stage exchange execution matches single-node `collect`. -#[tokio::test] -async fn group_by_data_crosses_the_exchange_between_two_stages() { - let ctx = context_with_partitioned_table().await; - let df = ctx - .sql("SELECT a, count(*) FROM t GROUP BY a") - .await - .unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - - let expected = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); - - // The split must produce a producer stage feeding a final stage over the - // exchange — i.e. a genuine 2-stage plan, not a single re-executed tree. - let probe_exchange = datafusion_scheduler::InMemoryExchange::new(); - let graph = create_stages(plan.clone(), &probe_exchange).unwrap(); - assert_eq!( - graph.stages.len(), - 2, - "GROUP BY should split into a producer stage and a final stage" - ); - assert_ne!( - graph.final_stage_id, 0, - "there must be a producer stage (id 0) feeding the final stage" - ); - - // Execute end-to-end: the producer stage's partial-aggregate output only - // reaches the final stage by crossing the in-memory exchange as IPC frames. - let config = SchedulerConfig::in_memory(&ctx); - let actual = run_distributed(&ctx, plan, config).await.unwrap(); - - let expected_rows: usize = expected.iter().map(|b| b.num_rows()).sum(); - let actual_rows: usize = actual.iter().map(|b| b.num_rows()).sum(); - assert_eq!( - expected_rows, actual_rows, - "row count across the exchange must match single-node collect" - ); -} From 0d629b58979e5523b63af393cd32a758a78fa30a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 16:21:33 -0600 Subject: [PATCH 3/4] chore: prettier format datafusion-examples/README.md --- datafusion-examples/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md index 3a99734cc44f0..ed1fa9b0827da 100644 --- a/datafusion-examples/README.md +++ b/datafusion-examples/README.md @@ -205,9 +205,9 @@ cargo run --example dataframe -- dataframe #### Category: Distributed -| Subcommand | File Path | Description | -| --------------------- | ----------------------------------------------------- | ---------------------------------------------------------------------------------------------------- | -| distributed_pipeline | [`scheduler/main.rs`](examples/scheduler/main.rs) | In-process model of stage-based distributed execution (stage splitting, per-task plan serialization) | +| Subcommand | File Path | Description | +| -------------------- | ------------------------------------------------- | ---------------------------------------------------------------------------------------------------- | +| distributed_pipeline | [`scheduler/main.rs`](examples/scheduler/main.rs) | In-process model of stage-based distributed execution (stage splitting, per-task plan serialization) | ## SQL Ops Examples From f8ab0d3cbc68946141574f8f789299e0307e46d7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 16:51:15 -0600 Subject: [PATCH 4/4] fix(examples): make scheduler README auto-generation match the check --- datafusion-examples/examples/scheduler/main.rs | 4 +--- datafusion-examples/src/utils/example_metadata/model.rs | 6 +++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/datafusion-examples/examples/scheduler/main.rs b/datafusion-examples/examples/scheduler/main.rs index 8ebbedd1e6330..f3ea4a2fbf7f0 100644 --- a/datafusion-examples/examples/scheduler/main.rs +++ b/datafusion-examples/examples/scheduler/main.rs @@ -45,9 +45,7 @@ //! - `all` — run all examples included in this module //! //! - `distributed_pipeline` -//! (desc: Run a multi-stage `GROUP BY ... ORDER BY` query through the -//! stage-graph executor and print the result alongside single-node -//! `collect` for comparison) +//! (file: main.rs, desc: In-process model of stage-based distributed execution (stage splitting, per-task plan serialization)) //! //! ## Tests //! diff --git a/datafusion-examples/src/utils/example_metadata/model.rs b/datafusion-examples/src/utils/example_metadata/model.rs index 11416d141eb74..79749a319e336 100644 --- a/datafusion-examples/src/utils/example_metadata/model.rs +++ b/datafusion-examples/src/utils/example_metadata/model.rs @@ -142,7 +142,7 @@ impl Category { /// Determines the category for a group by name. pub fn for_group(name: &str) -> Self { match name { - "flight" => Category::Distributed, + "flight" | "scheduler" => Category::Distributed, _ => Category::SingleProcess, } } @@ -185,6 +185,10 @@ mod tests { Category::for_group("flight"), Category::Distributed )); + assert!(matches!( + Category::for_group("scheduler"), + Category::Distributed + )); assert!(matches!( Category::for_group("anything_else"), Category::SingleProcess