adriangb opened a new pull request, #25458: URL: https://github.com/apache/datafusion/pull/25458
## Which issue does this PR close? - Part of https://github.com/apache/datafusion/issues/25329 and https://github.com/apache/datafusion/issues/23655. It does not fix them. It adds the check that finds them, and finds eight more shapes. ## Rationale for this change A query must not evaluate a volatile function more times than the query text says. `random()` gives a different answer each time, so a second evaluation returns wrong rows. An expensive function must not run two times either, because the second run costs and gives the same answer. Five guards in the logical optimizer exist only to keep that property. Each one was added after a bug report, and each one covers one call site: | Guard | Added by | |---|---| | `would_duplicate_volatile` in `extract_leaf_expressions`, 3 call sites | https://github.com/apache/datafusion/issues/24678 | | the volatile check in the `Projection` branch of `push_down_filter` | https://github.com/apache/datafusion/issues/20239 | | the volatile check in the `Aggregate` branch of `push_down_filter` | https://github.com/apache/datafusion/issues/25415 | | `merge_would_duplicate_kept_expr` in `extract_leaf_expressions` | https://github.com/apache/datafusion/issues/23655 | | the "referenced more than one time" guard in `optimize_projections` | https://github.com/apache/datafusion/issues/8296 | The next shape is always open. https://github.com/apache/datafusion/issues/25329 is the shape none of the five covers. A plan invariant covers the class instead of one call site. `assert_valid_optimization` already runs after every rule and checks the schema, so it is the place for it. ## Measured results The check ran over the full sqllogictest suite, 520 files, with both switches on. It reports **7 files, 14 failing queries and 10 distinct shapes. There are no false positives.** Volatile, 2 shapes: | Rule | Query | Sites | |---|---|---| | `simplify_expressions` | `SELECT random() BETWEEN 0.0 AND 1.0, random() = random()`, `expr.slt:1005` | `random()` 3 -> 4 | | `optimize_projections` | `SELECT file_row_index(), file_row_index() + 1, column1 FROM parquet_table ORDER BY column1`, `file_row_index.slt:46` | `file_row_index()` 1 -> 2 | The first one is a wrong results bug, filed as https://github.com/apache/datafusion/issues/25457. `SELECT count(*) FROM v WHERE random() BETWEEN 0.4 AND 0.6` over 100000 rows returns 35751 rows, against the 20000 that one draw gives. The second one is a logical plan duplication only. The physical planner resolves both calls to the same `__datafusion_file_row_index` column, so the answer stays right. `KeepInPlace`, 8 shapes: | Rule | Query | Sites | |---|---|---| | `push_down_filter` | `WITH test AS (SELECT substr(md5(i::text)::text, 1, 32) AS needle FROM generate_series(1, 100000) t(i)) SELECT count(*) FROM test WHERE needle IN (...)`, `array/array_has.slt:496` | `md5(...)` and `substr(...)` 1 -> 2 | | `push_down_filter` | `SELECT column1, column2 FROM any_op_test WHERE 5 > ANY(column2)`, `array/array_has.slt:687` | `array_min(column2)` 1 -> 3 | | `push_down_filter` | `SELECT t2.server['c3'] ... FROM (SELECT struct(time, load1, load2, host) AS server FROM t1) t2 WHERE t2.server['c3'] IS NOT NULL`, `expr.slt:2444` | `struct(...)` 1 -> 2 | | `extract_leaf_expressions` | `SELECT named_struct('a', a, 'b', b) AS s FROM ordered ORDER BY s['a']`, `order.slt:2001` | `named_struct(...)` 1 -> 2 | | `push_down_leaf_projections` | `SELECT s['b'] FROM struct_ctor_view WHERE s IS NOT NULL`, `struct.slt:1773` | `named_struct(...)` 1 -> 2 | | `replace_distinct_aggregate` | `SELECT DISTINCT ON (upper(c1)) c1, sum(c3) FROM aggregate_test_100 GROUP BY c1 ORDER BY upper(c1)`, `distinct_on.slt:223` | `upper(c1)` 2 -> 3 | | `simplify_expressions` | `SELECT abs(c1) BETWEEN 0 AND log(c1 * 100) FROM select_between_data`, `select.slt:943` | `abs(c1)` 1 -> 2 | | `simplify_expressions` | `SELECT column1, COALESCE(CAST(array_has(column1, column2) AS VARCHAR), 'null') FROM array_has_table_empty`, `array/array_has.slt:166` | `array_has(...)` 1 -> 2 | The two `simplify_expressions` shapes are the `BETWEEN` and `COALESCE` expansions, which clone the operand. The `push_down_filter` and leaf rule shapes are the trade-off described in https://github.com/apache/datafusion/issues/25329: the rule pays one more evaluation for a chance to prune at the scan. The `replace_distinct_aggregate` shape repeats the `DISTINCT ON` key as a sort key. The check adds about 12% to the CPU time of a sqllogictest run, measured with one binary and an interleaved A/B over 5 pairs. Child CPU time, median: 74.2 s off, 85.6 s on, minimum 73.1 s off, 81.7 s on. Wall clock on 12 threads is inside the noise of this machine. The switches are off in this PR, so the cost on `main` is zero. `EvaluationSites::of` returns an empty map and walks nothing. Note for whoever turns the switches on: `[profile.ci.package."*"]` in the workspace `Cargo.toml` sets `debug-assertions = false`. So `#[cfg(debug_assertions)]` code inside `datafusion-optimizer` does **not** run when the crate is built as a dependency with `--profile ci`. The existing `check_invariants(InvariantLevel::Executable)` call in the optimizer loop is in that state today: the sqllogictest suite does not exercise it. That is why this check uses a plain `const` and not `cfg(debug_assertions)`. ## What changes are included in this PR? A new crate private module, `datafusion/optimizer/src/evaluation_sites.rs`, and one call from the optimizer loop next to the schema invariant. A *site* is one `Expr::ScalarFunction` node that is volatile, or that has placement `ExpressionPlacement::KeepInPlace`. Three rules keep the count honest: 1. **Row lineage, not the whole plan.** `sites(node) = own_sites(node) + max over inputs of sites(input)`. The maximum, not the sum, because a row that enters the left input of a join never meets the nodes of the right input. This is what makes union branches, join sides and a filter pushed below a join free of reports. A subquery is another lineage. 2. **The class of a call ignores its arguments.** The placement of a call is asked for with every argument reduced to a literal or to a column. Without this, inlining `f(x)` into `get_field(c, 'k')` turns a `MoveTowardsLeafNodes` `get_field` into a `KeepInPlace` one and reports a duplication that did not happen. A nested call is a site of its own. 3. **Only a call that was already there is checked**, keyed by the text of the call. A rule is free to rewrite `concat_ws` into `concat`, or one Spark function into two DataFusion functions. That is a new call, not a second evaluation of an old one. A call with no argument, or on literals only, is not a site at all, because it is one value for the whole query. Each rule removed a class of false positive that the earlier drafts reported. Without rule 1 every union and join reports. Without rule 2, `struct.slt` reports five `get_field` calls over a literal struct. Without rule 3, five Spark function lowerings report. The error message names the call, the counts and the rule: ```text Check optimizer-specific invariants after optimizer rule: extract_leaf_expressions caused by Internal error: Optimizer rule added evaluation sites: KeepInPlace expression `arrow_field(t.a)` is evaluated at 2 sites, against 1 before the rule. A rule must not make a query evaluate a volatile or an expensive expression more times than the plan it was given does. Counts are per row lineage, see `EvaluationSites`. ``` ## What is the testing strategy for this PR? Eight unit tests in the new module cover the counting rules directly, so they run whatever the switches are set to: one node, a chain, union branches, join sides, a subquery, an inlined definition, a constant call, and the class of a call over a costly argument. The measured results above come from `cargo test --profile ci -p datafusion-sqllogictest --test sqllogictests` with both switches flipped to `true`. With the switches off, as this PR ships them, `cargo test --profile ci -p datafusion-optimizer` passes 890 + 26 + 5 tests and the full sqllogictest suite is green over all 520 files. ## Are there any user-facing changes? No. Both switches are off, so there is no behaviour change and no cost. There is no new public API. 🤖 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]
