Repository navigation
fix: lazy evaluation for coalesce - #17357
Conversation
| # due to the reason describe in https://github.com/apache/datafusion/issues/8927, | ||
| # the following queries will fail |
There was a problem hiding this comment.
Since the test below now works, I think we could move this comment to the test below this one.
alamb
left a comment
There was a problem hiding this comment.
Thanks @chenkovsky and @nuno-faria -- I think this PR is quite good and probably can be merged. My only potential concern is that we may mess up comet. Let's see if we get any more comments
| 02)--TableScan: t projection=[x, y] | ||
| physical_plan | ||
| 01)ProjectionExec: expr=[coalesce(1, CAST(y@1 / x@0 AS Int64)) as coalesce(Int64(1),t.y / t.x), coalesce(2, CAST(y@1 / x@0 AS Int64)) as coalesce(Int64(2),t.y / t.x)] | ||
| 01)ProjectionExec: expr=[CASE WHEN true THEN 1 ELSE CAST(y@1 / x@0 AS Int64) END as coalesce(Int64(1),t.y / t.x), CASE WHEN true THEN 2 ELSE CAST(y@1 / x@0 AS Int64) END as coalesce(Int64(2),t.y / t.x)] |
There was a problem hiding this comment.
It seems like this should be able to be simplifed even more -- CASE WHEN true THEN ... should go to 1
I filed a ticket to track this idea:
| )); | ||
| } | ||
|
|
||
| let n = args.len(); |
There was a problem hiding this comment.
This is quite a clever implementation.
However, I worry it may cause problems for comet which uses physical evaluation directly
@comphead or @andygrove do you know if comet uses the COALESCE implementation directly?
Comet was actually just thinking about coalesce: Maybe @coderfender has some thoughts about this? |
It seems like this PR does the same thing described in the comet PR Thus it seems like a good thing to merge Let's wait a while to see if anyone else has comments, otherwise I'll plan to merge this PR |
|
@alamb , @mbutrovich I made changes to comet to fallback to CASE statement to replicate |
|
Thanks everyone! |
|
I believe this PR caused the extended tests to fail (as more queries now start woring): |
This reverts commit e5dcc8c.
…ion (#18567) Note this targets the branch-51 release branch ## Which issue does this PR close? - part of #17558 - resolves #17801 in the 51 release branch ## Rationale for this change - We merged some clever rewrites for `coalesce` and `nvl2` to use `CASE` which are faster and more correct (👏 @chenkovsky @kosiew ) - However, these rewrites cause subtle schema mismatches in some cases planning (b/c the CASE simplification nullability logic can't determine the correct nullability in some cases - see #17801) - @pepijnve has some heroic efforts to fix the schema mismatch in #17813 (comment), but it is non trivial and I am worried about merging it so close to the 51 release and introducing new edge cases ## What changes are included in this PR? 1. Revert #17357 / e5dcc8c 3. Revert #17991 / ea83c26 2. Revert #18191 / 22c4214 2. Cherry-pick 6202254, a test that reproduces the schema mismatch issue (from #18536) 3. Cherry-pick 735cacf, a fix for the benchmarks that regressed due to the revert (from #17833) 4. Update datafusion-testing (see separate PR here) for extended tests (see apache/datafusion-testing#15) ## Are these changes tested? Yes I added a new test ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. -->
## Which issue does this PR close? - Closes #17801 - Obviates (contains) and thus Closes #17833 - Obviates (contains) and thus Closes #18536 ## Rationale for this change #17357 introduced a change that replaces `coalesce` function calls with `case` expressions. In the current implementation these two differ in the way they report their nullability. `coalesce` is more precise than `case` all will report itself as not nullable in situations where the equivalent `case` does report being nullable. The rest of the codebase currently does not expect the nullability property of an expression to change as a side effect of expression simplification. This PR is a first attempt to align the nullability of `coalesce` and `case`. ## What changes are included in this PR? Tweaks to the `nullable` logic for the logical and physical `case` expression code to report `case` as being not nullable in more situations. - For logical `case`, a best effort const evaluation of 'when' expressions is done to determine 'then' reachability. The code errs on the conservative side wrt nullability. - For physical `case`, const evaluation of 'when' expressions using a placeholder record batch is attempted to determine 'then' reachability. Again if const evaluation is not possible, the code errs on the conservative side. - The optimizer schema check has been relaxed slightly to allow nullability to be removed by optimizer passes without having to disable the schema check entirely - The panic'ing benchmark has been reenabled ## Are these changes tested? Additional unit tests have been added to test the new logic. ## Are there any user-facing changes? No --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
…e#17813) ## Which issue does this PR close? - Closes apache#17801 - Obviates (contains) and thus Closes apache#17833 - Obviates (contains) and thus Closes apache#18536 ## Rationale for this change apache#17357 introduced a change that replaces `coalesce` function calls with `case` expressions. In the current implementation these two differ in the way they report their nullability. `coalesce` is more precise than `case` all will report itself as not nullable in situations where the equivalent `case` does report being nullable. The rest of the codebase currently does not expect the nullability property of an expression to change as a side effect of expression simplification. This PR is a first attempt to align the nullability of `coalesce` and `case`. ## What changes are included in this PR? Tweaks to the `nullable` logic for the logical and physical `case` expression code to report `case` as being not nullable in more situations. - For logical `case`, a best effort const evaluation of 'when' expressions is done to determine 'then' reachability. The code errs on the conservative side wrt nullability. - For physical `case`, const evaluation of 'when' expressions using a placeholder record batch is attempted to determine 'then' reachability. Again if const evaluation is not possible, the code errs on the conservative side. - The optimizer schema check has been relaxed slightly to allow nullability to be removed by optimizer passes without having to disable the schema check entirely - The panic'ing benchmark has been reenabled ## Are these changes tested? Additional unit tests have been added to test the new logic. ## Are there any user-facing changes? No --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
`CoalesceFunc::simplify` rewrites `coalesce(a, b)` into `CASE WHEN a IS NOT NULL THEN a ELSE b END`. That rewrite names every non-final argument *twice* -- once in the `WHEN` predicate and once in the `THEN` result -- so for a volatile argument the two mentions are two independent draws. Since apache#17357 removed the runtime kernel (`invoke_with_args` was left as `internal_err!("coalesce should have been simplified to case")`), there was no other path and no volatility guard, so every volatile `coalesce` went through the duplicating rewrite. `nvl` delegates `simplify`/`invoke_with_args`/`short_circuits`/ `conditional_arguments` to `CoalesceFunc`, so it had the identical bug. Two observable symptoms, both reproduced in the new sqllogictests against the unpatched code: 1. Wrong results. Over 100000 rows, `count(coalesce(nullif(floor(random()*2), 0), -1))` returned 75112 instead of 100000 -- 24888 impossible NULLs, exactly the 0.5*0.5 two-draw rate, from rows where the `WHEN` draw was non-null but the independent `THEN` draw was null. 2. A hard failure. `return_field_from_args` marks the result non-nullable when any argument is non-nullable (here the `-1` fallback), so those impossible NULLs also trip `Arrow error: Invalid argument error: Column 'c' is declared as non-nullable but contains null values`. Fix (two parts, one file): * Guard `simplify`: if any argument `is_volatile()`, return `ExprSimplifyResult::Original` and leave the `coalesce` call intact. Single-argument `coalesce` is still unwrapped, volatile or not, because it is named only once. * Restore the runtime kernel in `invoke_with_args`, recovered from e5dcc8c (`new_null_array` + `is_not_null`/`is_null` + `zip`). It walks the already-evaluated `ColumnarValue`s and therefore evaluates each argument exactly once. Why this rather than a new physical expression (apache#25476): this is ~60 lines in a single file with no new public API, no new physical expression and no protobuf change, which answers the complexity objection raised on that PR directly. Nothing outside `coalesce.rs` moves. Trade-off -- the restored kernel is eager, which is exactly what apache#17357 removed, so this is deliberately confined to the volatile case: * Non-volatile `coalesce` is untouched and keeps the lazy `CASE` rewrite. `select.slt:1686` (`select coalesce(1, y/x)` with `x = 0`) still plans to `Projection: Int64(1)` and still never divides by zero. * Volatile `coalesce` is now eager, so `coalesce(random(), y/x)` will evaluate `y/x` and can raise divide-by-zero where it previously did not. That is accepted: today the same query silently returns wrong answers (or aborts with the Arrow nullability error above), and a surfaced error beats a silently wrong result. It also only restores the pre-apache#17357 behaviour, and only for the narrow slice of calls that actually contain a volatile argument. There is no way to get both per-row laziness and single-evaluation out of the `CASE` rewrite, because the rewrite is what duplicates the argument. `short_circuits()` stays `true` and `conditional_arguments` is unchanged. Its contract is "some subexpressions *may* not be evaluated", so an eager kernel does not violate it -- common-subexpression elimination simply declines to hoist out of the lazy arguments, which costs an optimization and never correctness. Tests: unit tests for the restored kernel (arrays, array + scalar fallback, all scalars, all-null scalars, all-null arrays, empty args) and for the simplify guard (non-volatile still expands to `CASE`; direct, nested and single-argument volatile cases); sqllogictests asserting the deterministic `count(c) = count(*)` row count, that no value comes from outside the two operands, and `EXPLAIN` output pinning that the volatile shape stays `coalesce(...)`/`nvl(...)` and does not expand to `CASE`, while the non-volatile shape still does. Closes apache#25477. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`CoalesceFunc::simplify` rewrites `coalesce(a, b)` into `CASE WHEN a IS NOT NULL THEN a ELSE b END`. That rewrite names every non-final argument *twice* -- once in the `WHEN` predicate and once in the `THEN` result -- so for a volatile argument the two mentions are two independent draws. Since apache#17357 removed the runtime kernel (`invoke_with_args` was left as `internal_err!("coalesce should have been simplified to case")`), there was no other path and no volatility guard, so every volatile `coalesce` went through the duplicating rewrite. `nvl` delegates `simplify`/`invoke_with_args`/`short_circuits`/ `conditional_arguments` to `CoalesceFunc`, so it had the identical bug. Two observable symptoms, both reproduced in the new sqllogictests against the unpatched code: 1. Wrong results. Over 100000 rows, `count(coalesce(nullif(floor(random()*2), 0), -1))` returned 75112 instead of 100000 -- 24888 impossible NULLs, exactly the 0.5*0.5 two-draw rate, from rows where the `WHEN` draw was non-null but the independent `THEN` draw was null. 2. A hard failure. `return_field_from_args` marks the result non-nullable when any argument is non-nullable (here the `-1` fallback), so those impossible NULLs also trip `Arrow error: Invalid argument error: Column 'c' is declared as non-nullable but contains null values`. Fix (two parts, one file): * Guard `simplify`: if any argument `is_volatile()`, return `ExprSimplifyResult::Original` and leave the `coalesce` call intact. Single-argument `coalesce` is still unwrapped, volatile or not, because it is named only once. * Restore the runtime kernel in `invoke_with_args`, recovered from e5dcc8c (`new_null_array` + `is_not_null`/`is_null` + `zip`). It walks the already-evaluated `ColumnarValue`s and therefore evaluates each argument exactly once. Why this rather than a new physical expression (apache#25476): this is ~60 lines in a single file with no new public API, no new physical expression and no protobuf change, which answers the complexity objection raised on that PR directly. Nothing outside `coalesce.rs` moves. Trade-off -- the restored kernel is eager, which is exactly what apache#17357 removed, so this is deliberately confined to the volatile case: * Non-volatile `coalesce` is untouched and keeps the lazy `CASE` rewrite. `select.slt:1686` (`select coalesce(1, y/x)` with `x = 0`) still plans to `Projection: Int64(1)` and still never divides by zero. * Volatile `coalesce` is now eager, so `coalesce(random(), y/x)` will evaluate `y/x` and can raise divide-by-zero where it previously did not. That is accepted: today the same query silently returns wrong answers (or aborts with the Arrow nullability error above), and a surfaced error beats a silently wrong result. It also only restores the pre-apache#17357 behaviour, and only for the narrow slice of calls that actually contain a volatile argument. There is no way to get both per-row laziness and single-evaluation out of the `CASE` rewrite, because the rewrite is what duplicates the argument. `short_circuits()` stays `true` and `conditional_arguments` is unchanged. Its contract is "some subexpressions *may* not be evaluated", so an eager kernel does not violate it -- common-subexpression elimination simply declines to hoist out of the lazy arguments, which costs an optimization and never correctness. Tests: unit tests for the restored kernel (arrays, array + scalar fallback, all scalars, all-null scalars, all-null arrays, empty args) and for the simplify guard (non-volatile still expands to `CASE`; direct, nested and single-argument volatile cases); sqllogictests asserting the deterministic `count(c) = count(*)` row count, that no value comes from outside the two operands, and `EXPLAIN` output pinning that the volatile shape stays `coalesce(...)`/`nvl(...)` and does not expand to `CASE`, while the non-volatile shape still does. Closes apache#25477. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…e#25561) ## Which issue does this PR close? - Closes apache#25477. ## Rationale for this change `coalesce` is rewritten by `simplify` into `CASE WHEN a IS NOT NULL THEN a ELSE b END`, which names every argument but the last one twice. For a volatile argument those two mentions are two independent draws, so the null test and the returned value disagree: ```sql -- both draws are independent, so COALESCE can return NULL, -- which a single evaluation can never produce SELECT count(*), count(c) FROM (SELECT coalesce(nullif(floor(random() * 2), 0), neg) AS c FROM w); +--------+----------+ | 100000 | 75117 | +--------+----------+ ``` 24,883 NULLs, which is exactly `P(draw1 != 0) * P(draw2 == 0) = 0.5 * 0.5`. With a non-null literal as the last argument the planner marks the output non-nullable and it fails outright: ``` Arrow error: Invalid argument error: Column 'c' is declared as non-nullable but contains null values ``` `nvl` delegates to `CoalesceFunc`, so it has the same bug. ## What changes are included in this PR? `simplify` still rewrites `coalesce` into the fast `CASE`, but it stops at the first volatile argument that is not the last one and leaves the remaining arguments as a `coalesce` call, so each argument is named once: ```sql coalesce(a, v, c) --> CASE WHEN a IS NOT NULL THEN a ELSE coalesce(v, c) END coalesce(v, b) --> unchanged ``` That leftover call is evaluated by `invoke_with_args`, which gets back the runtime kernel removed in e5dcc8c (apache#17357) and evaluates each argument exactly once. Calls with no volatile argument before the last plan exactly as before and never reach the kernel. A volatile *last* argument is named once by the `ELSE`, so it needs no guard. `nvl` delegates to `CoalesceFunc` and is fixed by the same change. `nvl2` names each argument once and is unaffected. As discussed in review, a scalar binding function (`apply(v, x -> CASE WHEN x IS NOT NULL THEN x ELSE coalesce(b, c) END)`) built on the new lambda support would be a cleaner long-term replacement for the kernel. It needs a new always-registered higher-order function and has Substrait implications, so I have kept this PR to the narrow fix. ## What is the testing strategy for this PR? `sqllogictest` cases in `coalesce.slt`. To get exact expected rows from a volatile argument they use `case when random() >= 0 then x end`, which is volatile but always equals `x`. They cover array arguments, a scalar fallback, all-null input, strings, and a volatile middle argument. The 100k-row `random()` reproducer above is kept as one query covering both `coalesce` and `nvl`, asserting `count(c) = count(*)`. ## Are there any user-facing changes? The bug fix itself (no more NULLs or non-nullable errors from a volatile `coalesce`/`nvl`). One behaviour change: in a call with a volatile non-final argument, that argument and the ones after it are evaluated eagerly by the kernel rather than lazily by `CASE`. For example, `coalesce(random(), y / x)` now evaluates `y / x`. Arguments before the first volatile one, and all non-volatile calls, keep the lazy `CASE` behaviour. ## AI disclosure Following the [DataFusion AI policy](https://datafusion.apache.org/contributor-guide/ai-policy.html) and the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html): I used Claude Code to help write the code, tests and this description, and I reviewed all of it. The restored kernel is DataFusion's own code from before apache#17357; no third-party code is included. Things I would like a reviewer to check: - Whether the eager evaluation of the arguments after a volatile one is acceptable, or whether it should be documented on `coalesce`/`nvl`. - The `case when random() >= 0 then x end` trick in the tests relies on `CASE` with a volatile condition staying volatile through optimization. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Which issue does this PR close?
Rationale for this change
arguments of coalesce are evaluated eagerly.
What changes are included in this PR?
simplify coalesce to case when expr.
Are these changes tested?
UT
Are there any user-facing changes?
No