feat: Allow struct field access projections to be pushed down into scans

Author: adriangb

Cargo.lock

@@ -41,7 +41,10 @@ pub mod dyn_eq;
 pub mod groups_accumulator;
 pub mod interval_arithmetic;
 pub mod operator;
+pub mod placement;
 pub mod signature;
 pub mod sort_properties;
 pub mod statistics;
 pub mod type_coercion;
+
+pub use placement::ExpressionPlacement;

datafusion/expr-common/src/lib.rs

@@ -0,0 +1,68 @@
+// 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.
+
+//! Expression placement classification for scalar functions.
+//!
+//! This module determines where in the query plan expressions should be placed
+//! to optimize data flow:
+//!
+//! - **Leaf placement**: Cheap expressions (field accessors, column references)
+//!   are pushed down to leaf nodes (near data sources) to reduce data volume early.
+//! - **Root placement**: Expensive expressions (computations, aggregates) are kept
+//!   at root nodes (after filtering) to operate on less data.
+
+/// Classification of expression placement for scalar functions.
+///
+/// This enum is used by [`ScalarUDFImpl::placement`] to allow
+/// functions to make context-dependent decisions about where they should
+/// be placed in the query plan based on the nature of their arguments.
+///
+/// For example, `get_field(struct_col, 'field_name')` is
+/// leaf-pushable (static field lookup), but `string_col like '%foo%'`
+/// performs expensive per-row computation and should be placed
+/// as further up the tree so that it can be run after filtering, sorting, etc.
+///
+/// # Why not pass in expressions directly to decide placement?
+///
+/// There are two reasons for using this enum instead of passing in the full expressions:
+///
+/// 1. **ScalarUDFImpl cannot reference PhysicalExpr**: The trait is defined in datafusion-expr,
+///    which cannot reference datafusion-physical-expr since the latter depends on the former
+///    (it would create a circular dependency).
+/// 2. **Simplicity**: Without this enum abstracting away logical / physical distinctions,
+///    we would need two distinct methods on ScalarUDFImpl: one for logical expression placement
+///    and one for physical expression placement. This would require implementors to duplicate logic
+///    and increases complexity for UDF authors.
+///
+/// [`ScalarUDFImpl::placement`]: https://docs.rs/datafusion-expr/latest/datafusion_expr/trait.ScalarUDFImpl.html#tymethod.placement
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ExpressionPlacement {
+    /// Argument is a literal constant value or an expression that can be
+    /// evaluated to a constant at planning time.
+    Literal,
+    /// Argument is a simple column reference.
+    Column,
+    /// Argument is a complex expression that can be safely placed at leaf nodes.
+    /// For example, if `get_field(struct_col, 'field_name')` is implemented as a
+    /// leaf-pushable expression, then it would return this variant.
+    /// Then `other_leaf_function(get_field(...), 42)` could also be classified as
+    /// leaf-pushable using the knowledge that `get_field(...)` is leaf-pushable.
+    PlaceAtLeafs,
+    /// Argument is a complex expression that should be placed at root nodes.
+    /// For example, `min(col1 + col2)` is not leaf-pushable because it requires per-row computation.
+    PlaceAtRoot,
+}

datafusion/expr-common/src/placement.rs

@@ -27,7 +27,7 @@ use std::sync::Arc;
 use crate::expr_fn::binary_expr;
 use crate::function::WindowFunctionSimplification;
 use crate::logical_plan::Subquery;
-use crate::{AggregateUDF, Volatility};
+use crate::{AggregateUDF, ExpressionPlacement, Volatility};
 use crate::{ExprSchemable, Operator, Signature, WindowFrame, WindowUDF};
 
 use arrow::datatypes::{DataType, Field, FieldRef};
@@ -1933,6 +1933,28 @@ impl Expr {
         }
     }
 
+    /// Returns the placement classification of this expression.
+    ///
+    /// This tells us if optimizers should preferentially
+    /// move this expression towards the leafs of the execution plan
+    /// tree (for cheap expressions or expressions that reduce the data size)
+    /// or towards the root of the execution plan tree (for expensive expressions
+    /// that should be run after filtering or parallelization, or expressions that increase the data size).
+    pub fn placement(&self) -> ExpressionPlacement {
+        match self {
+            Expr::Column(_) => ExpressionPlacement::Column,
+            Expr::Literal(_, _) => ExpressionPlacement::Literal,
+            Expr::ScalarFunction(func) => {
+                // Classify each argument's placement for context-aware decision making
+                let arg_placements: Vec<ExpressionPlacement> =
+                    func.args.iter().map(|arg| arg.placement()).collect();
+
+                func.func.placement_with_args(&arg_placements)
+            }
+            _ => ExpressionPlacement::PlaceAtRoot,
+        }
+    }
+
     /// Return all references to columns in this expression.
     ///
     /// # Example

