Repository navigation
fix: preserve field metadata through CASE - #27143
james-willis wants to merge 5 commits into
Conversation
The logical Expr::Case field and the physical CaseExpr field now merge the metadata of the value branches with coerced_fields_metadata, the same rules COALESCE and NVL2 use: the earlier branch wins for each key, an extension name stays with its parameters, and differing extension names are a planning error. CaseExpr::try_new keeps a typed NULL ELSE that carries metadata, and a NULL-literal ELSE still evaluates like no ELSE.
|
@paleolimbot @alamb @osipovartem This is my alternative solution for preserving metadata in those cases. |
paleolimbot
left a comment
There was a problem hiding this comment.
This approach makes sense to me but I'll defer to those more familiar with the case/when optimizations as to whether the other approach is preferable!
Codex complained to me that this approach requires that logical and physical representations of an expr have to match for this to work and pointed out a few places they don't at the moment which I filed as separate issues. Its noted test gap was "No test compares Expr::Case::to_field with CaseExpr::return_field for the same expression."
One more gap it suggested was to check where the CASE test where branches have different storage types and coercion inserts a cast. This matters because type-only casts strip extension metadata, and logical and physical cast behavior must agree.
Compare Expr::Case::to_field with the physical return_field, including a branch cast to the common type by type coercion. Use the built-in JSON extension type in the CASE metadata.slt cases, add a coercion case, and move the CASE WHEN false case next to the other CASE cases.
|
Thanks Dewey. Integrated all your feedback. Logical vs physical. Agreed, this relies on each branch's physical field matching its logical one, which is the property #27156 describes for every expression. CASE derives its field only from its branches in both layers, so it matches whenever its branches do, and a branch that doesn't (NOT in #27155, Alias in #27156) shows through. #26158 avoids that by carrying the logical field into the physical CASE, at the cost of storing it and a protobuf change. In practice the mismatch doesn't change the result schema, which takes its metadata from the logical plan. Only a function that takes the CASE as an argument sees the physical field. in the NOT case probably this doesnt really matter - its probably just extra metadata. In the Alias case it can be an issue but this is probably exceeding rate to call alias on a subexpression that doesnt persist past the operator.
Coercion casts. Logical and physical agree: the type-only cast that coercion inserts drops the extension type and keeps other keys on both sides. The unit test above covers it, and a new |
jayzhan211
left a comment
There was a problem hiding this comment.
@james-willis thanks for this! In expr_simplifier.rs around line 1485, the other rewrites in this arm still drop the merged metadata: ext_only(CASE WHEN 1 = 0 THEN g ELSE NULL END) (the UDF from the #24860 repro) fails planning with "argument lost its extension type", while CASE WHEN 1 = 0 THEN g END works. SELECT CASE WHEN 1 = 0 THEN g ELSE plain END AS x also reports the extension in df.schema() but not in the batches. Could the arm compare the field metadata before and after and keep the CASE when it changes? I ran this diff: every ext_only probe in the table below passes and the optimizer lib tests still pass. It does make the arm read the ELSE's field, so from reading the code (I didn't run this one) a simplifier without a schema would now error on CASE WHEN true THEN 1 ELSE col END.
- // CASE WHEN false THEN A END --> NULL
- // The NULL keeps the CASE's field metadata, such as an extension type
- if else_expr.is_none()
- && when_then_expr
- .iter()
- .all(|(when, _)| is_false(when.as_ref()))
- {
- let (_, field) = Expr::Case(Case::new(None, when_then_expr, None))
- .to_field(info.schema().as_ref())?;
- let metadata = FieldMetadata::new_from_field(&field);
- let null = Expr::Literal(
- ScalarValue::try_new_null(field.data_type())?,
- (!metadata.is_empty()).then_some(metadata),
- );
- return Ok(Transformed::yes(null));
- }
+ // Dropping a branch must not change the field of the CASE, whose
+ // metadata merges all of its value branches.
+ let original = Expr::Case(Case::new(
+ None,
+ when_then_expr.clone(),
+ else_expr.clone(),
+ ));
+ let schema = info.schema().as_ref();
+ let (_, field) = original.to_field(schema)?;
+ let metadata = FieldMetadata::new_from_field(&field);
@@
// Exclude CASE statement altogether if there are no when-then expressions left
- // CASE WHEN false THEN A ELSE B END --> B
- if new_when_then_expr.is_empty()
- && let Some(else_expr) = else_expr
- {
- return Ok(Transformed::yes(*else_expr));
- }
+ let simplified = match (new_when_then_expr.is_empty(), else_expr) {
+ // CASE WHEN false THEN A END --> NULL
+ (true, None) => {
+ return Ok(Transformed::yes(Expr::Literal(
+ ScalarValue::try_new_null(field.data_type())?,
+ (!metadata.is_empty()).then_some(metadata),
+ )));
+ }
+ // CASE WHEN false THEN A ELSE B END --> B
+ (true, Some(else_expr)) => *else_expr,
+ (false, else_expr) => Expr::Case(Case {
+ expr: None,
+ when_then_expr: new_when_then_expr,
+ else_expr,
+ }),
+ };
- Transformed::yes(Expr::Case(Case {
- expr: None,
- when_then_expr: new_when_then_expr,
- else_expr,
- }))
+ let (_, new_field) = simplified.to_field(schema)?;
+ if FieldMetadata::new_from_field(&new_field) == metadata {
+ Transformed::yes(simplified)
+ } else {
+ Transformed::no(original)
+ }Queries to add to the tests, and the results with this diff
Each arrow_metadata query should return arrow.json, the same as the non-constant CASE (expected results; I ran the same shapes through ext_only below, not these exact queries):
select arrow_metadata(case when 1 = 0 then cast(name as json) else NULL end, 'ARROW:extension:name') from table_with_metadata;
select arrow_metadata(case when 1 = 0 then cast(name as json) else upper(l_name) end, 'ARROW:extension:name') from table_with_metadata;
select arrow_metadata(case when 1 = 1 then upper(l_name) else cast(name as json) end, 'ARROW:extension:name') from table_with_metadata;
select arrow_metadata(case when id > 1 then upper(l_name) when 1 = 0 then cast(name as json) end, 'ARROW:extension:name') from table_with_metadata;
select arrow_metadata(coalesce(nonnull_name, cast(name as json)), 'ARROW:extension:name') from table_with_metadata;With ext_only requiring the extension on its argument (g has one; plain and nn don't, and nn is non-nullable):
| Argument | Now | With the diff |
|---|---|---|
CASE WHEN 1 = 0 THEN g ELSE NULL END |
plan error | ok |
CASE WHEN 1 = 0 THEN g ELSE plain END |
plan error | ok |
CASE WHEN 1 = 1 THEN plain ELSE g END |
plan error | ok |
CASE WHEN k = 1 THEN plain WHEN 1 = 0 THEN g END |
plan error | ok |
COALESCE(nn, g) |
plan error | ok |
NVL2(nn, plain, g) |
plan error | ok |
In expr_simplifier.rs around line 1555, the boolean lowering loses the merged metadata too: for CASE WHEN p THEN f END, with f a Boolean column carrying metadata, df.schema() reports it but the batches don't, and a metadata-checking UDF fails to plan. Could this arm skip the rewrite when a value branch has metadata, with a test on that query? Line 1613 looks like the same rewrite for literal THENs, so it likely needs the same guard; that's from reading the code, not a probe, and I haven't compiled this snippet.
/// Whether a value branch of a CASE has field metadata, which AND/OR cannot keep
fn case_branch_has_metadata(
when_then_expr: &[(Box<Expr>, Box<Expr>)],
else_expr: Option<&Expr>,
info: &SimplifyContext,
) -> Result<bool> {
let branches = when_then_expr
.iter()
.map(|(_, then)| then.as_ref())
.chain(else_expr);
for expr in branches {
let (_, field) = expr.to_field(info.schema().as_ref())?;
if !field.metadata().is_empty() {
return Ok(true);
}
}
Ok(false)
}- && can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref()) =>
+ && can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref())
+ && !case_branch_has_metadata(&when_then_expr, else_expr.as_deref(), info)? =>Optional cleanup: the "Metadata Handling" list on to_field in expr_schema.rs around line 509 could get a CASE entry now that CASE has its own rule.
/// - **Binary expressions**: field metadata is empty
+ /// - **CASE expressions**: Merges the metadata of the THEN and ELSE branches; the operand and WHEN conditions do not contribute| } | ||
| } | ||
|
|
||
| fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> { |
There was a problem hiding this comment.
The NOT mismatch from #27155 can be a planning error here, not only extra metadata: this query returns rows on main and now fails with "Cannot merge field metadata with differing values". It's fine to handle this in #27155 since the fix belongs in NOT; could you add this query there and adjust the note in the PR description?
-- t(id INT, b BOOLEAN)
select case when id > 1
then not with_metadata(b, 'ARROW:extension:name', 'a.type')
else with_metadata(b, 'ARROW:extension:name', 'b.type') end
from t;There was a problem hiding this comment.
I wish this just wouldn't plan but I didn't make the change in this or my prior PR (#24865) because of some unpleasant implications.
In my prior planning PR we made mismatched "ARROW:extension:name" an error but the pair of a present and a null one goes to the present value.
IMO this should be a planning error because ARROW:extension:name really defines the type and types should match between branches but it leads to some unpleasant situations especially around literals, for example COALESCE(json_col, '{}') would fail.
This pain is just kind of a side effect of ARROW:extension:name being an insufficient type system. Here its limitation is that it cant describe what kinds of up/down casts are allowed
A CASE with a constant WHEN now keeps the CASE when dropping a branch would change its field metadata, and a literal result takes the CASE's metadata. A Boolean CASE whose branch has metadata is no longer rewritten to AND/OR, which cannot carry it.
|
Thanks @jayzhan211, great catches. Addressed in 39697af. Constant WHENs (~L1485): took your diff with one tweak: if the simplified result is a literal, it takes the CASE's metadata, so You're right that a simplifier without a schema now errors on AND/OR lowering (~L1555, L1613): added your Docs: added the CASE entry to |
jayzhan211
left a comment
There was a problem hiding this comment.
@james-willis thanks for the quick turnaround! Every query from my last review now works, and when I remove any of the new simplifier guards, simplify_expr_case_keeps_field_metadata fails. Two cost issues, both inline with fixes I ran locally:
- planning time for a CASE nested inside function calls now grows much faster than on main (about 50× slower at 12 levels);
- the per-batch
return_fieldgrows quadratically with nesting depth.
The logical CASE field reads its type from the branch fields instead of calling get_type, which walked the branches a second time and made planning a CASE nested in function calls grow as 3^depth. The physical CASE computes only the metadata of a nested CASE branch, not its name and nullability, so return_field no longer grows quadratically with nesting depth. Also write the conflicting extension name error in metadata.slt as plain text.
|
thanks for all the help @jayzhan211 I made the tweaks from your review. |
jayzhan211
left a comment
There was a problem hiding this comment.
@james-willis thanks for applying both changes! One more nesting shape is left; details and a fix I ran are inline.
| Ok(Arc::new(metadata.add_to_field(Field::new( | ||
| &schema_name, | ||
| data_type, | ||
| self.nullable(schema)?, |
There was a problem hiding this comment.
One nesting shape still blows up: with the CASE in the ELSE and a non-nullable THEN, 20 levels of upper(CASE WHEN a > k THEN 'x' ELSE <inner> END) take 77 s to EXPLAIN here vs 1.9 s on main (debug build). to_field walks every branch, then self.nullable walks the ELSE again because 'x' can't be null, so the cost is 2^depth, where main walked the ELSE once. Taking each branch's nullability from the fields you already collected makes it linear: 0.07 s at depth 20 with the same plan. The expr, optimizer and physical-expr lib tests and the full sqllogictest suite pass with it.
// to_field, CASE arm
case_nullable(case, schema, |i| Ok(fields[i].is_nullable()))?,
// nullable, CASE arm
Expr::Case(case) => case_nullable(case, input_schema, |i| match case.when_then_expr.get(i) {
Some((_, then)) => then.nullable(input_schema),
None => case
.else_expr
.as_ref()
.map_or(Ok(true), |e| e.nullable(input_schema)),
}),Full diff (moves the existing nullable logic into a helper), query generator and timings
use crate::expr::{
- AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Cast, InList,
+ AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Case, Cast, InList,
@@ fn nullable
- Expr::Case(case) => {
- let nullable_then = case.when_then_expr.iter().find_map(|(w, t)| {
- let is_nullable = match t.nullable(input_schema) {
- ... (the rest of the current arm, moved unchanged into `case_nullable`)
- }
+ Expr::Case(case) => {
+ case_nullable(case, input_schema, |i| match case.when_then_expr.get(i) {
+ Some((_, then)) => then.nullable(input_schema),
+ None => case
+ .else_expr
+ .as_ref()
+ .map_or(Ok(true), |e| e.nullable(input_schema)),
+ })
+ }
@@ fn to_field, CASE arm
- self.nullable(schema)?,
+ case_nullable(case, schema, |i| Ok(fields[i].is_nullable()))?,
@@
+/// The nullability of `case`, given the nullability of each value branch:
+/// the THEN expressions in order, then the ELSE.
+fn case_nullable(
+ case: &Case,
+ input_schema: &dyn ExprSchema,
+ mut branch_nullable: impl FnMut(usize) -> Result<bool>,
+) -> Result<bool> {
+ let nullable_then = case
+ .when_then_expr
+ .iter()
+ .enumerate()
+ .find_map(|(i, (w, t))| {
+ let is_nullable = match branch_nullable(i) {
+ Err(e) => return Some(Err(e)),
+ Ok(n) => n,
+ };
+ // ... unchanged from the current arm ...
+ });
+
+ if let Some(nullable_then) = nullable_then {
+ nullable_then.map(|_| true)
+ } else if case.else_expr.is_some() {
+ branch_nullable(case.when_then_expr.len())
+ } else {
+ Ok(true)
+ }
+}python3 - 20 > q.sql <<'PY'
import sys
d = int(sys.argv[1]); e = "s"
for k in range(d):
e = f"upper(CASE WHEN a > {k} THEN 'x' ELSE {e} END)"
print("CREATE TABLE t AS VALUES (1, 'x');")
print("CREATE VIEW v AS SELECT column1 AS a, column2 AS s FROM t;")
print(f"EXPLAIN SELECT {e} AS r FROM v;")
PY
time datafusion-cli -q -f q.sql > /dev/null| depth | main | this PR | with the change |
|---|---|---|---|
| 16 | 0.04 s | 4.5–6.0 s | 0.08 s |
| 18 | 0.43 s | 22 s | 0.07 s |
| 20 | 1.9 s | 77 s | 0.07 s |
Which issue does this PR close?
CASEexpressions drop field metadata, silently stripping Arrow extension types and breaking metadata-aware UDFs #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. It also includes the simplifier fix from #26951, which depends on the CASE fix here.
Rationale for this change
CASEdrops 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 aCASE, or aCOALESCE,NVL,IFNULLorNVL2, since those are rewritten toCASEbefore execution. One example isvariant_getover a Parquet Variant column, reported on #26145; another is geometry and raster functions in SedonaDB.The metadata is dropped in three places:
Expr::Casefield, which goes through the metadata-lessto_fieldcatch-all;CaseExpr, which has noreturn_fieldoverride;CASE WHEN false THEN x ENDbecomes a NULL literal without metadata, a constant WHEN drops the other branches, and a Boolean CASE becomes AND/OR.What changes are included in this PR?
Expr::Case::to_fieldmerges the fields of its value branches (the THEN branches and the ELSE) withcoerced_fields_metadata, the helperCOALESCE/NVL2use since fix: propagate field metadata through coalesce, nvl, and nvl2 #24865.CaseExpr::return_fielddoes the same with its branches'return_field.CaseExpr::try_newno 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.CASE WHEN false THEN x END, the literal takes the field metadata of the CASE.So CASE follows the same rules as
COALESCE:NULL, does not block another branch's metadata.The operand of a simple
CASEand the WHEN conditions do not contribute.Differences from #26145 + #26158
About 150 lines of non-test code including the simplifier changes, against about 700 there.
NOTpasses its input's metadata through while logicalNOTdrops it (Expr::NotandPhysicalExpr::NotExprhave different metadata propagation rules #27155), and an alias's metadata can be lost when the physical expression is created (Alias metadata may be lost when the physical expression is created #27156). WithNOT, this can be a planning error thatmaindoesn't raise: when the other branch has a different extension name, the physical CASE can't merge the two. Those mismatches are inNOTandAliasand are not addressed here.Performance
These are local Criterion runs of the benchmarks from #26145/#26158 on a laptop, comparing this PR with
main, so treat them as indicative.CASE evaluation (the existing
case_whenbenchmarks) is unchanged.Physical planning is unchanged.
return_fieldis slower, because the merged metadata is rebuilt on every call: 954 ns vs 274 ns for a column CASE with 4 metadata keys.ScalarFunctionExprasks 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.A nested CASE branch only computes its metadata, not its name and nullability. For a CASE nested 32 deep in its ELSE, the extra cost is 7% of the CASE's evaluation time for 8,192 rows (measured by @jayzhan211).
Logical
to_fieldfor a CASE without any metadata takes 2–3x as long. It reads the type from the branch fields, so a CASE nested in function calls is planned about as fast as onmain.These can be addressed later with fast paths: reuse the metadata when only one branch has any, and 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.sltcases, checking the metadata at execution witharrow_metadata. Most use the built-inJSONextension type:coalesce,nvl,ifnullandnvl2, including a non-nullable first argument;Each metadata check fails without this change. Each
arrow_metadatacheck of a CASE that is still a CASE at execution also fails without the physicalCaseExpr::return_field. Each query uses a given CASE only once: a repeated expression is computed once in a lower projection, whose output field comes from the logical plan, so it would not test the physical field.A unit test in the physical planner checks that the physical CASE reports the same field metadata as the logical CASE. It covers extension types, merged keys, simple and nested CASE, a typed NULL ELSE with metadata, and branches cast to a common type.
A unit test for the CASE simplifications. It also covers the rewrite of a Boolean CASE with literal THENs to
NOT(CASE ...): in SQL, physicalNOTpasses the metadata through (Expr::NotandPhysicalExpr::NotExprhave different metadata propagation rules #27155), soarrow_metadatacan't see it being lost.The existing tests: the
datafusion-expr,datafusion-physical-expranddatafusion-optimizerlib tests pass.The full sqllogictest suite and the
core_integrationsqltests show no new failures. My checkout doesn't have the test-data submodules, so the files that need that data fail the same way onmain.The
metadata.sltfrom Preserve field metadata through physical CASE expressions #26158, which includes the cases from test: Add sqllogictests for field metadata through coalesce, nvl2 and CASE #26895, also passes with this change. I checked that locally with the current head of Preserve field metadata through physical CASE expressions #26158.Are there any user-facing changes?
CASEresults keep their field metadata, including extension types, also when the optimizer simplifies the CASE. So doCOALESCE,NVL,IFNULLandNVL2at execution.CASEbranches with differing extension names are now a planning error, as they already are forCOALESCE.No public API changes.