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

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)
}