alamb commented on code in PR #2226: URL: https://github.com/apache/arrow-datafusion/pull/2226#discussion_r856535693
########## 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: Those two links seem to point at a list called `waiters` (rather than `wait_list`) so not quite what is in this PR either :) I don't feel strongly about this -- 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]
