diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index bf92faa9a1104..f51e1c100934d 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -400,13 +400,7 @@ impl FileSource for ArrowSource { &Arc, ) -> Result, ) -> Result { - datafusion_physical_plan::apply_expression_roots( - self.projection - .source - .iter() - .map(|proj_expr| &proj_expr.expr), - f, - ) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) } /// Emit an `ArrowScan` node wrapping the shared base config. diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index 3956a7318d5ca..fcc50b559f00b 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/datafusion/datasource-avro/src/source.rs @@ -176,13 +176,7 @@ impl FileSource for AvroSource { &Arc, ) -> Result, ) -> Result { - datafusion_physical_plan::apply_expression_roots( - self.projection - .source - .iter() - .map(|proj_expr| &proj_expr.expr), - f, - ) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) } /// Emit an `AvroScan` node wrapping the shared base config. diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index d5fc6288eaaa3..08e4607498e62 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -316,13 +316,7 @@ impl FileSource for CsvSource { &Arc, ) -> Result, ) -> Result { - datafusion_physical_plan::apply_expression_roots( - self.projection - .source - .iter() - .map(|proj_expr| &proj_expr.expr), - f, - ) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) } /// Emit a `CsvScan` node wrapping the shared base config and CSV options. diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index c6d420bafb2f7..47241c9d99ab5 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -239,13 +239,7 @@ impl FileSource for JsonSource { &Arc, ) -> Result, ) -> Result { - datafusion_physical_plan::apply_expression_roots( - self.projection - .source - .iter() - .map(|proj_expr| &proj_expr.expr), - f, - ) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) } /// Emit a `JsonScan` node wrapping the shared base config. diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index e3fd6ddf744a9..d62ef76d8980b 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -69,6 +69,14 @@ impl PartialEq for ProjectionExpr { impl Eq for ProjectionExpr {} +/// Enables [`ProjectionExpr`] to be treated as a reference to its wrapped +/// [`Arc`] using [`AsRef::as_ref`]. +impl AsRef> for ProjectionExpr { + fn as_ref(&self) -> &Arc { + &self.expr + } +} + impl std::fmt::Display for ProjectionExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { if self.expr.to_string() == self.alias { diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 763ec2f7dfcd3..b1d7c32882b0d 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -35,13 +35,13 @@ pub use datafusion_common::utils::project_schema; pub use datafusion_common::{ColumnStatistics, Statistics, internal_err}; pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; pub use datafusion_expr::{Accumulator, ColumnarValue}; +use datafusion_physical_expr::projection::ProjectionExpr; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ Distribution, Partitioning, PhysicalExpr, expressions, }; use std::any::Any; -use std::borrow::Borrow; use std::collections::HashSet; use std::fmt::Debug; use std::sync::{Arc, LazyLock}; @@ -936,6 +936,42 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } +/// Allows a type to be treated as a reference to an +/// [`Arc`]. +/// +/// Used by [`apply_expression_roots`]. +pub trait AsPhysicalExprRef { + /// Returns the referenced physical expression. + fn as_physical_expr_ref(&self) -> &Arc; +} + +/// Allows an [`Arc`] to be treated as a reference to itself. +/// +/// This is needed because `Arc` does not implement +/// `AsRef>`. +impl AsPhysicalExprRef for Arc { + fn as_physical_expr_ref(&self) -> &Arc { + self + } +} + +/// Allows a [`ProjectionExpr`] to be treated as a reference to its +/// [`Arc`]. +impl AsPhysicalExprRef for ProjectionExpr { + fn as_physical_expr_ref(&self) -> &Arc { + self.as_ref() + } +} + +impl AsPhysicalExprRef for &T +where + T: AsPhysicalExprRef + ?Sized, +{ + fn as_physical_expr_ref(&self) -> &Arc { + (*self).as_physical_expr_ref() + } +} + /// Applies `f` to a shallow sequence of physical expression roots. /// /// [`TreeNodeRecursion::Stop`] stops iteration and is returned immediately. @@ -947,10 +983,10 @@ pub fn apply_expression_roots( ) -> Result where I: IntoIterator, - I::Item: Borrow>, + I::Item: AsPhysicalExprRef, { for root in roots { - match f(root.borrow())? { + match f(root.as_physical_expr_ref())? { TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop), TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {} } diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 6e1df1f840af0..941b4e561bb12 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -45,9 +45,10 @@ pub use crate::distribution_requirements::{ ChildSatisfactionOptions, InputDistributionRequirements, }; pub use crate::execution_plan::{ - ExecutionPlan, ExecutionPlanProperties, PlanProperties, apply_expression_roots, - collect, collect_partitioned, displayable, execute_input_stream, execute_stream, - execute_stream_partitioned, get_plan_string, with_new_children_if_necessary, + AsPhysicalExprRef, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + apply_expression_roots, collect, collect_partitioned, displayable, + execute_input_stream, execute_stream, execute_stream_partitioned, get_plan_string, + with_new_children_if_necessary, }; pub use crate::metrics::Metric; pub use crate::ordering::InputOrderMode; diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 32d444fba2b03..ecdd78cc2acc1 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -334,14 +334,7 @@ impl ExecutionPlan for ProjectionExec { &self, f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - crate::apply_expression_roots( - self.projector - .projection() - .as_ref() - .iter() - .map(|proj_expr| &proj_expr.expr), - f, - ) + crate::apply_expression_roots(self.projector.projection().as_ref().iter(), f) } fn with_new_children( diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index de684857f4446..3a67e55ea69a7 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -179,10 +179,7 @@ mod file_scan_config_serde { -> Result, ) -> Result { datafusion_physical_plan::apply_expression_roots( - self.projection - .iter() - .flatten() - .map(|proj_expr| &proj_expr.expr), + self.projection.iter().flatten(), f, ) }