Skip to content

fix: preserve field metadata through CASE - #27143

Open
james-willis wants to merge 5 commits into
apache:mainfrom
james-willis:case-field-metadata
Open

james-willis wants to merge 5 commits into
apache:mainfrom
james-willis:case-field-metadata

Conversation

@james-willis

@james-willis james-willis commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

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

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 three 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;
  • the simplifier, whose CASE rewrites can drop the branch that carries the metadata: CASE WHEN false THEN x END becomes 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?

  • 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 fix: propagate field metadata through coalesce, nvl, and nvl2 #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.
  • CASE simplifications keep the CASE's metadata.
    • When a constant WHEN leaves a literal, such as the NULL from CASE WHEN false THEN x END, the literal takes the field metadata of the CASE.
    • When it leaves another expression with different metadata, the CASE is not simplified.
    • A Boolean CASE with metadata on a branch is not rewritten to AND/OR.

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 150 lines of non-test code including the simplifier changes, against about 700 there.

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_when benchmarks) is unchanged.

  • Physical planning is unchanged.

  • return_field is slower, because the merged metadata is rebuilt on every call: 954 ns vs 274 ns for a column CASE with 4 metadata keys.

    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.

    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_field for 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 on main.

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.slt cases, checking the metadata at execution with arrow_metadata. Most use the built-in JSON extension type:

    • 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, including when no row reaches a THEN branch;
    • branches coerced to a common type: the cast drops the extension type and keeps other metadata;
    • nested CASE;
    • the CASE rewrites of coalesce, nvl, ifnull and nvl2, including a non-nullable first argument;
    • simplification of a constant WHEN, to a NULL or by dropping a branch;
    • a Boolean CASE that would otherwise be rewritten to AND/OR.

    Each metadata check fails without this change. Each arrow_metadata check of a CASE that is still a CASE at execution also fails without the physical CaseExpr::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, physical NOT passes the metadata through (Expr::Not and PhysicalExpr::NotExpr have different metadata propagation rules #27155), so arrow_metadata can't see it being lost.

  • 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 from 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?

  • CASE results keep their field metadata, including extension types, also when the optimizer simplifies the CASE. 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.

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.
@github-actions github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates sqllogictest SQL Logic Tests (.slt) labels Nov 5, 2026
@james-willis
james-willis marked this pull request as ready for review November 5, 2026 19:52
@james-willis

Copy link
Copy Markdown
Contributor Author

@paleolimbot @alamb @osipovartem This is my alternative solution for preserving metadata in those cases.

@paleolimbot paleolimbot left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

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.
@github-actions github-actions Bot added the optimizer Optimizer rules label Nov 6, 2026
@james-willis

Copy link
Copy Markdown
Contributor Author

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.

test_case_lowering_preserves_field_metadata in physical-expr/src/planner.rs (8ef3402) now compares Expr::Case::to_field with the physical return_field for column, literal, cast and nested CASE branches, including a typed NULL ELSE with metadata.

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 metadata.slt case checks that CASE ... THEN cast(name AS JSON) ELSE arrow_cast(l_name, 'LargeUtf8') END is LargeUtf8 without an extension type.

@jayzhan211 jayzhan211 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.

@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> {

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.

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;

@james-willis james-willis Nov 10, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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.
@james-willis

Copy link
Copy Markdown
Contributor Author

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 ELSE NULL and CASE WHEN true THEN 'x' ELSE g END still fold instead of keeping the whole CASE. Anything else only simplifies if its metadata matches the CASE's. Your queries, plus an nvl2 one, are in metadata.slt.

You're right that a simplifier without a schema now errors on CASE WHEN true THEN 1 ELSE col END. I left it since CASE WHEN false THEN col END and CASE WHEN c THEN col END already error without a schema today.

AND/OR lowering (~L1555, L1613): added your case_branch_has_metadata guard to both arms. The NOT(CASE ...) one can't be shown from SQL yet because physical NOT passes the metadata through (#27155), so simplify_expr_case_keeps_field_metadata covers it instead. Your CASE WHEN p THEN f END query is in metadata.slt.

Docs: added the CASE entry to to_field.

@jayzhan211 jayzhan211 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.

@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_field grows quadratically with nesting depth.

Comment thread datafusion/expr/src/expr_schema.rs Resolved
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.
@james-willis

Copy link
Copy Markdown
Contributor Author

thanks for all the help @jayzhan211 I made the tweaks from your review.

@jayzhan211 jayzhan211 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.

@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)?,

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.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

logical-expr Logical plan and expressions optimizer Optimizer rules physical-expr Changes to the physical-expr crates sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CASE expressions drop field metadata, silently stripping Arrow extension types and breaking metadata-aware UDFs

3 participants