apache / apache/datafusion

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

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

Description

## Describe the bug

Field metadata — including the `ARROW:extension:name` / `ARROW:extension:metadata` keys that encode Arrow extension types — is not propagated through `CASE` expressions. `COALESCE` is affected identically because it simplifies into `CASE` (and its own `return_field_from_args` also drops metadata).

Two user-visible symptoms:

1. **Silent extension-type stripping.** Projecting a `CASE` over an extension-typed column returns a plain storage-typed column with empty metadata and no error. Downstream consumers (FFI, IPC, Parquet writers, other engines) see e.g. plain `Binary` where the input was `geoarrow.wkb`.

2. **Metadata-aware UDFs fail to plan.** A `ScalarUDFImpl` that dispatches on its argument `Field`s via `return_field_from_args` (the mechanism added for exactly this use case) cannot see the extension type of a `CASE` argument, so type resolution fails even though every branch of the `CASE` is properly typed.

## To Reproduce

Run the program below against `datafusion = "55.0.0"` (the drops are also present on `main` as of 2026-09-01). It registers a table with an extension-typed `Binary` column `g` and a UDF `ext_only(x)` that requires its argument field to carry the extension metadata:

```sql
SELECT g FROM t
-> output field metadata: {"ARROW:extension:name": "geoarrow.wkb"}

SELECT CASE WHEN k = 1 THEN g END AS g FROM t
-> output field metadata: {}

SELECT ext_only(g) FROM t
-> OK

SELECT ext_only(CASE WHEN k = 1 THEN g END) FROM t
-> PLAN ERR: ext_only(): argument lost its extension type; field = Field { name: "CASE WHEN t.k = Int64(1) THEN t.g END", data_type: Binary, nullable: true }

SELECT ext_only(CASE WHEN 1 = 0 THEN g END) FROM t
-> PLAN ERR: ext_only(): argument lost its extension type; field = Field { name: "CASE WHEN Int64(1) = Int64(0) THEN t.g END", data_type: Binary, nullable: true }

SELECT ext_only(COALESCE(g, g)) FROM t
-> PLAN ERR: ext_only(): argument lost its extension type; field = Field { name: "coalesce(t.g,t.g)", data_type: Binary, nullable: true }
```

Full repro (main.rs, no dependencies beyond datafusion and tokio)

```rust
use std::collections::HashMap;
use std::sync::Arc;

use datafusion::arrow::array::{BinaryArray, Int32Array, RecordBatch};
use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema};
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::logical_expr::{
ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl,
Signature, Volatility,
};
use datafusion::prelude::*;

const EXT_KEY: &str = "ARROW:extension:name";
const EXT_NAME: &str = "geoarrow.wkb";

/// A UDF that requires its argument's Field to carry the extension-type
/// metadata, the way geoarrow/sedona-db/uuid/json extension-aware UDFs do
/// via `return_field_from_args` / `invoke_with_args`.
#[derive(Debug, PartialEq, Eq, Hash)]
struct ExtOnly {
signature: Signature,
}

impl ScalarUDFImpl for ExtOnly {
fn name(&self) -> &str {
"ext_only"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _args: &[DataType]) -> Result {
exec_err!("return_field_from_args should be called")
}
fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result {
let f = &args.arg_fields[0];
if f.metadata().get(EXT_KEY).map(String::as_str) != Some(EXT_NAME) {
return exec_err!(
"ext_only(): argument lost its extension type; field = {f:?}"
);
}
Ok(Arc::new(Field::new(self.name(), DataType::Boolean, true)))
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result {
Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(
args.arg_fields[0].metadata().get(EXT_KEY).map(String::as_str)
== Some(EXT_NAME),
))))
}
}

#[tokio::main]
async fn main() -> Result<()> {
let ext_metadata: HashMap =
[(EXT_KEY.to_string(), EXT_NAME.to_string())].into();
let schema = Arc::new(Schema::new(vec![
Field::new("g", DataType::Binary, true).with_metadata(ext_metadata),
Field::new("k", DataType::Int32, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(BinaryArray::from_opt_vec(vec![Some(b"\x00".as_ref())])),
Arc::new(Int32Array::from(vec![1])),
],
)?;

let ctx = SessionContext::new();
ctx.register_batch("t", batch)?;
ctx.register_udf(ScalarUDF::from(ExtOnly {
signature: Signature::any(1, Volatility::Immutable),
}));

// 1. Silent stripping: projecting CASE drops the extension metadata.
for sql in [
"SELECT g FROM t",
"SELECT CASE WHEN k = 1 THEN g END AS g FROM t",
] {
let df = ctx.sql(sql).await?;
let batches = df.collect().await?;
let field = batches[0].schema().field(0).clone();
println!("{sql}\n -> output field metadata: {:?}\n", field.metadata());
}

// 2. Metadata-aware UDFs fail to plan over CASE / COALESCE.
for sql in [
"SELECT ext_only(g) FROM t",
"SELECT ext_only(CASE WHEN k = 1 THEN g END) FROM t",
"SELECT ext_only(CASE WHEN 1 = 0 THEN g END) FROM t",
"SELECT ext_only(COALESCE(g, g)) FROM t",
] {
let out = match ctx.sql(sql).await {
Ok(df) => match df.collect().await {
Ok(b) => format!("OK: {:?}", b[0].column(0)),
Err(e) => format!("EXEC ERR: {}", first_line(&e.to_string())),
},
Err(e) => format!("PLAN ERR: {}", first_line(&e.to_string())),
};
println!("{sql}\n -> {out}\n");
}
Ok(())
}

fn first_line(s: &str) -> &str {
s.lines().next().unwrap_or(s)
}
```

## Expected behavior

`CASE` (and `COALESCE`) preserve the field metadata of their branches: when the branches that determine the result type agree on metadata, the output field carries it. `SELECT CASE WHEN k = 1 THEN g END` has the same extension type as `g`, and `return_field_from_args`-based UDFs see extension-typed arg fields for `CASE` arguments.

## Additional context

Real-world instance: Apache SedonaDB models geometry (`geoarrow.wkb`) and rasters as Arrow extension types and resolves kernels through `return_field_from_args` / `invoke_with_args` arg fields. Any conditional spatial SQL fails, e.g.

```sql
SELECT ST_X(CASE WHEN k = 1 THEN geom END) FROM t
-- Error during planning: st_x(binary): No kernel matching arguments
SELECT RS_BandNoDataValue(CASE WHEN 1 = 0 THEN rast END, 1) FROM r
-- rs_bandnodatavalue(struct, int64): No kernel matching arguments
```

while the same SQL works on engines whose type systems carry the user-defined type through `CASE` (e.g. Sedona on Spark). Cataloged as xfails in apache/sedona-db#1203.

Note the metadata is dropped independently at several layers (logical `to_field`, physical `CaseExpr`, the simplifier's `CASE WHEN false THEN a END → NULL` fold, and `coalesce`'s return field), so a fix at one layer alone doesn't resolve the end-to-end behavior. Happy to work on a PR.

Related:

- #6886 (added `to_field` metadata propagation for type-preserving expressions in 2023; `CASE` was left in the empty-metadata catch-all)
- #12644 (extension types epic)
- #22079 (cast metadata propagation)
- #21982 (`make_array` / `array_agg` inner-field metadata)
- #22108 (NTH_VALUE/FIRST_VALUE/LAST_VALUE metadata, fixed)
- #23529 (projection metadata during physical planning, fixed)

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.