Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 22 additions & 8 deletions datafusion/expr/src/expr_schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -177,15 +178,28 @@ impl ExprSchemable for Expr {
Expr::ScalarVariable(field, _) => Ok(field.data_type().clone()),
Expr::Literal(l, _) => Ok(l.data_type()),
Expr::Case(case) => {
for (_, then_expr) in &case.when_then_expr {
let then_type = then_expr.get_type(schema)?;
if !then_type.is_null() {
return Ok(then_type);
}
}
case.else_expr
let then_types = case
.when_then_expr
.iter()
.map(|(_, then_expr)| then_expr.get_type(schema))
.collect::<Result<Vec<_>>>()?;
let else_type = case
.else_expr
.as_ref()
.map_or(Ok(DataType::Null), |e| e.get_type(schema))
.map(|expr| expr.get_type(schema))
.transpose()?;
// Leave unsupported combinations for the analyzer, including
// custom rules that may resolve the branches before coercion.
Ok(
get_coerce_type_for_case_expression(&then_types, else_type.as_ref())
.unwrap_or_else(|| {
then_types
.into_iter()
.find(|data_type| !data_type.is_null())
.or(else_type)
.unwrap_or(DataType::Null)
}),
)
}
Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => {
Ok(field.data_type().clone())
Expand Down
5 changes: 3 additions & 2 deletions datafusion/expr/src/logical_plan/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
103 changes: 76 additions & 27 deletions datafusion/expr/src/logical_plan/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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<LogicalPlan> {
match self {
// Build `LogicalPlan::Values` from the values for type inference.
Expand All @@ -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(),
}
Expand Down Expand Up @@ -3541,8 +3547,9 @@ impl Union {
}

/// Constructs new Union instance deriving schema from inputs.
/// Inputs do not have to have matching types and produced schema will
/// take type from the first input.
/// Infers common output types without inserting casts into the inputs.
/// If inference fails, retains the first input's type until analysis.
/// The analyzer coerces the inputs after their expressions are resolved.
// TODO (https://github.com/apache/datafusion/issues/14380): Avoid creating uncoerced union at all.
pub fn try_new_with_loose_types(inputs: Vec<Arc<LogicalPlan>>) -> Result<Self> {
let schema = Self::derive_schema_from_inputs(&inputs, true, false)?;
Expand Down Expand Up @@ -3585,9 +3592,9 @@ impl Union {
);
}
}
wrapped_inputs.push(Arc::new(LogicalPlan::Projection(
Projection::try_new_with_schema(expr, input, Arc::clone(schema))?,
)));
wrapped_inputs.push(Arc::new(LogicalPlan::Projection(Projection::try_new(
expr, input,
)?)));
}

Ok(wrapped_inputs)
Expand All @@ -3596,8 +3603,8 @@ impl Union {
/// Constructs new Union instance deriving schema from inputs.
///
/// If `loose_types` is true, inputs do not need to have matching types and
/// the produced schema will use the type from the first input.
/// TODO (<https://github.com/apache/datafusion/issues/14380>): 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.
Expand All @@ -3621,14 +3628,38 @@ impl Union {
inputs: &[Arc<LogicalPlan>],
loose_types: bool,
) -> Result<DFSchemaRef> {
type FieldData<'a> = (&'a DataType, bool, Vec<&'a Metadata>, usize);
// Keep the first type as a fallback if any input prevents inference.
type FieldData<'a> = (
&'a DataType,
Option<DataType>,
bool,
Vec<&'a Metadata>,
usize,
);
let mut cols: Vec<(&str, FieldData)> = Vec::new();
for input in inputs.iter() {
for (input_index, input) in inputs.iter().enumerate() {
// Alignment projections supply NULL for missing columns. Include
// those inputs in order, since NULL can unwrap dictionary types.
if loose_types {
for (name, (_, common_type, ..)) in &mut cols {
if !input.schema().has_column_with_unqualified_name(name) {
*common_type = common_type.as_ref().and_then(|data_type| {
type_union_coercion(data_type, &DataType::Null)
});
}
}
}
for field in input.schema().fields() {
if let Some((_, (data_type, is_nullable, metadata, occurrences))) =
cols.iter_mut().find(|(name, _)| name == field.name())
if let Some((
_,
(first_type, common_type, is_nullable, metadata, occurrences),
)) = cols.iter_mut().find(|(name, _)| name == field.name())
{
if !loose_types && *data_type != field.data_type() {
if loose_types {
*common_type = common_type.as_ref().and_then(|data_type| {
type_union_coercion(data_type, field.data_type())
});
} else if *first_type != field.data_type() {
return plan_err!(
"Found different types for field {}",
field.name()
Expand All @@ -3645,6 +3676,11 @@ impl Union {
field.name(),
(
field.data_type(),
if loose_types && input_index > 0 {
type_union_coercion(&DataType::Null, field.data_type())
} else {
Some(field.data_type().clone())
},
field.is_nullable(),
vec![field.metadata()],
1,
Expand All @@ -3657,7 +3693,16 @@ impl Union {
let union_fields = cols
.into_iter()
.map(
|(name, (data_type, is_nullable, unmerged_metadata, occurrences))| {
|(
name,
(
first_type,
common_type,
is_nullable,
unmerged_metadata,
occurrences,
),
)| {
// If the final number of occurrences of the field is less
// than the number of inputs (i.e. the field is missing from
// one or more inputs), then it must be treated as nullable.
Expand All @@ -3667,8 +3712,8 @@ impl Union {
true
};

let mut field =
Field::new(name, data_type.clone(), final_is_nullable);
let data_type = common_type.unwrap_or_else(|| first_type.clone());
let mut field = Field::new(name, data_type, final_is_nullable);
field.set_metadata(merge_metadata_for_union(unmerged_metadata));

(None, Arc::new(field))
Expand Down Expand Up @@ -3715,14 +3760,18 @@ impl Union {
let base_name = first_field.name().to_string();

let data_type = if loose_types {
// TODO apply type coercion here, or document why it's better to defer
// temporarily use the data type from the left input and later rely on the analyzer to
// coerce the two schemas into a common one.
first_field.data_type()
fields
.iter()
.skip(1)
.try_fold(first_field.data_type().clone(), |acc, field| {
type_union_coercion(&acc, field.data_type())
})
.unwrap_or_else(|| first_field.data_type().clone())
} else {
fields.iter().skip(1).try_fold(
first_field.data_type(),
|acc, field| {
fields
.iter()
.skip(1)
.try_fold(first_field.data_type(), |acc, field| {
if acc != field.data_type() {
return plan_err!(
"UNION field {i} have different type in inputs: \
Expand All @@ -3732,8 +3781,8 @@ impl Union {
);
}
Ok(acc)
},
)?
})?
.clone()
};
let nullable = fields.iter().any(|field| field.is_nullable());

Expand All @@ -3746,7 +3795,7 @@ impl Union {
base_name
};

let mut field = Field::new(&name, data_type.clone(), nullable);
let mut field = Field::new(&name, data_type, nullable);
let field_metadata =
merge_metadata_for_union(fields.iter().map(|field| field.metadata()));
field.set_metadata(field_metadata);
Expand Down
14 changes: 10 additions & 4 deletions datafusion/optimizer/src/analyzer/type_coercion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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<LogicalPlan>]) -> Result<DFSchema> {
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<LogicalPlan>],
base_schema: &DFSchemaRef,
) -> Result<DFSchema> {
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())
Expand Down
12 changes: 12 additions & 0 deletions datafusion/sqllogictest/test_files/case.slt
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,18 @@ SELECT CASE WHEN false THEN 1 ELSE 0 END FROM foo
0

# test then value type coercion
# DESCRIBE checks the inferred schema before analyzer coercion.
query TTT
DESCRIBE SELECT
CASE WHEN p THEN 1::INT ELSE 0.5 END AS else_widens,
CASE WHEN p THEN 1::INT WHEN NOT p THEN 0.5 END AS then_widens,
CASE WHEN p THEN 1::SMALLINT WHEN NOT p THEN 2::BIGINT ELSE true END AS deferred
FROM (VALUES (true), (false)) AS t(p);
----
else_widens Float64 NO
then_widens Float64 YES
deferred Int16 NO

# List(Utf8) will be casted to List(Int64)
query ?
SELECT CASE 1 WHEN 1 THEN ['1', '2', '3'] WHEN 2 THEN [1, 2, 3] ELSE null END;
Expand Down
30 changes: 30 additions & 0 deletions datafusion/sqllogictest/test_files/union.slt
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,36 @@
## UNION Tests
##########

# Infer the common schema before analysis, including timestamp precision/timezone
# and string representations (https://github.com/apache/datafusion/issues/27023).
query TTT
DESCRIBE SELECT
arrow_cast(1, 'Timestamp(Microsecond, Some("UTC"))') AS t,
arrow_cast('a', 'Utf8') AS large,
arrow_cast('a', 'Utf8') AS view
UNION ALL SELECT
CAST(NULL AS TIMESTAMP),
arrow_cast('b', 'LargeUtf8'),
arrow_cast('b', 'Utf8View');
----
t Timestamp(ns, "UTC") YES
large LargeUtf8 NO
view Utf8View NO

# Failed inference keeps the first type for analyzer rules to resolve or reject.
query TTT
DESCRIBE SELECT 1::SMALLINT AS x UNION ALL SELECT true;
----
x Int16 NO

# Infer dictionary coercion from the actual inputs, without folding the inferred
# UNION type back into the inputs and restoring the dictionary representation.
query T
SELECT DISTINCT arrow_typeof(x)
FROM (SELECT NULL AS x UNION ALL SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x);
----
Utf8

statement ok
CREATE TABLE t1(
id INT,
Expand Down
30 changes: 30 additions & 0 deletions datafusion/sqllogictest/test_files/union_by_name.slt
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,36 @@ INSERT INTO t2 VALUES (2, 2), (4, 4);


# Test binding
# Infer common types by name while preserving missing-column nullability.
query TTT
DESCRIBE SELECT 1::INT AS x
UNION ALL BY NAME SELECT 'a' AS y, 2::BIGINT AS x;
----
x Int64 NO
y Utf8 YES

# Retain the first type if a later input prevents inference.
query TTT
DESCRIBE SELECT 1::SMALLINT AS x
UNION ALL BY NAME SELECT true AS x;
----
x Int16 NO

# Missing columns contribute NULL to type inference, including dictionaries.
query TTT
DESCRIBE SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x
UNION ALL BY NAME SELECT arrow_cast('b', 'Dictionary(Int32, Utf8)') AS y;
----
x Utf8 YES
y Utf8 YES

query TT
SELECT DISTINCT arrow_typeof(x), arrow_typeof(y)
FROM (SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x
UNION ALL BY NAME SELECT arrow_cast('b', 'Dictionary(Int32, Utf8)') AS y);
----
Utf8 Utf8

query I
SELECT t1.x FROM t1 UNION BY NAME SELECT x FROM t1 ORDER BY x;
----
Expand Down
8 changes: 8 additions & 0 deletions docs/source/library-user-guide/upgrading/56.0.0.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading