#17972 Restore case expr/expr optimisation while ensuring lazy evaluation #17973

pepijnve · 2025-10-08T17:13:53Z

Which issue does this PR close?

Closes Restore expr/expr optimisation for case expressions #17972

Rationale for this change

For non-trivial case when ... then ... else ... end expressions the general code path is currently taken which is not very efficient. The original optimisation for these expressions was reverted in #15384. This PR restores the optimisation while avoiding the mistake that was made in the original implementation.

What changes are included in this PR?

The is_cheap_and_infallible precondition for the ExpressionOrExpression code path has been removed
Two guard clauses have been added for the when_value.true_count() == 0 and when_value.true_count() == batch.num_rows() cases. In these situations calling evaluate_selection for the branch that will never be taken is avoided.

Are these changes tested?

Covered by tests added in #15384

Are there any user-facing changes?

No

…evaluation

alamb

Thanks @pepijnve -- I think the idea is great -- thank you. I have one question about how it works, but I did verify test coverage and things look ok

I also queued up a benchmark run for this PR to gather some more information as well

FYI @findepi and @aweltsch

alamb · 2025-10-09T18:32:29Z

datafusion/physical-expr/src/expressions/case.rs

+            // Avoid evaluate_selection when all rows are false/null
+            return self.else_expr.as_ref().unwrap().evaluate(batch);
+        }
+


My reading of this code is that it will still evaluate the then expression as long as there is at least one true value in when -- if evaluate_selection is a problem, shouldn't we fix evaluate_selection?

My reading of this code is that it will still evaluate the then expression as long as there is at least one true value in when

Yes, that's correct. There's no way to avoid that.

This particular bit of code is both an optimisation and a correctness thing.

From a performance point of view, we already know the selection vector is redundant, so there's really no point in calling evaluate_selection.

For correctness, what's being avoid here is calling either then or else with a selection vector that will result in an empty record batch after filtering. We could add similar checks in evaluate_selection to prevent evaluating the downstream expression for empty record batches as well. Its current contract requires it to return an array with the same length as the unfiltered input batch though. You can't avoid having to create an all-nulls array then.

My reading of this code is that it will still evaluate the then expression as long as there is at least one true value in when

Yes, that's correct. There's no way to avoid that.

isn't this reintroducing the bug that was fixed in #15384, just in a big more complex wrapping?

isn't this reintroducing the bug that was fixed in #15384, just in a big more complex wrapping?

No, why do you think that's the case? If you write case foo is not null then foo else 1/0 end and foo happens to be NULL, what do you expect to happen?

The original issue was that in the example above for a single row with a non-null foo, the code was evaluating the then branch with [true] as selection vector and the else branch with [false]. The latter was passed to evaluate_selection which then filters the record batch down to an empty record batch and then calls the else expression with that record batch. For an expression like 1/0, you end up getting executing that division anyway even though the result would be discarded.

There are two ways to fix this:

don't evaluate binary expressions and literals for empty batches (as you had suggested in a comment earlier I believe)

don't call evaluate with empty input batches

The earlier fix had the effect of 2. as well, just in a less explicit way. The fix here does the same but adds the necessary checks in the explicit expr/expr code path. The code that's being seen as an optimisation is intended to prevent calling evaluate_selection with an all-false selection vector.

I woke up early this morning to the realisation that the actual bug was a subtlety in the implementation of evaluate_selection. There's a difference between calling it with the empty set vs a non-empty set and a false selection vector. The implementation was actually treating both cases identically which can cause a spurious row to get materialised. I've pushed a correction for this and tweaked the comments in the code a bit.

I believe this properly addresses the original evaluation problem. All SQL logic tests pass even when commenting out the optimisation for true and false selection vectors in expr_or_expr.

I agree with the change in evaluate_selection
Now that evaluate_selection is changed, do we need those lines here?

(They look nice but my only concern is additional code complexity which is harder to cover with SLT tests. Now that we have branching here, we should have a bunch of SLT cases that clearly exercise all-true, all-false, some-true situations.)

Yes, more SLTs! I will add some for the various cases you mentioned.

Now that evaluate_selection is changed, do we need those lines here?

@findepi Yes I would prefer to keep these early outs since they're pretty trivial and I think they're appropriate for the expr_or_expr function since the knowledge that it's "then or else" is located here. evaluate_selection cannot be implemented with awareness of this particular usage pattern.

Calling evaluate_selection still has to produce a value –it can't return None– so you pay at least a non-zero cost for calling it unnecessarily. If it's obvious that it will not perform any useful work, it makes sense to avoid it IMO.

Just for context, the queries I'm working on are quite case heavy. Any work we can save in the inner loop of the queries seems worthwhile.

alamb · 2025-10-09T19:11:24Z

