From 854cf98258174e15b9dcafd732477358e3cfe0e7 Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Wed, 4 Nov 2026 20:47:12 -0500 Subject: [PATCH 1/4] Infer CASE result types from all result branches --- datafusion/core/tests/dataframe/case.rs | 48 +++++++++++++++++++++ datafusion/core/tests/dataframe/mod.rs | 1 + datafusion/expr/src/expr_schema.rs | 22 ++++++---- datafusion/sqllogictest/test_files/case.slt | 4 +- 4 files changed, 65 insertions(+), 10 deletions(-) create mode 100644 datafusion/core/tests/dataframe/case.rs diff --git a/datafusion/core/tests/dataframe/case.rs b/datafusion/core/tests/dataframe/case.rs new file mode 100644 index 0000000000000..361a0d3b5707e --- /dev/null +++ b/datafusion/core/tests/dataframe/case.rs @@ -0,0 +1,48 @@ +// 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 + +use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::error::Result; +use datafusion::prelude::SessionContext; + +#[tokio::test] +async fn case_schema_matches_execution() -> Result<()> { + let ctx = SessionContext::new(); + let cases = [ + // The ELSE branch widens the first THEN branch. + ("CASE WHEN p THEN 1::INT ELSE 0.5 END", DataType::Float64), + // A later THEN branch widens the result even without ELSE. + ( + "CASE WHEN p THEN 1::INT WHEN NOT p THEN 0.5 END", + DataType::Float64, + ), + // The inferred type must propagate through nested expressions. + ( + "CASE WHEN p THEN 2::INT ELSE (CASE WHEN p THEN 1::INT ELSE 0.5 END) END + 1", + DataType::Float64, + ), + // No concrete result type is available. + ("CASE WHEN p THEN NULL ELSE NULL END", DataType::Null), + // Use the same timestamp precision rule as the analyzer. + ( + "CASE WHEN p THEN arrow_cast(1, 'Timestamp(Microsecond, None)') ELSE arrow_cast(1, 'Timestamp(Nanosecond, None)') END", + DataType::Timestamp(TimeUnit::Nanosecond, None), + ), + ]; + for (expr, expected_type) in cases { + let df = ctx + .sql(&format!( + "SELECT {expr} AS x FROM (VALUES (true), (false)) AS t(p)" + )) + .await?; + assert_eq!(df.schema().field(0).data_type(), &expected_type, "{expr}"); + let declared = df.schema().as_arrow().clone(); + let batches = df.collect().await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + for batch in batches { + assert_eq!(batch.schema().as_ref(), &declared, "{expr}"); + } + } + Ok(()) +} diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 49635ed110e7f..e93566acaa415 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -3,6 +3,7 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements // Include tests in dataframe_functions +mod case; mod dataframe_functions; mod describe; diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index a1db445dd90c5..b880757b7e80e 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,20 @@ 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()?; + get_coerce_type_for_case_expression(&then_types, else_type.as_ref()) + .ok_or_else(|| plan_datafusion_err!( + "Cannot find a common type for CASE result branches {then_types:?} and ELSE {else_type:?}" + )) } Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { Ok(field.data_type().clone()) diff --git a/datafusion/sqllogictest/test_files/case.slt b/datafusion/sqllogictest/test_files/case.slt index 4727dd19d112f..87a39a35c8eb4 100644 --- a/datafusion/sqllogictest/test_files/case.slt +++ b/datafusion/sqllogictest/test_files/case.slt @@ -173,7 +173,7 @@ SELECT CASE 1 WHEN 1 THEN arrow_cast([1, 2], 'FixedSizeList(2, Int64)') WHEN 2 T ---- [1, 2] -query error DataFusion error: type_coercion +query error DataFusion error: Error during planning: Cannot find a common type for CASE result branches SELECT CASE 1 WHEN 1 THEN [1,2,3] WHEN 2 THEN 'test' ELSE null END; # test case when type coercion @@ -411,7 +411,7 @@ FROM t; # coerce structs with subset of fields - field count mismatch causes type coercion failure # column3 has 2 fields but column4 has only 1 field -query error DataFusion error: type_coercion\ncaused by\nError during planning: Failed to coerce then .* and else .* to common types in CASE WHEN expression +query error DataFusion error: Error during planning: Cannot find a common type for CASE result branches .* and ELSE .* SELECT case when column1 > 0 then column3 From 78fc261b1c404a3d8cf34805fed9adc8df38015a Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Wed, 4 Nov 2026 21:04:16 -0500 Subject: [PATCH 2/4] Infer UNION output types before analysis --- datafusion/core/tests/dataframe/mod.rs | 1 + datafusion/core/tests/dataframe/union.rs | 125 +++++++++++++++++++ datafusion/expr/src/logical_plan/builder.rs | 5 +- datafusion/expr/src/logical_plan/plan.rs | 70 +++++++---- datafusion/sqllogictest/test_files/union.slt | 2 +- 5 files changed, 176 insertions(+), 27 deletions(-) create mode 100644 datafusion/core/tests/dataframe/union.rs diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index e93566acaa415..6ce0d01c27949 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -6,6 +6,7 @@ mod case; mod dataframe_functions; mod describe; +mod union; use arrow::array::{ Array, ArrayRef, BooleanArray, DictionaryArray, FixedSizeListArray, diff --git a/datafusion/core/tests/dataframe/union.rs b/datafusion/core/tests/dataframe/union.rs new file mode 100644 index 0000000000000..fe4848e075f54 --- /dev/null +++ b/datafusion/core/tests/dataframe/union.rs @@ -0,0 +1,125 @@ +// 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 + +use datafusion::arrow::array::Float64Array; +use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::common::ScalarValue; +use datafusion::dataframe::DataFrame; +use datafusion::error::Result; +use datafusion::prelude::SessionContext; + +async fn check_schema(df: DataFrame, expected: DataType) -> Result<()> { + assert_eq!(df.schema().field(0).data_type(), &expected); + let schema = df.schema().as_arrow().clone(); + let batches = df.collect().await?; + assert!(!batches.is_empty()); + for batch in batches { + assert_eq!(batch.schema().as_ref(), &schema); + } + Ok(()) +} + +#[tokio::test] +async fn union_schema_matches_execution() -> Result<()> { + let ctx = SessionContext::new(); + let cases = [ + ( + "SELECT t FROM (SELECT arrow_cast(arrow_cast(1, 'Timestamp(Microsecond, None)'), 'Timestamp(Microsecond, Some(\"UTC\"))') AS t UNION ALL SELECT CAST(NULL AS TIMESTAMP) AS t) u", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + ), + ( + "SELECT arrow_cast('a', 'Utf8') AS x UNION ALL SELECT arrow_cast('b', 'LargeUtf8') AS x", + DataType::LargeUtf8, + ), + ( + "SELECT arrow_cast('a', 'Utf8') AS x UNION ALL SELECT arrow_cast('b', 'Utf8View') AS x", + DataType::Utf8View, + ), + ( + "SELECT x FROM (SELECT 1::INT AS x UNION ALL SELECT 2::BIGINT AS x) u UNION ALL SELECT 3::SMALLINT AS x", + DataType::Int64, + ), + ( + "SELECT * FROM (SELECT 1::INT AS x) a UNION ALL SELECT * FROM (SELECT 2::BIGINT AS x) b", + DataType::Int64, + ), + ( + "SELECT 1::INT AS x UNION SELECT 2::BIGINT AS x", + DataType::Int64, + ), + ( + "SELECT 1::INT AS x UNION ALL BY NAME SELECT 2::BIGINT AS x", + DataType::Int64, + ), + ( + "SELECT 1::INT AS x UNION ALL BY NAME SELECT 'a' AS y, 2::BIGINT AS x", + DataType::Int64, + ), + ]; + for (sql, expected) in cases { + check_schema(ctx.sql(sql).await?, expected).await?; + } + Ok(()) +} + +#[tokio::test] +async fn union_schema_after_binding() -> Result<()> { + let ctx = SessionContext::new(); + for sql in [ + "SELECT $1 AS x UNION ALL SELECT 2::BIGINT AS x", + "SELECT CASE WHEN false THEN 1::INT ELSE $1 END AS x UNION ALL SELECT 2::BIGINT AS x", + "SELECT x FROM (SELECT $1 AS x UNION ALL SELECT 2::BIGINT AS x) u", + "SELECT $1 AS x UNION ALL BY NAME SELECT 2::BIGINT AS x", + ] { + let df = ctx + .sql(sql) + .await? + .with_param_values(vec![ScalarValue::Float64(Some(0.5))])?; + check_fractional_values(df).await?; + } + Ok(()) +} + +#[tokio::test] +async fn dataframe_union_schema() -> Result<()> { + let ctx = SessionContext::new(); + let left = ctx.sql("SELECT 1::INT AS x").await?; + let right = ctx.sql("SELECT 2::BIGINT AS x").await?; + check_schema(left.union(right)?, DataType::Int64).await +} + +// These expressions previously made early UNION casts truncate the fractional +// result. Both inference and execution must retain the common floating type. +#[tokio::test] +async fn union_preserves_fractional_results() -> Result<()> { + let ctx = SessionContext::new(); + for sql in [ + "SELECT CASE WHEN false THEN 1::INT ELSE 0.5 END AS x UNION ALL SELECT 2::BIGINT AS x", + "WITH t(v) AS (VALUES (2::INT), (NULL::INT)) SELECT nvl(v, 0.5) AS x FROM t UNION ALL SELECT NULL AS x WHERE false", + ] { + check_fractional_values(ctx.sql(sql).await?).await?; + } + Ok(()) +} + +async fn check_fractional_values(df: DataFrame) -> Result<()> { + assert_eq!(df.schema().field(0).data_type(), &DataType::Float64); + let mut values = Vec::new(); + for batch in df.collect().await? { + // Removing the empty NULL input can tighten nullability. + assert_eq!(batch.schema().field(0).data_type(), &DataType::Float64); + assert_eq!(batch.column(0).null_count(), 0); + values.extend( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + } + values.sort_by(f64::total_cmp); + assert_eq!(values, vec![0.5, 2.0]); + Ok(()) +} 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..5c3202b4ba016 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,8 @@ 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. + /// 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 +3591,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 +3602,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. @@ -3617,18 +3623,32 @@ impl Union { } } + fn coerce_field_type(data_type: &DataType, field: &Field) -> Result { + 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() + ) + }) + } + fn derive_schema_from_inputs_by_name( inputs: &[Arc], loose_types: bool, ) -> Result { - type FieldData<'a> = (&'a DataType, bool, Vec<&'a Metadata>, usize); + type FieldData<'a> = (DataType, bool, Vec<&'a Metadata>, usize); let mut cols: Vec<(&str, FieldData)> = Vec::new(); for input in inputs.iter() { for field in input.schema().fields() { if let Some((_, (data_type, is_nullable, metadata, occurrences))) = cols.iter_mut().find(|(name, _)| name == field.name()) { - if !loose_types && *data_type != field.data_type() { + if loose_types { + *data_type = Self::coerce_field_type(data_type, field)?; + } else if data_type != field.data_type() { return plan_err!( "Found different types for field {}", field.name() @@ -3644,7 +3664,7 @@ impl Union { cols.push(( field.name(), ( - field.data_type(), + field.data_type().clone(), field.is_nullable(), vec![field.metadata()], 1, @@ -3667,8 +3687,7 @@ impl Union { true }; - let mut field = - Field::new(name, data_type.clone(), final_is_nullable); + 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 +3734,17 @@ 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| { + Self::coerce_field_type(&acc, field) + })? } 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 +3754,8 @@ impl Union { ); } Ok(acc) - }, - )? + })? + .clone() }; let nullable = fields.iter().any(|field| field.is_nullable()); @@ -3746,7 +3768,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/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 1a8d88ee149bb..22557ce17f1e4 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -916,7 +916,7 @@ SELECT NULL WHERE FALSE; 1 # Test Union of List Types. Issue: https://github.com/apache/datafusion/issues/12291 -query error DataFusion error: type_coercion\ncaused by\nError during planning: Incompatible inputs for Union: Previous inputs were of type List(.*), but got incompatible type List(.*) on column 'x' +query error DataFusion error: Error during planning: Incompatible inputs for Union: Previous inputs were of type List(.*), but got incompatible type List(.*) on column 'x' SELECT make_array(2) x UNION ALL SELECT make_array(now()) x; query ? rowsort From dea0561008aecbbf535be97ff9b509efa9d4cd1c Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Thu, 5 Nov 2026 12:54:54 -0500 Subject: [PATCH 3/4] test: cover CASE and UNION schemas with SQL --- datafusion/core/tests/dataframe/case.rs | 48 ------- datafusion/core/tests/dataframe/mod.rs | 2 - datafusion/core/tests/dataframe/union.rs | 125 ------------------ datafusion/sqllogictest/test_files/case.slt | 10 ++ datafusion/sqllogictest/test_files/union.slt | 16 +++ .../sqllogictest/test_files/union_by_name.slt | 8 ++ 6 files changed, 34 insertions(+), 175 deletions(-) delete mode 100644 datafusion/core/tests/dataframe/case.rs delete mode 100644 datafusion/core/tests/dataframe/union.rs diff --git a/datafusion/core/tests/dataframe/case.rs b/datafusion/core/tests/dataframe/case.rs deleted file mode 100644 index 361a0d3b5707e..0000000000000 --- a/datafusion/core/tests/dataframe/case.rs +++ /dev/null @@ -1,48 +0,0 @@ -// 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 - -use datafusion::arrow::datatypes::{DataType, TimeUnit}; -use datafusion::error::Result; -use datafusion::prelude::SessionContext; - -#[tokio::test] -async fn case_schema_matches_execution() -> Result<()> { - let ctx = SessionContext::new(); - let cases = [ - // The ELSE branch widens the first THEN branch. - ("CASE WHEN p THEN 1::INT ELSE 0.5 END", DataType::Float64), - // A later THEN branch widens the result even without ELSE. - ( - "CASE WHEN p THEN 1::INT WHEN NOT p THEN 0.5 END", - DataType::Float64, - ), - // The inferred type must propagate through nested expressions. - ( - "CASE WHEN p THEN 2::INT ELSE (CASE WHEN p THEN 1::INT ELSE 0.5 END) END + 1", - DataType::Float64, - ), - // No concrete result type is available. - ("CASE WHEN p THEN NULL ELSE NULL END", DataType::Null), - // Use the same timestamp precision rule as the analyzer. - ( - "CASE WHEN p THEN arrow_cast(1, 'Timestamp(Microsecond, None)') ELSE arrow_cast(1, 'Timestamp(Nanosecond, None)') END", - DataType::Timestamp(TimeUnit::Nanosecond, None), - ), - ]; - for (expr, expected_type) in cases { - let df = ctx - .sql(&format!( - "SELECT {expr} AS x FROM (VALUES (true), (false)) AS t(p)" - )) - .await?; - assert_eq!(df.schema().field(0).data_type(), &expected_type, "{expr}"); - let declared = df.schema().as_arrow().clone(); - let batches = df.collect().await?; - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); - for batch in batches { - assert_eq!(batch.schema().as_ref(), &declared, "{expr}"); - } - } - Ok(()) -} diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 6ce0d01c27949..49635ed110e7f 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -3,10 +3,8 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements // Include tests in dataframe_functions -mod case; mod dataframe_functions; mod describe; -mod union; use arrow::array::{ Array, ArrayRef, BooleanArray, DictionaryArray, FixedSizeListArray, diff --git a/datafusion/core/tests/dataframe/union.rs b/datafusion/core/tests/dataframe/union.rs deleted file mode 100644 index fe4848e075f54..0000000000000 --- a/datafusion/core/tests/dataframe/union.rs +++ /dev/null @@ -1,125 +0,0 @@ -// 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 - -use datafusion::arrow::array::Float64Array; -use datafusion::arrow::datatypes::{DataType, TimeUnit}; -use datafusion::common::ScalarValue; -use datafusion::dataframe::DataFrame; -use datafusion::error::Result; -use datafusion::prelude::SessionContext; - -async fn check_schema(df: DataFrame, expected: DataType) -> Result<()> { - assert_eq!(df.schema().field(0).data_type(), &expected); - let schema = df.schema().as_arrow().clone(); - let batches = df.collect().await?; - assert!(!batches.is_empty()); - for batch in batches { - assert_eq!(batch.schema().as_ref(), &schema); - } - Ok(()) -} - -#[tokio::test] -async fn union_schema_matches_execution() -> Result<()> { - let ctx = SessionContext::new(); - let cases = [ - ( - "SELECT t FROM (SELECT arrow_cast(arrow_cast(1, 'Timestamp(Microsecond, None)'), 'Timestamp(Microsecond, Some(\"UTC\"))') AS t UNION ALL SELECT CAST(NULL AS TIMESTAMP) AS t) u", - DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), - ), - ( - "SELECT arrow_cast('a', 'Utf8') AS x UNION ALL SELECT arrow_cast('b', 'LargeUtf8') AS x", - DataType::LargeUtf8, - ), - ( - "SELECT arrow_cast('a', 'Utf8') AS x UNION ALL SELECT arrow_cast('b', 'Utf8View') AS x", - DataType::Utf8View, - ), - ( - "SELECT x FROM (SELECT 1::INT AS x UNION ALL SELECT 2::BIGINT AS x) u UNION ALL SELECT 3::SMALLINT AS x", - DataType::Int64, - ), - ( - "SELECT * FROM (SELECT 1::INT AS x) a UNION ALL SELECT * FROM (SELECT 2::BIGINT AS x) b", - DataType::Int64, - ), - ( - "SELECT 1::INT AS x UNION SELECT 2::BIGINT AS x", - DataType::Int64, - ), - ( - "SELECT 1::INT AS x UNION ALL BY NAME SELECT 2::BIGINT AS x", - DataType::Int64, - ), - ( - "SELECT 1::INT AS x UNION ALL BY NAME SELECT 'a' AS y, 2::BIGINT AS x", - DataType::Int64, - ), - ]; - for (sql, expected) in cases { - check_schema(ctx.sql(sql).await?, expected).await?; - } - Ok(()) -} - -#[tokio::test] -async fn union_schema_after_binding() -> Result<()> { - let ctx = SessionContext::new(); - for sql in [ - "SELECT $1 AS x UNION ALL SELECT 2::BIGINT AS x", - "SELECT CASE WHEN false THEN 1::INT ELSE $1 END AS x UNION ALL SELECT 2::BIGINT AS x", - "SELECT x FROM (SELECT $1 AS x UNION ALL SELECT 2::BIGINT AS x) u", - "SELECT $1 AS x UNION ALL BY NAME SELECT 2::BIGINT AS x", - ] { - let df = ctx - .sql(sql) - .await? - .with_param_values(vec![ScalarValue::Float64(Some(0.5))])?; - check_fractional_values(df).await?; - } - Ok(()) -} - -#[tokio::test] -async fn dataframe_union_schema() -> Result<()> { - let ctx = SessionContext::new(); - let left = ctx.sql("SELECT 1::INT AS x").await?; - let right = ctx.sql("SELECT 2::BIGINT AS x").await?; - check_schema(left.union(right)?, DataType::Int64).await -} - -// These expressions previously made early UNION casts truncate the fractional -// result. Both inference and execution must retain the common floating type. -#[tokio::test] -async fn union_preserves_fractional_results() -> Result<()> { - let ctx = SessionContext::new(); - for sql in [ - "SELECT CASE WHEN false THEN 1::INT ELSE 0.5 END AS x UNION ALL SELECT 2::BIGINT AS x", - "WITH t(v) AS (VALUES (2::INT), (NULL::INT)) SELECT nvl(v, 0.5) AS x FROM t UNION ALL SELECT NULL AS x WHERE false", - ] { - check_fractional_values(ctx.sql(sql).await?).await?; - } - Ok(()) -} - -async fn check_fractional_values(df: DataFrame) -> Result<()> { - assert_eq!(df.schema().field(0).data_type(), &DataType::Float64); - let mut values = Vec::new(); - for batch in df.collect().await? { - // Removing the empty NULL input can tighten nullability. - assert_eq!(batch.schema().field(0).data_type(), &DataType::Float64); - assert_eq!(batch.column(0).null_count(), 0); - values.extend( - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values(), - ); - } - values.sort_by(f64::total_cmp); - assert_eq!(values, vec![0.5, 2.0]); - Ok(()) -} diff --git a/datafusion/sqllogictest/test_files/case.slt b/datafusion/sqllogictest/test_files/case.slt index 87a39a35c8eb4..c986f7c36518b 100644 --- a/datafusion/sqllogictest/test_files/case.slt +++ b/datafusion/sqllogictest/test_files/case.slt @@ -132,6 +132,16 @@ 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 +FROM (VALUES (true), (false)) AS t(p); +---- +else_widens Float64 NO +then_widens Float64 YES + # 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 22557ce17f1e4..0070f1423118f 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -6,6 +6,22 @@ ## 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 + 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..5cc49b5450ce5 100644 --- a/datafusion/sqllogictest/test_files/union_by_name.slt +++ b/datafusion/sqllogictest/test_files/union_by_name.slt @@ -40,6 +40,14 @@ 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 + query I SELECT t1.x FROM t1 UNION BY NAME SELECT x FROM t1 ORDER BY x; ---- From 21290b2768ac733137c082de07160e805b1a266c Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Thu, 5 Nov 2026 15:57:07 -0500 Subject: [PATCH 4/4] stealing ideas from @kumarUjjawal Incorporate deferred CASE/UNION coercion errors, dictionary coercion from actual inputs, and an empty-input guard, informed by: https://github.com/apache/datafusion/pull/27139 Keep regression coverage in SQL and document the dictionary output-type change for UNION BY NAME. --- datafusion/expr/src/expr_schema.rs | 16 +++-- datafusion/expr/src/logical_plan/plan.rs | 71 +++++++++++++------ .../optimizer/src/analyzer/type_coercion.rs | 14 ++-- datafusion/sqllogictest/test_files/case.slt | 8 ++- datafusion/sqllogictest/test_files/union.slt | 16 ++++- .../sqllogictest/test_files/union_by_name.slt | 22 ++++++ .../library-user-guide/upgrading/56.0.0.md | 8 +++ 7 files changed, 121 insertions(+), 34 deletions(-) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index b880757b7e80e..3a0f70b0dc3c2 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -188,10 +188,18 @@ impl ExprSchemable for Expr { .as_ref() .map(|expr| expr.get_type(schema)) .transpose()?; - get_coerce_type_for_case_expression(&then_types, else_type.as_ref()) - .ok_or_else(|| plan_datafusion_err!( - "Cannot find a common type for CASE result branches {then_types:?} and ELSE {else_type:?}" - )) + // 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/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 5c3202b4ba016..518df4a4ffe1c 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -3548,6 +3548,7 @@ impl Union { /// Constructs new Union instance deriving schema from inputs. /// 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 { @@ -3623,32 +3624,42 @@ impl Union { } } - fn coerce_field_type(data_type: &DataType, field: &Field) -> Result { - 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() - ) - }) - } - fn derive_schema_from_inputs_by_name( inputs: &[Arc], loose_types: bool, ) -> Result { - type FieldData<'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 = Self::coerce_field_type(data_type, field)?; - } else if data_type != field.data_type() { + *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() @@ -3664,7 +3675,12 @@ impl Union { cols.push(( field.name(), ( - field.data_type().clone(), + 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, @@ -3677,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. @@ -3687,6 +3712,7 @@ impl Union { true }; + 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)); @@ -3738,8 +3764,9 @@ impl Union { .iter() .skip(1) .try_fold(first_field.data_type().clone(), |acc, field| { - Self::coerce_field_type(&acc, field) - })? + type_union_coercion(&acc, field.data_type()) + }) + .unwrap_or_else(|| first_field.data_type().clone()) } else { fields .iter() 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 c986f7c36518b..9f6711a24f7fd 100644 --- a/datafusion/sqllogictest/test_files/case.slt +++ b/datafusion/sqllogictest/test_files/case.slt @@ -136,11 +136,13 @@ SELECT CASE WHEN false THEN 1 ELSE 0 END FROM foo 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::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 ? @@ -183,7 +185,7 @@ SELECT CASE 1 WHEN 1 THEN arrow_cast([1, 2], 'FixedSizeList(2, Int64)') WHEN 2 T ---- [1, 2] -query error DataFusion error: Error during planning: Cannot find a common type for CASE result branches +query error DataFusion error: type_coercion SELECT CASE 1 WHEN 1 THEN [1,2,3] WHEN 2 THEN 'test' ELSE null END; # test case when type coercion @@ -421,7 +423,7 @@ FROM t; # coerce structs with subset of fields - field count mismatch causes type coercion failure # column3 has 2 fields but column4 has only 1 field -query error DataFusion error: Error during planning: Cannot find a common type for CASE result branches .* and ELSE .* +query error DataFusion error: type_coercion\ncaused by\nError during planning: Failed to coerce then .* and else .* to common types in CASE WHEN expression SELECT case when column1 > 0 then column3 diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 0070f1423118f..5096f94fef911 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -22,6 +22,20 @@ 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, @@ -932,7 +946,7 @@ SELECT NULL WHERE FALSE; 1 # Test Union of List Types. Issue: https://github.com/apache/datafusion/issues/12291 -query error DataFusion error: Error during planning: Incompatible inputs for Union: Previous inputs were of type List(.*), but got incompatible type List(.*) on column 'x' +query error DataFusion error: type_coercion\ncaused by\nError during planning: Incompatible inputs for Union: Previous inputs were of type List(.*), but got incompatible type List(.*) on column 'x' SELECT make_array(2) x UNION ALL SELECT make_array(now()) x; query ? rowsort diff --git a/datafusion/sqllogictest/test_files/union_by_name.slt b/datafusion/sqllogictest/test_files/union_by_name.slt index 5cc49b5450ce5..c531863001650 100644 --- a/datafusion/sqllogictest/test_files/union_by_name.slt +++ b/datafusion/sqllogictest/test_files/union_by_name.slt @@ -48,6 +48,28 @@ 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