fix: one inlining policy for the rules that inline a projection column

[adriangb]

Which issue does this PR close?

• Closes CommonSubexprEliminate + leaf-expression pushdown duplicate a UDF call wrapped in get_field #25329

Rationale for this change

A function that keeps its place, for example a struct-returning UDF, is evaluated two times when the query hoists it into a common sub-expression and then reads one field of it inside a filter.

MRE:

CREATE TABLE t(a INT) AS VALUES (1), (2);
EXPLAIN SELECT a FROM t WHERE CASE WHEN arrow_field(a) IS NOT NULL THEN arrow_field(a)['nullable'] IS NULL END;

Before, on main. arrow_field(a) is in the plan two times, so it runs two times for each row.

logical_plan
01)Projection: t.a
02)--Filter: CASE WHEN __common_expr_7 IS NOT NULL THEN __datafusion_extracted_8 IS NULL END
03)----Projection: arrow_field(t.a) AS __common_expr_7, t.a, get_field(arrow_field(t.a), Utf8("nullable")) AS __datafusion_extracted_8
04)------TableScan: t projection=[a]
physical_plan
01)FilterExec: CASE WHEN __common_expr_7@0 IS NOT NULL THEN __datafusion_extracted_8@2 IS NULL END, projection=[a@1]
02)--ProjectionExec: expr=[arrow_field(a@0) as __common_expr_7, a@0 as a, get_field(arrow_field(a@0), nullable) as __datafusion_extracted_8]
03)----DataSourceExec: partitions=1, partition_sizes=[1]

After, with this PR. arrow_field(a) is in the plan one time:

logical_plan
01)Projection: t.a
02)--Filter: CASE WHEN __common_expr_1 IS NOT NULL THEN get_field(__common_expr_1, Utf8("nullable")) IS NULL END
03)----Projection: arrow_field(t.a) AS __common_expr_1, t.a
04)------TableScan: t projection=[a]
physical_plan
01)FilterExec: CASE WHEN __common_expr_1@0 IS NOT NULL THEN get_field(__common_expr_1@0, nullable) IS NULL END, projection=[a@1]
02)--ProjectionExec: expr=[arrow_field(a@0) as __common_expr_1, a@0 as a]
03)----DataSourceExec: partitions=1, partition_sizes=[1]

The same statement in a SELECT list runs the function one time since #23691. That fix counts the reference sites inside one projection. In the filter shape the two sites are in two different nodes, the filter and the extracted expression, so the count is 1 and the guard does not fire.

CommonSubexprEliminate puts an expression into its own column so that the plan evaluates it one time. Four rules then replace a column reference with the expression that defines it. Each rule got its own guard against the duplication that this causes, one guard for each bug report:

Guard	Rule	Added by
would_duplicate_volatile (3 call sites)	ExtractLeafExpressions, PushDownLeafProjections	#24720
merge_would_duplicate_kept_expr (written by tohuya6, reviewed by kosiew; subsumed, not wrong, and its two regression tests stay)	PushDownLeafProjections	#23691
the volatile and placement partition in rewrite_projection	PushDownFilter	-
the referral count in merge_consecutive_projections_one_level	OptimizeProjections	#8296

Each guard sees only its own rule and its own node. This PR puts the decision in one place, so a rule can not miss a reference site again.

What changes are included in this PR?

A crate-private ProjectionInliner (datafusion/optimizer/src/projection_inliner.rs) with one policy:

1. Classify each definition of the projection by the cost of one more evaluation: free (a column, a literal or a MoveTowardsLeafNodes expression), cheap (arithmetic, CAST, CASE and similar kernels), expensive (contains a scalar function that reports KeepInPlace), or forbidden (volatile, or contains a subquery).
2. Count the evaluations that the plan has after the rewrite. This counts the references in the expressions that the rule inlines, and the columns that the other consumers keep.
3. Report the definitions that must stay where they are: a forbidden definition, and a definition above the cost that the caller accepts that the plan would evaluate more than one time.

The four rules call it instead of building their own replacement map:

• OptimizeProjections accepts free definitions only, which is the behaviour of Conflicting optimization rules: common_sub_expression_eliminate and push_down_projection #8296. The reference count now resolves unqualified names too, which the old HashMap<&Column, usize> did not, so the new code can pin a definition the old code merged. No plan in the test suites changes; a case where it would is a merge the old code should not have done.
• PushDownFilter accepts cheap definitions. A predicate stays above the projection when it references a volatile definition, a MoveTowardsLeafNodes definition (as before), or an expensive definition that the predicates reference more than one time. A predicate with one reference is still pushed, so a cheap cast or an arithmetic expression still reaches the scan.
• ExtractLeafExpressions and PushDownLeafProjections accept free definitions only, and count the extracted expressions together with the columns that the nodes above reference directly. This is the fix for the issue.