datafusion/expr/src/expr.rs

@@ -92,6 +92,7 @@ pub use datafusion_doc::{
     DocSection, Documentation, DocumentationBuilder, aggregate_doc_sections,
     scalar_doc_sections, window_doc_sections,
 };
+pub use datafusion_expr_common::ExpressionPlacement;
 pub use datafusion_expr_common::accumulator::Accumulator;
 pub use datafusion_expr_common::columnar_value::ColumnarValue;
 pub use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};

datafusion/expr/src/lib.rs

@@ -31,6 +31,7 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_common::{ExprSchema, Result, ScalarValue, not_impl_err};
 use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
 use datafusion_expr_common::interval_arithmetic::Interval;
+use datafusion_expr_common::placement::ExpressionPlacement;
 use std::any::Any;
 use std::cmp::Ordering;
 use std::fmt::Debug;
@@ -123,6 +124,22 @@ impl ScalarUDF {
         Self { inner: fun }
     }
 
+    /// Returns the placement classification of this function given its arguments' placement.
+    ///
+    /// This allows functions to make context-dependent decisions about where they should
+    /// be placed in the query plan. For example, `get_field(struct_col, 'field_name')` is
+    /// leaf-pushable (static field lookup), but `string_col like '%foo%'`
+    /// performs expensive per-row computation and should be placed
+    /// as further up the tree so that it can be run after filtering, sorting, etc.
+    ///
+    /// See [`ScalarUDFImpl::placement`] for more details.
+    pub fn placement_with_args(
+        &self,
+        args: &[ExpressionPlacement],
+    ) -> ExpressionPlacement {
+        self.inner.placement(args)
+    }
+
     /// Return the underlying [`ScalarUDFImpl`] trait object for this function
     pub fn inner(&self) -> &Arc<dyn ScalarUDFImpl> {
         &self.inner
@@ -885,6 +902,37 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     fn documentation(&self) -> Option<&Documentation> {
         None
     }
+
+    /// Returns the placement classification of this function given its arguments' placement.
+    ///
+    /// This method allows functions to make context-dependent decisions about
+    /// where they should be placed in the query plan. The default implementation
+    /// returns [`ExpressionPlacement::PlaceAtRoot`] (conservative default).
+    ///
+    /// Leaf-pushable functions are lightweight accessor functions like `get_field`
+    /// (struct field access) that simply access nested data within a column
+    /// without significant computation.
+    /// These can be pushed down to leaf nodes near data sources to reduce data volume early in the plan.
+    ///
+    /// [`ExpressionPlacement::PlaceAtRoot`] represents expressions that should be kept after filtering,
+    /// such as expensive computations or aggregates that benefit from operating
+    /// on fewer rows.
+    ///
+    /// # Example
+    ///
+    /// - `get_field(struct_col, 'field_name')` with a literal key is leaf-pushable as it
+    ///    performs metadata only (cheap) extraction of a sub-array from a struct column.
+    ///    Thus, it can be placed near the data source to minimize data early.
+    /// - `string_col like '%foo%'` performs expensive per-row computation and should be placed
+    ///    further up the tree so that it can be run after filtering, sorting, etc.
+    ///
+    /// # Arguments
+    ///
+    /// * `args` - Classification of each argument's placement, collected from the expression tree
+    ///    by the caller.
+    fn placement(&self, _args: &[ExpressionPlacement]) -> ExpressionPlacement {
+        ExpressionPlacement::PlaceAtRoot
+    }
 }
 
 /// ScalarUDF that adds an alias to the underlying function. It is better to
@@ -1012,6 +1060,10 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
     fn documentation(&self) -> Option<&Documentation> {
         self.inner.documentation()
     }
+
+    fn placement(&self, args: &[ExpressionPlacement]) -> ExpressionPlacement {
+        self.inner.placement(args)
+    }
 }
 
 #[cfg(test)]

datafusion/expr/src/udf.rs

