feat(push-scheduler): port PR #2226 push scheduler onto tokio + crossbeam

Cargo.toml

@@ -46,6 +46,7 @@ members = [
     "datafusion/physical-optimizer",
     "datafusion/pruning",
     "datafusion/physical-plan",
+    "datafusion/push-scheduler",
     "datafusion/proto",
     "datafusion/proto/gen",
     "datafusion/proto-common",
@@ -146,6 +147,7 @@ datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter",
 datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "53.0.0", default-features = false }
 datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "53.0.0" }
 datafusion-physical-plan = { path = "datafusion/physical-plan", version = "53.0.0" }
+datafusion-push-scheduler = { path = "datafusion/push-scheduler", version = "53.0.0" }
 datafusion-proto = { path = "datafusion/proto", version = "53.0.0" }
 datafusion-proto-common = { path = "datafusion/proto-common", version = "53.0.0" }
 datafusion-pruning = { path = "datafusion/pruning", version = "53.0.0" }

benchmarks/Cargo.toml

@@ -45,6 +45,7 @@ bytes = { workspace = true }
 clap = { version = "4.5.60", features = ["derive"] }
 datafusion = { workspace = true, default-features = true }
 datafusion-common = { workspace = true, default-features = true }
+datafusion-push-scheduler = { workspace = true }
 env_logger = { workspace = true }
 futures = { workspace = true }
 libmimalloc-sys = { version = "0.1", optional = true }

benchmarks/src/clickbench.rs

@@ -19,7 +19,10 @@ use std::fs;
 use std::io::ErrorKind;
 use std::path::{Path, PathBuf};
 
-use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats};
+use crate::util::{
+    BenchmarkRun, CommonOpt, QueryResult, collect_sql_via_push_scheduler,
+    print_memory_stats,
+};
 use clap::Args;
 use datafusion::logical_expr::{ExplainFormat, ExplainOption};
 use datafusion::{
@@ -255,7 +258,11 @@ impl RunOpt {
         let mut query_results = vec![];
         for i in 0..self.iterations() {
             let start = Instant::now();
-            let results = ctx.sql(sql).await?.collect().await?;
+            let results = if self.common.push_scheduler {
+                collect_sql_via_push_scheduler(ctx, sql).await?
+            } else {
+                ctx.sql(sql).await?.collect().await?
+            };
             let elapsed = start.elapsed();
             let ms = elapsed.as_secs_f64() * 1000.0;
             millis.push(ms);

benchmarks/src/util/mod.rs

@@ -19,8 +19,10 @@
 pub mod latency_object_store;
 mod memory;
 mod options;
+pub mod push_scheduler;
 mod run;
 
 pub use memory::print_memory_stats;
 pub use options::CommonOpt;
+pub use push_scheduler::collect_sql_via_push_scheduler;
 pub use run::{BenchQuery, BenchmarkRun, QueryResult};

benchmarks/src/util/options.rs

@@ -70,6 +70,19 @@ pub struct CommonOpt {
     /// Adds random latency in the range 20-200ms to each object store operation.
     #[arg(long = "simulate-latency")]
     pub simulate_latency: bool,
+
+    /// Execute queries via `datafusion-push-scheduler` — the push-based
+    /// morsel-driven scheduler ported from apache/datafusion#2226. On by
+    /// default; pass `--push-scheduler=false` to fall back to the
+    /// pull-based path.
+    #[arg(
+        long = "push-scheduler",
+        default_value_t = true,
+        num_args = 0..=1,
+        default_missing_value = "true",
+        action = clap::ArgAction::Set,
+    )]
+    pub push_scheduler: bool,
 }
 
 impl CommonOpt {
@@ -190,6 +203,7 @@ mod tests {
             sort_spill_reservation_bytes: None,
             debug: false,
             simulate_latency: false,
+            push_scheduler: true,
         };
 
         // With env var set, builder should succeed and have a memory pool

benchmarks/src/util/push_scheduler.rs

@@ -0,0 +1,46 @@
+// 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.
+
+//! Benchmark helpers for driving queries through
+//! [`datafusion_push_scheduler::Scheduler`].
+
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::error::Result;
+use datafusion::prelude::SessionContext;
+use datafusion_push_scheduler::Scheduler;
+use futures::StreamExt;
+
+/// Compile a SQL query via `ctx`, then execute the resulting physical plan
+/// through [`Scheduler`] and collect all output batches. The worker count
+/// is taken from the session's `target_partitions`, matching how the
+/// default path fans out.
+pub async fn collect_sql_via_push_scheduler(
+    ctx: &SessionContext,
+    sql: &str,
+) -> Result<Vec<RecordBatch>> {
+    let df = ctx.sql(sql).await?;
+    let plan = df.create_physical_plan().await?;
+    let workers = ctx.state().config().target_partitions().max(1);
+    let scheduler = Scheduler::new(workers)?;
+    let mut stream = scheduler.schedule(plan, ctx.task_ctx())?.stream();
+
+    let mut out = Vec::new();
+    while let Some(batch) = stream.next().await {
+        out.push(batch?);
+    }
+    Ok(out)
+}

datafusion/push-scheduler/Cargo.toml

@@ -0,0 +1,51 @@
+# 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-push-scheduler"
+description = "Push-based morsel scheduler for DataFusion (port of apache/datafusion#2226 onto plain tokio + crossbeam)"
+keywords = ["arrow", "query", "sql", "scheduler", "morsel", "push"]
+version = { workspace = true }
+edition = { workspace = true }
+homepage = { workspace = true }
+repository = { workspace = true }
+license = { workspace = true }
+authors = { workspace = true }
+rust-version = { workspace = true }
+
+[lints]
+workspace = true
+
+[lib]
+name = "datafusion_push_scheduler"
+
+[dependencies]
+ahash = { version = "0.8", default-features = false, features = ["runtime-rng"] }
+arrow = { workspace = true }
+crossbeam-deque = "0.8"
+datafusion-common = { workspace = true, default-features = true }
+datafusion-execution = { workspace = true, default-features = true }
+datafusion-physical-expr = { workspace = true, default-features = true }
+datafusion-physical-plan = { workspace = true }
+futures = { workspace = true }
+log = { workspace = true }
+parking_lot = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt", "sync"] }
+
+[dev-dependencies]
+datafusion = { workspace = true, default-features = true }
+tokio = { workspace = true, features = ["rt-multi-thread", "time", "macros"] }

datafusion/push-scheduler/src/lib.rs

@@ -0,0 +1,79 @@
+// 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.
+
+//! Push-based, morsel-driven scheduler for Apache DataFusion.
+//!
+//! This crate is a port of the experimental scheduler from
+//! [apache/datafusion#2226](https://github.com/apache/datafusion/pull/2226)
+//! onto plain `tokio` + `crossbeam-deque` (the original used `rayon`).
+//!
+//! # Architecture
+//!
+//! * A [`PipelinePlanner`] compiles an
+//!   [`ExecutionPlan`](datafusion_physical_plan::ExecutionPlan) into a flat
+//!   [`PipelinePlan`] — a vector of [`RoutablePipeline`]s plus their
+//!   [`OutputLink`]s. Linear chains of pull-based operators are grouped
+//!   into an [`ExecutionPipeline`](pipelines::execution::ExecutionPipeline);
+//!   `RepartitionExec`, `CoalescePartitionsExec`, `SortExec`, and partial
+//!   `AggregateExec` become dedicated breaker pipelines.
+//!
+//! * A [`Scheduler`] owns a pool of worker OS threads, each running a tokio
+//!   `current_thread` runtime + `LocalSet` and pulling tasks from a
+//!   `crossbeam_deque` work-stealing queue. Idle workers steal from peers;
+//!   external submissions go through a shared `Injector`.
+//!
+//! * A `Task = (pipeline_idx, partition)` is scheduled per output partition.
+//!   Each worker iteration calls
+//!   [`Pipeline::poll_partition`](pipeline::Pipeline::poll_partition) on the
+//!   task's pipeline and routes the result to the downstream pipeline's
+//!   `push` / `close`. `Poll::Pending` parks the task on a
+//!   [`ArcWake`](futures::task::ArcWake) that re-enqueues it when the
+//!   underlying future wakes.
+//!
+//! # Entry point
+//!
+//! ```ignore
+//! use std::sync::Arc;
+//! use datafusion_push_scheduler::Scheduler;
+//!
+//! # async fn run(plan: std::sync::Arc<dyn datafusion_physical_plan::ExecutionPlan>,
+//! #              ctx: std::sync::Arc<datafusion_execution::TaskContext>)
+//! #   -> datafusion_common::Result<()> {
+//! let scheduler = Scheduler::new(num_cpus::get());
+//! let results = scheduler.schedule(plan, ctx)?;
+//! let mut stream = results.stream();
+//! while let Some(batch) = futures::StreamExt::next(&mut stream).await {
+//!     let _ = batch?;
+//! }
+//! # Ok(()) }
+//! ```
+
+#![deny(clippy::clone_on_ref_ptr)]
+
+pub mod pipeline;
+pub mod pipelines;
+pub mod plan;
+pub mod scheduler;
+pub mod task;
+pub mod worker_pool;
+
+pub use pipeline::Pipeline;
+pub use pipelines::execution::ExecutionPipeline;
+pub use pipelines::repartition::RepartitionPipeline;
+pub use plan::{OutputLink, PipelinePlan, PipelinePlanner, RoutablePipeline};
+pub use scheduler::Scheduler;
+pub use task::{ExecutionResults, spawn_plan};

datafusion/push-scheduler/src/pipeline.rs

@@ -0,0 +1,57 @@
+// 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.
+
+//! Push-based [`Pipeline`] trait. Ported verbatim from PR apache/datafusion#2226.
+
+use std::fmt::Debug;
+use std::task::{Context, Poll};
+
+use arrow::record_batch::RecordBatch;
+use datafusion_common::Result;
+
+/// A push-based operator. The scheduler drives progress by polling one of the
+/// pipeline's output partitions via [`Pipeline::poll_partition`]; when a batch
+/// is produced the scheduler routes it to the next pipeline's
+/// [`Pipeline::push`]. End-of-input is signalled by [`Pipeline::close`].
+///
+/// The split between eager (push-time) and lazy (poll-time) work is entirely
+/// the pipeline's choice — e.g. a repartition breaker does most work in
+/// `push`, while a streaming adapter forwards from a wrapped
+/// `SendableRecordBatchStream` inside `poll_partition`.
+pub trait Pipeline: Debug + Send + Sync {
+    /// Push a batch into `child` of the pipeline for the given `partition`.
+    fn push(&self, input: RecordBatch, child: usize, partition: usize) -> Result<()>;
+
+    /// Signal that `child` / `partition` will receive no further batches.
+    fn close(&self, child: usize, partition: usize);
+
+    /// Number of output partitions this pipeline produces.
+    fn output_partitions(&self) -> usize;
+
+    /// Attempt to drive progress for one output partition.
+    ///
+    /// * `Poll::Ready(Some(Ok(batch)))` — next batch for the partition.
+    /// * `Poll::Ready(Some(Err(e)))` — query failed.
+    /// * `Poll::Ready(None)` — partition exhausted.
+    /// * `Poll::Pending` — pipeline has stored `cx.waker()` and will wake
+    ///   the scheduler when it can make progress.
+    fn poll_partition(
+        &self,
+        cx: &mut Context<'_>,
+        partition: usize,
+    ) -> Poll<Option<Result<RecordBatch>>>;
+}