You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/16 01:19:11 UTC
(datafusion) branch gh-readonly-queue/main/pr-23691-767d1422d9a6459e6251c8b6cc0515abf07d08a2 created (now 2c1f74f10b)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23691-767d1422d9a6459e6251c8b6cc0515abf07d08a2
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 2c1f74f10b fix: evaluate struct-returning UDFs once across repeated field accesses (#23691)
This branch includes the following new commits:
new 2c1f74f10b fix: evaluate struct-returning UDFs once across repeated field accesses (#23691)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: evaluate struct-returning UDFs once across repeated field accesses (#23691)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23691-767d1422d9a6459e6251c8b6cc0515abf07d08a2
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 2c1f74f10b2d6bd3506fad3fc37e5361621f6a82
Author: Varun <to...@gmail.com>
AuthorDate: Wed Sep 16 01:18:47 2026 +0000
fix: evaluate struct-returning UDFs once across repeated field accesses (#23691)
## Which issue does this PR close?
- Closes #23655.
## Rationale for this change
`SELECT f(c)['a'], f(c)['b']`, where `f` is a struct-returning UDF,
evaluates `f(c)` twice in the physical `ProjectionExec` — once per
`get_field` — although the two calls are identical.
`CommonSubexprEliminate` already hoists the shared `f(c)` into a
`__common_expr_N` column so it runs once.
[`PushDownLeafProjections`](https://github.com/apache/datafusion/blob/main/datafusion/optimizer/src/extract_leaf_expressions.rs)
then undoes that, one pass at a time:
- After CSE, each field access reads its base from the hoist column, so
`get_field`'s placement flips to `MoveTowardsLeafNodes` and the pass
extracts it.
- Merging the extraction into the hoist projection inlines that column's
definition back into every access (via `build_projection_replace_map`),
re-duplicating the `KeepInPlace` UDF.
- `optimize_projections` drops the now-unused hoist column, and the two
passes oscillate pass after pass — settling on a double evaluation.
This is the instability the `is_ignored` note in
`common_subexpr_eliminate` already warns about.
## What changes are included in this PR?
Adds a guard `merge_would_duplicate_kept_expr` and calls it from
`split_and_push_projection`, which now bails out — leaving the
projection untouched — when merging the extractions into the input
projection would inline a `KeepInPlace` definition that more than one
extraction pair references. This mirrors the existing
`optimize_projections` merge guard (#8296): a compute-once expression
referenced more than once is kept in place instead of duplicated.
With the guard, CSE's hoist survives, `optimize_projections` keeps it,
and the pipeline converges to a single evaluation:
Before:
```text
Projection: get_field(f(c), 'a'), get_field(f(c), 'b') -- f(c) twice
TableScan
```
After:
```text
Projection: get_field(__common_expr_1, 'a'), get_field(__common_expr_1, 'b')
Projection: f(c) AS __common_expr_1 -- f(c) once
TableScan
```
## Are these changes tested?
Yes.
- `test_struct_returning_udf_evaluated_once` (in
`extract_leaf_expressions.rs`) runs the four interacting rules — CSE →
extraction → pushdown → projection pruning — and asserts no field access
wraps a fresh UDF call, plus an `insta` snapshot of the single-eval
plan. It runs with the default pass budget, so any residual oscillation
would surface as a double evaluation.
- `get_field_like` — a small test UDF added to `test/udfs.rs` —
reproduces `GetFieldFunc`'s conditional placement
(`MoveTowardsLeafNodes` only when the base is pushable and the keys are
literals). This lets the optimizer crate model struct-field access
without depending on `datafusion-functions`, which it intentionally does
not.
`cargo fmt --all` and `cargo clippy --all-targets --all-features -- -D
warnings` are clean.
## Are there any user-facing changes?
Yes, but only in the sense that a redundantly-evaluated path becomes
single-evaluation. No API or query-result changes: affected plans
evaluate a struct-returning UDF once instead of once per field access, a
saving that scales with the number of distinct field accesses on the
same UDF.
---------
Co-authored-by: tohuya6 <20...@users.noreply.github.com>
---
.../optimizer/src/extract_leaf_expressions.rs | 170 ++++++++++++++++++++-
datafusion/optimizer/src/test/udfs.rs | 65 +++++++-
2 files changed, 233 insertions(+), 2 deletions(-)
diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs
index dfaf93bc1e..8a4abfcb48 100644
--- a/datafusion/optimizer/src/extract_leaf_expressions.rs
+++ b/datafusion/optimizer/src/extract_leaf_expressions.rs
@@ -837,6 +837,66 @@ fn passthrough_column(expr: &Expr) -> Option<&Column> {
}
}
+/// Returns `true` if merging the extracted `pairs` (and the pass-through
+/// `columns_needed`) into `child` would inline a `KeepInPlace` definition into
+/// more than one reference site — e.g. a struct-returning UDF `f(c)` behind both
+/// `f(c)['a']` and `f(c)['b']`, or behind `f(c)['a']` and a bare `f(c)`.
+///
+/// `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)],
+ columns_needed: &IndexSet<Column>,
+ child: &Projection,
+) -> bool {
+ // Columns whose `KeepInPlace` definition must stay put; inlining one would
+ // duplicate a compute-once expression. Plain columns, pushable exprs, and
+ // literals are all cheap to duplicate, so only `KeepInPlace` counts.
+ let kept_columns: std::collections::HashSet<String> = child
+ .schema
+ .iter()
+ .zip(child.expr.iter())
+ .filter_map(|((qualifier, field), expr)| {
+ if matches!(expr.placement(), ExpressionPlacement::KeepInPlace) {
+ Some(Column::from((qualifier, field)).flat_name())
+ } else {
+ None
+ }
+ })
+ .collect();
+ if kept_columns.is_empty() {
+ return false;
+ }
+
+ // Count every site that references a kept column. 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;
+ }
+ }
+ }
+ // Pass-through / pre-existing-alias inputs are additional reference sites:
+ // a bare `f(c)` alongside `f(c)['a']` also pins the compute-once column, so
+ // counting only `pairs` under-counted and re-duplicated it (#23655).
+ for col in columns_needed {
+ let name = col.flat_name();
+ if kept_columns.contains(&name) {
+ *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.
///
@@ -990,6 +1050,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, columns_needed, child)
+ {
+ return Ok(None);
+ }
+
// ── Phase 2: Push down ──────────────────────────────────────────────
let proj_input = Arc::clone(&proj.input);
let pushed = push_extraction_pairs(
@@ -1346,8 +1414,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;
@@ -3431,6 +3500,105 @@ mod tests {
alias_generator.next(EXTRACTED_EXPR_PREFIX),
"__datafusion_extracted_8"
);
+
+ 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(())
+ }
+
+ #[test]
+ fn test_struct_returning_udf_evaluated_once_with_bare_reference() -> Result<()> {
+ // Regression for the `SELECT f(c)['a'], f(c)` shape: the struct UDF is
+ // referenced once through a field access and once as a bare value. The
+ // bare reference is tracked in `columns_needed`, not `extraction_pairs`,
+ // so a guard that counted only pairs would re-inline `f(c)` and evaluate
+ // it twice (#23655).
+ 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"), f_c])?
+ .build()?;
+
+ 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, |_, _| {})?;
+
+ // The field access must read the UDF back through `__common_expr_1`, not
+ // wrap 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")), __common_expr_1 AS keep_in_place_udf(test.c)
+ 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 72b1a63f67..e139039957 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, Volatility,
+ ScalarUDFImpl, Signature, TypeSignature, Volatility, lit,
};
/// A configurable test UDF for optimizer tests.
@@ -103,3 +103,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)])
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org