Amogh-2404 commented on code in PR #26201:
URL: https://github.com/apache/datafusion/pull/26201#discussion_r4476251434
##########
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:
Good point. Checking only the first child would hide a short input further
along. The old schema is kept only when every input still has the same width
and a type does not match yet. A short later input, or a union with one input,
returns the error. There is a case for the short input on the right, and
another where the types line up afterwards and the metadata and nullability
come back.
##########
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:
Done. Same query without the aggregate. `x` still has `PARQUET:field_id = 1`
on the optimized plan and on the batch that comes back.
##########
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:
Added a same-width union where the stored field is `NOT NULL` and one child
is nullable. After recompute the name, type, and metadata are the same, and the
field is nullable.
##########
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:
`with_new_exprs` still keeps the stored schema when the width is unchanged.
The metadata and nullability get fixed in `recompute_schema`, which is what
runs when a child schema changes. The new test does the rebuild first, metadata
still empty, then recompute, metadata back.
--
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]
