fix: preserve schema consistency across union and cast

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;
 

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 <https://github.com/apache/datafusion/issues/26178>.
+
+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<RecordBatch> {
+    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
+}

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(())
+    }
 }