Deleted: would_duplicate_volatile, volatile_output_columns, merge_would_duplicate_kept_expr, build_projection_replace_map and the contain helper of push_down_filter.

The commits are one for the new module, one for each rule, and one for the tests.

This PR is complementary to the draft #25388. That draft stops a push-down whenever the projection expression is KeepInPlace, which is a wider rule: its own snapshots show that a cast and an arithmetic expression no longer reach partial_filters and full_filters. This PR only stops a rewrite that evaluates a definition more times than before, which is a bug in every case. cc @pepijnve

What is the testing strategy for this PR?

• datafusion/sqllogictest/test_files/cse.slt: both end-to-end shapes, the projection shape of Optimize redundant calls to UDFs producing StructArray #23655 and the filter shape of this issue. Both plans must show arrow_field(a) one time.
• projection_inliner.rs: unit tests for the cost classification, the evaluation count, the substitution and the volatile case.
• push_down_filter.rs: filter_with_repeated_expensive_reference_not_pushed_through_projection, filter_with_single_expensive_reference_pushed_through_projection and filter_not_pushed_through_nested_computed_projection (the shape of the draft above, with the plan that the narrow policy gives).
• extract_leaf_expressions.rs: test_struct_returning_udf_in_filter_evaluated_once.

No existing snapshot and no existing .slt expectation changed.

Are there any user-facing changes?

Plans that evaluated a KeepInPlace expression two times now evaluate it one time. There is no API change: ProjectionInliner is crate-private.

Known remaining duplication sites

With the evaluation-site invariant from #25458 switched on, the arrow_field duplication of #25329 no longer appears. Two leaf-rule shapes still evaluate a KeepInPlace definition at two sites and are not changed by this PR: named_struct(...) referenced from an ORDER BY field access (order.slt, through extract_leaf_expressions) and named_struct(...) referenced from a field access above the projection (struct.slt, through push_down_leaf_projections). They are tracked in the EPIC and are the next step for the policy.

Behaviour change to note: a projection definition that contains a subquery is now classed as forbidden to duplicate, so a predicate over such a column is no longer pushed below the projection. On main, WHERE m > 0 over SELECT (SELECT max(a) FROM t) AS m pushed the subquery below the projection while the projection still computed it.

Volatility is now asked about before placement. DuplicationCost::of returned Free as soon as a definition reported a placement of MoveTowardsLeafNodes, before it walked the expression for volatile nodes, so a volatile function with that placement was classed free to inline at two sites. file_row_index() is such a function, and file_row_index.slt shows it: the extraction column is now kept and both references read it, so the call is evaluated one time. A new EXPLAIN in that file pins the plan, and a unit test next to the other DuplicationCost tests asserts Forbidden for a volatile MoveTowardsLeafNodes function.

pinned_for_known_consumers is renamed to pinned_with_unknown_consumers. The old name said the opposite of its own doc comment: the method is for a projection that can stay in the plan for consumers that the caller can not see.

Part of the leaf-pushdown EPIC: #25459

🤖 Generated with Claude Code

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.67857% with 41 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.36%. Comparing base (3a647e4) to head (67f541c).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
...tafusion/optimizer/src/extract_leaf_expressions.rs	88.26%	11 Missing and 12 partials ⚠️
datafusion/optimizer/src/push_down_filter.rs	78.08%	2 Missing and 14 partials ⚠️
datafusion/optimizer/src/projection_inliner.rs	99.30%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25456      +/-   ##
==========================================
+ Coverage   82.33%   82.36%   +0.02%     
==========================================
  Files        1137     1138       +1     
  Lines      432498   433234     +736     
  Branches   432498   433234     +736     
==========================================
+ Hits       356115   356822     +707     
- Misses      54844    54863      +19     
- Partials    21539    21549      +10     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[pepijnve]

I haven't gone through the code changes in detail yet, but the description sounds like a good way forward. Implementing the assessment once rather than trying to patch things ad hoc per optimiser is the way to go.

An open question I still had related to keep-in-place was if this should be refined a bit. On #25388 I tweaked Expr::placement so that for simple expressions (unary and binary ops mostly) it is based on the operands placement rather than being keep-in-place. Is that something we would want to take along here? Or should I make a separate PR out of that?

[pepijnve]

This enumeration and the way DuplicationCost::of calculates it seems to overlap largely with the existing expression placement enumeration and Expr::placement. Are those two trying to express the same thing or is there a fundamental semantic difference?