apache / apache/datafusion

Simplifier: push immutable scalar functions into CASE branches with literal outputs

Open
#24,477 1 comment 0 reactions 1 assignee Claimed by @radmirnovii View on GitHub
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

## Is your feature request related to a problem or challenge?

Applying an expensive deterministic function to a `CASE` whose branches are all literals evaluates
the function once per row at execution time, even though only one evaluation per branch is ever
needed:

```sql
CREATE TABLE t AS
SELECT * FROM (VALUES ('a', true), ('b', false)) AS v(s, flag);

EXPLAIN SELECT
to_timestamp(CASE WHEN flag
THEN '2024-03-01T00:00:00Z'
ELSE '2024-09-01T00:00:00Z' END)
FROM t;
```

Nothing in the current simplifier touches this expression — observed output on current main
(`EXPLAIN FORMAT indent`):

```
logical_plan Projection: to_timestamp(CASE WHEN t.flag THEN Utf8("2024-03-01T00:00:00Z") ELSE Utf8("2024-09-01T00:00:00Z") END)
physical_plan ProjectionExec: expr=[to_timestamp(CASE WHEN flag@0 THEN 2024-03-01T00:00:00Z ELSE 2024-09-01T00:00:00Z END) ...]
```

So the CASE materializes a string array at execution time and `to_timestamp` string-parses it per
row. Two controls confirm the blocker is precisely the column reference: the comparison form
`(CASE WHEN flag THEN 'x' ELSE 'y' END) = 'x'` already simplifies (all the way to `flag`) via the
existing Eq/NotEq rule, and replacing `WHEN flag` with `WHEN 1 = 1` lets `ConstEvaluator` fold the
whole expression to a timestamp constant.

The hand-rewritten form already folds today — writing the pushdown result manually,
`CASE WHEN flag THEN to_timestamp('2024-03-01T00:00:00Z') ELSE to_timestamp('2024-09-01T00:00:00Z') END`,
plans as

```
Projection: CASE WHEN ... THEN TimestampNanosecond(1709251200000000000, None) ELSE TimestampNanosecond(1725148800000000000, None) END
```

— both `to_timestamp` calls evaluated once at plan time. So everything after the rewrite already
works; only the rewrite itself is missing. The per-row cost is visible in `EXPLAIN ANALYZE` over
`generate_series(1, 2000000)`: the operator evaluating the current form reports
`elapsed_compute=3.65s` vs `348ms` for the folded form (debug build, illustration only — a
criterion bench would accompany the PR).

The same shape appears with regex/JSON-path compilation inside UDFs, interval parsing, and other
parse-heavy functions over flag- or category-driven literal choices.

## Describe the solution you'd like

Generalize the existing Eq/NotEq literal pushdown
(`datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs:1440-1473`, added in #17743) from
`{=, !=} × literal` to any immutable scalar function:

```
f(a_1, ..., CASE WHEN w_i THEN t_i ... [ELSE e] END, ..., a_k)
--> CASE WHEN w_i THEN f(a_1, ..., t_i, ..., a_k) ... ELSE f(a_1, ..., e_or_NULL, ..., a_k) END
```

Guards: `f` is `Volatility::Immutable`; exactly one argument is the CASE and every other argument is
a literal; the CASE passes the existing `is_case_with_literal_outputs` guard; a missing ELSE is
materialized as an explicit `ELSE f(..., NULL, ...)` (never assume `f(NULL)` is NULL). `Cast` is
excluded (its fold errors fail fast at plan time by design).

On the next simplifier cycle `ConstEvaluator` folds each `f(literal)` branch to a constant, so the
per-row cost disappears and the result even hits the physical `ScalarOrScalar` fast path.
Correctness follows from CASE evaluating branches only on selected rows before and after the
rewrite, and from `ConstEvaluator` preserving expressions whose fold errors ("to allow
short-circuit evaluation at execution time") — an invalid literal in a never-taken branch stays
unfolded and never executes. I am happy to submit a PR with unit tests, `case.slt` coverage, and a
criterion bench.

## Describe alternatives you've considered

- **Per-UDF `simplify()` overrides**: each parse-heavy UDF could implement the pushdown itself, but
that duplicates the same rewrite across functions and misses third-party UDFs.
- **A separate analyzer pass**: runs once, outside the simplifier's ConstEvaluator↔Simplifier
cycle, so the folded-literal payoff would need a re-run; the simplifier arm gets folding for free
on the next cycle.
- **Do nothing**: users can hand-rewrite queries, but the shape is common in generated SQL where the
CASE comes from a parameter/flag expansion.

## Additional context

- Precedent: #17743 (merged) contains this exact rewrite for comparisons, with the guard helpers the
general rule would reuse (`is_case_with_literal_outputs`, `is_lit`).
- Cautionary precedent: PR #19732 (closed unmerged) was rejected because *removing* a CASE changed
error semantics. This proposal keeps the CASE structure and its WHEN guards intact; only branch
bodies change, and only when they are literals.
- No existing issue or PR proposes this generalization (searched issues/PRs for simplify/push/
distribute/inline CASE variants).

**Prior art in other engines**

- Spark Catalyst has exactly this rule: [`PushFoldableIntoBranches`](https://github.com/apache/spark/blob/cd687e6259bbbeb4ab3ea35f31d96e3883c9b39f/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala#L716-L801) (SPARK-33798), paired with the same preserve-on-error constant folding for conditional branches ([`ConstantFolding.tryFold`](https://github.com/apache/spark/blob/cd687e6259bbbeb4ab3ea35f31d96e3883c9b39f/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala#L63-L80)).
- Calcite applies the boolean-call subset: [`ReduceExpressionsRule.pushPredicateIntoCase`](https://github.com/apache/calcite/blob/7925800cb86892e32183959776cf476c4add1244/core/src/main/java/org/apache/calcite/rel/rules/ReduceExpressionsRule.java#L855-L908).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.