kosiew commented on code in PR #26201:
URL: https://github.com/apache/datafusion/pull/26201#discussion_r4453251324
##########
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) => {
Review Comment:
Could we restrict this fallback to expected pre-coercion type mismatches?
Checking only the first child's width can suppress a later child's width
mismatch or other reconstruction errors. Please validate all child widths and
propagate unrelated errors, preferably using an explicit compatibility
predicate rather than error strings. Add tests for loose-type compatibility,
later-child width mismatch, and metadata/nullability handling, including when
consistency is restored if the old schema must be retained.
##########
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
Review Comment:
Optional: could we also test the non-aggregate UNION and assert that `x`
retains `PARQUET:field_id=1` in both the optimized logical schema and emitted
batch schema? The existing aggregate regression checks planning and results,
but an end-to-end metadata assertion would protect against silently stripping
metadata.
##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -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()?;
Review Comment:
Optional: please consider a same-width UNION test where the stored schema is
non-nullable but a current child is nullable. Verify that recomputation updates
nullability while preserving the expected name, type, and metadata, since the
existing fixture exercises metadata refresh but not this nullability transition.
##########
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()) {
Review Comment:
Optional follow-up: could we compare this path with rebuilding the same
rewritten children through `with_new_exprs`? That UNION arm still retains the
old schema when width is unchanged. A focused test and clarification of whether
recomputation must follow rebuilding would establish the intended contract
without requiring a broader refactor.
--
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]
