Skip to content

fix: lazy evaluation for coalesce - #17357

Merged
alamb merged 2 commits into
apache:mainfrom
chenkovsky:fix/lazy_coalesce
Sep 8, 2025
Merged

alamb merged 2 commits into
apache:mainfrom
chenkovsky:fix/lazy_coalesce

Conversation

@chenkovsky

@chenkovsky chenkovsky commented Aug 30, 2025 •

Copy link
Copy Markdown
Contributor

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

@github-actions github-actions Bot added sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation labels Aug 30, 2025

@nuno-faria nuno-faria left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Comment on lines 1689 to 1690
# due to the reason describe in https://github.com/apache/datafusion/issues/8927,
# the following queries will fail

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the test below now works, I think we could move this comment to the test below this one.

@alamb alamb left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@mbutrovich

Copy link
Copy Markdown
Contributor

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

Comet was actually just thinking about coalesce:

apache/datafusion-comet#2270

Maybe @coderfender has some thoughts about this?

@alamb

alamb commented Sep 6, 2025

Copy link
Copy Markdown
Contributor

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

Comet was actually just thinking about coalesce:

apache/datafusion-comet#2270

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

@coderfender

coderfender commented Sep 6, 2025 •

Copy link
Copy Markdown
Contributor

@alamb , @mbutrovich I made changes to comet to fallback to CASE statement to replicate lazy evaluation with coalesce (and then plan to work on this PR). Glad to see that the changes are looking good and once we update DataFusion's version in comet, we should be able to undo my comet changes to leverage DF's coalesce function . I will create a github issue on the comet side for this

@alamb
alamb merged commit e5dcc8c into apache:main Sep 8, 2025
@alamb

alamb commented Sep 8, 2025

Copy link
Copy Markdown
Contributor

Thanks everyone!

@alamb

alamb commented Sep 9, 2025

Copy link
Copy Markdown
Contributor

I believe this PR caused the extended tests to fail (as more queries now start woring):

alamb added a commit to alamb/datafusion that referenced this pull request Nov 9, 2025
alamb added a commit that referenced this pull request Nov 10, 2025
…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.
-->
github-merge-queue Bot pushed a commit that referenced this pull request Nov 20, 2025
## 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>
logan-keede pushed a commit to logan-keede/datafusion that referenced this pull request Nov 23, 2025
…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>
DeviousCardi added a commit to DeviousCardi/datafusion that referenced this pull request Sep 21, 2026
`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>
DeviousCardi added a commit to DeviousCardi/datafusion that referenced this pull request Oct 23, 2026
`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>
adriangb pushed a commit to pydantic/datafusion that referenced this pull request Oct 30, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

functions Changes to functions implementation sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

COALESCE expr in datafusion should perform lazy evaluation of the operands Evaluates COALESCE arguments past first non-NULL value

5 participants