fix: evaluate struct-returning UDFs once across repeated field accesses

datafusion/optimizer/src/extract_leaf_expressions.rs

@@ -757,6 +757,51 @@ fn try_push_input(
     split_and_push_projection(proj, alias_generator)
 }
 
+/// Returns `true` if merging the extracted `pairs` into `child` would inline a
+/// `KeepInPlace` definition into more than one extracted expression — e.g. a
+/// struct-returning UDF `f(c)` behind both `f(c)['a']` and `f(c)['b']`.
+///
+/// `CommonSubexprEliminate` hoists such an expression into a shared column so it
+/// runs once; [`build_extraction_projection_impl`] then resolves that column
+/// back to its definition and re-duplicates it. Mirrors the
+/// `optimize_projections` merge guard (#8296): a compute-once expression used
+/// more than once stays in place.
+fn merge_would_duplicate_kept_expr(pairs: &[(Expr, String)], child: &Projection) -> bool {
+    // Columns whose `KeepInPlace` definition must stay put; inlining one would
+    // duplicate a compute-once expression. Plain and pushable definitions are
+    // cheap to duplicate.
+    let kept_columns: std::collections::HashSet<String> = child
+        .schema
+        .iter()
+        .zip(child.expr.iter())
+        .filter_map(|((qualifier, field), expr)| {
+            if expr.placement().should_push_to_leaves() {
+                None
+            } else {
+                Some(Column::from((qualifier, field)).flat_name())
+            }
+        })
+        .collect();
+    if kept_columns.is_empty() {
+        return false;
+    }
+
+    // A pair that references the same kept column twice still inlines one copy,
+    // so count each pair at most once per column.
+    let mut usage: HashMap<String, usize> = HashMap::new();
+    for (expr, _alias) in pairs {
+        let mut seen: std::collections::HashSet<String> =
+            std::collections::HashSet::new();
+        for col in expr.column_refs() {
+            let name = col.flat_name();
+            if kept_columns.contains(&name) && seen.insert(name.clone()) {
+                *usage.entry(name).or_insert(0) += 1;
+            }
+        }
+    }
+    usage.values().any(|&count| count > 1)
+}
+
 /// Splits a projection into extractable pieces, pushes them towards leaf
 /// nodes, and adds a recovery projection if needed.
 ///
@@ -899,6 +944,14 @@ fn split_and_push_projection(
         return Ok(None);
     }
 
