diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 7dec89c61c0..32577773110 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/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> { 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 { +fn merge_branch_types(left_type: &DataType, right_type: &DataType) -> Option { use arrow::datatypes::FieldRef; - fn field(then_field: &FieldRef, else_field: &FieldRef) -> Option { + fn field(left: &FieldRef, right: &FieldRef) -> Option { 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::>>()? .into(), @@ -133,7 +137,7 @@ fn if_common_type(then_type: &DataType, else_type: &DataType) -> Option { 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::(); 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, batch: &RecordBatch) -> ArrayRef { + fn evaluate_conditional(expr: &Arc, 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::(); 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}"); + } } } diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql b/spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql new file mode 100644 index 00000000000..f8106cb2a25 --- /dev/null +++ b/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) + 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 diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql new file mode 100644 index 00000000000..45596f3cc26 --- /dev/null +++ b/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