diff --git a/datafusion/core/tests/sql/mod.rs b/datafusion/core/tests/sql/mod.rs index 8c0b328c46c52..cf9ab2498a64e 100644 --- a/datafusion/core/tests/sql/mod.rs +++ b/datafusion/core/tests/sql/mod.rs @@ -61,6 +61,7 @@ mod sql_api; mod tpcds_planning; mod union_comparison; mod union_nullable; +mod union_schema; mod unparser; async fn register_aggregate_csv_by_sql(ctx: &SessionContext) { diff --git a/datafusion/core/tests/sql/union_nullable.rs b/datafusion/core/tests/sql/union_nullable.rs index f8850f6dadfa2..1fedc7455bb0f 100644 --- a/datafusion/core/tests/sql/union_nullable.rs +++ b/datafusion/core/tests/sql/union_nullable.rs @@ -61,12 +61,8 @@ async fn union_all_mismatched_nullable( .await } -/// The schema DataFusion actually commits to for a query: the logical plan -/// after the `Analyzer` (which includes the `UNION` nullability/type -/// coercion this test targets) and `Optimizer` have run. `DataFrame::schema` -/// alone is not enough here -- it reflects the raw, pre-`Analyzer` plan (see -/// `SessionState::create_logical_plan`), which for a `UNION` still has the -/// first leg's un-coerced type. +/// This test checks the schema after the Analyzer and Optimizer run, +/// because they can change nullability. fn analyzed_schema(df: &DataFrame) -> Result { Ok(df .clone() diff --git a/datafusion/core/tests/sql/union_schema.rs b/datafusion/core/tests/sql/union_schema.rs new file mode 100644 index 0000000000000..65b6bddae306a --- /dev/null +++ b/datafusion/core/tests/sql/union_schema.rs @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! The pre-analysis UNION schema must describe the output types (#27023). + +use arrow::datatypes::{DataType, TimeUnit}; +use datafusion::prelude::*; +use datafusion_common::{Result, ScalarValue}; + +fn assert_union_schema(df: &DataFrame, expected: &[(&str, DataType)]) { + let fields = df.schema().fields(); + assert_eq!(fields.len(), expected.len()); + for (field, (name, data_type)) in fields.iter().zip(expected) { + assert_eq!(field.name(), name); + assert_eq!(field.data_type(), data_type); + } +} + +#[tokio::test] +async fn union_infers_types_before_analysis() -> Result<()> { + let ctx = SessionContext::new(); + let cases = [ + ( + r#"arrow_cast(arrow_cast(1, 'Timestamp(Microsecond, None)'), 'Timestamp(Microsecond, Some("UTC"))')"#, + "CAST(NULL AS TIMESTAMP)", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + ), + ( + "arrow_cast('a', 'Utf8')", + "arrow_cast('b', 'LargeUtf8')", + DataType::LargeUtf8, + ), + ( + "arrow_cast('a', 'Utf8')", + "arrow_cast('b', 'Utf8View')", + DataType::Utf8View, + ), + ("CAST(1 AS SMALLINT)", "2.5", DataType::Float64), + ( + "CASE WHEN false THEN 1 ELSE 2.5 END", + "NULL", + DataType::Float64, + ), + ( + "CASE WHEN false THEN 1 WHEN true THEN 2.5 END", + "NULL", + DataType::Float64, + ), + ]; + for (left, right, expected) in cases { + for (left, right) in [(left, right), (right, left)] { + for operation in ["ALL", "DISTINCT", "ALL BY NAME", "BY NAME"] { + let sql = format!( + "SELECT x FROM (SELECT {left} AS x UNION {operation} SELECT {right} AS x) u" + ); + assert_union_schema(&ctx.sql(&sql).await?, &[("x", expected.clone())]); + } + } + } + Ok(()) +} + +#[tokio::test] +async fn union_infers_dictionary_types_before_analysis() -> Result<()> { + let ctx = SessionContext::new(); + for (sql, expected) in [ + ( + "SELECT NULL AS x UNION ALL SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x", + DataType::Utf8, + ), + ( + "SELECT CAST(1 AS SMALLINT) AS x UNION ALL SELECT arrow_cast(2, 'Dictionary(Int32, Int32)') AS x", + DataType::Int32, + ), + ] { + assert_union_schema(&ctx.sql(sql).await?, &[("x", expected)]); + } + for sql in [ + "SELECT 1 AS y UNION ALL BY NAME SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x", + "SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x UNION ALL BY NAME SELECT 1 AS y", + ] { + let df = ctx.sql(sql).await?.select_columns(&["x", "y"])?; + assert_union_schema(&df, &[("x", DataType::Utf8), ("y", DataType::Int64)]); + } + Ok(()) +} + +#[tokio::test] +async fn union_dataframe_and_by_name_schemas() -> Result<()> { + let ctx = SessionContext::new(); + for variant in 0..4 { + let left = ctx.sql("SELECT CAST(1 AS SMALLINT) AS x").await?; + let right = ctx.sql("SELECT 2.5 AS x").await?; + let df = match variant { + 0 => left.union(right)?, + 1 => left.union_distinct(right)?, + 2 => left.union_by_name(right)?, + _ => left.union_by_name_distinct(right)?, + }; + assert_union_schema(&df.select_columns(&["x"])?, &[("x", DataType::Float64)]); + } + for sql in [ + "SELECT CAST(1 AS SMALLINT) AS x, 1 AS y UNION ALL BY NAME SELECT 2 AS y, 2.5 AS x", + "SELECT CAST(1 AS SMALLINT) AS x UNION ALL BY NAME SELECT 2.5 AS x, 2 AS y", + ] { + assert_union_schema( + &ctx.sql(sql).await?, + &[("x", DataType::Float64), ("y", DataType::Int64)], + ); + } + let df = ctx + .sql( + "SELECT CAST(NULL AS SMALLINT) AS x WHERE false \ + UNION ALL SELECT CAST(NULL AS BIGINT) AS x WHERE false", + ) + .await?; + assert_union_schema(&df, &[("x", DataType::Int64)]); + let df = ctx + .sql( + "SELECT max(x) AS x FROM (SELECT 1 AS x UNION ALL SELECT 2.5 AS x \ + UNION ALL SELECT 3 AS x) u", + ) + .await?; + assert_union_schema(&df, &[("x", DataType::Float64)]); + Ok(()) +} + +#[tokio::test] +async fn union_bound_parameters_keep_inferred_type() -> Result<()> { + let ctx = SessionContext::new(); + let df = ctx + .sql("SELECT CAST($1 AS INT) AS x UNION ALL SELECT 2.5 AS x") + .await? + .with_param_values(vec![ScalarValue::Int32(Some(1))])?; + assert_union_schema(&df, &[("x", DataType::Float64)]); + Ok(()) +} diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index a1db445dd90c5..ddabb95740afa 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, @@ -33,6 +34,10 @@ use std::sync::Arc; /// Trait to allow expr to typable with respect to a schema pub trait ExprSchemable { /// Given a schema, return the type of the expr + /// + /// `CASE` expressions infer a common type from all `THEN` and `ELSE` + /// branches without inserting casts. Incompatible branches retain a + /// provisional type until analysis. fn get_type(&self, schema: &dyn ExprSchema) -> Result; /// Given a schema, return the nullability of the expr @@ -177,15 +182,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(|(_, expr)| 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()?; + // If the branches have no common type, keep the old provisional + // type. The analyzer reports the error. + let common_type = + get_coerce_type_for_case_expression(&then_types, else_type.as_ref()); + let result_type = common_type.unwrap_or_else(|| { + then_types + .into_iter() + .find(|t| !t.is_null()) + .or(else_type) + .unwrap_or(DataType::Null) + }); + Ok(result_type) } Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { Ok(field.data_type().clone()) @@ -1684,4 +1702,33 @@ mod tests { } } } + + #[test] + fn case_infers_all_result_types_without_casting() -> Result<()> { + use crate::expr::Case; + let schema = DFSchema::empty(); + let nested = Expr::Case(Case::new( + None, + vec![(Box::new(lit(false)), Box::new(lit(1_i64)))], + Some(Box::new(lit(2.5))), + )); + for (then_values, else_value, expected) in [ + (vec![lit(1_i64)], Some(lit(2.5)), DataType::Float64), + (vec![lit(1_i64), lit(2.5)], None, DataType::Float64), + (vec![lit(2.5), lit(1_i64)], None, DataType::Float64), + (vec![lit(2.5)], Some(lit(1_i64)), DataType::Float64), + (vec![lit(3_i64)], Some(nested), DataType::Float64), + (vec![lit(ScalarValue::Null)], None, DataType::Null), + // Incompatible raw types retain a provisional type for the analyzer. + (vec![lit(true)], Some(lit("a")), DataType::Boolean), + ] { + let when_then = then_values + .into_iter() + .map(|value| (Box::new(lit(true)), Box::new(value))) + .collect(); + let expr = Expr::Case(Case::new(None, when_then, else_value.map(Box::new))); + assert_eq!(expr.get_type(&schema)?, expected); + } + Ok(()) + } } diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index c9575e7a556fa..2a3939031cad2 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -2189,8 +2189,8 @@ 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 common output types, but does not insert casts into the inputs. +/// The analyzer must still resolve input expressions and insert any required casts. /// /// 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 0fac69e7a0ec4..22e2deae97a9c 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -29,6 +29,7 @@ use crate::logical_plan::materialized_cte::{ MaterializedCte, MaterializedCtePlan, MaterializedCteScan, }; use crate::logical_plan::{DmlStatement, Statement, WriteOp}; +use crate::type_coercion::other::coerce_union_type; use crate::utils::{ check_aggregate_and_window_nesting, check_no_window_functions, enumerate_grouping_sets, expr_to_columns, exprlist_to_fields, @@ -45,7 +46,7 @@ use crate::{ use crate::logical_plan::information_schema::ColumnsDetail; use crate::statistics::StatisticsRequest; use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, FieldRef, Metadata, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datafusion_common::cse::{NormalizeEq, Normalizeable}; use datafusion_common::format::{ExplainAnalyzeCategories, ExplainFormat, MetricType}; use datafusion_common::metadata::{ @@ -3541,9 +3542,11 @@ 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. - // TODO (https://github.com/apache/datafusion/issues/14380): Avoid creating uncoerced union at all. + /// Infers common output types without inserting casts into the inputs. + /// If a common type cannot be inferred, keeps the first input's type so + /// the analyzer can resolve or reject the inputs later. + // TODO (https://github.com/apache/datafusion/issues/14380): Inputs are still + // uncast until the analyzer runs. pub fn try_new_with_loose_types(inputs: Vec>) -> Result { let schema = Self::derive_schema_from_inputs(&inputs, true, false)?; Ok(Union { inputs, schema }) @@ -3552,6 +3555,11 @@ impl Union { /// Constructs a new Union instance that combines rows from different tables by name, /// instead of by position. This means that the specified inputs need not have schemas /// that are all the same width. + /// + /// Like [`Self::try_new_with_loose_types`], this infers common output types + /// without inserting casts. Missing columns contribute `NULL` at their input + /// position. For example, a missing `Dictionary(Int32, Utf8)` column and a + /// present one produce `Utf8` in either input order. pub fn try_new_by_name(inputs: Vec>) -> Result { let schema = Self::derive_schema_from_inputs(&inputs, true, true)?; let inputs = Self::rewrite_inputs_from_schema(&schema, inputs)?; @@ -3595,9 +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. + /// If `loose_types` is true, infers common types where possible without + /// inserting casts. Otherwise, input types must match exactly. /// /// If `by_name` is `true`, input schemas need not be the same width. That is, /// the constructed schema follows `UNION BY NAME` semantics. @@ -3621,60 +3628,62 @@ impl Union { inputs: &[Arc], loose_types: bool, ) -> Result { - type FieldData<'a> = (&'a DataType, bool, Vec<&'a Metadata>, usize); - let mut cols: Vec<(&str, FieldData)> = Vec::new(); - for input in inputs.iter() { + let mut cols: Vec<(&str, Vec>)> = Vec::new(); + for (input_index, input) in inputs.iter().enumerate() { for field in input.schema().fields() { - if let Some((_, (data_type, is_nullable, metadata, occurrences))) = + if let Some((_, fields)) = cols.iter_mut().find(|(name, _)| name == field.name()) { - if !loose_types && *data_type != field.data_type() { - return plan_err!( - "Found different types for field {}", - field.name() - ); - } - - metadata.push(field.metadata()); - // If the field is nullable in any one of the inputs, - // then the field in the final schema is also nullable. - *is_nullable |= field.is_nullable(); - *occurrences += 1; + fields[input_index] = Some(field.as_ref()); } else { - cols.push(( - field.name(), - ( - field.data_type(), - field.is_nullable(), - vec![field.metadata()], - 1, - ), - )); + let mut fields = vec![None; inputs.len()]; + fields[input_index] = Some(field.as_ref()); + cols.push((field.name(), fields)); } } } + // The analyzer sees missing fields as NULL literals in the alignment + // projections. Include them so both passes infer the same type. + let null_field = Field::new("", DataType::Null, true); let union_fields = cols .into_iter() - .map( - |(name, (data_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. - let final_is_nullable = if occurrences == inputs.len() { - is_nullable - } else { - true - }; - - let mut field = - Field::new(name, data_type.clone(), final_is_nullable); - field.set_metadata(merge_metadata_for_union(unmerged_metadata)); - - (None, Arc::new(field)) - }, - ) - .collect::, _)>>(); + .map(|(name, fields)| { + let first_type = fields + .iter() + .flatten() + .next() + .expect("column has at least one input field") + .data_type(); + let data_type = if loose_types { + coerce_union_type( + fields.iter().map(|field| field.unwrap_or(&null_field)), + ) + .unwrap_or_else(|_| first_type.clone()) + } else { + if fields + .iter() + .flatten() + .any(|field| field.data_type() != first_type) + { + return plan_err!("Found different types for field {name}"); + } + first_type.clone() + }; + let nullable = fields + .iter() + .any(|field| field.is_none_or(|field| field.is_nullable())); + let metadata = merge_metadata_for_union( + fields.iter().flatten().map(|field| field.metadata()), + ); + Ok(( + None, + Arc::new( + Field::new(name, data_type, nullable).with_metadata(metadata), + ), + )) + }) + .collect::, _)>>>()?; let union_schema_metadata = merge_metadata_for_union( inputs.iter().map(|input| input.schema().metadata()), @@ -3715,14 +3724,13 @@ 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() + coerce_union_type(fields.iter().map(|field| field.as_ref())) + .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 +3740,8 @@ impl Union { ); } Ok(acc) - }, - )? + })? + .clone() }; let nullable = fields.iter().any(|field| field.is_nullable()); @@ -3746,7 +3754,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); @@ -5518,6 +5526,7 @@ mod tests { GroupingSet, binary_expr, col, exists, in_subquery, lit, placeholder, scalar_subquery, }; + use arrow::datatypes::Metadata; use datafusion_common::metadata::ScalarAndMetadata; use datafusion_common::tree_node::{ TransformedResult, TreeNodeRewriter, TreeNodeVisitor, @@ -7760,4 +7769,51 @@ mod tests { Ok(()) } + + #[test] + fn union_infers_types_without_rewriting_inputs() -> Result<()> { + let input = |data_type, nullable, source: &str| { + table_scan( + Some(source), + &Schema::new(vec![ + Field::new("x", data_type, nullable) + .with_metadata(Metadata::new().with("source", source)), + ]), + None, + )? + .build() + .map(Arc::new) + }; + let left = input(DataType::Int16, false, "left")?; + let right = input(DataType::Int64, true, "right")?; + let inputs = vec![Arc::clone(&left), Arc::clone(&right)]; + let union = Union::try_new_with_loose_types(inputs.clone())?; + assert_eq!(union.inputs, inputs); + assert_eq!(union.schema.field(0).name(), "x"); + assert_eq!(union.schema.field(0).data_type(), &DataType::Int64); + assert!(union.schema.field(0).is_nullable()); + assert_eq!( + union + .schema + .field(0) + .metadata() + .get("source") + .map(String::as_str), + Some("right") + ); + assert!(Union::try_new(inputs.clone()).is_err()); + assert!(Union::try_new_with_loose_types(vec![Arc::clone(&left)]).is_err()); + + // Keep the original type when a later input prevents builtin inference. + // A custom analyzer can still resolve the raw inputs without earlier casts. + let incompatible = input(DataType::Boolean, false, "third")?; + let inputs = vec![left, right, incompatible]; + for union in [ + Union::try_new_with_loose_types(inputs.clone())?, + Union::try_new_by_name(inputs)?, + ] { + assert_eq!(union.schema.field(0).data_type(), &DataType::Int16); + } + Ok(()) + } } diff --git a/datafusion/expr/src/type_coercion/other.rs b/datafusion/expr/src/type_coercion/other.rs index a605951d1e6b8..033e51af1c2fb 100644 --- a/datafusion/expr/src/type_coercion/other.rs +++ b/datafusion/expr/src/type_coercion/other.rs @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use arrow::datatypes::DataType; +use arrow::datatypes::{DataType, Field}; +use datafusion_common::{Result, plan_datafusion_err, plan_err}; use super::binary::{comparison_coercion, type_union_coercion}; @@ -19,6 +20,30 @@ fn fold_coerce( }) } +/// Infer a common output type for one UNION column, in input order. +/// +/// Field names are used only to report incompatible inputs. This does not +/// insert casts or require input expressions to have been analyzed. +pub fn coerce_union_type<'a>( + fields: impl IntoIterator, +) -> Result { + let mut fields = fields.into_iter(); + let Some(first) = fields.next() else { + return plan_err!("UNION requires at least one input field"); + }; + fields.try_fold(first.data_type().clone(), |data_type, field| { + type_union_coercion(&data_type, field.data_type()).ok_or_else(|| { + plan_datafusion_err!( + "Incompatible inputs for Union: Previous inputs were \ + of type {}, but got incompatible type {} on column '{}'", + data_type, + field.data_type(), + field.name() + ) + }) + }) +} + /// Attempts to coerce the types of `list_types` to be comparable with the /// `expr_type` for IN list predicates. /// Returns the common data type for `expr_type` and `list_types`. diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 296fccec4fb19..6c3e2109c6398 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -6,7 +6,7 @@ use arrow::compute::can_cast_types; use datafusion_expr::binary::BinaryTypeCoercer; -use itertools::{Itertools as _, izip}; +use itertools::Itertools as _; use std::sync::{Arc, LazyLock}; use crate::analyzer::AnalyzerRule; @@ -34,12 +34,13 @@ use datafusion_expr::expr_schema::{cast_subquery, cast_subquery_columns}; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::type_coercion::binary::{ comparison_coercion_with_session_timezone, like_coercion, regex_coercion, - type_union_coercion, }; use datafusion_expr::type_coercion::functions::{ UDFCoercionExt, fields_with_udf, value_fields_with_higher_order_udf_and_lambdas, }; -use datafusion_expr::type_coercion::other::get_coerce_type_for_case_expression; +use datafusion_expr::type_coercion::other::{ + coerce_union_type, get_coerce_type_for_case_expression, +}; use datafusion_expr::type_coercion::{ is_datetime, is_interval, is_signed_numeric, is_timestamp, }; @@ -1649,7 +1650,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 the input order, which can matter for dictionary and nested types. /// /// **Nullability merging**: Nullability is accumulated using logical OR (`||`). /// Once any input field is nullable, the result field becomes nullable permanently. @@ -1671,81 +1672,56 @@ fn coerce_case_expression( /// ``` /// /// **Precedence Summary**: -/// - **Datatypes**: Determined by `type_union_coercion()` rules, not input order +/// - **Datatypes**: Folded in input order using `type_union_coercion()` /// - **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 - .fields() - .iter() - .map(|f| f.data_type().clone()) - .collect::>(); - let mut union_nullabilities = base_schema - .fields() - .iter() - .map(|f| f.is_nullable()) - .collect::>(); - for (i, plan) in inputs.iter().enumerate() { - let plan_schema = plan.schema(); - - if plan_schema.fields().len() != base_schema.fields().len() { + if plan.schema().fields().len() != base_schema.fields().len() { return plan_err!( "Union schemas have different number of fields: \ query 1 has {} fields whereas query {} has {} fields", base_schema.fields().len(), i + 1, - plan_schema.fields().len() + plan.schema().fields().len() ); } - - // coerce data type and nullability for each field - for (union_datatype, union_nullable, plan_field) in izip!( - union_datatypes.iter_mut(), - union_nullabilities.iter_mut(), - plan_schema.fields().iter() - ) { - let coerced_type = - type_union_coercion(union_datatype, plan_field.data_type()).ok_or_else( - || { - plan_datafusion_err!( - "Incompatible inputs for Union: Previous inputs were \ - of type {}, but got incompatible type {} on column '{}'", - union_datatype, - plan_field.data_type(), - plan_field.name() - ) - }, - )?; - - *union_datatype = coerced_type; - *union_nullable = *union_nullable || plan_field.is_nullable(); - } } let all_schemas = std::iter::once(base_schema) .chain(inputs.iter().map(|p| p.schema())) .collect::>(); - let union_qualified_fields = izip!( - base_schema.fields().iter().enumerate(), - union_datatypes.into_iter(), - union_nullabilities - ) - .map(|((i, field), datatype, nullable)| { - let mut field = Field::new(field.name().clone(), datatype, nullable); - field.set_metadata(merge_metadata_for_union( - all_schemas.iter().map(|schema| schema.field(i).metadata()), - )); - (None, field.into()) - }) - .collect::>(); + let union_qualified_fields = base_schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| { + // The stored schema may already contain an inferred common type. Feeding + // it back into coercion can change the result, notably for dictionaries. + let datatype = coerce_union_type( + inputs.iter().map(|plan| plan.schema().field(i).as_ref()), + )?; + let nullable = field.is_nullable() + || inputs + .iter() + .any(|plan| plan.schema().field(i).is_nullable()); + let mut field = Field::new(field.name().clone(), datatype, nullable); + field.set_metadata(merge_metadata_for_union( + all_schemas.iter().map(|schema| schema.field(i).metadata()), + )); + Ok((None, field.into())) + }) + .collect::>>()?; let metadata = merge_metadata_for_union(all_schemas.iter().map(|schema| schema.metadata())); @@ -1969,6 +1945,18 @@ mod test { ) } + #[test] + fn union_schema_validates_input_shapes() -> Result<()> { + use super::coerce_union_schema; + assert!(coerce_union_schema(&[]).is_err()); + let input = empty_with_type(DataType::Int32); + let schema = coerce_union_schema(&[Arc::clone(&input)])?; + assert_eq!(schema.field(0).data_type(), &DataType::Int32); + let incompatible_width = empty(); + assert!(coerce_union_schema(&[input, incompatible_width]).is_err()); + Ok(()) + } + #[test] fn test_coerce_union() -> Result<()> { let left_plan = Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { diff --git a/datafusion/sqllogictest/test_files/prepare.slt b/datafusion/sqllogictest/test_files/prepare.slt index 2b5af841a3031..e71f94c53b776 100644 --- a/datafusion/sqllogictest/test_files/prepare.slt +++ b/datafusion/sqllogictest/test_files/prepare.slt @@ -694,3 +694,15 @@ DEALLOCATE my_plan statement ok drop table keys_t + +statement ok +PREPARE union_coercion(INT) AS SELECT CAST($1 AS INT) AS x UNION ALL SELECT 2.5 AS x; + +query R rowsort +EXECUTE union_coercion(1); +---- +1 +2.5 + +statement ok +DEALLOCATE union_coercion; diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 1a8d88ee149bb..c44b657ab2a7c 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -1148,3 +1148,139 @@ UNION ALL # Config reset statement ok RESET datafusion.execution.batch_size; + +# Type inference must preserve timestamp precision in either input order. +query TI rowsort +SELECT arrow_typeof(x), arrow_cast(x, 'Int64') FROM ( + SELECT arrow_cast(1, 'Timestamp(Microsecond, None)') AS x + UNION ALL + SELECT arrow_cast(1001, 'Timestamp(Nanosecond, None)') AS x +) u; +---- +Timestamp(ns) 1000 +Timestamp(ns) 1001 + +query TI rowsort +SELECT arrow_typeof(x), arrow_cast(x, 'Int64') FROM ( + SELECT arrow_cast(1001, 'Timestamp(Nanosecond, None)') AS x + UNION DISTINCT + SELECT arrow_cast(1, 'Timestamp(Microsecond, None)') AS x +) u; +---- +Timestamp(ns) 1000 +Timestamp(ns) 1001 + +# https://github.com/apache/datafusion/issues/27023 +query TI rowsort +SELECT arrow_typeof(x), arrow_cast(x, 'Int64') FROM ( + SELECT arrow_cast(1, 'Timestamp(Microsecond, Some("UTC"))') AS x + UNION ALL + SELECT CAST(NULL AS TIMESTAMP) AS x +) u; +---- +Timestamp(ns, "UTC") 1000 +Timestamp(ns, "UTC") NULL + +query TI rowsort +SELECT arrow_typeof(x), arrow_cast(x, 'Int64') FROM ( + SELECT CAST(NULL AS TIMESTAMP) AS x + UNION DISTINCT + SELECT arrow_cast(1, 'Timestamp(Microsecond, Some("UTC"))') AS x +) u; +---- +Timestamp(ns, "UTC") 1000 +Timestamp(ns, "UTC") NULL + +query TT rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT arrow_cast('a', 'Utf8') AS x + UNION ALL + SELECT arrow_cast('b', 'LargeUtf8') AS x +) u; +---- +LargeUtf8 a +LargeUtf8 b + +query TT rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT arrow_cast('b', 'Utf8View') AS x + UNION DISTINCT + SELECT arrow_cast('a', 'Utf8') AS x +) u; +---- +Utf8View a +Utf8View b + +query TR rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT CAST(1 AS SMALLINT) AS x + UNION ALL + SELECT 2.5 AS x +) u; +---- +Float64 1 +Float64 2.5 + +query TR rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT 2.5 AS x + UNION DISTINCT + SELECT CAST(1 AS SMALLINT) AS x +) u; +---- +Float64 1 +Float64 2.5 + +query TR rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT CASE WHEN false THEN 1 ELSE 2.5 END AS x + UNION ALL + SELECT CAST(3 AS INT) AS x +) u; +---- +Float64 2.5 +Float64 3 + +query TR rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT NULL AS x + UNION DISTINCT + SELECT CASE WHEN false THEN 1 WHEN true THEN 2.5 END AS x +) u; +---- +Float64 2.5 +Float64 NULL + +# The inferred output type must not seed coercion of dictionary inputs again. +query TT rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT NULL AS x + UNION ALL + SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x +) u; +---- +Utf8 NULL +Utf8 a + +query TI rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT CAST(1 AS SMALLINT) AS x + UNION ALL + SELECT arrow_cast(2, 'Dictionary(Int32, Int32)') AS x +) u; +---- +Int32 1 +Int32 2 + +query TR +SELECT arrow_typeof(max(x)), max(x) FROM ( + SELECT 1 AS x UNION ALL SELECT 2.5 AS x UNION ALL SELECT 3 AS x +) u; +---- +Float64 3 + +query I +SELECT CAST(NULL AS SMALLINT) AS x WHERE false +UNION ALL +SELECT CAST(NULL AS BIGINT) AS x WHERE false; +---- diff --git a/datafusion/sqllogictest/test_files/union_by_name.slt b/datafusion/sqllogictest/test_files/union_by_name.slt index 81f7317b2a39f..ddff4ba243667 100644 --- a/datafusion/sqllogictest/test_files/union_by_name.slt +++ b/datafusion/sqllogictest/test_files/union_by_name.slt @@ -326,3 +326,84 @@ select x, y, z from t3 union all by name select z, y, x, 'd' as zz from t3; ---- a b c NULL a b c d + +# Missing columns participate as NULL, including during dictionary coercion. +query TTIT rowsort +SELECT arrow_typeof(x), x, y, arrow_typeof(y) FROM ( + SELECT 1 AS y + UNION ALL BY NAME + SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x +) u; +---- +Utf8 NULL 1 Int64 +Utf8 a NULL Int64 + +query TTIT rowsort +SELECT arrow_typeof(x), x, y, arrow_typeof(y) FROM ( + SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x + UNION ALL BY NAME + SELECT 1 AS y +) u; +---- +Utf8 NULL 1 Int64 +Utf8 a NULL Int64 + +query TRIT rowsort +SELECT arrow_typeof(x), x, y, arrow_typeof(y) FROM ( + SELECT CAST(1 AS SMALLINT) AS x, 1 AS y + UNION ALL BY NAME + SELECT 2 AS y, 2.5 AS x +) u; +---- +Float64 1 1 Int64 +Float64 2.5 2 Int64 + +query TRIT rowsort +SELECT arrow_typeof(x), x, y, arrow_typeof(y) FROM ( + SELECT CAST(1 AS SMALLINT) AS x + UNION BY NAME + SELECT 2.5 AS x, 2 AS y +) u; +---- +Float64 1 NULL Int64 +Float64 2.5 2 Int64 + +query TT rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT arrow_cast('b', 'LargeUtf8') AS x + UNION BY NAME + SELECT arrow_cast('a', 'Utf8') AS x +) u; +---- +LargeUtf8 a +LargeUtf8 b + +query TT rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT arrow_cast('a', 'Utf8') AS x + UNION ALL BY NAME + SELECT arrow_cast('b', 'Utf8View') AS x +) u; +---- +Utf8View a +Utf8View b + +query TR rowsort +SELECT arrow_typeof(x), x FROM ( + SELECT CASE WHEN false THEN 1 ELSE 2.5 END AS x + UNION BY NAME + SELECT NULL AS x +) u; +---- +Float64 2.5 +Float64 NULL + +query TI rowsort +SELECT arrow_typeof(x), arrow_cast(x, 'Int64') FROM ( + SELECT CAST(NULL AS TIMESTAMP) AS x + UNION ALL BY NAME + SELECT arrow_cast(1, 'Timestamp(Microsecond, Some("UTC"))') AS x +) u; +---- +Timestamp(ns, "UTC") 1000 +Timestamp(ns, "UTC") NULL