@@ -33,8 +33,8 @@ use datafusion_common::{
 use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::simplify::ExprSimplifyResult;
 use datafusion_expr::{
-    ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF,
-    ScalarUDFImpl, Signature, Volatility,
+    ColumnarValue, Documentation, Expr, ExpressionPlacement, ReturnFieldArgs,
+    ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
 };
 use datafusion_macros::user_doc;
 
@@ -499,6 +499,37 @@ impl ScalarUDFImpl for GetFieldFunc {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn placement(&self, args: &[ExpressionPlacement]) -> ExpressionPlacement {
+        // get_field is leaf-pushable if:
+        // 1. The struct/map argument is a Column or PlaceAtLeafs (not Literal or PlaceAtRoot)
+        // 2. All key arguments are literals (static field access, not dynamic per-row lookup)
+        //
+        // Literal base is not considered leaf-pushable because it would be constant-folded anyway.
+        if args.is_empty() {
+            return ExpressionPlacement::PlaceAtRoot;
+        }
+
+        // Check if the base (struct/map) argument is Column or PlaceAtLeafs
+        if !matches!(
+            args[0],
+            ExpressionPlacement::Column | ExpressionPlacement::PlaceAtLeafs
+        ) {
+            return ExpressionPlacement::PlaceAtRoot;
+        }
+
+        // All key arguments (after the first) must be literals for static field access
+        let keys_literal = args
+            .iter()
+            .skip(1)
+            .all(|a| *a == ExpressionPlacement::Literal);
+
+        if keys_literal {
+            ExpressionPlacement::PlaceAtLeafs
+        } else {
+            ExpressionPlacement::PlaceAtRoot
+        }
+    }
 }
 
 #[cfg(test)]
@@ -542,4 +573,80 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_placement_with_args_literal_key() {
+        let func = GetFieldFunc::new();
+
+        // get_field(col, 'literal') -> leaf-pushable (static field access)
+        let args = vec![ExpressionPlacement::Column, ExpressionPlacement::Literal];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtLeafs);
+
+        // get_field(col, 'a', 'b') -> leaf-pushable (nested static field access)
+        let args = vec![
+            ExpressionPlacement::Column,
+            ExpressionPlacement::Literal,
+            ExpressionPlacement::Literal,
+        ];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtLeafs);
+
+        // get_field(get_field(col, 'a'), 'b') represented as PlaceAtLeafs for base
+        let args = vec![
+            ExpressionPlacement::PlaceAtLeafs,
+            ExpressionPlacement::Literal,
+        ];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtLeafs);
+    }
+
+    #[test]
+    fn test_placement_with_args_column_key() {
+        let func = GetFieldFunc::new();
+
+        // get_field(col, other_col) -> NOT leaf-pushable (dynamic per-row lookup)
+        let args = vec![ExpressionPlacement::Column, ExpressionPlacement::Column];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtRoot);
+
+        // get_field(col, 'a', other_col) -> NOT leaf-pushable (dynamic nested lookup)
+        let args = vec![
+            ExpressionPlacement::Column,
+            ExpressionPlacement::Literal,
+            ExpressionPlacement::Column,
+        ];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtRoot);
+    }
+
+    #[test]
+    fn test_placement_with_args_root() {
+        let func = GetFieldFunc::new();
+
+        // get_field(root_expr, 'literal') -> NOT leaf-pushable
+        let args = vec![
+            ExpressionPlacement::PlaceAtRoot,
+            ExpressionPlacement::Literal,
+        ];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtRoot);
+
+        // get_field(col, root_expr) -> NOT leaf-pushable
+        let args = vec![
+            ExpressionPlacement::Column,
+            ExpressionPlacement::PlaceAtRoot,
+        ];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtRoot);
+    }
+
+    #[test]
+    fn test_placement_with_args_edge_cases() {
+        let func = GetFieldFunc::new();
+
+        // Empty args -> NOT leaf-pushable
+        assert_eq!(func.placement(&[]), ExpressionPlacement::PlaceAtRoot);
+
+        // Just base, no key -> PlaceAtLeafs (not a valid call but should handle gracefully)
+        let args = vec![ExpressionPlacement::Column];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtLeafs);
+
+        // Literal base with literal key -> NOT leaf-pushable (would be constant-folded)
+        let args = vec![ExpressionPlacement::Literal, ExpressionPlacement::Literal];
+        assert_eq!(func.placement(&args), ExpressionPlacement::PlaceAtRoot);
+    }
 }

datafusion/functions/src/core/getfield.rs

@@ -30,8 +30,8 @@ use datafusion_common::{
 };
 use datafusion_expr::expr::Alias;
 use datafusion_expr::{
-    Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScan, Unnest, Window,
-    logical_plan::LogicalPlan,
+    Aggregate, Distinct, EmptyRelation, Expr, ExpressionPlacement, Projection, TableScan,
+    Unnest, Window, logical_plan::LogicalPlan,
 };
 
 use crate::optimize_projections::required_indices::RequiredIndices;
@@ -530,9 +530,11 @@ fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Project
     // For details, see: https://github.com/apache/datafusion/issues/8296
     if column_referral_map.into_iter().any(|(col, usage)| {
         usage > 1
-            && !is_expr_trivial(
-                &prev_projection.expr
-                    [prev_projection.schema.index_of_column(col).unwrap()],
+            && matches!(
+                prev_projection.expr
+                    [prev_projection.schema.index_of_column(col).unwrap()]
+                .placement(),
+                ExpressionPlacement::PlaceAtRoot
             )
     }) {
         // no change
@@ -586,11 +588,6 @@ 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(_, _))
-}
-
 /// Rewrites a projection expression using the projection before it (i.e. its input)
 /// This is a subroutine to the `merge_consecutive_projections` function.
 ///