diff --git a/datafusion/expr-common/src/type_coercion/binary.rs b/datafusion/expr-common/src/type_coercion/binary.rs index 3195218ea28e9..f2d273ac40b17 100644 --- a/datafusion/expr-common/src/type_coercion/binary.rs +++ b/datafusion/expr-common/src/type_coercion/binary.rs @@ -614,7 +614,11 @@ pub fn try_type_union_resolution_with_struct( let mut keys_string: Option = None; for data_type in data_types { if let DataType::Struct(fields) = data_type { - let keys = fields.iter().map(|f| f.name().to_owned()).join(","); + let keys = fields + .iter() + .map(|f| f.name().to_owned()) + .sorted() + .join(","); if let Some(ref k) = keys_string { if *k != keys { return exec_err!("Expect same keys for struct type but got mismatched pair {} and {}", *k, keys); @@ -671,7 +675,6 @@ pub fn try_type_union_resolution_with_struct( } final_struct_types.push(DataType::Struct(new_fields.into())) } - Ok(final_struct_types) } diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 446ae94108b1c..3d09afb1b45c8 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -58,6 +58,7 @@ use datafusion_common::{ FunctionalDependencies, ParamValues, Result, ScalarValue, TableReference, UnnestOptions, }; +use datafusion_expr_common::type_coercion::binary::try_type_union_resolution_with_struct; use indexmap::IndexSet; // backwards compatibility @@ -2695,7 +2696,15 @@ impl Union { // TODO apply type coercion here, or document why it's better to defer // temporarily use the data type from the left input and later rely on the analyzer to // coerce the two schemas into a common one. - first_field.data_type() + let first_type = first_field.data_type(); + if matches!(first_type, DataType::Struct(_)) { + let types: Vec = + fields.iter().map(|f| f.data_type().clone()).collect(); + &try_type_union_resolution_with_struct(&types) + .map(|types| types[0].clone())? + } else { + first_type + } } else { fields.iter().skip(1).try_fold( first_field.data_type(), @@ -2728,7 +2737,6 @@ impl Union { // Functional Dependencies doesn't preserve after UNION operation let schema = DFSchema::new_with_metadata(union_fields, union_schema_metadata)?; let schema = Arc::new(schema); - Ok(schema) } } diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 5097ff37643d6..b5f1815cdb59a 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -19,7 +19,7 @@ use std::sync::Arc; -use datafusion_expr::binary::BinaryTypeCoercer; +use datafusion_expr::binary::{try_type_union_resolution_with_struct, BinaryTypeCoercer}; use itertools::izip; use arrow::datatypes::{DataType, Field, IntervalUnit, Schema}; @@ -843,7 +843,7 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { .as_ref() .map(|expr| expr.get_type(schema)) .transpose()?; - let then_types = case + let mut then_types = case .when_then_expr .iter() .map(|(_when, then)| then.get_type(schema)) @@ -853,7 +853,6 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { .as_ref() .map(|expr| expr.get_type(schema)) .transpose()?; - // find common coercible types let case_when_coerce_type = case_type .as_ref() @@ -873,6 +872,21 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { }) }) .transpose()?; + // do checks + if then_types.iter().any(|t| matches!(t, DataType::Struct(_))) + || else_type + .as_ref() + .is_some_and(|t| matches!(t, DataType::Struct(_))) + { + if let Some(ref else_t) = else_type { + then_types.push(else_t.clone()); + try_type_union_resolution_with_struct(&then_types).map_err(|_| { + DataFusionError::Execution("failed to do coercsion for case".to_string()) + })?; + then_types.pop(); + } + } + let then_else_coerce_type = get_coerce_type_for_case_expression(&then_types, else_type.as_ref()).ok_or_else( || { @@ -882,7 +896,6 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { ) }, )?; - // do cast if found common coercible types let case_expr = case .expr diff --git a/datafusion/sqllogictest/test_files/case.slt b/datafusion/sqllogictest/test_files/case.slt index 46e9c86c7591c..a1dc4473852f5 100644 --- a/datafusion/sqllogictest/test_files/case.slt +++ b/datafusion/sqllogictest/test_files/case.slt @@ -408,7 +408,7 @@ FROM t; {xxx: c, foo: d} # coerce structs with subset of fields -query error Failed to coerce then +query error failed to do coercsion for case SELECT case when column1 > 0 then column3 @@ -416,5 +416,40 @@ SELECT end FROM t; + + +statement ok +drop table t + +statement ok +create table t as values +( + { 'foo': 'baz' }, + { 'xxx': arrow_cast('blarg', 'Utf8View') } +); + +query error +select CASE WHEN 1=2 THEN column1 ELSE column2 END from t ; +---- +DataFusion error: type_coercion +caused by +Execution error: failed to do coercsion for case + + +statement ok +drop table t + +statement ok +create table t as values +( + { 'name': 'Alice', 'age': 25 }, + { 'age': 30, 'name': 'Bob' } +); + +query ? +select CASE WHEN 1=2 THEN column1 ELSE column2 END from t; +---- +{age: 30, name: Bob} + statement ok drop table t diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index cbd19bf3806fe..1373f530e6f7f 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -851,3 +851,27 @@ FROM ( ---- NULL false foo true + +statement ok +drop table t + +statement ok +create table t as values +( + { 'foo': 'baz' }, + { 'xxx': arrow_cast('blarg', 'Utf8View') }, + { 'name': 'Alice', 'age': 20 }, + { 'age': 30, 'name': 'Bob' } +); + +query error DataFusion error: Execution error: Expect same keys for struct type but got mismatched pair foo and xxx +select column1 from t UNION ALL select column2 from t; + +query ? +select column3 from t UNION ALL select column4 from t order by column3; +---- +{name: 30, age: Bob} +{name: Alice, age: 20} + +statement ok +drop table t