From f7203aeaf500b7e90c5356c260cbebb0c4336fe3 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 10 Aug 2026 14:43:39 +0000 Subject: [PATCH 1/2] refactor: make apply_expression_roots more ergonomic --- datafusion/datasource-arrow/src/source.rs | 8 +---- datafusion/datasource-avro/src/source.rs | 8 +---- datafusion/datasource-csv/src/source.rs | 8 +---- datafusion/datasource-json/src/source.rs | 8 +---- datafusion/physical-expr/src/projection.rs | 6 ++++ .../physical-plan/src/execution_plan.rs | 33 +++++++++++++++++-- datafusion/physical-plan/src/lib.rs | 7 ++-- datafusion/physical-plan/src/projection.rs | 9 +---- datafusion/proto/src/physical_plan/mod.rs | 5 +-- 9 files changed, 46 insertions(+), 46 deletions(-) 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..c52e94c385e91 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -69,6 +69,12 @@ impl PartialEq for ProjectionExpr { impl Eq for ProjectionExpr {} +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..435207e6326c7 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,33 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } +/// A value that contains a physical expression root. +pub trait PhysicalExprRoot { + /// Returns the physical expression at this root. + fn as_physical_expr_root(&self) -> &Arc; +} + +impl PhysicalExprRoot for Arc { + fn as_physical_expr_root(&self) -> &Arc { + self + } +} + +impl PhysicalExprRoot for ProjectionExpr { + fn as_physical_expr_root(&self) -> &Arc { + self.as_ref() + } +} + +impl PhysicalExprRoot for &T +where + T: PhysicalExprRoot + ?Sized, +{ + fn as_physical_expr_root(&self) -> &Arc { + (*self).as_physical_expr_root() + } +} + /// Applies `f` to a shallow sequence of physical expression roots. /// /// [`TreeNodeRecursion::Stop`] stops iteration and is returned immediately. @@ -947,10 +974,10 @@ pub fn apply_expression_roots( ) -> Result where I: IntoIterator, - I::Item: Borrow>, + I::Item: PhysicalExprRoot, { for root in roots { - match f(root.borrow())? { + match f(root.as_physical_expr_root())? { 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..c85f43e0f9002 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, + ExecutionPlan, ExecutionPlanProperties, PhysicalExprRoot, 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, ) } From 5d6146ca3fc0c0c601c11bd04ed55e6eb4327dfe Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 14:54:41 +0000 Subject: [PATCH 2/2] fix docs --- datafusion/physical-expr/src/projection.rs | 2 + .../physical-plan/src/execution_plan.rs | 37 ++++++++++++------- datafusion/physical-plan/src/lib.rs | 2 +- 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index c52e94c385e91..d62ef76d8980b 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -69,6 +69,8 @@ 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 diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 435207e6326c7..b1d7c32882b0d 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -936,30 +936,39 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } -/// A value that contains a physical expression root. -pub trait PhysicalExprRoot { - /// Returns the physical expression at this root. - fn as_physical_expr_root(&self) -> &Arc; +/// 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; } -impl PhysicalExprRoot for Arc { - fn as_physical_expr_root(&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 } } -impl PhysicalExprRoot for ProjectionExpr { - fn as_physical_expr_root(&self) -> &Arc { +/// 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 PhysicalExprRoot for &T +impl AsPhysicalExprRef for &T where - T: PhysicalExprRoot + ?Sized, + T: AsPhysicalExprRef + ?Sized, { - fn as_physical_expr_root(&self) -> &Arc { - (*self).as_physical_expr_root() + fn as_physical_expr_ref(&self) -> &Arc { + (*self).as_physical_expr_ref() } } @@ -974,10 +983,10 @@ pub fn apply_expression_roots( ) -> Result where I: IntoIterator, - I::Item: PhysicalExprRoot, + I::Item: AsPhysicalExprRef, { for root in roots { - match f(root.as_physical_expr_root())? { + 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 c85f43e0f9002..941b4e561bb12 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -45,7 +45,7 @@ pub use crate::distribution_requirements::{ ChildSatisfactionOptions, InputDistributionRequirements, }; pub use crate::execution_plan::{ - ExecutionPlan, ExecutionPlanProperties, PhysicalExprRoot, PlanProperties, + 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,