alamb commented on code in PR #27139:
URL: https://github.com/apache/datafusion/pull/27139#discussion_r4483796536
##########
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
Review Comment:
If we want to document the behavior of how schema resolution works, it
should be in the code documentation I think
##########
datafusion/expr/src/expr_schema.rs:
##########
@@ -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::<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()?;
+ // If the branches have no common type, keep the old
provisional
+ // type. The analyzer reports the error.
+ Ok(
Review Comment:
we could probably make this more concse / less indented with some
intermediate variables
##########
datafusion/optimizer/src/analyzer/type_coercion.rs:
##########
@@ -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<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");
Review Comment:
this is a nice cleanup
##########
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
Review Comment:
This is a bug fix. Why does it need a upgrade guide entry?
##########
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,
Review Comment:
thank you for trimming the AI verbosity
##########
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<Vec<RecordBatch>> {
+ 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::<usize>(),
+ 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<()> {
Review Comment:
I feel like most of these tests can be written as .slt tests (but better). I
recommend leaving some tests here that cover the expected schema case but don't
run the query -- instead run the queries as a .slt tests
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