🤖 ./gh_compare_branch_bench.sh Benchmark Script Running
Linux aal-dev 6.14.0-1016-gcp #17~24.04.1-Ubuntu SMP Wed Sep 3 01:55:36 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing issue_17972 (cf78119) to 07a7eb2 diff
BENCH_NAME=case_when
BENCH_COMMAND=cargo bench --bench case_when
BENCH_FILTER=
BENCH_BRANCH_NAME=issue_17972
Results will be posted here when complete

alamb · 2025-10-09T19:18:49Z

🤖: Benchmark completed

Details

group                          issue_17972                            main
-----                          -----------                            ----
case_when: CASE expr           1.00     24.0±0.26µs        ? ?/sec    1.01     24.4±0.19µs        ? ?/sec
case_when: column or null      1.00   1409.5±2.43ns        ? ?/sec    1.01   1421.6±3.52ns        ? ?/sec
case_when: expr or expr        1.00     31.0±0.29µs        ? ?/sec    1.01     31.3±0.16µs        ? ?/sec
case_when: scalar or scalar    1.00      7.9±0.02µs        ? ?/sec    1.00      7.9±0.03µs        ? ?/sec

alamb · 2025-10-09T19:18:51Z

🤖 ./gh_compare_branch_bench.sh Benchmark Script Running
Linux aal-dev 6.14.0-1016-gcp #17~24.04.1-Ubuntu SMP Wed Sep 3 01:55:36 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing issue_17972 (cf78119) to 07a7eb2 diff
BENCH_NAME=case_when
BENCH_COMMAND=cargo bench --bench case_when
BENCH_FILTER=
BENCH_BRANCH_NAME=issue_17972
Results will be posted here when complete

alamb · 2025-10-09T19:25:28Z

🤖: Benchmark completed

Details

group                          issue_17972                            main
-----                          -----------                            ----
case_when: CASE expr           1.00     24.4±0.38µs        ? ?/sec    1.02     24.9±0.25µs        ? ?/sec
case_when: column or null      1.00  1406.5±11.79ns        ? ?/sec    1.01   1414.3±8.58ns        ? ?/sec
case_when: expr or expr        1.00     31.3±0.27µs        ? ?/sec    1.01     31.7±0.30µs        ? ?/sec
case_when: scalar or scalar    1.00      7.9±0.03µs        ? ?/sec    1.00      7.9±0.18µs        ? ?/sec

findepi · 2025-10-09T20:10:58Z

datafusion/physical-expr/src/expressions/case.rs

+        if true_count == batch.num_rows() {
+            // Avoid evaluate_selection when all rows are true
+            return self.when_then_expr[0].1.evaluate(batch);
+        } else if true_count == 0 {
+            // Avoid evaluate_selection when all rows are false/null
+            return self.else_expr.as_ref().unwrap().evaluate(batch);
+        }


per @alamb comment https://github.com/apache/datafusion/pull/17973/files#r2417644396, those lines look like an optimization (and a reasonable one).
but the code should also work if we take them out, just a bit slower. Do all the test pass if we comment this block out?

No, you're back to the original issue for which you added the guard clauses then. The fix back then was local, so I applied a local fix in this PR as well.

I've pushed an extra commit that adds an extra guard clause in evaluate_selection as well. With that change you can comment out this block and all SQL logic tests still pass.

…election` for empty selections.

…sets and all false filters

…path.

findepi · 2025-10-10T09:04:16Z

datafusion/physical-expr-common/src/physical_expr.rs

+        let selection_count = selection.true_count();

-        let tmp_result = self.evaluate(&tmp_batch)?;
+        if batch.num_rows() == 0 || selection_count == batch.num_rows() {


Suggested change

if batch.num_rows() == 0 || selection_count == batch.num_rows() {

if selection_count == batch.num_rows() {

(i'd assume batch.num_rows() == 0 implies that also selection_count == 0. If it is not the case, we have a more permissive if condition, but requires a code comment)

There was no assertion for this in place and as far as I can tell (but I did not verify yet), you can call evaluate_selection today with a mismatch between batch.num_rows() and selection.len() and it will do something. I haven't tested this yet, so I'm not 100% sure what the outcome would be.
I'll write an extra unit test to cover this.

The idea here is to avoid any extra work whatsoever in the trivial cases. Not sure what a useful, non-trivial comment for that would be.

@alamb @findepi I've added some unit tests to help define the behaviour of evaluate_selection. Some are currently still failing. Could you guys take a look at the tests to see if what they're asserting is correct?

The issues are all related to record batch / selection vector size mismatches. I don't think the current behaviour makes sense tbh (which is also present on main). I would expect to either get an error or as many output rows as there were input rows.

I've taken the liberty of adding the size mismatch check and returning an execution error in case of mismatch. Within the DataFusion library case is the only client of this API and for that this change should be fine. The behaviour in case of mismatch was pretty weird, so I doubt people would be making active use of this. You never know of course. I'll add this to the user visible changes list.

I've added a couple of tests guide by code coverage.expr_or_expr is definitely well covered.

findepi · 2025-10-10T09:07:30Z

datafusion/physical-expr-common/src/physical_expr.rs

-        let tmp_batch = filter_record_batch(batch, selection)?;
+        let selection_count = selection.true_count();

-        let tmp_result = self.evaluate(&tmp_batch)?;


There is a lot of new code which repeats logic of filter_record_batch.
What if we just changed this line only?

let tmp_result = if tmp_batch.is_empty { // Do not call `evaluate` when the selection is empty. // When `evaluate_selection` is being used for conditional, lazy evaluation, // evaluating an expression for a false selection vector may end up unintentionally // evaluating a fallible expression. let datatype = self.data_type(batch.schema_ref().as_ref())?; ColumnarValue::Array(make_builder(&datatype, 0).finish()) } else { self.evaluate(&tmp_batch)?; }

Note how this does not inspect selection / selection_count directly, leveraging the work done by filter_record_batch.

I might be missing it, but I don't see the overlap with filter_record_batch. AFAICT there are no checks to avoid creating a new record batch. What this code is doing is preparing an empty result value while filter_record_batch has optimised code to an empty record batch if the filter is all-false.

findepi · 2025-10-10T09:10:16Z

datafusion/physical-expr/src/expressions/case.rs

+            // Avoid evaluate_selection when all rows are false/null
+            return self.else_expr.as_ref().unwrap().evaluate(batch);
+        }
+


I agree with the change in evaluate_selection
Now that evaluate_selection is changed, do we need those lines here?

(They look nice but my only concern is additional code complexity which is harder to cover with SLT tests. Now that we have branching here, we should have a bunch of SLT cases that clearly exercise all-true, all-false, some-true situations.)

…_selection

- Add extra comments - Use match for the scatter paragraph - Validate that the size of selection and batch match

alamb · 2025-10-14T21:08:51Z

CI failure seems unrelated to this PR:

CI is failing on main / current_date() = cast(now() as date); #18062

pepijnve · 2025-10-15T09:23:21Z

I updated this PR again since the failing test was reverted on main

alamb · 2025-10-15T10:36:01Z

Thanks again @pepijnve 🙏 -

apache#17972 Restore case expr/expr optimisation while ensuring lazy …

cf78119

…evaluation

github-actions bot added the physical-expr Changes to the physical-expr crates label Oct 8, 2025

alamb reviewed Oct 9, 2025

View reviewed changes

findepi reviewed Oct 9, 2025

View reviewed changes

pepijnve added 6 commits October 9, 2025 23:01

Avoid calling PhysicalExpr::evaluate from `PhysicalExpr::evaluate_s…

b782628

…election` for empty selections.

Make PhysicalExpr::evaluate_selection correctly handle empty input …

eed9a34

…sets and all false filters

Reoragnize code to avoid scatter codepath when using evaluate fast …

c8524d3

…path.

Clarify comments in case

f9f67c5

Move null handling after true count check.

b849738

Tweaking comments

efbd205

findepi reviewed Oct 10, 2025

View reviewed changes

findepi approved these changes Oct 10, 2025

View reviewed changes

pepijnve added 4 commits October 10, 2025 13:18

Add unit tests to help define the boundary case behaviour of evaluate…

480a747

…_selection

Code polishing

c8186ec

- Add extra comments - Use match for the scatter paragraph - Validate that the size of selection and batch match

Fix clippy errors

aff3f18

Add additional case SLTs

99a32fc

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Oct 10, 2025

Merge branch 'main' into issue_17972

e1a4f79

Merge branch 'main' into issue_17972

7e85be2

alamb added this pull request to the merge queue Oct 15, 2025

Merged via the queue into apache:main with commit 6d3854f Oct 15, 2025
28 checks passed

alamb added the performance Make DataFusion faster label Oct 15, 2025

This was referenced Oct 15, 2025

[EPIC] A collection of items to improve CASE performance #18075

Open

extended tests failures on main #18084

Closed

Update extended tests with new results apache/datafusion-testing#14

Merged

	if batch.num_rows() == 0 \|\| selection_count == batch.num_rows() {
	if selection_count == batch.num_rows() {

#17972 Restore case expr/expr optimisation while ensuring lazy evaluation #17973

#17972 Restore case expr/expr optimisation while ensuring lazy evaluation #17973

Uh oh!

Conversation

pepijnve commented Oct 8, 2025

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

findepi Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pepijnve Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pepijnve Oct 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Oct 9, 2025

Uh oh!

alamb commented Oct 9, 2025

Uh oh!

alamb commented Oct 9, 2025

Uh oh!

alamb commented Oct 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Oct 14, 2025

Uh oh!

pepijnve commented Oct 15, 2025

Uh oh!

alamb commented Oct 15, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

findepi Oct 9, 2025 •

edited

Loading

pepijnve Oct 9, 2025 •

edited

Loading

pepijnve Oct 10, 2025 •

edited

Loading