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

[james-willis]

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 Fields 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:

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)

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<DataType> {
        exec_err!("return_field_from_args should be called")
    }
    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
        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<ColumnarValue> {
        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<String, String> =
        [(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.

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:

• Preserve field metadata across expressions in logical plans #6886 (added to_field metadata propagation for type-preserving expressions in 2023; CASE was left in the empty-metadata catch-all)
• Support Extension Types / User Defined Types in DataFusion #12644 (extension types epic)
• Propagation of metadata through casts can strips extension type from the destination field and can result in invalid extension types #22079 (cast metadata propagation)
• make_array / array_agg drop inner-Field metadata when constructing List<T> #21982 (make_array / array_agg inner-field metadata)
• NTH_VALUE, FIRST_VALUE, and LAST_VALUE drop metadata #22108 (NTH_VALUE/FIRST_VALUE/LAST_VALUE metadata, fixed)
• Missing field metadata for projections during physical planning #23529 (projection metadata during physical planning, fixed)

[james-willis]

take