diff --git a/datafusion/core/src/physical_plan/mod.rs b/datafusion/core/src/physical_plan/mod.rs index 66254ee6f5f8..cda0b19b4f52 100644 --- a/datafusion/core/src/physical_plan/mod.rs +++ b/datafusion/core/src/physical_plan/mod.rs @@ -17,6 +17,7 @@ //! Traits for physical query plan, supporting parallel execution for partitioned relations. +mod visitor; pub use self::metrics::Metric; use self::metrics::MetricsSet; use self::{ @@ -26,6 +27,7 @@ use crate::datasource::physical_plan::FileScanConfig; use crate::physical_plan::expressions::PhysicalSortExpr; use datafusion_common::Result; pub use datafusion_common::{ColumnStatistics, Statistics}; +pub use visitor::{accept, visit_execution_plan, ExecutionPlanVisitor}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -331,82 +333,6 @@ pub fn displayable(plan: &dyn ExecutionPlan) -> DisplayableExecutionPlan<'_> { DisplayableExecutionPlan::new(plan) } -/// Visit all children of this plan, according to the order defined on `ExecutionPlanVisitor`. -// Note that this would be really nice if it were a method on -// ExecutionPlan, but it can not be because it takes a generic -// parameter and `ExecutionPlan` is a trait -pub fn accept( - plan: &dyn ExecutionPlan, - visitor: &mut V, -) -> Result<(), V::Error> { - visitor.pre_visit(plan)?; - for child in plan.children() { - visit_execution_plan(child.as_ref(), visitor)?; - } - visitor.post_visit(plan)?; - Ok(()) -} - -/// Trait that implements the [Visitor -/// pattern](https://en.wikipedia.org/wiki/Visitor_pattern) for a -/// depth first walk of `ExecutionPlan` nodes. `pre_visit` is called -/// before any children are visited, and then `post_visit` is called -/// after all children have been visited. -//// -/// To use, define a struct that implements this trait and then invoke -/// ['accept']. -/// -/// For example, for an execution plan that looks like: -/// -/// ```text -/// ProjectionExec: id -/// FilterExec: state = CO -/// CsvExec: -/// ``` -/// -/// The sequence of visit operations would be: -/// ```text -/// visitor.pre_visit(ProjectionExec) -/// visitor.pre_visit(FilterExec) -/// visitor.pre_visit(CsvExec) -/// visitor.post_visit(CsvExec) -/// visitor.post_visit(FilterExec) -/// visitor.post_visit(ProjectionExec) -/// ``` -pub trait ExecutionPlanVisitor { - /// The type of error returned by this visitor - type Error; - - /// Invoked on an `ExecutionPlan` plan before any of its child - /// inputs have been visited. If Ok(true) is returned, the - /// recursion continues. If Err(..) or Ok(false) are returned, the - /// recursion stops immediately and the error, if any, is returned - /// to `accept` - fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result; - - /// Invoked on an `ExecutionPlan` plan *after* all of its child - /// inputs have been visited. The return value is handled the same - /// as the return value of `pre_visit`. The provided default - /// implementation returns `Ok(true)`. - fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result { - Ok(true) - } -} - -/// Recursively calls `pre_visit` and `post_visit` for this node and -/// all of its children, as described on [`ExecutionPlanVisitor`] -pub fn visit_execution_plan( - plan: &dyn ExecutionPlan, - visitor: &mut V, -) -> Result<(), V::Error> { - visitor.pre_visit(plan)?; - for child in plan.children() { - visit_execution_plan(child.as_ref(), visitor)?; - } - visitor.post_visit(plan)?; - Ok(()) -} - /// Execute the [ExecutionPlan] and collect the results in memory pub async fn collect( plan: Arc, diff --git a/datafusion/core/src/physical_plan/visitor.rs b/datafusion/core/src/physical_plan/visitor.rs new file mode 100644 index 000000000000..573e4f8b02be --- /dev/null +++ b/datafusion/core/src/physical_plan/visitor.rs @@ -0,0 +1,94 @@ +// 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 super::ExecutionPlan; + +/// Visit all children of this plan, according to the order defined on `ExecutionPlanVisitor`. +// Note that this would be really nice if it were a method on +// ExecutionPlan, but it can not be because it takes a generic +// parameter and `ExecutionPlan` is a trait +pub fn accept( + plan: &dyn ExecutionPlan, + visitor: &mut V, +) -> Result<(), V::Error> { + visitor.pre_visit(plan)?; + for child in plan.children() { + visit_execution_plan(child.as_ref(), visitor)?; + } + visitor.post_visit(plan)?; + Ok(()) +} + +/// Trait that implements the [Visitor +/// pattern](https://en.wikipedia.org/wiki/Visitor_pattern) for a +/// depth first walk of `ExecutionPlan` nodes. `pre_visit` is called +/// before any children are visited, and then `post_visit` is called +/// after all children have been visited. +//// +/// To use, define a struct that implements this trait and then invoke +/// ['accept']. +/// +/// For example, for an execution plan that looks like: +/// +/// ```text +/// ProjectionExec: id +/// FilterExec: state = CO +/// CsvExec: +/// ``` +/// +/// The sequence of visit operations would be: +/// ```text +/// visitor.pre_visit(ProjectionExec) +/// visitor.pre_visit(FilterExec) +/// visitor.pre_visit(CsvExec) +/// visitor.post_visit(CsvExec) +/// visitor.post_visit(FilterExec) +/// visitor.post_visit(ProjectionExec) +/// ``` +pub trait ExecutionPlanVisitor { + /// The type of error returned by this visitor + type Error; + + /// Invoked on an `ExecutionPlan` plan before any of its child + /// inputs have been visited. If Ok(true) is returned, the + /// recursion continues. If Err(..) or Ok(false) are returned, the + /// recursion stops immediately and the error, if any, is returned + /// to `accept` + fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result; + + /// Invoked on an `ExecutionPlan` plan *after* all of its child + /// inputs have been visited. The return value is handled the same + /// as the return value of `pre_visit`. The provided default + /// implementation returns `Ok(true)`. + fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result { + Ok(true) + } +} + +/// Recursively calls `pre_visit` and `post_visit` for this node and +/// all of its children, as described on [`ExecutionPlanVisitor`] +pub fn visit_execution_plan( + plan: &dyn ExecutionPlan, + visitor: &mut V, +) -> Result<(), V::Error> { + visitor.pre_visit(plan)?; + for child in plan.children() { + visit_execution_plan(child.as_ref(), visitor)?; + } + visitor.post_visit(plan)?; + Ok(()) +}