From 287009597f7e33a1f613df7d818bfc7722717cb0 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Thu, 5 Nov 2026 22:41:37 +0530 Subject: [PATCH 1/2] fix: infer UNION output types before analysis --- datafusion/core/tests/sql/mod.rs | 1 + datafusion/core/tests/sql/union_nullable.rs | 8 +- datafusion/core/tests/sql/union_schema.rs | 220 ++++++++++++++++++ datafusion/expr/src/expr_schema.rs | 59 ++++- datafusion/expr/src/logical_plan/builder.rs | 4 +- datafusion/expr/src/logical_plan/plan.rs | 173 +++++++++----- datafusion/expr/src/type_coercion/other.rs | 27 ++- .../optimizer/src/analyzer/type_coercion.rs | 102 ++++---- .../library-user-guide/upgrading/56.0.0.md | 28 +++ 9 files changed, 487 insertions(+), 135 deletions(-) create mode 100644 datafusion/core/tests/sql/union_schema.rs 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..e9f7af5a2e731 --- /dev/null +++ b/datafusion/core/tests/sql/union_schema.rs @@ -0,0 +1,220 @@ +// 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::array::{AsArray, RecordBatch}; +use arrow::datatypes::{DataType, Float64Type, Schema, TimeUnit}; +use datafusion::assert_batches_sorted_eq; +use datafusion::prelude::*; +use datafusion_common::{Result, ScalarValue}; + +async fn assert_union_schema( + df: DataFrame, + expected: &[(&str, DataType)], + rows: usize, +) -> Result> { + let assert_fields = |schema: &Schema| { + assert_eq!(schema.fields().len(), expected.len()); + for (field, (name, data_type)) in schema.fields().iter().zip(expected) { + assert_eq!(field.name(), name); + assert_eq!(field.data_type(), data_type); + } + }; + assert_fields(df.schema().as_arrow()); + // Empty results must also have the advertised type. + assert_fields(&df.clone().create_physical_plan().await?.schema()); + let batches = df.collect().await?; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + rows + ); + for batch in &batches { + assert_fields(&batch.schema()); + } + Ok(batches) +} + +#[tokio::test] +async fn union_declared_types_match_execution() -> 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())], 2) + .await?; + } + } + } + Ok(()) +} + +#[tokio::test] +async fn union_dictionary_types_use_actual_inputs() -> Result<()> { + let ctx = SessionContext::new(); + // Reseeding coercion with the inferred output can select Dictionary again. + 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)], 2).await?; + } + // A missing left column previously preserved Dictionary. Both orders now + // match positional UNION and return Utf8. + 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)], 2) + .await?; + } + 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)], 2) + .await?; + } + 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)], + 2, + ) + .await?; + } + 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)], 0).await?; + Ok(()) +} + +#[tokio::test] +async fn union_coercion_preserves_values() -> Result<()> { + let ctx = SessionContext::new(); + let batches = ctx + .sql( + "SELECT arrow_cast(x, 'Int64') AS x FROM (\ + SELECT arrow_cast(1, 'Timestamp(Microsecond, None)') AS x UNION ALL \ + SELECT arrow_cast(1001, 'Timestamp(Nanosecond, None)') AS x) u", + ) + .await? + .collect() + .await?; + assert_batches_sorted_eq!( + [ + "+------+", "| x |", "+------+", "| 1000 |", "| 1001 |", "+------+" + ], + &batches + ); + for (sql, expected) in [ + ( + "SELECT nvl(v1, 0.5) AS x FROM (VALUES (CAST(1 AS INT)), (NULL)) t(v1) \ + UNION ALL SELECT NULL WHERE false", + vec![0.5, 1.0], + ), + ( + "SELECT CASE WHEN false THEN 1 ELSE 2.5 END AS x \ + UNION ALL SELECT CAST(3 AS INT)", + vec![2.5, 3.0], + ), + ] { + let batches = + assert_union_schema(ctx.sql(sql).await?, &[("x", DataType::Float64)], 2) + .await?; + let mut values = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect::>(); + values.sort_by(f64::total_cmp); + assert_eq!(values, expected); + } + 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)], 1).await?; + 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))])?; + let batches = assert_union_schema(df, &[("x", DataType::Float64)], 2).await?; + assert_batches_sorted_eq!( + [ + "+-----+", "| x |", "+-----+", "| 1.0 |", "| 2.5 |", "+-----+" + ], + &batches + ); + Ok(()) +} diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index a1db445dd90c5..c963bbfc765a2 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(|(_, 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. + Ok( + get_coerce_type_for_case_expression(&then_types, else_type.as_ref()) + .unwrap_or_else(|| { + then_types + .into_iter() + .find(|t| !t.is_null()) + .or(else_type) + .unwrap_or(DataType::Null) + }), + ) } Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { Ok(field.data_type().clone()) @@ -1684,4 +1698,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 c5b5b23ba8cc9..4f9aefc511b0d 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 }) @@ -3595,9 +3598,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 +3623,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 +3719,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 +3735,8 @@ impl Union { ); } Ok(acc) - }, - )? + })? + .clone() }; let nullable = fields.iter().any(|field| field.is_nullable()); @@ -3746,7 +3749,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); @@ -5520,6 +5523,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, @@ -7762,4 +7766,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 fc798449b0b2a..ff18f79ebd11a 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/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index b567e505e27e4..1f24c1c7b3bda 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,34 @@ 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 and CASE infer common output types before analysis + +`Union::try_new_with_loose_types` and `logical_plan::builder::union` now infer +common output types from the input schemas. Input casts remain deferred until +analysis. Logical schemas before analysis, for example `DataFrame::schema()`, +previously used the first input's type. +Unsupported combinations keep a provisional type until analysis. + +`ExprSchemable::get_type` now considers all `CASE` result branches, including +`ELSE`. It can return a wider type and report branch errors earlier. Deeply +nested `CASE` expressions can require more planning work. + +Low-level physical expression construction still requires coerced inputs. +`SessionState::create_physical_expr` applies coercion automatically. Direct +callers of `datafusion_physical_expr::create_physical_expr` must coerce +mixed-type `CASE` expressions first. + +`datafusion_expr::type_coercion::other::coerce_union_type` is a new helper +that infers one UNION column type from input fields. + +For `UNION BY NAME`, missing columns participate as `NULL` during type +inference. With two inputs, a missing `Dictionary(Int32, Utf8)` column now +produces `Utf8` in either input order. Previously, a missing left column preserved the +dictionary representation. Decoded values remain the same. + +Consumers that depend on the Arrow representation must update their schema +expectations or cast each UNION input to the required type. + ### Session time zone now applies to mixed timestamp operations When `datafusion.execution.time_zone` is set, DataFusion now uses it to From 7f6236673d3a885af7f964a239ef5b87e59ee52f Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Fri, 6 Nov 2026 09:42:37 +0530 Subject: [PATCH 2/2] test: move UNION execution regressions to SLT --- datafusion/core/tests/sql/union_schema.rs | 118 +++------------ datafusion/expr/src/expr_schema.rs | 24 ++-- datafusion/expr/src/logical_plan/plan.rs | 5 + .../sqllogictest/test_files/prepare.slt | 12 ++ datafusion/sqllogictest/test_files/union.slt | 136 ++++++++++++++++++ .../sqllogictest/test_files/union_by_name.slt | 81 +++++++++++ .../library-user-guide/upgrading/56.0.0.md | 28 ---- 7 files changed, 266 insertions(+), 138 deletions(-) diff --git a/datafusion/core/tests/sql/union_schema.rs b/datafusion/core/tests/sql/union_schema.rs index e9f7af5a2e731..65b6bddae306a 100644 --- a/datafusion/core/tests/sql/union_schema.rs +++ b/datafusion/core/tests/sql/union_schema.rs @@ -4,40 +4,21 @@ //! The pre-analysis UNION schema must describe the output types (#27023). -use arrow::array::{AsArray, RecordBatch}; -use arrow::datatypes::{DataType, Float64Type, Schema, TimeUnit}; -use datafusion::assert_batches_sorted_eq; +use arrow::datatypes::{DataType, TimeUnit}; use datafusion::prelude::*; use datafusion_common::{Result, ScalarValue}; -async fn assert_union_schema( - df: DataFrame, - expected: &[(&str, DataType)], - rows: usize, -) -> Result> { - let assert_fields = |schema: &Schema| { - assert_eq!(schema.fields().len(), expected.len()); - for (field, (name, data_type)) in schema.fields().iter().zip(expected) { - assert_eq!(field.name(), name); - assert_eq!(field.data_type(), data_type); - } - }; - assert_fields(df.schema().as_arrow()); - // Empty results must also have the advertised type. - assert_fields(&df.clone().create_physical_plan().await?.schema()); - let batches = df.collect().await?; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - rows - ); - for batch in &batches { - assert_fields(&batch.schema()); +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); } - Ok(batches) } #[tokio::test] -async fn union_declared_types_match_execution() -> Result<()> { +async fn union_infers_types_before_analysis() -> Result<()> { let ctx = SessionContext::new(); let cases = [ ( @@ -73,8 +54,7 @@ async fn union_declared_types_match_execution() -> Result<()> { 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())], 2) - .await?; + assert_union_schema(&ctx.sql(&sql).await?, &[("x", expected.clone())]); } } } @@ -82,9 +62,8 @@ async fn union_declared_types_match_execution() -> Result<()> { } #[tokio::test] -async fn union_dictionary_types_use_actual_inputs() -> Result<()> { +async fn union_infers_dictionary_types_before_analysis() -> Result<()> { let ctx = SessionContext::new(); - // Reseeding coercion with the inferred output can select Dictionary again. for (sql, expected) in [ ( "SELECT NULL AS x UNION ALL SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x", @@ -95,17 +74,14 @@ async fn union_dictionary_types_use_actual_inputs() -> Result<()> { DataType::Int32, ), ] { - assert_union_schema(ctx.sql(sql).await?, &[("x", expected)], 2).await?; + assert_union_schema(&ctx.sql(sql).await?, &[("x", expected)]); } - // A missing left column previously preserved Dictionary. Both orders now - // match positional UNION and return Utf8. 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)], 2) - .await?; + assert_union_schema(&df, &[("x", DataType::Utf8), ("y", DataType::Int64)]); } Ok(()) } @@ -122,19 +98,16 @@ async fn union_dataframe_and_by_name_schemas() -> Result<()> { 2 => left.union_by_name(right)?, _ => left.union_by_name_distinct(right)?, }; - assert_union_schema(df.select_columns(&["x"])?, &[("x", DataType::Float64)], 2) - .await?; + 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?, + &ctx.sql(sql).await?, &[("x", DataType::Float64), ("y", DataType::Int64)], - 2, - ) - .await?; + ); } let df = ctx .sql( @@ -142,63 +115,14 @@ async fn union_dataframe_and_by_name_schemas() -> Result<()> { UNION ALL SELECT CAST(NULL AS BIGINT) AS x WHERE false", ) .await?; - assert_union_schema(df, &[("x", DataType::Int64)], 0).await?; - Ok(()) -} - -#[tokio::test] -async fn union_coercion_preserves_values() -> Result<()> { - let ctx = SessionContext::new(); - let batches = ctx - .sql( - "SELECT arrow_cast(x, 'Int64') AS x FROM (\ - SELECT arrow_cast(1, 'Timestamp(Microsecond, None)') AS x UNION ALL \ - SELECT arrow_cast(1001, 'Timestamp(Nanosecond, None)') AS x) u", - ) - .await? - .collect() - .await?; - assert_batches_sorted_eq!( - [ - "+------+", "| x |", "+------+", "| 1000 |", "| 1001 |", "+------+" - ], - &batches - ); - for (sql, expected) in [ - ( - "SELECT nvl(v1, 0.5) AS x FROM (VALUES (CAST(1 AS INT)), (NULL)) t(v1) \ - UNION ALL SELECT NULL WHERE false", - vec![0.5, 1.0], - ), - ( - "SELECT CASE WHEN false THEN 1 ELSE 2.5 END AS x \ - UNION ALL SELECT CAST(3 AS INT)", - vec![2.5, 3.0], - ), - ] { - let batches = - assert_union_schema(ctx.sql(sql).await?, &[("x", DataType::Float64)], 2) - .await?; - let mut values = batches - .iter() - .flat_map(|b| { - b.column(0) - .as_primitive::() - .values() - .iter() - .copied() - }) - .collect::>(); - values.sort_by(f64::total_cmp); - assert_eq!(values, expected); - } + 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)], 1).await?; + assert_union_schema(&df, &[("x", DataType::Float64)]); Ok(()) } @@ -209,12 +133,6 @@ async fn union_bound_parameters_keep_inferred_type() -> Result<()> { .sql("SELECT CAST($1 AS INT) AS x UNION ALL SELECT 2.5 AS x") .await? .with_param_values(vec![ScalarValue::Int32(Some(1))])?; - let batches = assert_union_schema(df, &[("x", DataType::Float64)], 2).await?; - assert_batches_sorted_eq!( - [ - "+-----+", "| x |", "+-----+", "| 1.0 |", "| 2.5 |", "+-----+" - ], - &batches - ); + 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 c963bbfc765a2..ddabb95740afa 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -34,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 @@ -190,16 +194,16 @@ impl ExprSchemable for Expr { .transpose()?; // If the branches have no common type, keep the old provisional // type. The analyzer reports the error. - Ok( - get_coerce_type_for_case_expression(&then_types, else_type.as_ref()) - .unwrap_or_else(|| { - then_types - .into_iter() - .find(|t| !t.is_null()) - .or(else_type) - .unwrap_or(DataType::Null) - }), - ) + 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()) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 4f9aefc511b0d..74a107349a255 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -3555,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)?; 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 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 1f24c1c7b3bda..b567e505e27e4 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -12,34 +12,6 @@ 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 and CASE infer common output types before analysis - -`Union::try_new_with_loose_types` and `logical_plan::builder::union` now infer -common output types from the input schemas. Input casts remain deferred until -analysis. Logical schemas before analysis, for example `DataFrame::schema()`, -previously used the first input's type. -Unsupported combinations keep a provisional type until analysis. - -`ExprSchemable::get_type` now considers all `CASE` result branches, including -`ELSE`. It can return a wider type and report branch errors earlier. Deeply -nested `CASE` expressions can require more planning work. - -Low-level physical expression construction still requires coerced inputs. -`SessionState::create_physical_expr` applies coercion automatically. Direct -callers of `datafusion_physical_expr::create_physical_expr` must coerce -mixed-type `CASE` expressions first. - -`datafusion_expr::type_coercion::other::coerce_union_type` is a new helper -that infers one UNION column type from input fields. - -For `UNION BY NAME`, missing columns participate as `NULL` during type -inference. With two inputs, a missing `Dictionary(Int32, Utf8)` column now -produces `Utf8` in either input order. Previously, a missing left column preserved the -dictionary representation. Decoded values remain the same. - -Consumers that depend on the Arrow representation must update their schema -expectations or cast each UNION input to the required type. - ### Session time zone now applies to mixed timestamp operations When `datafusion.execution.time_zone` is set, DataFusion now uses it to