fix: preserve CASE WHEN struct field names from the first THEN branch

native/spark-expr/src/conditional_funcs/case_when.rs

@@ -31,7 +31,6 @@ use datafusion::common::{
     exec_err, internal_datafusion_err, internal_err, DataFusionError, Result, ScalarValue,
 };
 use datafusion::logical_expr::type_coercion::binary::type_union_coercion;
-use datafusion::logical_expr::type_coercion::other::get_coerce_type_for_case_expression;
 use datafusion::logical_expr::{ColumnarValue, Operator};
 use datafusion::physical_expr::expressions::{
     BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr,
@@ -60,8 +59,17 @@ pub fn create_case_when(
         .transpose()?
         .unwrap_or(DataType::Null);
 
-    let Some(coerce_type) = get_coerce_type_for_case_expression(&then_types, Some(&else_type))
-    else {
+    // Spark merges THEN types from left to right, with ELSE last, retaining the first THEN's
+    // struct field names. DataFusion's CASE coercion starts with ELSE and can match fields by name.
+    let coerce_type = then_types.split_first().and_then(|(first, remaining)| {
+        remaining
+            .iter()
+            .chain([&else_type])
+            .try_fold(first.clone(), |common, next| {
+                merge_branch_types(&common, next)
+            })
+    });
+    let Some(coerce_type) = coerce_type else {
         return Ok(Arc::new(CaseWhenExpr::try_new(when_then, else_expr)?));
     };
     let when_then = when_then
@@ -86,7 +94,7 @@ pub fn create_if_expr(
 ) -> Result<Arc<dyn PhysicalExpr>> {
     let true_type = true_expr.data_type(input_schema)?;
     let false_type = false_expr.data_type(input_schema)?;
-    let Some(common_type) = if_common_type(&true_type, &false_type) else {
+    let Some(common_type) = merge_branch_types(&true_type, &false_type) else {
         return Ok(Arc::new(IfExpr::new(if_expr, true_expr, false_expr)));
     };
     Ok(Arc::new(IfExpr::new(
@@ -96,34 +104,30 @@ pub fn create_if_expr(
     )))
 }
 
-/// Reconciles Spark IF branches positionally, retaining THEN names and merging nullability.
+/// Merges conditional branch types positionally, retaining left names and merging nullability.
 /// Spark has already coerced the branches to the same SQL type. DataFusion's struct union may
 /// instead match by name, pairing different positions when names differ only in case.
-fn if_common_type(then_type: &DataType, else_type: &DataType) -> Option<DataType> {
+fn merge_branch_types(left_type: &DataType, right_type: &DataType) -> Option<DataType> {
     use arrow::datatypes::FieldRef;
 
-    fn field(then_field: &FieldRef, else_field: &FieldRef) -> Option<FieldRef> {
+    fn field(left: &FieldRef, right: &FieldRef) -> Option<FieldRef> {
         Some(Arc::new(
-            then_field
-                .as_ref()
+            left.as_ref()
                 .clone()
-                .with_data_type(if_common_type(
-                    then_field.data_type(),
-                    else_field.data_type(),
-                )?)
-                .with_nullable(then_field.is_nullable() || else_field.is_nullable()),
+                .with_data_type(merge_branch_types(left.data_type(), right.data_type())?)
+                .with_nullable(left.is_nullable() || right.is_nullable()),
         ))
     }
 
-    match (then_type, else_type) {
-        (DataType::Struct(then_fields), DataType::Struct(else_fields)) => {
-            if then_fields.len() != else_fields.len() {
+    match (left_type, right_type) {
+        (DataType::Struct(left_fields), DataType::Struct(right_fields)) => {
+            if left_fields.len() != right_fields.len() {
                 return None;
             }
             Some(DataType::Struct(
-                then_fields
+                left_fields
                     .iter()
-                    .zip(else_fields)
+                    .zip(right_fields)
                     .map(|(t, e)| field(t, e))
                     .collect::<Option<Vec<_>>>()?
                     .into(),
@@ -133,7 +137,7 @@ fn if_common_type(then_type: &DataType, else_type: &DataType) -> Option<DataType
         (DataType::Map(t, t_sorted), DataType::Map(e, e_sorted)) => {
             Some(DataType::Map(field(t, e)?, *t_sorted && *e_sorted))
         }
-        _ => type_union_coercion(then_type, else_type),
+        _ => type_union_coercion(left_type, right_type),
     }
 }
 
@@ -1338,11 +1342,17 @@ mod tests {
 
     /// CASE branches that share a Spark timestamp type but carry different Arrow timezone labels
     /// are reconciled by relabelling them, which used to panic because the cast had no timezone.
+    /// Like IF, CASE retains the first THEN's label when present.
     #[test]
     fn case_reconciles_timestamp_timezone_labels() {
         let utc = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()));
         for then_label in [Some("Etc/UTC"), None] {
             let then_type = DataType::Timestamp(TimeUnit::Microsecond, then_label.map(Into::into));
+            let expected_type = if then_label.is_some() {
+                then_type.clone()
+            } else {
+                utc.clone()
+            };
             let schema = Schema::new(vec![
                 Field::new("b", DataType::Boolean, true),
                 Field::new("t", then_type, true),
@@ -1362,16 +1372,17 @@ mod tests {
             let c = |name: &str| col(name, &schema).unwrap();
             let case = create_case_when(vec![(c("b"), c("t"))], Some(c("e")), &schema).unwrap();
             let result = case.evaluate(&batch).unwrap().into_array(2).unwrap();
-            assert_eq!(result.data_type(), &utc, "{then_label:?}");
+            assert_eq!(case.data_type(&schema).unwrap(), expected_type);
+            assert_eq!(result.data_type(), &expected_type, "{then_label:?}");
             let values = result.as_primitive::<TimestampMicrosecondType>();
             assert_eq!(values.values().to_vec(), vec![1, 20], "{then_label:?}");
         }
     }
 
-    /// Evaluates an IF over `batch`, whose first row takes the THEN branch and second row the ELSE
-    /// branch, and over each row alone, and checks that every result has the type the IF reports.
+    /// Evaluates a conditional over `batch`, whose first row takes THEN and second row ELSE,
+    /// and over each row alone, checking that every result has the type the expression reports.
     /// Returns the result for the whole batch.
-    fn evaluate_if(expr: &Arc<dyn PhysicalExpr>, batch: &RecordBatch) -> ArrayRef {
+    fn evaluate_conditional(expr: &Arc<dyn PhysicalExpr>, batch: &RecordBatch) -> ArrayRef {
         let data_type = expr.data_type(&batch.schema()).unwrap();
         // Both branches, only THEN, and only ELSE, whose array `IfExpr` returns unchanged
         for (offset, len) in [(0, 2), (0, 1), (1, 1)] {
@@ -1389,7 +1400,7 @@ mod tests {
     /// Spark matches struct fields by position even when their case-distinct names are reordered.
     /// Name-based union pairs the first INT with the second DOUBLE and silently widens it.
     #[test]
-    fn if_reconciles_case_variant_fields_positionally() {
+    fn conditional_branches_reconcile_case_variant_fields_positionally() {
         use arrow::array::Float64Array;
 
         let then_fields: arrow::datatypes::Fields = vec![
@@ -1436,12 +1447,6 @@ mod tests {
         )
         .unwrap();
         let c = |name: &str| col(name, &schema).unwrap();
-        let expr = create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap();
-        assert_eq!(
-            expr.data_type(&schema).unwrap(),
-            DataType::Struct(expected_fields.clone())
-        );
-        let result = evaluate_if(&expr, &batch);
         let expected = StructArray::new(
             expected_fields,
             vec![
@@ -1450,7 +1455,14 @@ mod tests {
             ],
             None,
         );
-        assert_eq!(result.as_ref(), &expected);
+        for expr in [
+            create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap(),
+            create_case_when(vec![(c("b"), c("t"))], Some(c("e")), &schema).unwrap(),
+        ] {
+            assert_eq!(expr.data_type(&schema).unwrap(), *expected.data_type());
+            let result = evaluate_conditional(&expr, &batch);
+            assert_eq!(result.as_ref(), &expected, "{expr}");
+        }
     }
 
     /// IF branches that share a Spark timestamp type but carry different Arrow timezone labels are
@@ -1493,17 +1505,17 @@ mod tests {
                 timestamp(label),
                 "{labels}"
             );
-            let result = evaluate_if(&expr, &batch);
+            let result = evaluate_conditional(&expr, &batch);
             let values = result.as_primitive::<TimestampMicrosecondType>();
             assert_eq!(values.values().to_vec(), vec![1, 20], "{labels}");
         }
     }
 
-    /// IF branches whose struct field differs only in whether it can be NULL, or in the case of its
+    /// Conditional branches whose struct field differs only in nullability or in the case of its
     /// name, which Spark adds no cast for. The result's field can be NULL if either branch's can,
-    /// and has the THEN branch's name, as in Spark's `If.dataType`.
+    /// and has the THEN branch's name, as in Spark's `If.dataType` and `CaseWhen.dataType`.
     #[test]
-    fn if_reconciles_struct_field_nullability_and_names() {
+    fn conditional_branches_reconcile_struct_field_nullability_and_names() {
         let field =
             |name: &str, nullable: bool| Arc::new(Field::new(name, DataType::Int32, nullable));
         let struct_type = |field: &FieldRef| DataType::Struct(vec![Arc::clone(field)].into());
@@ -1542,16 +1554,20 @@ mod tests {
             )
             .unwrap();
             let c = |name: &str| col(name, &schema).unwrap();
-            let expr = create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap();
             let fields = format!("{then_field:?} {else_field:?}");
-            assert_eq!(
-                expr.data_type(&schema).unwrap(),
-                struct_type(&result_field),
-                "{fields}"
-            );
-            let result = evaluate_if(&expr, &batch);
             let expected = column(&result_field, vec![Some(1), else_values[1]]);
-            assert_eq!(result.as_ref(), expected.as_ref(), "{fields}");
+            for expr in [
+                create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap(),
+                create_case_when(vec![(c("b"), c("t"))], Some(c("e")), &schema).unwrap(),
+            ] {
+                assert_eq!(
+                    expr.data_type(&schema).unwrap(),
+                    struct_type(&result_field),
+                    "{fields} {expr}"
+                );
+                let result = evaluate_conditional(&expr, &batch);
+                assert_eq!(result.as_ref(), expected.as_ref(), "{fields} {expr}");
+            }
         }
     }
 

spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql

@@ -0,0 +1,67 @@
+-- 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.
+
+-- Spark preserves the first THEN's struct field names when branches differ only in case.
+-- Native to_json exposes Arrow field names that the row comparison would otherwise ignore.
+-- https://github.com/apache/datafusion-comet/issues/6482
+-- Config: spark.sql.caseSensitive=false
+-- Config: spark.comet.expression.StructsToJson.allowIncompatible=true
+
+statement
+CREATE TABLE test_case_field_names(p boolean, q boolean, i int, d double) USING parquet
+
+statement
+INSERT INTO test_case_field_names VALUES
+  (true, false, 1, 1.5), (false, true, 2, 2.5),
+  (NULL, false, 3, 3.5), (false, NULL, NULL, NULL)
+
+-- The result keeps the THEN name for both selected branches, in either name order.
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN p THEN named_struct('x', i) ELSE named_struct('X', i) END),
+  to_json(CASE WHEN p THEN named_struct('X', i) ELSE named_struct('x', i) END)
+FROM test_case_field_names
+
+-- Preserve the first THEN's names at both levels of a nested struct, in either branch order.
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN p THEN named_struct('a', named_struct('x', i))
+    ELSE named_struct('A', named_struct('X', i)) END),
+  to_json(CASE WHEN p THEN named_struct('A', named_struct('X', i))
+    ELSE named_struct('a', named_struct('x', i)) END)
+FROM test_case_field_names
+
+-- Later THEN branches and an omitted ELSE must also preserve the first THEN's name.
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN p THEN named_struct('x', i)
+    WHEN q THEN named_struct('X', i + 10) ELSE named_struct('X', i + 20) END),
+  to_json(CASE WHEN p THEN named_struct('x', i)
+    WHEN q THEN named_struct('X', i + 10) END)
+FROM test_case_field_names
+
+-- A typed NULL in the first THEN still determines the result's struct field name.
+query expect_native(casewhen)
+SELECT to_json(CASE WHEN p THEN CAST(NULL AS STRUCT<x: INT>)
+  WHEN q THEN named_struct('x', i) ELSE named_struct('X', i) END)
+FROM test_case_field_names
+
+-- Fields align by position. Matching these names would incorrectly combine INT with DOUBLE.
+query expect_native(casewhen)
+SELECT to_json(CASE WHEN p THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))
+  ELSE named_struct('X', 0, 'x', d) END)
+FROM test_case_field_names

spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql

@@ -0,0 +1,36 @@
+-- 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.
+
+-- COALESCE lowers to CASE WHEN with its last argument as ELSE, but must retain the first
+-- argument's struct field names. Native to_json exposes names that row comparison ignores.
+-- https://github.com/apache/datafusion-comet/issues/6482
+-- Config: spark.sql.caseSensitive=false
+-- Config: spark.comet.expression.StructsToJson.allowIncompatible=true
+
+statement
+CREATE TABLE test_coalesce_field_names(p boolean, q boolean, i int) USING parquet
+
+statement
+INSERT INTO test_coalesce_field_names VALUES
+  (true, false, 1), (false, true, 2), (NULL, false, 3), (false, NULL, NULL)
+
+query expect_native(coalesce)
+SELECT to_json(coalesce(
+  CASE WHEN p THEN named_struct('x', i) END,
+  CASE WHEN q THEN named_struct('X', i + 100) END,
+  named_struct('X', 7)))
+FROM test_coalesce_field_names