diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index a1db445dd90c5..3a0f70b0dc3c2 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -13,6 +13,7 @@ use crate::expr::{FieldMetadata, LambdaVariable}; use crate::higher_order_function::HigherOrderReturnFieldArgs; use crate::type_coercion::functions::value_fields_with_higher_order_udf_and_lambdas; use crate::type_coercion::functions::{UDFCoercionExt, fields_with_udf}; +use crate::type_coercion::other::get_coerce_type_for_case_expression; use crate::udf::ReturnFieldArgs; use crate::{ LogicalPlan, Operator, Projection, Subquery, WindowFunctionDefinition, utils, @@ -177,15 +178,28 @@ impl ExprSchemable for Expr { Expr::ScalarVariable(field, _) => Ok(field.data_type().clone()), Expr::Literal(l, _) => Ok(l.data_type()), Expr::Case(case) => { - for (_, then_expr) in &case.when_then_expr { - let then_type = then_expr.get_type(schema)?; - if !then_type.is_null() { - return Ok(then_type); - } - } - case.else_expr + let then_types = case + .when_then_expr + .iter() + .map(|(_, then_expr)| then_expr.get_type(schema)) + .collect::>>()?; + let else_type = case + .else_expr .as_ref() - .map_or(Ok(DataType::Null), |e| e.get_type(schema)) + .map(|expr| expr.get_type(schema)) + .transpose()?; + // Leave unsupported combinations for the analyzer, including + // custom rules that may resolve the branches before coercion. + Ok( + get_coerce_type_for_case_expression(&then_types, else_type.as_ref()) + .unwrap_or_else(|| { + then_types + .into_iter() + .find(|data_type| !data_type.is_null()) + .or(else_type) + .unwrap_or(DataType::Null) + }), + ) } Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { Ok(field.data_type().clone()) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index c9575e7a556fa..26f6caede28a4 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -2189,8 +2189,9 @@ pub fn validate_unique_names<'a>( /// Union two [`LogicalPlan`]s. /// -/// Constructs the UNION plan, but does not perform type-coercion. Therefore the -/// subtree expressions will not be properly typed until the optimizer pass. +/// Infers the common output schema, but leaves input casts to the analyzer. +/// Input expressions and unbound parameters may still change type during analysis +/// or parameter binding. /// /// If a properly typed UNION plan is needed, refer to [`TypeCoercionRewriter::coerce_union`] /// or alternatively, merge the union input schema using [`coerce_union_schema`] and diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c5b5b23ba8cc9..518df4a4ffe1c 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -44,6 +44,7 @@ use crate::{ use crate::logical_plan::information_schema::ColumnsDetail; use crate::statistics::StatisticsRequest; +use crate::type_coercion::binary::type_union_coercion; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef, Metadata, Schema, SchemaRef}; use datafusion_common::cse::{NormalizeEq, Normalizeable}; @@ -1878,8 +1879,8 @@ impl LogicalPlan { /// Recompute schema fields' data type after replacing params, ensuring fields data type can be /// updated according to the new parameters. /// - /// Unlike `recompute_schema()`, this method rebuilds VALUES plans entirely to properly infer - /// types from literal values after placeholder substitution. + /// Unlike `recompute_schema()`, this also rebuilds VALUES and UNION schemas + /// to infer types from literal values after placeholder substitution. fn update_schema_data_type(self) -> Result { match self { // Build `LogicalPlan::Values` from the values for type inference. @@ -1888,6 +1889,11 @@ impl LogicalPlan { LogicalPlan::Values(Values { values, schema: _ }) => { LogicalPlanBuilder::values(values)?.build() } + // UNION output types can change when previously untyped parameters + // acquire their bound values. Do not preserve the provisional schema. + LogicalPlan::Union(Union { inputs, .. }) => { + Ok(LogicalPlan::Union(Union::try_new_with_loose_types(inputs)?)) + } // other plans can just use `recompute_schema` directly. plan => plan.recompute_schema(), } @@ -3541,8 +3547,9 @@ impl Union { } /// Constructs new Union instance deriving schema from inputs. - /// Inputs do not have to have matching types and produced schema will - /// take type from the first input. + /// Infers common output types without inserting casts into the inputs. + /// If inference fails, retains the first input's type until analysis. + /// The analyzer coerces the inputs after their expressions are resolved. // TODO (https://github.com/apache/datafusion/issues/14380): Avoid creating uncoerced union at all. pub fn try_new_with_loose_types(inputs: Vec>) -> Result { let schema = Self::derive_schema_from_inputs(&inputs, true, false)?; @@ -3585,9 +3592,9 @@ impl Union { ); } } - wrapped_inputs.push(Arc::new(LogicalPlan::Projection( - Projection::try_new_with_schema(expr, input, Arc::clone(schema))?, - ))); + wrapped_inputs.push(Arc::new(LogicalPlan::Projection(Projection::try_new( + expr, input, + )?))); } Ok(wrapped_inputs) @@ -3596,8 +3603,8 @@ impl Union { /// Constructs new Union instance deriving schema from inputs. /// /// If `loose_types` is true, inputs do not need to have matching types and - /// the produced schema will use the type from the first input. - /// TODO (): This is not necessarily reasonable behavior. + /// the produced schema will use their common types. Input casts are deferred + /// until analysis, so unresolved parameters can still change these types. /// /// If `by_name` is `true`, input schemas need not be the same width. That is, /// the constructed schema follows `UNION BY NAME` semantics. @@ -3621,14 +3628,38 @@ impl Union { inputs: &[Arc], loose_types: bool, ) -> Result { - type FieldData<'a> = (&'a DataType, bool, Vec<&'a Metadata>, usize); + // Keep the first type as a fallback if any input prevents inference. + type FieldData<'a> = ( + &'a DataType, + Option, + bool, + Vec<&'a Metadata>, + usize, + ); let mut cols: Vec<(&str, FieldData)> = Vec::new(); - for input in inputs.iter() { + for (input_index, input) in inputs.iter().enumerate() { + // Alignment projections supply NULL for missing columns. Include + // those inputs in order, since NULL can unwrap dictionary types. + if loose_types { + for (name, (_, common_type, ..)) in &mut cols { + if !input.schema().has_column_with_unqualified_name(name) { + *common_type = common_type.as_ref().and_then(|data_type| { + type_union_coercion(data_type, &DataType::Null) + }); + } + } + } for field in input.schema().fields() { - if let Some((_, (data_type, is_nullable, metadata, occurrences))) = - cols.iter_mut().find(|(name, _)| name == field.name()) + if let Some(( + _, + (first_type, common_type, is_nullable, metadata, occurrences), + )) = cols.iter_mut().find(|(name, _)| name == field.name()) { - if !loose_types && *data_type != field.data_type() { + if loose_types { + *common_type = common_type.as_ref().and_then(|data_type| { + type_union_coercion(data_type, field.data_type()) + }); + } else if *first_type != field.data_type() { return plan_err!( "Found different types for field {}", field.name() @@ -3645,6 +3676,11 @@ impl Union { field.name(), ( field.data_type(), + if loose_types && input_index > 0 { + type_union_coercion(&DataType::Null, field.data_type()) + } else { + Some(field.data_type().clone()) + }, field.is_nullable(), vec![field.metadata()], 1, @@ -3657,7 +3693,16 @@ impl Union { let union_fields = cols .into_iter() .map( - |(name, (data_type, is_nullable, unmerged_metadata, occurrences))| { + |( + name, + ( + first_type, + common_type, + is_nullable, + unmerged_metadata, + occurrences, + ), + )| { // If the final number of occurrences of the field is less // than the number of inputs (i.e. the field is missing from // one or more inputs), then it must be treated as nullable. @@ -3667,8 +3712,8 @@ impl Union { true }; - let mut field = - Field::new(name, data_type.clone(), final_is_nullable); + let data_type = common_type.unwrap_or_else(|| first_type.clone()); + let mut field = Field::new(name, data_type, final_is_nullable); field.set_metadata(merge_metadata_for_union(unmerged_metadata)); (None, Arc::new(field)) @@ -3715,14 +3760,18 @@ impl Union { let base_name = first_field.name().to_string(); let data_type = if loose_types { - // 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() + fields + .iter() + .skip(1) + .try_fold(first_field.data_type().clone(), |acc, field| { + type_union_coercion(&acc, field.data_type()) + }) + .unwrap_or_else(|| first_field.data_type().clone()) } else { - fields.iter().skip(1).try_fold( - first_field.data_type(), - |acc, field| { + fields + .iter() + .skip(1) + .try_fold(first_field.data_type(), |acc, field| { if acc != field.data_type() { return plan_err!( "UNION field {i} have different type in inputs: \ @@ -3732,8 +3781,8 @@ impl Union { ); } Ok(acc) - }, - )? + })? + .clone() }; let nullable = fields.iter().any(|field| field.is_nullable()); @@ -3746,7 +3795,7 @@ impl Union { base_name }; - let mut field = Field::new(&name, data_type.clone(), nullable); + let mut field = Field::new(&name, data_type, nullable); let field_metadata = merge_metadata_for_union(fields.iter().map(|field| field.metadata())); field.set_metadata(field_metadata); diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index fc798449b0b2a..2d30c1c57bde4 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1649,7 +1649,7 @@ fn coerce_case_expression( /// /// **Type coercion precedence**: The coerced type is determined by iteratively applying /// `type_union_coercion()` between the accumulated type and each new input's type. The -/// result depends on type coercion rules, not input order. +/// result follows input order, which matters for dictionary types. /// /// **Nullability merging**: Nullability is accumulated using logical OR (`||`). /// Once any input field is nullable, the result field becomes nullable permanently. @@ -1671,17 +1671,23 @@ fn coerce_case_expression( /// ``` /// /// **Precedence Summary**: -/// - **Datatypes**: Determined by `type_union_coercion()` rules, not input order +/// - **Datatypes**: Determined by `type_union_coercion()` in input order /// - **Nullability**: Later inputs can add nullability but cannot remove it (logical OR) /// - **Metadata**: Later inputs take precedence for same keys (HashMap::extend semantics) pub fn coerce_union_schema(inputs: &[Arc]) -> Result { - coerce_union_schema_with_schema(&inputs[1..], inputs[0].schema()) + let Some(first) = inputs.first() else { + return plan_err!("UNION requires at least one input"); + }; + coerce_union_schema_with_schema(inputs, first.schema()) } fn coerce_union_schema_with_schema( inputs: &[Arc], base_schema: &DFSchemaRef, ) -> Result { - let mut union_datatypes = base_schema + // The output schema may already contain inferred types. Feeding those types + // back into coercion can change the result, notably for dictionaries. + let mut union_datatypes = inputs[0] + .schema() .fields() .iter() .map(|f| f.data_type().clone()) diff --git a/datafusion/sqllogictest/test_files/case.slt b/datafusion/sqllogictest/test_files/case.slt index 4727dd19d112f..9f6711a24f7fd 100644 --- a/datafusion/sqllogictest/test_files/case.slt +++ b/datafusion/sqllogictest/test_files/case.slt @@ -132,6 +132,18 @@ SELECT CASE WHEN false THEN 1 ELSE 0 END FROM foo 0 # test then value type coercion +# DESCRIBE checks the inferred schema before analyzer coercion. +query TTT +DESCRIBE SELECT + CASE WHEN p THEN 1::INT ELSE 0.5 END AS else_widens, + CASE WHEN p THEN 1::INT WHEN NOT p THEN 0.5 END AS then_widens, + CASE WHEN p THEN 1::SMALLINT WHEN NOT p THEN 2::BIGINT ELSE true END AS deferred +FROM (VALUES (true), (false)) AS t(p); +---- +else_widens Float64 NO +then_widens Float64 YES +deferred Int16 NO + # List(Utf8) will be casted to List(Int64) query ? SELECT CASE 1 WHEN 1 THEN ['1', '2', '3'] WHEN 2 THEN [1, 2, 3] ELSE null END; diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 1a8d88ee149bb..5096f94fef911 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -6,6 +6,36 @@ ## UNION Tests ########## +# Infer the common schema before analysis, including timestamp precision/timezone +# and string representations (https://github.com/apache/datafusion/issues/27023). +query TTT +DESCRIBE SELECT + arrow_cast(1, 'Timestamp(Microsecond, Some("UTC"))') AS t, + arrow_cast('a', 'Utf8') AS large, + arrow_cast('a', 'Utf8') AS view +UNION ALL SELECT + CAST(NULL AS TIMESTAMP), + arrow_cast('b', 'LargeUtf8'), + arrow_cast('b', 'Utf8View'); +---- +t Timestamp(ns, "UTC") YES +large LargeUtf8 NO +view Utf8View NO + +# Failed inference keeps the first type for analyzer rules to resolve or reject. +query TTT +DESCRIBE SELECT 1::SMALLINT AS x UNION ALL SELECT true; +---- +x Int16 NO + +# Infer dictionary coercion from the actual inputs, without folding the inferred +# UNION type back into the inputs and restoring the dictionary representation. +query T +SELECT DISTINCT arrow_typeof(x) +FROM (SELECT NULL AS x UNION ALL SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x); +---- +Utf8 + statement ok CREATE TABLE t1( id INT, diff --git a/datafusion/sqllogictest/test_files/union_by_name.slt b/datafusion/sqllogictest/test_files/union_by_name.slt index 81f7317b2a39f..c531863001650 100644 --- a/datafusion/sqllogictest/test_files/union_by_name.slt +++ b/datafusion/sqllogictest/test_files/union_by_name.slt @@ -40,6 +40,36 @@ INSERT INTO t2 VALUES (2, 2), (4, 4); # Test binding +# Infer common types by name while preserving missing-column nullability. +query TTT +DESCRIBE SELECT 1::INT AS x +UNION ALL BY NAME SELECT 'a' AS y, 2::BIGINT AS x; +---- +x Int64 NO +y Utf8 YES + +# Retain the first type if a later input prevents inference. +query TTT +DESCRIBE SELECT 1::SMALLINT AS x +UNION ALL BY NAME SELECT true AS x; +---- +x Int16 NO + +# Missing columns contribute NULL to type inference, including dictionaries. +query TTT +DESCRIBE SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x +UNION ALL BY NAME SELECT arrow_cast('b', 'Dictionary(Int32, Utf8)') AS y; +---- +x Utf8 YES +y Utf8 YES + +query TT +SELECT DISTINCT arrow_typeof(x), arrow_typeof(y) +FROM (SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x + UNION ALL BY NAME SELECT arrow_cast('b', 'Dictionary(Int32, Utf8)') AS y); +---- +Utf8 Utf8 + query I SELECT t1.x FROM t1 UNION BY NAME SELECT x FROM t1 ORDER BY x; ---- diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 769d3fdbeb1db..632b3dfe28b64 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -12,6 +12,14 @@ SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one in this section pertains to features and changes that have already been merged to the main branch and are awaiting release in this version. +### UNION BY NAME dictionary output types + +When a `Dictionary(Int32, Utf8)` column is missing from one input to +`UNION BY NAME`, the output type is now `Utf8` in either input order. +Previously, a missing left column preserved the dictionary representation. +Values are unchanged; applications inspecting Arrow schemas may need to update +their expectations. + ### Session time zone now applies to mixed timestamp operations When `datafusion.execution.time_zone` is set, DataFusion now uses it to