+    // Merging here would re-inline an expression `CommonSubexprEliminate` hoisted
+    // to run once, undoing it (issue #23655); leave the projection in place.
+    if let LogicalPlan::Projection(child) = input.as_ref()
+        && merge_would_duplicate_kept_expr(&extraction_pairs, child)
+    {
+        return Ok(None);
+    }
+
     // ── Phase 2: Push down ──────────────────────────────────────────────
     let proj_input = Arc::clone(&proj.input);
     let pushed = push_extraction_pairs(
@@ -1246,8 +1299,9 @@ fn try_push_into_inputs(
 mod tests {
 
     use super::*;
+    use crate::common_subexpr_eliminate::CommonSubexprEliminate;
     use crate::optimize_projections::OptimizeProjections;
-    use crate::test::udfs::PlacementTestUDF;
+    use crate::test::udfs::{PlacementTestUDF, get_field_like};
     use crate::test::*;
     use crate::{Optimizer, OptimizerContext};
     use datafusion_expr::expr::ScalarFunction;
@@ -3163,4 +3217,57 @@ mod tests {
 
         Ok(())
     }
+
+    /// Regression test for issue #23655: `f(c)['a'], f(c)['b']` on a
+    /// struct-returning UDF must evaluate `f(c)` once, not once per field access.
+    ///
+    /// `keep_in_place_udf` / `get_field_like` model `f(c)` and `f(c)[key]` — the
+    /// crate can't depend on the real `get_field` (see [`get_field_like`]).
+    /// Without the guard, `PushDownLeafProjections` re-inlines the shared UDF that
+    /// `CommonSubexprEliminate` hoisted and the passes oscillate to a double
+    /// evaluation.
+    #[test]
+    fn test_struct_returning_udf_evaluated_once() -> Result<()> {
+        let udf = ScalarUDF::new_from_impl(
+            PlacementTestUDF::new().with_placement(ExpressionPlacement::KeepInPlace),
+        );
+        let f_c = udf.call(vec![col("c")]);
+        let plan = LogicalPlanBuilder::from(test_table_scan()?)
+            .project(vec![
+                get_field_like(f_c.clone(), "a"),
+                get_field_like(f_c, "b"),
+            ])?
+            .build()?;
+
+        // The full four-rule interaction, default pass budget so any oscillation
+        // would surface here as a double evaluation.
+        let ctx = OptimizerContext::new();
+        let optimizer = Optimizer::with_rules(vec![
+            Arc::new(CommonSubexprEliminate::new()),
+            Arc::new(ExtractLeafExpressions::new()),
+            Arc::new(PushDownLeafProjections::new()),
+            Arc::new(OptimizeProjections::new()),
+        ]);
+        let optimized = optimizer.optimize(plan, &ctx, |_, _| {})?;
+
+        // Every field access reads the UDF back through `__common_expr_1`; none
+        // wraps a fresh `keep_in_place_udf(...)` call. Asserted explicitly so a
+        // careless snapshot re-accept can't quietly restore the double evaluation.
+        let formatted = format!("{optimized}");
+        assert_eq!(
+            formatted
+                .matches("get_field_like(keep_in_place_udf")
+                .count(),
+            0,
+            "struct-returning UDF must not be re-evaluated inside a field access; got:\n{formatted}"
+        );
+
+        insta::assert_snapshot!(formatted, @r#"
+        Projection: get_field_like(__common_expr_1 AS keep_in_place_udf(test.c), Utf8("a")), get_field_like(__common_expr_1 AS keep_in_place_udf(test.c), Utf8("b"))
+          Projection: keep_in_place_udf(test.c) AS __common_expr_1
+            TableScan: test projection=[c]
+        "#);
+
+        Ok(())
+    }
 }

datafusion/optimizer/src/test/udfs.rs

@@ -19,7 +19,7 @@ use arrow::datatypes::DataType;
 use datafusion_common::Result;
 use datafusion_expr::{
     ColumnarValue, Expr, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF,
-    ScalarUDFImpl, Signature, TypeSignature,
+    ScalarUDFImpl, Signature, TypeSignature, Volatility, lit,
 };
 
 /// A configurable test UDF for optimizer tests.
@@ -44,7 +44,7 @@ impl PlacementTestUDF {
             // The actual types don't matter since this UDF is not intended for execution.
             signature: Signature::new(
                 TypeSignature::OneOf(vec![TypeSignature::Any(1), TypeSignature::Any(2)]),
-                datafusion_expr::Volatility::Immutable,
+                Volatility::Immutable,
             ),
             placement: ExpressionPlacement::MoveTowardsLeafNodes,
             id: 0,
@@ -96,3 +96,66 @@ pub fn leaf_udf_expr(arg: Expr) -> Expr {
     );
     udf.call(vec![arg])
 }
+
+/// A `get_field`-like test UDF for modeling struct-field access (`f(c)['a']`),
+/// so tests need not depend on `datafusion-functions` (which this crate avoids —
+/// see its `Cargo.toml`).
+///
+/// Placement mirrors `GetFieldFunc`: `MoveTowardsLeafNodes` when the base is
+/// pushable and the keys are literals, else `KeepInPlace`. That flip drives the
+/// issue #23655 pushdown interaction: the access is `KeepInPlace` until
+/// `CommonSubexprEliminate` hoists the UDF into a column, then becomes pushable.
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct GetFieldLikeUDF {
+    signature: Signature,
+}
+
+impl Default for GetFieldLikeUDF {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl GetFieldLikeUDF {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
+        }
+    }
+}
+
+impl ScalarUDFImpl for GetFieldLikeUDF {
+    fn name(&self) -> &str {
+        "get_field_like"
+    }
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::UInt32)
+    }
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        panic!("GetFieldLikeUDF: not intended for execution")
+    }
+    fn placement(&self, args: &[ExpressionPlacement]) -> ExpressionPlacement {
+        // Mirror `GetFieldFunc`: cheap to push down only when the base is a
+        // column (or already pushable) and every field key is a literal.
+        match args.first() {
+            Some(
+                ExpressionPlacement::Column | ExpressionPlacement::MoveTowardsLeafNodes,
+            ) if args[1..]
+                .iter()
+                .all(|p| matches!(p, ExpressionPlacement::Literal)) =>
+            {
+                ExpressionPlacement::MoveTowardsLeafNodes
+            }
+            _ => ExpressionPlacement::KeepInPlace,
+        }
+    }
+}
+
+/// Create a `get_field_like(base, key)` expression that models `base[key]`.
+pub fn get_field_like(base: Expr, key: &str) -> Expr {
+    let udf = ScalarUDF::new_from_impl(GetFieldLikeUDF::new());
+    udf.call(vec![base, lit(key)])
+}