james-willis opened a new pull request, #27143:
URL: https://github.com/apache/datafusion/pull/27143
## Which issue does this PR close?
- Part of #24860.
This is a smaller alternative to #26145 + #26158, following the suggestion
on #24865 to look for a simpler way to reach the same goal.
## Rationale for this change
`CASE` drops the field metadata of its branches, so Arrow extension types
are lost. That breaks functions that check the extension type of their argument
when the argument is a `CASE`, or a `COALESCE`, `NVL`, `IFNULL` or `NVL2`,
since those are rewritten to `CASE` before execution. One example is
`variant_get` over a Parquet Variant column, reported on #26145; another is
geometry and raster functions in SedonaDB.
The metadata is dropped in two places:
- the logical `Expr::Case` field, which goes through the metadata-less
`to_field` catch-all;
- the physical `CaseExpr`, which has no `return_field` override.
## What changes are included in this PR?
- **Logical `Expr::Case::to_field`** merges the fields of its value branches
(the THEN branches and the ELSE) with `coerced_fields_metadata`, the helper
`COALESCE`/`NVL2` use since #24865.
- **Physical `CaseExpr::return_field`** does the same with its branches'
`return_field`.
- **A typed NULL ELSE that carries metadata.** `CaseExpr::try_new` no longer
drops such an ELSE, because its metadata is part of the result field. A
NULL-literal ELSE is still evaluated like no ELSE, so CASE keeps the same
evaluation paths.
So CASE follows the same rules as `COALESCE`:
- the earlier branch wins for each key;
- an extension name is kept together with its parameters;
- differing extension names are a planning error;
- a branch without metadata, or an untyped `NULL`, does not block another
branch's metadata.
The operand of a simple `CASE` and the WHEN conditions do not contribute.
### Differences from #26145 + #26158
About 60 lines of non-test code, against about 700 there.
- **No stored copy of the logical field, and no protobuf change.** The
physical field is derived from the physical branches.
- **No depth limit.** In #26158, a function branch nested more than 8 levels
deep loses its metadata.
- **A branch without metadata never clears the result**, in line with
`COALESCE`. In #26145, `CASE WHEN id > 1 THEN id > 2 ELSE with_metadata(id > 3,
'source', 'f') END` returns no metadata, while `coalesce(id > 2,
with_metadata(id > 3, 'source', 'f'))` returns `f`.
- **Each branch's field is used as that branch reports it.** One
consequence: physical `NOT` passes its input's metadata through while logical
`NOT` drops it. That inconsistency is in `NOT` and is not addressed here.
### Performance
These are local Criterion runs of the benchmarks from #26145/#26158 on a
laptop, so treat them as indicative.
- **CASE evaluation** (the existing `case_when` benchmarks) is unchanged.
- **Physical planning** is unchanged from `main`. With metadata present,
#26158's version takes 1.7–3.5x as long as `main`, for example 689 µs vs 197 µs
for a CASE nested 32 deep.
- **`return_field` is slower than in #26158**, because the merged metadata
is rebuilt on every call:
- 954 ns vs 282 ns for a column CASE with 4 metadata keys;
- 89 µs vs 3.8 µs for a CASE nested 32 deep, because nested branches build
names that are not used.
`ScalarFunctionExpr` asks its arguments for their field on every batch, so
this cost applies once per batch when a CASE is a function argument. Field plus
evaluation for 8,192 rows with 4 keys took 13.6 µs vs 13.1 µs.
- **Logical `to_field`** for a CASE without any metadata takes 2–3x as long
as `main`.
All of these can be addressed later with fast paths: skip building names for
nested CASE, column and literal branches; reuse the metadata when only one
branch has any; return early when no branch has metadata. I left them out to
keep this change small.
## What is the testing strategy for this PR?
- **New `metadata.slt` cases**, checking the metadata at execution with
`arrow_metadata`:
- precedence across branches, and simple CASE;
- conflicting extension names, which are a planning error;
- a branch without metadata, an untyped NULL ELSE, and no ELSE;
- a typed NULL ELSE with metadata, including when no row reaches a THEN
branch;
- nested CASE;
- the CASE rewrites of `coalesce`, `nvl`, `ifnull` and `nvl2`.
All of them fail without this change.
- **The existing tests**: the `datafusion-expr`, `datafusion-physical-expr`
and `datafusion-optimizer` lib tests pass.
The full sqllogictest suite and the `core_integration` `sql` tests show no
new failures. My checkout doesn't have the test-data submodules, so the files
that need that data fail the same way on `main`.
- **The `metadata.slt` cases from #26158**, which include those from #26895,
also pass with this change. I checked that locally against the base of #26158.
## Are there any user-facing changes?
- `CASE` results keep their field metadata, including extension types. So do
`COALESCE`, `NVL`, `IFNULL` and `NVL2` at execution.
- `CASE` branches with differing extension names are now a planning error,
as they already are for `COALESCE`.
No public API changes.
--
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]
