adriangb opened a new pull request, #25456: URL: https://github.com/apache/datafusion/pull/25456
## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/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: ```sql 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. ```text 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: ```text 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 https://github.com/apache/datafusion/pull/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` | https://github.com/apache/datafusion/pull/24720 | | `merge_would_duplicate_kept_expr` | `PushDownLeafProjections` | https://github.com/apache/datafusion/pull/23691 | | the volatile and placement partition in `rewrite_projection` | `PushDownFilter` | - | | the referral count in `merge_consecutive_projections_one_level` | `OptimizeProjections` | https://github.com/apache/datafusion/issues/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 https://github.com/apache/datafusion/issues/8296. The plans do not change. - `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 https://github.com/apache/datafusion/pull/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 https://github.com/apache/datafusion/issues/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. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
