refactor: make apply_expression_roots more ergonomic

datafusion/datasource-arrow/src/source.rs

@@ -400,13 +400,7 @@ impl FileSource for ArrowSource {
             &Arc<dyn datafusion_physical_plan::PhysicalExpr>,
         ) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
-        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.

datafusion/datasource-avro/src/source.rs

@@ -176,13 +176,7 @@ impl FileSource for AvroSource {
             &Arc<dyn datafusion_physical_plan::PhysicalExpr>,
         ) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
-        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.

datafusion/datasource-csv/src/source.rs

@@ -316,13 +316,7 @@ impl FileSource for CsvSource {
             &Arc<dyn datafusion_physical_plan::PhysicalExpr>,
         ) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
-        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.

datafusion/datasource-json/src/source.rs

@@ -239,13 +239,7 @@ impl FileSource for JsonSource {
             &Arc<dyn datafusion_physical_plan::PhysicalExpr>,
         ) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
-        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.

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<dyn PhysicalExpr>`] using [`AsRef::as_ref`].
+impl AsRef<Arc<dyn PhysicalExpr>> for ProjectionExpr {
+    fn as_ref(&self) -> &Arc<dyn PhysicalExpr> {
+        &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 {

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<dyn PhysicalExpr>`].
+///
+/// Used by [`apply_expression_roots`].
+pub trait AsPhysicalExprRef {
+    /// Returns the referenced physical expression.
+    fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr>;
+}
+
+/// Allows an [`Arc<dyn PhysicalExpr>`] to be treated as a reference to itself.
+///
+/// This is needed because `Arc<dyn PhysicalExpr>` does not implement
+/// `AsRef<Arc<dyn PhysicalExpr>>`.
+impl AsPhysicalExprRef for Arc<dyn PhysicalExpr> {
+    fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> {
+        self
+    }
+}
+
+/// Allows a [`ProjectionExpr`] to be treated as a reference to its
+/// [`Arc<dyn PhysicalExpr>`].
+impl AsPhysicalExprRef for ProjectionExpr {
+    fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> {
+        self.as_ref()
+    }
+}
+
+impl<T> AsPhysicalExprRef for &T
+where
+    T: AsPhysicalExprRef + ?Sized,
+{
+    fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> {
+        (*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<I>(
 ) -> Result<TreeNodeRecursion>
 where
     I: IntoIterator,
-    I::Item: Borrow<Arc<dyn PhysicalExpr>>,
+    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 => {}
         }

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;

datafusion/physical-plan/src/projection.rs

@@ -334,14 +334,7 @@ impl ExecutionPlan for ProjectionExec {
         &self,
         f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
-        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(

datafusion/proto/src/physical_plan/mod.rs

@@ -179,10 +179,7 @@ mod file_scan_config_serde {
                 -> Result<datafusion_common::tree_node::TreeNodeRecursion>,
         ) -> Result<datafusion_common::tree_node::TreeNodeRecursion> {
             datafusion_physical_plan::apply_expression_roots(
-                self.projection
-                    .iter()
-                    .flatten()
-                    .map(|proj_expr| &proj_expr.expr),
+                self.projection.iter().flatten(),
                 f,
             )
         }