diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index b855f224c420b..d2f6160c4142f 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/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 = 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 = HashMap::new(); + for (expr, _alias) in pairs { + let mut seen: std::collections::HashSet = + 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(()) + } } diff --git a/datafusion/optimizer/src/test/udfs.rs b/datafusion/optimizer/src/test/udfs.rs index ba71b6a04a7a2..504c9806ba512 100644 --- a/datafusion/optimizer/src/test/udfs.rs +++ b/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 { + Ok(DataType::UInt32) + } + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + 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)]) +}