fix: preserve field metadata for simple CASE branches

datafusion/common/src/dfschema.rs

@@ -1244,6 +1244,14 @@ pub trait ExprSchema: std::fmt::Debug {
 
     // Return the column's field
     fn field_from_column(&self, col: &Column) -> Result<&FieldRef>;
+
+    /// Return the field and column nullability together.
+    ///
+    /// The default preserves custom implementations of [`Self::nullable`].
+    fn field_and_nullable(&self, col: &Column) -> Result<(&FieldRef, bool)> {
+        let field = self.field_from_column(col)?;
+        Ok((field, self.nullable(col)?))
+    }
 }
 
 // Implement `ExprSchema` for `Arc<DFSchema>`
@@ -1267,6 +1275,10 @@ impl<P: AsRef<DFSchema> + std::fmt::Debug> ExprSchema for P {
     fn field_from_column(&self, col: &Column) -> Result<&FieldRef> {
         self.as_ref().field_from_column(col)
     }
+
+    fn field_and_nullable(&self, col: &Column) -> Result<(&FieldRef, bool)> {
+        self.as_ref().field_and_nullable(col)
+    }
 }
 
 impl ExprSchema for DFSchema {
@@ -1276,6 +1288,11 @@ impl ExprSchema for DFSchema {
             None => self.field_with_unqualified_name(&col.name),
         }
     }
+
+    fn field_and_nullable(&self, col: &Column) -> Result<(&FieldRef, bool)> {
+        let field = self.field_from_column(col)?;
+        Ok((field, field.is_nullable()))
+    }
 }
 
 /// DataFusion-specific extensions to [`Schema`].

datafusion/core/tests/sql/select.rs

@@ -18,6 +18,7 @@
 use std::collections::HashMap;
 
 use super::*;
+use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY;
 use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata};
 use insta::assert_snapshot;
 
@@ -387,6 +388,81 @@ async fn test_query_parameters_with_metadata() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_sql_case_field_metadata() -> Result<()> {
+    let ctx = SessionContext::new();
+    let metadata = HashMap::from([
+        ("structured_type".to_string(), "ARRAY(NUMBER)".to_string()),
+        (
+            EXTENSION_TYPE_NAME_KEY.to_string(),
+            "example.case".to_string(),
+        ),
+    ]);
+    let other_metadata = HashMap::from([
+        ("structured_type".to_string(), "OBJECT".to_string()),
+        (
+            EXTENSION_TYPE_NAME_KEY.to_string(),
+            "example.other".to_string(),
+        ),
+    ]);
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("flag", DataType::Boolean, false),
+        Field::new("a", DataType::Int32, false).with_metadata(metadata.clone()),
+        Field::new("b", DataType::Int32, false).with_metadata(metadata.clone()),
+        Field::new("c", DataType::Int32, false).with_metadata(other_metadata),
+    ]));
+    ctx.register_batch(
+        "t",
+        RecordBatch::try_new(
+            schema,
+            vec![
+                Arc::new(BooleanArray::from(vec![true, false])),
+                Arc::new(Int32Array::from(vec![1, 2])),
+                Arc::new(Int32Array::from(vec![3, 4])),
+                Arc::new(Int32Array::from(vec![5, 6])),
+            ],
+        )?,
+    )?;
+
+    for sql in [
+        "SELECT CASE WHEN flag THEN a ELSE b END AS value FROM t",
+        "SELECT CASE flag WHEN TRUE THEN a ELSE b END AS value FROM t",
+        "SELECT CASE WHEN flag THEN a ELSE NULL END AS value FROM t",
+        "SELECT CASE WHEN flag THEN NULL ELSE a END AS value FROM t",
+    ] {
+        let df = ctx.sql(sql).await?;
+        let optimized = ctx.state().optimize(df.logical_plan())?;
+        for field in [df.schema().field(0), optimized.schema().field(0)] {
+            assert_eq!(field.metadata(), &metadata);
+        }
+        if sql == "SELECT CASE WHEN flag THEN a ELSE b END AS value FROM t" {
+            let batches = df.collect().await?;
+            datafusion::assert_batches_eq!(
+                [
+                    "+-------+",
+                    "| value |",
+                    "+-------+",
+                    "| 1     |",
+                    "| 4     |",
+                    "+-------+"
+                ],
+                &batches
+            );
+        }
+    }
+
+    let conflicting = ctx
+        .sql("SELECT CASE WHEN flag THEN a ELSE c END AS value FROM t")
+        .await;
+    let error = match conflicting {
+        Ok(df) => ctx.state().optimize(df.logical_plan()).unwrap_err(),
+        Err(error) => error,
+    };
+    assert!(error.to_string().contains("example.case"), "{error}");
+    assert!(error.to_string().contains("example.other"), "{error}");
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_limit_offset_parameters_named() -> Result<()> {
     let ctx = SessionContext::new();

datafusion/expr/Cargo.toml

@@ -78,3 +78,7 @@ insta = { workspace = true }
 [[bench]]
 name = "normalize_columns"
 harness = false
+
+[[bench]]
+name = "case_field_metadata"
+harness = false

datafusion/expr/benches/case_field_metadata.rs

@@ -0,0 +1,81 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::HashMap;
+use std::hint::black_box;
+
+use arrow::datatypes::{DataType, Field};
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_common::DFSchema;
+use datafusion_expr::{Expr, ExprSchemable, col, lit, when};
+
+fn literal_case(width: usize) -> Expr {
+    let mut builder = when(lit(true), lit(0_i32));
+    for i in 1..width {
+        builder = builder.when(lit(false), lit(i32::try_from(i).unwrap()));
+    }
+    builder.end().unwrap()
+}
+
+fn column_case(width: usize, marked_last: bool) -> (DFSchema, Expr) {
+    let fields = (0..width)
+        .map(|i| {
+            let field = Field::new(format!("c{i}"), DataType::Int32, false);
+            if marked_last && i == width - 1 {
+                field.with_metadata(HashMap::from([(
+                    "source".to_owned(),
+                    "marked".to_owned(),
+                )]))
+            } else {
+                field
+            }
+        })
+        .collect::<Vec<_>>();
+    let schema =
+        DFSchema::from_unqualified_fields(fields.into(), HashMap::new()).unwrap();
+    let mut builder = when(lit(true), col("c0"));
+    for i in 1..width {
+        builder = builder.when(lit(false), col(format!("c{i}")));
+    }
+    (schema, builder.end().unwrap())
+}
+
+fn case_field_metadata(c: &mut Criterion) {
+    let mut group = c.benchmark_group("case_field_metadata");
+    for width in [10, 100, 1000] {
+        let literal = literal_case(width);
+        let empty_schema = DFSchema::empty();
+        group.bench_with_input(
+            BenchmarkId::new("plain_literals", width),
+            &width,
+            |b, _| {
+                b.iter(|| black_box(literal.to_field(black_box(&empty_schema)).unwrap()))
+            },
+        );
+
+        for (name, marked_last) in [("plain_columns", false), ("marked_last", true)] {
+            let (schema, expression) = column_case(width, marked_last);
+            group.bench_with_input(BenchmarkId::new(name, width), &width, |b, _| {
+                b.iter(|| black_box(expression.to_field(black_box(&schema)).unwrap()))
+            });
+        }
+    }
+    group.finish();
+}
+
+criterion_group!(benches, case_field_metadata);
+criterion_main!(benches);