diff --git a/datafusion/core/tests/sql/mod.rs b/datafusion/core/tests/sql/mod.rs index afed2f82d57a8..f8aea80631d55 100644 --- a/datafusion/core/tests/sql/mod.rs +++ b/datafusion/core/tests/sql/mod.rs @@ -71,6 +71,7 @@ mod runtime_config; pub mod select; mod sql_api; mod union_comparison; +mod union_field_metadata; mod union_nullable; mod unparser; diff --git a/datafusion/core/tests/sql/union_field_metadata.rs b/datafusion/core/tests/sql/union_field_metadata.rs new file mode 100644 index 0000000000000..e6c435f919386 --- /dev/null +++ b/datafusion/core/tests/sql/union_field_metadata.rs @@ -0,0 +1,164 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Regression tests for field-metadata consistency under `UNION ALL` when +//! expression simplification rewrites a cast child from a metadata-free +//! function result into a metadata-carrying column. +//! +//! See . + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion::assert_batches_eq; +use datafusion::prelude::*; +use datafusion_common::Result; + +fn table_batch(with_field_metadata: bool) -> Result { + let mut field = Field::new("b", DataType::Int64, false); + if with_field_metadata { + field = field.with_metadata(HashMap::from([( + "PARQUET:field_id".to_string(), + "1".to_string(), + )])); + } + let schema = Arc::new(Schema::new(vec![field])); + Ok(RecordBatch::try_new( + schema, + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?) +} + +fn register_t(ctx: &SessionContext, with_field_metadata: bool) -> Result<()> { + ctx.register_batch("t", table_batch(with_field_metadata)?)?; + Ok(()) +} + +async fn assert_max_count( + ctx: &SessionContext, + sql: &str, + expected_count: i64, +) -> Result<()> { + let batches = ctx.sql(sql).await?.collect().await?; + let expected = [ + "+---+---+".to_string(), + "| m | c |".to_string(), + "+---+---+".to_string(), + format!("| 3 | {expected_count} |"), + "+---+---+".to_string(), + ]; + let expected_refs: Vec<&str> = expected.iter().map(String::as_str).collect(); + assert_batches_eq!(&expected_refs, &batches); + Ok(()) +} + +const COALESCE_CAST_UNION_AGG: &str = "\ +SELECT max(x) AS m, count(x) AS c FROM (\ + SELECT CAST(coalesce(b, 7) AS BIGINT) AS x FROM t \ + UNION ALL \ + SELECT CAST(coalesce(b, 7) AS BIGINT) AS x FROM t\ +) q"; + +const CAST_COLUMN_UNION_AGG: &str = "\ +SELECT max(x) AS m, count(x) AS c FROM (\ + SELECT CAST(b AS BIGINT) AS x FROM t \ + UNION ALL \ + SELECT CAST(b AS BIGINT) AS x FROM t\ +) q"; + +const COALESCE_CAST_AGG_NO_UNION: &str = "\ +SELECT max(x) AS m, count(x) AS c FROM (\ + SELECT CAST(coalesce(b, 7) AS BIGINT) AS x FROM t\ +) q"; + +const DIRECT_COLUMN_UNION_AGG: &str = "\ +SELECT max(x) AS m, count(x) AS c FROM (\ + SELECT b AS x FROM t \ + UNION ALL \ + SELECT b AS x FROM t\ +) q"; + +/// Controls that isolate the failure mode reported in #26178. +#[tokio::test] +async fn coalesce_cast_union_aggregate_controls() -> Result<()> { + let ctx = SessionContext::new(); + register_t(&ctx, false)?; + assert_max_count(&ctx, COALESCE_CAST_UNION_AGG, 6).await?; + + let ctx = SessionContext::new(); + register_t(&ctx, true)?; + assert_max_count(&ctx, CAST_COLUMN_UNION_AGG, 6).await?; + // Metadata-bearing direct column (no coalesce rewrite): Union schema has + // field metadata from construction and must keep agreeing with physical. + assert_max_count(&ctx, DIRECT_COLUMN_UNION_AGG, 6).await?; + // Without UNION ALL the three input rows are not doubled. + assert_max_count(&ctx, COALESCE_CAST_AGG_NO_UNION, 3).await?; + Ok(()) +} + +/// `CAST(coalesce(non_null_col, lit))` under `UNION ALL` + aggregate must plan +/// when the column carries Arrow field metadata. +/// +/// `b` is `NOT NULL`, so `coalesce(b, 7)` simplifies to `b`. The Union's +/// logical schema must be refreshed from the simplified children so it stays +/// consistent with the physical plan's inherited column metadata. +#[tokio::test] +async fn coalesce_cast_union_aggregate_preserves_field_metadata() -> Result<()> { + let ctx = SessionContext::new(); + register_t(&ctx, true)?; + assert_max_count(&ctx, COALESCE_CAST_UNION_AGG, 6).await +} + +/// One leg with field metadata and one without: empty metadata is skipped in +/// the Union intersection, so planning under an aggregate still succeeds after +/// coalesce simplification on both legs. +#[tokio::test] +async fn coalesce_cast_union_aggregate_mixed_field_metadata() -> Result<()> { + let ctx = SessionContext::new(); + + let with_meta = Field::new("b", DataType::Int64, false).with_metadata(HashMap::from( + [("PARQUET:field_id".to_string(), "1".to_string())], + )); + let without_meta = Field::new("b", DataType::Int64, false); + ctx.register_batch( + "t1", + RecordBatch::try_new( + Arc::new(Schema::new(vec![with_meta])), + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?, + )?; + ctx.register_batch( + "t2", + RecordBatch::try_new( + Arc::new(Schema::new(vec![without_meta])), + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?, + )?; + + assert_max_count( + &ctx, + "SELECT max(x) AS m, count(x) AS c FROM (\ + SELECT CAST(coalesce(b, 7) AS BIGINT) AS x FROM t1 \ + UNION ALL \ + SELECT CAST(coalesce(b, 7) AS BIGINT) AS x FROM t2\ + ) q", + 6, + ) + .await +} diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9d360d052f835..b02de32f70c07 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -752,19 +752,30 @@ impl LogicalPlan { })) } LogicalPlan::Union(Union { inputs, schema }) => { - let first_input_schema = inputs[0].schema(); - if schema.fields().len() == first_input_schema.fields().len() { - // If inputs are not pruned do not change schema - Ok(LogicalPlan::Union(Union { inputs, schema })) - } else { - // A note on `Union`s constructed via `try_new_by_name`: - // - // At this point, the schema for each input should have - // the same width. Thus, we do not need to save whether a - // `Union` was created `BY NAME`, and can safely rely on the - // `try_new` initializer to derive the new schema based on - // column positions. - Ok(LogicalPlan::Union(Union::try_new(inputs)?)) + // Re-derive from the current children so field metadata and + // nullability stay aligned after expression rewrites (for + // example `CAST(coalesce(non_null_col, lit))` simplifying to + // `CAST(col)` under `UNION ALL` — see #26178). + // + // Previously this arm kept the stored schema whenever the + // field count was unchanged, which left stale empty metadata + // after child projections were refreshed. + // + // When width is unchanged but `try_new` fails (loose-typed + // children not yet coerced), keep the previous schema. When + // width changed, still surface the `try_new` error — that was + // the prior behavior for pruned inputs. For `Union`s built + // `BY NAME`, inputs should already share width here, so + // positional `try_new` is sufficient. + match Union::try_new(inputs.clone()) { + Ok(union) => Ok(LogicalPlan::Union(union)), + Err(e) => { + if schema.fields().len() == inputs[0].schema().fields().len() { + Ok(LogicalPlan::Union(Union { inputs, schema })) + } else { + Err(e) + } + } } } LogicalPlan::Distinct(distinct) => { @@ -7509,4 +7520,132 @@ mod tests { Ok(()) } + + /// `recompute_schema` must refresh Union field metadata from children even + /// when the field count is unchanged (#26178). + #[test] + fn recompute_schema_refreshes_union_field_metadata() -> Result<()> { + use crate::LogicalPlanBuilder; + use arrow::datatypes::{DataType, Field, Schema}; + use std::collections::HashMap; + + let meta = HashMap::from([("PARQUET:field_id".to_string(), "1".to_string())]); + let schema = Schema::new(vec![ + Field::new("b", DataType::Int64, false).with_metadata(meta.clone()), + ]); + let scan = table_scan(Some("t"), &schema, None)?.build()?; + let proj = LogicalPlanBuilder::from(scan) + .project([col("b").alias("x")])? + .build()?; + + let fresh = Union::try_new(vec![Arc::new(proj.clone()), Arc::new(proj.clone())])?; + assert_eq!(fresh.schema.field(0).metadata(), &meta); + + // Simulate a stale Union that kept a pre-rewrite schema with empty + // metadata while its children already advertise field metadata. + let stale_schema = Arc::new(DFSchema::from_unqualified_fields( + vec![Field::new("x", DataType::Int64, false)].into(), + HashMap::new(), + )?); + let stale = LogicalPlan::Union(Union { + inputs: vec![Arc::new(proj.clone()), Arc::new(proj)], + schema: stale_schema, + }); + assert!(stale.schema().field(0).metadata().is_empty()); + + let refreshed = stale.recompute_schema()?; + assert_eq!( + refreshed.schema().field(0).metadata(), + &meta, + "recompute_schema must refresh Union field metadata from children" + ); + Ok(()) + } + + /// Loose-typed unions that cannot satisfy `Union::try_new` keep their + /// previous schema instead of erroring out of `recompute_schema`. + #[test] + fn recompute_schema_preserves_loose_union_when_types_differ() -> Result<()> { + use crate::LogicalPlanBuilder; + use arrow::datatypes::{DataType, Field, Schema}; + + let left = table_scan( + Some("l"), + &Schema::new(vec![Field::new("x", DataType::Int32, true)]), + None, + )? + .build()?; + let right = table_scan( + Some("r"), + &Schema::new(vec![Field::new("x", DataType::Int64, true)]), + None, + )? + .build()?; + // Project to unqualified `x` so positional union can be formed. + let left = LogicalPlanBuilder::from(left) + .project([col("x")])? + .build()?; + let right = LogicalPlanBuilder::from(right) + .project([col("x")])? + .build()?; + + let loose = LogicalPlan::Union(Union::try_new_with_loose_types(vec![ + Arc::new(left), + Arc::new(right), + ])?); + let before = Arc::clone(loose.schema()); + let after = loose.recompute_schema()?; + assert_eq!( + after.schema().as_ref(), + before.as_ref(), + "loose unions should keep their schema when try_new cannot succeed" + ); + assert_eq!(after.schema().field(0).data_type(), &DataType::Int32); + Ok(()) + } + + /// Width changes must still surface `Union::try_new` errors instead of + /// silently keeping a stale schema. + #[test] + fn recompute_schema_errors_when_union_width_mismatches() -> Result<()> { + use crate::LogicalPlanBuilder; + use arrow::datatypes::{DataType, Field, Schema}; + + let wide = table_scan( + Some("w"), + &Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + ]), + None, + )? + .build()?; + let narrow = table_scan( + Some("n"), + &Schema::new(vec![Field::new("a", DataType::Int64, true)]), + None, + )? + .build()?; + let wide = LogicalPlanBuilder::from(wide) + .project([col("a"), col("b")])? + .build()?; + let narrow = LogicalPlanBuilder::from(narrow) + .project([col("a")])? + .build()?; + + // Build a Union with a 2-field schema, then replace one child with a + // 1-field plan and ask recompute_schema to notice the mismatch. + let union = Union::try_new(vec![Arc::new(wide.clone()), Arc::new(wide)])?; + let stale = LogicalPlan::Union(Union { + inputs: vec![Arc::new(narrow), Arc::clone(&union.inputs[1])], + schema: union.schema, + }); + let err = stale.recompute_schema().expect_err("width mismatch"); + assert!( + err.to_string().contains("different number of columns") + || err.to_string().contains("UNION"), + "unexpected error: {err}" + ); + Ok(()) + } }