alamb commented on code in PR #2226: URL: https://github.com/apache/arrow-datafusion/pull/2226#discussion_r856178334
########## datafusion/scheduler/src/lib.rs: ########## @@ -0,0 +1,381 @@ +// 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 futures::stream::BoxStream; +use log::{debug, error}; + +use datafusion::error::Result; +use datafusion::execution::context::TaskContext; +use datafusion::physical_plan::ExecutionPlan; + +use crate::query::{Query, QueryBuilder, RoutablePipeline}; +use crate::task::{spawn_query, Task}; + +use rayon::{ThreadPool, ThreadPoolBuilder}; + +pub use task::QueryResults; + +mod pipeline; +mod query; +mod task; + +/// Builder for a [`Scheduler`] +#[derive(Debug)] +pub struct SchedulerBuilder { + inner: ThreadPoolBuilder, +} + +impl SchedulerBuilder { + /// Create a new [`SchedulerConfig`] with the provided number of threads + pub fn new(num_threads: usize) -> Self { + let builder = ThreadPoolBuilder::new() + .num_threads(num_threads) + .panic_handler(|p| error!("{}", format_worker_panic(p))) + .thread_name(|idx| format!("df-worker-{}", idx)); + + Self { inner: builder } + } + + /// Registers a custom panic handler + /// + /// Used by tests + #[allow(dead_code)] + fn panic_handler<H>(self, panic_handler: H) -> Self + where + H: Fn(Box<dyn std::any::Any + Send>) + Send + Sync + 'static, + { + Self { + inner: self.inner.panic_handler(panic_handler), + } + } + + /// Build a new [`Scheduler`] + fn build(self) -> Scheduler { + Scheduler { + pool: Arc::new(self.inner.build().unwrap()), + } + } +} + +/// A [`Scheduler`] maintains a pool of dedicated worker threads on which +/// query execution can be scheduled. This is based on the idea of [Morsel-Driven Parallelism] +/// and is designed to decouple the execution parallelism from the parallelism expressed in +/// the physical plan as partitions. +/// +/// # Implementation +/// +/// When provided with an [`ExecutionPlan`] the [`Scheduler`] first breaks it up into smaller +/// chunks called pipelines. Each pipeline may consist of one or more nodes from the +/// [`ExecutionPlan`] tree. +/// +/// The scheduler then maintains a list of pending [`Task`], that identify a partition within +/// a particular pipeline that may be able to make progress on some "morsel" of data. These +/// [`Task`] are then scheduled on the worker pool, with a preference for scheduling work +/// on a given "morsel" on the same thread that produced it. +/// +/// # Rayon +/// +/// Under-the-hood these [`Task`] are scheduled by [rayon], which is a lightweight, work-stealing +/// scheduler optimised for CPU-bound workloads. Pipelines may exploit this fact, and use [rayon]'s +/// structured concurrency primitives to express additional parallelism that may be exploited +/// if there are idle threads available at runtime +/// +/// # Shutdown +/// +/// Queries scheduled on a [`Scheduler`] will run to completion even if the +/// [`Scheduler`] is dropped +/// +/// [Morsel-Driven Parallelism]: https://db.in.tum.de/~leis/papers/morsels.pdf +/// [rayon]: https://docs.rs/rayon/latest/rayon/ +/// +pub struct Scheduler { Review Comment: Stylistically I recommend putting this structure (and its very nice documentation) at the top of this source file so people see it first. I also think a doc example showing how to use it, would be very very helpful (so people didn't have to look at the unit test). Something like: ```rust let scheduler = Scheduler::new(4); let config = SessionConfig::new().with_target_partitions(4); let context = SessionContext::with_config(config); context.register_csv("example", "tests/example.csv", CsvReadOptions::new()).await?; let plan = context.sql("SELECT MIN(b) FROM example") .await .unwrap() .create_physical_plan() .unwrap() let stream = scheduler.schedule(plan, task).unwrap(); let scheduled: Vec<_> = stream.try_collect().await.unwrap(); let results = query.collect().await.unwrap(); ``` ########## datafusion/scheduler/src/lib.rs: ########## @@ -0,0 +1,381 @@ +// 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 futures::stream::BoxStream; +use log::{debug, error}; + +use datafusion::error::Result; +use datafusion::execution::context::TaskContext; +use datafusion::physical_plan::ExecutionPlan; + +use crate::query::{Query, QueryBuilder, RoutablePipeline}; +use crate::task::{spawn_query, Task}; + +use rayon::{ThreadPool, ThreadPoolBuilder}; + +pub use task::QueryResults; + +mod pipeline; +mod query; +mod task; + +/// Builder for a [`Scheduler`] +#[derive(Debug)] +pub struct SchedulerBuilder { + inner: ThreadPoolBuilder, +} + +impl SchedulerBuilder { + /// Create a new [`SchedulerConfig`] with the provided number of threads + pub fn new(num_threads: usize) -> Self { + let builder = ThreadPoolBuilder::new() + .num_threads(num_threads) + .panic_handler(|p| error!("{}", format_worker_panic(p))) + .thread_name(|idx| format!("df-worker-{}", idx)); + + Self { inner: builder } + } + + /// Registers a custom panic handler + /// + /// Used by tests + #[allow(dead_code)] + fn panic_handler<H>(self, panic_handler: H) -> Self + where + H: Fn(Box<dyn std::any::Any + Send>) + Send + Sync + 'static, + { + Self { + inner: self.inner.panic_handler(panic_handler), + } + } + + /// Build a new [`Scheduler`] + fn build(self) -> Scheduler { + Scheduler { + pool: Arc::new(self.inner.build().unwrap()), + } + } +} + +/// A [`Scheduler`] maintains a pool of dedicated worker threads on which +/// query execution can be scheduled. This is based on the idea of [Morsel-Driven Parallelism] +/// and is designed to decouple the execution parallelism from the parallelism expressed in +/// the physical plan as partitions. +/// +/// # Implementation +/// +/// When provided with an [`ExecutionPlan`] the [`Scheduler`] first breaks it up into smaller +/// chunks called pipelines. Each pipeline may consist of one or more nodes from the +/// [`ExecutionPlan`] tree. +/// +/// The scheduler then maintains a list of pending [`Task`], that identify a partition within +/// a particular pipeline that may be able to make progress on some "morsel" of data. These +/// [`Task`] are then scheduled on the worker pool, with a preference for scheduling work +/// on a given "morsel" on the same thread that produced it. +/// +/// # Rayon +/// +/// Under-the-hood these [`Task`] are scheduled by [rayon], which is a lightweight, work-stealing +/// scheduler optimised for CPU-bound workloads. Pipelines may exploit this fact, and use [rayon]'s +/// structured concurrency primitives to express additional parallelism that may be exploited +/// if there are idle threads available at runtime +/// +/// # Shutdown +/// +/// Queries scheduled on a [`Scheduler`] will run to completion even if the +/// [`Scheduler`] is dropped +/// +/// [Morsel-Driven Parallelism]: https://db.in.tum.de/~leis/papers/morsels.pdf +/// [rayon]: https://docs.rs/rayon/latest/rayon/ +/// +pub struct Scheduler { + pool: Arc<ThreadPool>, +} + +impl Scheduler { + /// Create a new [`Scheduler`] with `num_threads` new threads in a dedicated thread pool + pub fn new(num_threads: usize) -> Self { + SchedulerBuilder::new(num_threads).build() + } + + /// Schedule the provided [`ExecutionPlan`] on this [`Scheduler`]. + /// + /// Returns a [`BoxStream`] that can be used to receive results as they are produced Review Comment: ```suggestion /// Returns a [`QueryResults`] (stream of [`RecordBatch`]es) that can be used to receive results as they are produced ``` ########## datafusion/scheduler/Cargo.toml: ########## @@ -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. + +[package] +name = "datafusion-scheduler" +description = "Scheduling for DataFusion query engine" +version = "7.0.0" +homepage = "https://github.com/apache/arrow-datafusion"; +repository = "https://github.com/apache/arrow-datafusion"; +readme = "../README.md" +authors = ["Apache Arrow <[email protected]>"] +license = "Apache-2.0" +keywords = ["arrow", "query", "sql"] +edition = "2021" +rust-version = "1.58" + +[lib] +name = "datafusion_scheduler" +path = "src/lib.rs" + +[features] + +[dependencies] +ahash = { version = "0.7", default-features = false } +arrow = { version = "12" } +async-trait = "0.1" +datafusion = { path = "../core", version = "7.0.0" } Review Comment: by structuring the crate this way it means `datafusions-scheduler` has to be brought in as its own crate. I think it would be more consistent if datafusion-scheduler was an optional dependency of `datafusion` (aka `datafusion/core` like the `jit` module. That way users of the scheduler do not have to import the scheduler crate itself. ########## datafusion/scheduler/src/query.rs: ########## @@ -0,0 +1,276 @@ +// 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::context::TaskContext; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; + +use crate::pipeline::{ + execution::ExecutionPipeline, repartition::RepartitionPipeline, Pipeline, +}; + +/// Identifies the [`Pipeline`] within the [`Query`] to route output to +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct OutputLink { + /// The index of the [`Pipeline`] in [`Query`] to route output to + pub pipeline: usize, + + /// The child of the [`Pipeline`] to route output to + pub child: usize, +} + +/// Combines a [`Pipeline`] with an [`OutputLink`] identifying where to send its output +#[derive(Debug)] +pub struct RoutablePipeline { + /// The pipeline that produces data + pub pipeline: Box<dyn Pipeline>, + + /// Where to send output the output of `pipeline` + /// + /// If `None`, the output should be sent to the query output + pub output: Option<OutputLink>, +} + +/// [`Query`] is the scheduler's representation of the [`ExecutionPlan`] passed to +/// [`super::Scheduler::schedule`]. It combines the list of [Pipeline`] with the information +/// necessary to route output from one stage to the next +#[derive(Debug)] +pub struct Query { + pub pipelines: Vec<RoutablePipeline>, +} + +/// When converting [`ExecutionPlan`] to [`Pipeline`] we may wish to group +/// together multiple operators, [`OperatorGroup`] stores this state +struct OperatorGroup { + /// Where to route the output of the eventual [`Pipeline`] + output: Option<OutputLink>, + + /// The [`ExecutionPlan`] from which to start recursing + root: Arc<dyn ExecutionPlan>, + + /// The number of times to recurse into the [`ExecutionPlan`]'s children + depth: usize, +} + +/// A utility struct to assist converting from [`ExecutionPlan`] to [`Query`] +/// +/// The [`ExecutionPlan`] is visited in a depth-first fashion, gradually building +/// up the [`RoutablePipeline`] for the [`Query`]. As nodes are visited depth-first, +/// a node is visited only after its parent has been. +pub struct QueryBuilder { + task_context: Arc<TaskContext>, + + /// The current list of completed pipelines + in_progress: Vec<RoutablePipeline>, + + /// A list of [`ExecutionPlan`] still to visit, along with + /// where they should route their output + to_visit: Vec<(Arc<dyn ExecutionPlan>, Option<OutputLink>)>, + + /// Stores one or more operators to combine + /// together into a single [`ExecutionPipeline`] + execution_operators: Option<OperatorGroup>, +} + +impl QueryBuilder { + pub fn new(plan: Arc<dyn ExecutionPlan>, task_context: Arc<TaskContext>) -> Self { + Self { + in_progress: vec![], + to_visit: vec![(plan, None)], + task_context, + execution_operators: None, + } + } + + /// Flush the current group of operators stored in `execution_operators` + /// into a single [`ExecutionPipeline] + fn flush_exec(&mut self) -> Result<usize> { + let group = self.execution_operators.take().unwrap(); + let node_idx = self.in_progress.len(); + self.in_progress.push(RoutablePipeline { + pipeline: Box::new(ExecutionPipeline::new( + group.root, + self.task_context.clone(), + group.depth, + )?), + output: group.output, + }); + Ok(node_idx) + } + + /// Visit a non-special cased [`ExecutionPlan`] + fn visit_exec( + &mut self, + plan: Arc<dyn ExecutionPlan>, + parent: Option<OutputLink>, + ) -> Result<()> { + let children = plan.children(); + + // Add the operator to the current group of operators to be combined + // into a single [`ExecutionPipeline`]. + // + // TODO: More sophisticated policy, just because we can combine them doesn't mean we should + match self.execution_operators.as_mut() { + Some(buffer) => { + assert_eq!(parent, buffer.output, "QueryBuilder out of sync"); + buffer.depth += 1; + } + None => { + self.execution_operators = Some(OperatorGroup { + output: parent, + root: plan, + depth: 0, + }) + } + } + + match children.len() { + 1 => { + // Enqueue the children with the parent of the `OperatorGroup` + self.to_visit + .push((children.into_iter().next().unwrap(), parent)) + } + _ => { + // We can only recursively group through nodes with a single child, therefore + // if this node has multiple children, we now need to flush the buffer and + // enqueue its children with this new pipeline as its parent + let node = self.flush_exec()?; + self.enqueue_children(children, node); + } + } + + Ok(()) + } + + /// Add the given list of children to the stack of [`ExecutionPlan`] to visit + fn enqueue_children( + &mut self, + children: Vec<Arc<dyn ExecutionPlan>>, + parent_node_idx: usize, + ) { + for (child_idx, child) in children.into_iter().enumerate() { + self.to_visit.push(( + child, + Some(OutputLink { + pipeline: parent_node_idx, + child: child_idx, + }), + )) + } + } + + /// Push a new [`RoutablePipeline`] and enqueue its children to be visited + fn push_pipeline( + &mut self, + node: RoutablePipeline, + children: Vec<Arc<dyn ExecutionPlan>>, + ) { + let node_idx = self.in_progress.len(); + self.in_progress.push(node); + self.enqueue_children(children, node_idx) + } + + /// Push a new [`RepartitionPipeline`] first flushing any buffered [`OperatorGroup`] + fn push_repartition( + &mut self, + input: Partitioning, + output: Partitioning, + parent: Option<OutputLink>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<()> { + let parent = match &self.execution_operators { + Some(buffer) => { + assert_eq!(buffer.output, parent, "QueryBuilder out of sync"); + Some(OutputLink { + pipeline: self.flush_exec()?, + child: 0, // Must be the only child + }) + } + None => parent, + }; + + let node = Box::new(RepartitionPipeline::try_new(input, output)?); + self.push_pipeline( + RoutablePipeline { + pipeline: node, + output: parent, + }, + children, + ); + Ok(()) + } + + /// Visit an [`ExecutionPlan`] operator and add it to the [`Query`] being built + fn visit_operator( + &mut self, + plan: Arc<dyn ExecutionPlan>, + parent: Option<OutputLink>, + ) -> Result<()> { + if let Some(repartition) = plan.as_any().downcast_ref::<RepartitionExec>() { Review Comment: Is it worth some commentary here that `RepartitionExec` and `CoalscePartitonsExec` are handled natively by the scheduler and thus not directly added into a `Pipeline` ########## datafusion/scheduler/src/pipeline/execution.rs: ########## @@ -0,0 +1,330 @@ +// 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::any::Any; +use std::collections::VecDeque; +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, Waker}; + +use arrow::error::ArrowError; +use async_trait::async_trait; +use futures::{Stream, StreamExt, TryStreamExt}; +use parking_lot::Mutex; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::{error::Result as ArrowResult, record_batch::RecordBatch}; +use datafusion::error::Result; +use datafusion::execution::context::TaskContext; +use datafusion::physical_plan::expressions::PhysicalSortExpr; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ + displayable, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + RecordBatchStream, SendableRecordBatchStream, Statistics, +}; + +use crate::pipeline::Pipeline; +use crate::BoxStream; + +/// An [`ExecutionPipeline`] wraps a portion of an [`ExecutionPlan`] and Review Comment: This is very clear -- thank you for the writeup ########## datafusion/scheduler/src/lib.rs: ########## @@ -0,0 +1,381 @@ +// 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 futures::stream::BoxStream; +use log::{debug, error}; + +use datafusion::error::Result; +use datafusion::execution::context::TaskContext; +use datafusion::physical_plan::ExecutionPlan; + +use crate::query::{Query, QueryBuilder, RoutablePipeline}; +use crate::task::{spawn_query, Task}; + +use rayon::{ThreadPool, ThreadPoolBuilder}; + +pub use task::QueryResults; + +mod pipeline; +mod query; +mod task; + +/// Builder for a [`Scheduler`] +#[derive(Debug)] +pub struct SchedulerBuilder { + inner: ThreadPoolBuilder, +} + +impl SchedulerBuilder { + /// Create a new [`SchedulerConfig`] with the provided number of threads + pub fn new(num_threads: usize) -> Self { + let builder = ThreadPoolBuilder::new() + .num_threads(num_threads) + .panic_handler(|p| error!("{}", format_worker_panic(p))) + .thread_name(|idx| format!("df-worker-{}", idx)); + + Self { inner: builder } + } + + /// Registers a custom panic handler + /// + /// Used by tests + #[allow(dead_code)] + fn panic_handler<H>(self, panic_handler: H) -> Self + where + H: Fn(Box<dyn std::any::Any + Send>) + Send + Sync + 'static, + { + Self { + inner: self.inner.panic_handler(panic_handler), + } + } + + /// Build a new [`Scheduler`] + fn build(self) -> Scheduler { + Scheduler { + pool: Arc::new(self.inner.build().unwrap()), + } + } +} + +/// A [`Scheduler`] maintains a pool of dedicated worker threads on which +/// query execution can be scheduled. This is based on the idea of [Morsel-Driven Parallelism] +/// and is designed to decouple the execution parallelism from the parallelism expressed in +/// the physical plan as partitions. +/// +/// # Implementation +/// +/// When provided with an [`ExecutionPlan`] the [`Scheduler`] first breaks it up into smaller +/// chunks called pipelines. Each pipeline may consist of one or more nodes from the +/// [`ExecutionPlan`] tree. +/// +/// The scheduler then maintains a list of pending [`Task`], that identify a partition within +/// a particular pipeline that may be able to make progress on some "morsel" of data. These +/// [`Task`] are then scheduled on the worker pool, with a preference for scheduling work +/// on a given "morsel" on the same thread that produced it. +/// +/// # Rayon +/// +/// Under-the-hood these [`Task`] are scheduled by [rayon], which is a lightweight, work-stealing +/// scheduler optimised for CPU-bound workloads. Pipelines may exploit this fact, and use [rayon]'s +/// structured concurrency primitives to express additional parallelism that may be exploited +/// if there are idle threads available at runtime +/// +/// # Shutdown +/// +/// Queries scheduled on a [`Scheduler`] will run to completion even if the +/// [`Scheduler`] is dropped +/// +/// [Morsel-Driven Parallelism]: https://db.in.tum.de/~leis/papers/morsels.pdf +/// [rayon]: https://docs.rs/rayon/latest/rayon/ +/// +pub struct Scheduler { + pool: Arc<ThreadPool>, +} + +impl Scheduler { + /// Create a new [`Scheduler`] with `num_threads` new threads in a dedicated thread pool + pub fn new(num_threads: usize) -> Self { + SchedulerBuilder::new(num_threads).build() + } + + /// Schedule the provided [`ExecutionPlan`] on this [`Scheduler`]. + /// + /// Returns a [`BoxStream`] that can be used to receive results as they are produced + pub fn schedule( + &self, + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, + ) -> Result<QueryResults> { + let query = QueryBuilder::new(plan, context).build()?; + Ok(self.schedule_query(query)) + } + + /// Schedule the provided [`Query`] on this [`Scheduler`]. + pub(crate) fn schedule_query(&self, query: Query) -> QueryResults { + spawn_query(query, self.spawner()) + } + + fn spawner(&self) -> Spawner { + Spawner { + pool: self.pool.clone(), + } + } +} + +/// Formats a panic message for a worker +fn format_worker_panic(panic: Box<dyn std::any::Any + Send>) -> String { + let maybe_idx = rayon::current_thread_index(); + let worker: &dyn std::fmt::Display = match &maybe_idx { + Some(idx) => idx, + None => &"UNKNOWN", + }; + + let message = if let Some(msg) = panic.downcast_ref::<&str>() { Review Comment: I think i got in trouble in IOx by only handling &str and String -- what about trying to downcast to `&dyn Display` and `&dyn Debug` instead which would cover the &str and String cases as well ########## datafusion/scheduler/src/query.rs: ########## @@ -0,0 +1,276 @@ +// 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::context::TaskContext; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; + +use crate::pipeline::{ + execution::ExecutionPipeline, repartition::RepartitionPipeline, Pipeline, +}; + +/// Identifies the [`Pipeline`] within the [`Query`] to route output to +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct OutputLink { + /// The index of the [`Pipeline`] in [`Query`] to route output to + pub pipeline: usize, + + /// The child of the [`Pipeline`] to route output to + pub child: usize, +} + +/// Combines a [`Pipeline`] with an [`OutputLink`] identifying where to send its output +#[derive(Debug)] +pub struct RoutablePipeline { + /// The pipeline that produces data + pub pipeline: Box<dyn Pipeline>, + + /// Where to send output the output of `pipeline` + /// + /// If `None`, the output should be sent to the query output + pub output: Option<OutputLink>, +} + +/// [`Query`] is the scheduler's representation of the [`ExecutionPlan`] passed to +/// [`super::Scheduler::schedule`]. It combines the list of [Pipeline`] with the information +/// necessary to route output from one stage to the next +#[derive(Debug)] +pub struct Query { Review Comment: Editorially I find `Query` confusing in this context (as it is already used for other purposes in DataFusion and databases in general). It is confusing that we go from `sqlparser::Query --> LogicalPlan --> ExecutionPlan --> scheduler::Query` I suggest a name like `PipelineList`, `PipelinePlan` or`PushPipelines` would make this code easier to navigate for others ########## datafusion/scheduler/src/pipeline/execution.rs: ########## @@ -0,0 +1,330 @@ +// 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::any::Any; +use std::collections::VecDeque; +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, Waker}; + +use arrow::error::ArrowError; +use async_trait::async_trait; +use futures::{Stream, StreamExt, TryStreamExt}; +use parking_lot::Mutex; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::{error::Result as ArrowResult, record_batch::RecordBatch}; +use datafusion::error::Result; +use datafusion::execution::context::TaskContext; +use datafusion::physical_plan::expressions::PhysicalSortExpr; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ + displayable, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + RecordBatchStream, SendableRecordBatchStream, Statistics, +}; + +use crate::pipeline::Pipeline; +use crate::BoxStream; + +/// An [`ExecutionPipeline`] wraps a portion of an [`ExecutionPlan`] and +/// converts it to the push-based [`Pipeline`] interface +/// +/// Internally [`ExecutionPipeline`] is still pull-based which limits its parallelism +/// to that of its output partitioning, however, it provides full compatibility with +/// [`ExecutionPlan`] allowing full interoperability with the existing ecosystem +/// +/// Longer term we will likely want to introduce new traits that differentiate between +/// pipeline-able operators like filters, and pipeline-breakers like aggregations, and +/// are better aligned with a push-based execution model. +/// +/// This in turn will allow for [`Pipeline`] implementations that are able to introduce +/// parallelism beyond that expressed in their partitioning +pub struct ExecutionPipeline { + proxied: Arc<dyn ExecutionPlan>, + inputs: Vec<Vec<Arc<Mutex<InputPartition>>>>, + outputs: Vec<Mutex<BoxStream<'static, ArrowResult<RecordBatch>>>>, +} + +impl std::fmt::Debug for ExecutionPipeline { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let tree = debug_tree(self.proxied.as_ref()); + f.debug_tuple("ExecutionNode").field(&tree).finish() + } +} + +impl ExecutionPipeline { + pub fn new( + plan: Arc<dyn ExecutionPlan>, + task_context: Arc<TaskContext>, + depth: usize, + ) -> Result<Self> { + // The point in the plan at which to splice the plan graph + let mut splice_point = plan; + let mut parent_plans = Vec::with_capacity(depth.saturating_sub(1)); + for _ in 0..depth { + let children = splice_point.children(); + assert_eq!( + children.len(), + 1, + "can only group through nodes with a single child" + ); + parent_plans.push(splice_point); + splice_point = children.into_iter().next().unwrap(); + } + + // The children to replace with [`ProxyExecutionPlan`] + let children = splice_point.children(); + let mut inputs = Vec::with_capacity(children.len()); + + // The spliced plan with its children replaced with [`ProxyExecutionPlan`] + let spliced = if !children.is_empty() { + let mut proxies: Vec<Arc<dyn ExecutionPlan>> = + Vec::with_capacity(children.len()); + + for child in children { + let count = child.output_partitioning().partition_count(); + + let mut child_inputs = Vec::with_capacity(count); + for _ in 0..count { + child_inputs.push(Default::default()) + } + + inputs.push(child_inputs.clone()); + proxies.push(Arc::new(ProxyExecutionPlan { + inner: child, + inputs: child_inputs, + })); + } + + splice_point.with_new_children(proxies)? + } else { + splice_point.clone() + }; + + // Reconstruct the parent graph + let mut proxied = spliced; + for parent in parent_plans.into_iter().rev() { + proxied = parent.with_new_children(vec![proxied])? + } + + // Construct the output streams + let output_count = proxied.output_partitioning().partition_count(); + let outputs = (0..output_count) + .map(|x| { + let proxy_captured = proxied.clone(); + let task_captured = task_context.clone(); + let fut = async move { + proxy_captured + .execute(x, task_captured) + .await + .map_err(|e| ArrowError::ExternalError(Box::new(e))) + }; + + // Use futures::stream::once to handle operators that perform computation + // within `ExecutionPlan::execute`. If we evaluated these futures here + // we could potentially block indefinitely waiting for inputs that will + // never arrive as the query isn't scheduled yet + Mutex::new(futures::stream::once(fut).try_flatten().boxed()) + }) + .collect(); + + Ok(Self { + proxied, + inputs, + outputs, + }) + } +} + +impl Pipeline for ExecutionPipeline { + /// Push a [`RecordBatch`] to the given input partition + fn push(&self, input: RecordBatch, child: usize, partition: usize) -> Result<()> { + let mut partition = self.inputs[child][partition].lock(); + assert!(!partition.is_closed); + + partition.buffer.push_back(input); + for waker in partition.wait_list.drain(..) { + waker.wake() + } + Ok(()) + } + + fn close(&self, child: usize, partition: usize) { + let mut partition = self.inputs[child][partition].lock(); + assert!(!partition.is_closed); + + partition.is_closed = true; + for waker in partition.wait_list.drain(..) { + waker.wake() + } + } + + fn output_partitions(&self) -> usize { + self.outputs.len() + } + + /// Poll an output partition, attempting to get its output + fn poll_partition( + &self, + cx: &mut Context<'_>, + partition: usize, + ) -> Poll<Option<Result<RecordBatch>>> { + self.outputs[partition] + .lock() + .poll_next_unpin(cx) + .map(|opt| opt.map(|r| r.map_err(Into::into))) + } +} + +#[derive(Debug, Default)] +struct InputPartition { + buffer: VecDeque<RecordBatch>, + wait_list: Vec<Waker>, Review Comment: Maybe calling it `wake_list` would be better to align with the list of wakers ```suggestion wake_list: Vec<Waker>, ``` ########## datafusion/scheduler/src/task.rs: ########## @@ -0,0 +1,439 @@ +// 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 crate::query::Query; +use crate::{is_worker, spawn_local, spawn_local_fifo, RoutablePipeline, Spawner}; +use arrow::record_batch::RecordBatch; +use datafusion::error::{DataFusionError, Result}; +use futures::channel::mpsc; +use futures::task::ArcWake; +use futures::{Stream, StreamExt}; +use log::{debug, trace}; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Weak}; +use std::task::{Context, Poll}; + +/// Spawns a query using the provided [`Spawner`] +pub fn spawn_query(query: Query, spawner: Spawner) -> QueryResults { + debug!("Spawning query: {:#?}", query); + + let (sender, receiver) = mpsc::unbounded(); + let query = Arc::new(QueryTask { + spawner, + pipelines: query.pipelines, + output: sender, + }); + + for (pipeline_idx, query_pipeline) in query.pipelines.iter().enumerate() { + for partition in 0..query_pipeline.pipeline.output_partitions() { + query.spawner.spawn(Task { + query: query.clone(), + waker: Arc::new(TaskWaker { + query: Arc::downgrade(&query), + wake_count: AtomicUsize::new(1), + pipeline: pipeline_idx, + partition, + }), + }); + } + } + + QueryResults { + query, + inner: receiver, + } +} + +/// A [`Task`] identifies an output partition within a given pipeline that may be able to +/// make progress. The [`Scheduler`][super::Scheduler] maintains a list of outstanding +/// [`Task`] and distributes them amongst its worker threads. +/// +/// A [`Query`] is considered completed when it has no outstanding [`Task`] +pub struct Task { + /// Maintain a link to the [`QueryTask`] this is necessary to be able to + /// route the output of the partition to its destination, and also because + /// when [`QueryTask`] is dropped it signals completion of query execution + query: Arc<QueryTask>, + + /// A [`ArcWake`] that can be used to re-schedule this [`Task`] for execution + waker: Arc<TaskWaker>, +} + +impl std::fmt::Debug for Task { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let output = &self.query.pipelines[self.waker.pipeline].output; + + f.debug_struct("Task") + .field("pipeline", &self.waker.pipeline) + .field("partition", &self.waker.partition) + .field("output", &output) + .finish() + } +} + +impl Task { + fn handle_error(&self, routable: &RoutablePipeline, error: DataFusionError) { + self.query.send_query_output(Err(error)); + if let Some(link) = routable.output { + trace!( + "Closing pipeline: {:?}, partition: {}, due to error", + link, + self.waker.partition, + ); + + self.query.pipelines[link.pipeline] + .pipeline + .close(link.child, self.waker.partition); + } + } + /// Call [`Pipeline::poll_partition`] attempting to make progress on query execution + pub fn do_work(self) { + assert!(is_worker(), "Task::do_work called outside of worker pool"); + if self.query.is_cancelled() { + return; + } + + // Capture the wake count prior to calling [`Pipeline::poll_partition`] + // this allows us to detect concurrent wake ups and handle them correctly + let wake_count = self.waker.wake_count.load(Ordering::SeqCst); + + let node = self.waker.pipeline; + let partition = self.waker.partition; + + let waker = futures::task::waker_ref(&self.waker); + let mut cx = Context::from_waker(&*waker); + + let pipelines = &self.query.pipelines; + let routable = &pipelines[node]; + match routable.pipeline.poll_partition(&mut cx, partition) { + Poll::Ready(Some(Ok(batch))) => { + trace!("Poll {:?}: Ok: {}", self, batch.num_rows()); + match routable.output { + Some(link) => { + trace!( + "Publishing batch to pipeline {:?} partition {}", + link, + partition + ); + + let r = pipelines[link.pipeline] + .pipeline + .push(batch, link.child, partition); + + if let Err(e) = r { + self.handle_error(routable, e); + + // Return without rescheduling this output again + return; + } + } + None => { + trace!("Publishing batch to output"); + self.query.send_query_output(Ok(batch)) + } + } + + // Reschedule this pipeline again + // + // We want to prioritise running tasks triggered by the most recent + // batch, so reschedule with FIFO ordering + // + // Note: We must schedule after we have routed the batch, otherwise + // we introduce a potential ordering race where the newly scheduled + // task runs before this task finishes routing the output + spawn_local_fifo(self); + } + Poll::Ready(Some(Err(e))) => { + trace!("Poll {:?}: Error: {:?}", self, e); + self.handle_error(routable, e) + } + Poll::Ready(None) => { + trace!("Poll {:?}: None", self); + match routable.output { + Some(link) => { + trace!("Closing pipeline: {:?}, partition: {}", link, partition); + pipelines[link.pipeline] + .pipeline + .close(link.child, partition) + } + None => self.query.finish(), + } + } + Poll::Pending => { + trace!("Poll {:?}: Pending", self); + // Attempt to reset the wake count with the value obtained prior + // to calling [`Pipeline::poll_partition`]. + // + // If this fails it indicates a wakeup was received whilst executing + // [`Pipeline::poll_partition`] and we should reschedule the task + let reset = self.waker.wake_count.compare_exchange( + wake_count, + 0, + Ordering::SeqCst, + Ordering::SeqCst, + ); + + if reset.is_err() { + trace!("Wakeup triggered whilst polling: {:?}", self); + spawn_local(self); + } + } + } + } +} + +/// The result stream for a query +/// +/// # Cancellation +/// +/// Dropping this will cancel the inflight query +pub struct QueryResults { Review Comment: ```suggestion pub struct ExecutionResults { ``` Or something like ```suggestion pub struct PipelinePlanResults { ``` (related to comment above) ########## datafusion/scheduler/src/task.rs: ########## @@ -0,0 +1,439 @@ +// 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 crate::query::Query; +use crate::{is_worker, spawn_local, spawn_local_fifo, RoutablePipeline, Spawner}; +use arrow::record_batch::RecordBatch; +use datafusion::error::{DataFusionError, Result}; +use futures::channel::mpsc; +use futures::task::ArcWake; +use futures::{Stream, StreamExt}; +use log::{debug, trace}; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Weak}; +use std::task::{Context, Poll}; + +/// Spawns a query using the provided [`Spawner`] +pub fn spawn_query(query: Query, spawner: Spawner) -> QueryResults { + debug!("Spawning query: {:#?}", query); + + let (sender, receiver) = mpsc::unbounded(); + let query = Arc::new(QueryTask { + spawner, + pipelines: query.pipelines, + output: sender, + }); + + for (pipeline_idx, query_pipeline) in query.pipelines.iter().enumerate() { + for partition in 0..query_pipeline.pipeline.output_partitions() { + query.spawner.spawn(Task { + query: query.clone(), + waker: Arc::new(TaskWaker { + query: Arc::downgrade(&query), + wake_count: AtomicUsize::new(1), + pipeline: pipeline_idx, + partition, + }), + }); + } + } + + QueryResults { + query, + inner: receiver, + } +} + +/// A [`Task`] identifies an output partition within a given pipeline that may be able to +/// make progress. The [`Scheduler`][super::Scheduler] maintains a list of outstanding +/// [`Task`] and distributes them amongst its worker threads. +/// +/// A [`Query`] is considered completed when it has no outstanding [`Task`] +pub struct Task { + /// Maintain a link to the [`QueryTask`] this is necessary to be able to + /// route the output of the partition to its destination, and also because + /// when [`QueryTask`] is dropped it signals completion of query execution + query: Arc<QueryTask>, + + /// A [`ArcWake`] that can be used to re-schedule this [`Task`] for execution + waker: Arc<TaskWaker>, +} + +impl std::fmt::Debug for Task { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let output = &self.query.pipelines[self.waker.pipeline].output; + + f.debug_struct("Task") + .field("pipeline", &self.waker.pipeline) + .field("partition", &self.waker.partition) + .field("output", &output) + .finish() + } +} + +impl Task { + fn handle_error(&self, routable: &RoutablePipeline, error: DataFusionError) { + self.query.send_query_output(Err(error)); + if let Some(link) = routable.output { + trace!( + "Closing pipeline: {:?}, partition: {}, due to error", + link, + self.waker.partition, + ); + + self.query.pipelines[link.pipeline] + .pipeline + .close(link.child, self.waker.partition); + } + } + /// Call [`Pipeline::poll_partition`] attempting to make progress on query execution + pub fn do_work(self) { + assert!(is_worker(), "Task::do_work called outside of worker pool"); + if self.query.is_cancelled() { + return; + } + + // Capture the wake count prior to calling [`Pipeline::poll_partition`] + // this allows us to detect concurrent wake ups and handle them correctly + let wake_count = self.waker.wake_count.load(Ordering::SeqCst); + + let node = self.waker.pipeline; + let partition = self.waker.partition; + + let waker = futures::task::waker_ref(&self.waker); + let mut cx = Context::from_waker(&*waker); + + let pipelines = &self.query.pipelines; + let routable = &pipelines[node]; + match routable.pipeline.poll_partition(&mut cx, partition) { + Poll::Ready(Some(Ok(batch))) => { + trace!("Poll {:?}: Ok: {}", self, batch.num_rows()); + match routable.output { + Some(link) => { + trace!( + "Publishing batch to pipeline {:?} partition {}", + link, + partition + ); + + let r = pipelines[link.pipeline] + .pipeline + .push(batch, link.child, partition); + + if let Err(e) = r { + self.handle_error(routable, e); + + // Return without rescheduling this output again + return; + } + } + None => { + trace!("Publishing batch to output"); + self.query.send_query_output(Ok(batch)) + } + } + + // Reschedule this pipeline again + // + // We want to prioritise running tasks triggered by the most recent + // batch, so reschedule with FIFO ordering + // + // Note: We must schedule after we have routed the batch, otherwise + // we introduce a potential ordering race where the newly scheduled + // task runs before this task finishes routing the output + spawn_local_fifo(self); + } + Poll::Ready(Some(Err(e))) => { + trace!("Poll {:?}: Error: {:?}", self, e); + self.handle_error(routable, e) + } + Poll::Ready(None) => { + trace!("Poll {:?}: None", self); + match routable.output { + Some(link) => { + trace!("Closing pipeline: {:?}, partition: {}", link, partition); + pipelines[link.pipeline] + .pipeline + .close(link.child, partition) + } + None => self.query.finish(), + } + } + Poll::Pending => { + trace!("Poll {:?}: Pending", self); + // Attempt to reset the wake count with the value obtained prior + // to calling [`Pipeline::poll_partition`]. + // + // If this fails it indicates a wakeup was received whilst executing + // [`Pipeline::poll_partition`] and we should reschedule the task + let reset = self.waker.wake_count.compare_exchange( + wake_count, + 0, + Ordering::SeqCst, + Ordering::SeqCst, + ); + + if reset.is_err() { + trace!("Wakeup triggered whilst polling: {:?}", self); + spawn_local(self); + } + } + } + } +} + +/// The result stream for a query +/// +/// # Cancellation +/// +/// Dropping this will cancel the inflight query +pub struct QueryResults { + inner: mpsc::UnboundedReceiver<Option<Result<RecordBatch>>>, + + /// Keep a reference to the [`QueryTask`] so it isn't dropped early + #[allow(unused)] + query: Arc<QueryTask>, +} + +impl Stream for QueryResults { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.inner.poll_next_unpin(cx).map(Option::flatten) + } +} + +/// The shared state of all [`Task`] created from the same [`Query`] +#[derive(Debug)] +struct QueryTask { + /// Spawner for this query + spawner: Spawner, + + /// List of pipelines that belong to this query, pipelines are addressed + /// based on their index within this list + pipelines: Vec<RoutablePipeline>, + + /// The output stream for this query's execution + output: mpsc::UnboundedSender<Option<Result<RecordBatch>>>, +} + +impl Drop for QueryTask { + fn drop(&mut self) { + debug!("Query dropped"); + } +} + +impl QueryTask { + /// Returns `true` if this query has been dropped, specifically if the + /// stream returned by [`super::Scheduler::schedule`] has been dropped + fn is_cancelled(&self) -> bool { + self.output.is_closed() + } + + /// Sends `output` to this query's output stream + fn send_query_output(&self, output: Result<RecordBatch>) { + let _ = self.output.unbounded_send(Some(output)); + } + + /// Mark this query as finished + fn finish(&self) { + let _ = self.output.unbounded_send(None); + } +} + +struct TaskWaker { Review Comment: Since we already use the futures `Waker` elsewhere in this PR, I wonder if we can use the same here? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
