initial implementation of projection pushdown for nexted expressions

Author: adriangb

datafusion/common/src/column.rs

@@ -30,7 +30,7 @@ use std::fmt;
 pub struct Column {
     /// relation/table reference.
     pub relation: Option<TableReference>,
-    /// field/column name.
+    /// Field/column name.
     pub name: String,
     /// Original source code location, if known
     pub spans: Spans,

datafusion/expr/src/expr.rs

@@ -1883,6 +1883,28 @@ impl Expr {
         }
     }
 
+    /// Returns true if this expression is trivial (cheap to evaluate).
+    ///
+    /// Trivial expressions include column references, literals, and nested
+    /// field access via `get_field`.
+    ///
+    /// # Example
+    /// ```
+    /// # use datafusion_expr::col;
+    /// let expr = col("foo");
+    /// assert!(expr.is_trivial());
+    /// ```
+    pub fn is_trivial(&self) -> bool {
+        match self {
+            Expr::Column(_) | Expr::Literal(_, _) => true,
+            Expr::ScalarFunction(func) => {
+                func.func.is_trivial()
+                    && func.args.first().is_some_and(|arg| arg.is_trivial())
+            }
+            _ => false,
+        }
+    }
+
     /// Return all references to columns in this expression.
     ///
     /// # Example

datafusion/expr/src/udf.rs

@@ -122,6 +122,11 @@ impl ScalarUDF {
         Self { inner: fun }
     }
 
+    /// Returns true if this function is trivial (cheap to evaluate).
+    pub fn is_trivial(&self) -> bool {
+        self.inner.is_trivial()
+    }
+
     /// Return the underlying [`ScalarUDFImpl`] trait object for this function
     pub fn inner(&self) -> &Arc<dyn ScalarUDFImpl> {
         &self.inner
@@ -846,6 +851,18 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     fn documentation(&self) -> Option<&Documentation> {
         None
     }
+
+    /// Returns true if this function is trivial (cheap to evaluate).
+    ///
+    /// Trivial functions are lightweight accessor functions like `get_field`
+    /// (struct field access) that simply access nested data within a column
+    /// without significant computation.
+    ///
+    /// This is used to identify expressions that are cheap to duplicate or
+    /// don't benefit from caching/partitioning optimizations.
+    fn is_trivial(&self) -> bool {
+        false
+    }
 }
 
 /// ScalarUDF that adds an alias to the underlying function. It is better to

datafusion/functions/src/core/getfield.rs

@@ -499,6 +499,10 @@ impl ScalarUDFImpl for GetFieldFunc {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn is_trivial(&self) -> bool {
+        true
+    }
 }
 
 #[cfg(test)]

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -593,7 +593,7 @@ fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Project
 
 // Check whether `expr` is trivial; i.e. it doesn't imply any computation.
 fn is_expr_trivial(expr: &Expr) -> bool {
-    matches!(expr, Expr::Column(_) | Expr::Literal(_, _))
+    expr.is_trivial()
 }
 
 /// Rewrites a projection expression using the projection before it (i.e. its input)

datafusion/physical-expr-common/src/physical_expr.rs

@@ -430,6 +430,20 @@ pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash {
     fn is_volatile_node(&self) -> bool {
         false
     }
+
+    /// Returns true if this expression is trivial (cheap to evaluate).
+    ///
+    /// Trivial expressions include:
+    /// - Column references
+    /// - Literal values
+    /// - Struct field access via `get_field`
+    /// - Nested combinations of field accessors (e.g., `col['a']['b']`)
+    ///
+    /// This is used to identify expressions that are cheap to duplicate or
+    /// don't benefit from caching/partitioning optimizations.
+    fn is_trivial(&self) -> bool {
+        false
+    }
 }
 
 #[deprecated(

datafusion/physical-expr/src/expressions/column.rs

@@ -146,6 +146,10 @@ impl PhysicalExpr for Column {
     fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.name)
     }
+
+    fn is_trivial(&self) -> bool {
+        true
+    }
 }
 
 impl Column {

datafusion/physical-expr/src/expressions/literal.rs

@@ -134,6 +134,10 @@ impl PhysicalExpr for Literal {
     fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         std::fmt::Display::fmt(self, f)
     }
+
+    fn is_trivial(&self) -> bool {
+        true
+    }
 }
 
 /// Create a literal expression

datafusion/physical-expr/src/scalar_function.rs

@@ -366,6 +366,13 @@ impl PhysicalExpr for ScalarFunctionExpr {
     fn is_volatile_node(&self) -> bool {
         self.fun.signature().volatility == Volatility::Volatile
     }
+
+    fn is_trivial(&self) -> bool {
+        if !self.fun.is_trivial() {
+            return false;
+        }
+        self.args.iter().all(|arg| arg.is_trivial())
+    }
 }
 
 #[cfg(test)]

datafusion/physical-plan/src/projection.rs

@@ -20,7 +20,7 @@
 //! of a projection on table `t1` where the expressions `a`, `b`, and `a+b` are the
 //! projection expressions. `SELECT` without `FROM` will only evaluate expressions.
 
-use super::expressions::{Column, Literal};
+use super::expressions::Column;
 use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use super::{
     DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream,
@@ -253,18 +253,16 @@ impl ExecutionPlan for ProjectionExec {
     }
 
     fn benefits_from_input_partitioning(&self) -> Vec<bool> {
-        let all_simple_exprs =
-            self.projector
-                .projection()
-                .as_ref()
-                .iter()
-                .all(|proj_expr| {
-                    proj_expr.expr.as_any().is::<Column>()
-                        || proj_expr.expr.as_any().is::<Literal>()
-                });
-        // If expressions are all either column_expr or Literal, then all computations in this projection are reorder or rename,
-        // and projection would not benefit from the repartition, benefits_from_input_partitioning will return false.
-        vec![!all_simple_exprs]
+        let all_trivial_exprs = self
+            .projector
+            .projection()
+            .as_ref()
+            .iter()
+            .all(|proj_expr| proj_expr.expr.is_trivial());
+        // If expressions are all trivial (columns, literals, or field accessors),
+        // then all computations in this projection are reorder or rename,
+        // and projection would not benefit from the repartition.
+        vec![!all_trivial_exprs]
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
@@ -630,11 +628,10 @@ pub fn make_with_child(
         .map(|e| Arc::new(e) as _)
 }
 
-/// Returns `true` if all the expressions in the argument are `Column`s.
-pub fn all_columns(exprs: &[ProjectionExpr]) -> bool {
-    exprs
-        .iter()
-        .all(|proj_expr| proj_expr.expr.as_any().is::<Column>())
+/// Returns `true` if all the expressions in the argument are trivial
+/// (columns, literals, or field accessors).
+pub fn all_trivial(exprs: &[ProjectionExpr]) -> bool {
+    exprs.iter().all(|proj_expr| proj_expr.expr.is_trivial())
 }
 
 /// Updates the given lexicographic ordering according to given projected
@@ -990,10 +987,9 @@ fn new_columns_for_join_on(
 }
 
 /// Checks if the given expression is trivial.
-/// An expression is considered trivial if it is either a `Column` or a `Literal`.
+/// An expression is considered trivial if it is a `Column`, `Literal`, or field accessor.
 fn is_expr_trivial(expr: &Arc<dyn PhysicalExpr>) -> bool {
-    expr.as_any().downcast_ref::<Column>().is_some()
-        || expr.as_any().downcast_ref::<Literal>().is_some()
+    expr.is_trivial()
 }
 
 #[cfg(test)]