fix: preserve schema consistency across union and cast by Amogh-2404 · Pull Request #26201 · apache/datafusion

Amogh-2404 · 2026-10-11T08:22:13Z

Which issue does this PR close?

Closes Internal error: CAST over an expression simplified to a metadata-carrying column, under UNION ALL + aggregate (physical vs logical field metadata) #26178.

Rationale for this change

CAST(coalesce(non_null_col, lit) AS T) under UNION ALL + aggregate fails planning when the column carries Arrow field metadata (PARQUET:field_id):

Internal error: Physical input schema should be the same as the one converted
from logical input schema. Differences:
  - field metadata at index 0 [x]: (physical) {"PARQUET:field_id": "1"} vs (logical) {}.

coalesce on a NOT NULL column simplifies to the column. Child projections then advertise that metadata; the Union's stored schema still has {} from the pre-simplify coalesce field. Physical planning derives from the simplified exprs, so the aggregate schema check fails.

What changes are included in this PR?

LogicalPlan::recompute_schema for Union no longer keeps the old schema whenever field count is unchanged. It re-derives via Union::try_new so metadata/nullability track the current children. If width is unchanged but try_new fails (loose-typed children before coercion), the previous schema is kept; width mismatches still return the try_new error.

What is the testing strategy for this PR?

Integration: datafusion/core/tests/sql/union_field_metadata.rs — issue query (m=3, c=6), controls (no metadata / CAST(b) / direct b AS x / no union), mixed-metadata legs.
Unit: recompute_schema_refreshes_union_field_metadata, loose-union preserve, width-mismatch error.

Verified fail-before / pass-after on coalesce_cast_union_aggregate_preserves_field_metadata by temporarily reverting the plan.rs change.

Commands run:

cargo test -p datafusion-expr --lib logical_plan::plan::tests::recompute_schema
cargo test -p datafusion --test core_integration coalesce_cast_union_aggregate
cargo test -p datafusion --test core_integration union_
cargo test -p datafusion-optimizer --lib
cargo test -p datafusion --lib physical_planner
cargo test -p datafusion-sqllogictest --test sqllogictests metadata.slt
cargo clippy -p datafusion-expr -p datafusion --tests -- -D warnings
cargo fmt --all -- --check

Not run: full ./dev/rust_lint.sh, full extended CI suite, full sqllogictest corpus.

Are there any user-facing changes?

No public API change. Queries that previously hit this internal error now plan. Union output may correctly carry field metadata that a stale schema had dropped after expression rewrites. CAST metadata inheritance is unchanged.

kosiew

@Amogh-2404,

Thanks for the fix. Recomputing the logical UNION schema from rewritten children addresses the reported metadata mismatch. I independently verified that the b * 1 and CASE WHEN true THEN b ELSE 7 END reproducers fail before the fix and pass afterward, supporting the normal-path correction.

The logical UNION is the right boundary for maintaining this contract. Its constructor owns field naming, type validation, metadata intersection, and nullability, so reusing that policy is preferable to compensating during physical planning or changing CAST metadata behavior.

I'm requesting changes because the compatibility fallback can suppress structural errors unrelated to the intended pre-coercion type mismatch. Please constrain that exception and verify its schema-consistency behavior. This concern is about the recovery policy, not a demonstrated regression in the normal execution path. The additional tests are non-blocking suggestions.

kosiew · 2026-11-03T10:06:32Z

+                // positional `try_new` is sufficient.
+                match Union::try_new(inputs.clone()) {
+                    Ok(union) => Ok(LogicalPlan::Union(union)),
+                    Err(e) => {


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.

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.

kosiew · 2026-11-03T10:06:32Z

+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


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.

Done. Same query without the aggregate. x still has PARQUET:field_id = 1 on the optimized plan and on the batch that comes back.

kosiew · 2026-11-03T10:06:32Z

+        });
+        assert!(stale.schema().field(0).metadata().is_empty());
+
+        let refreshed = stale.recompute_schema()?;


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.

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.

kosiew · 2026-11-03T10:06:32Z

+                // 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()) {


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.

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.

akupchinskiy · 2026-11-04T17:18:36Z

Another case this fixes: DISTINCT ON below a UNION. The rewrite makes a nullable (first_value is always nullable) while the Union keeps the old NOT NULL schema.

CREATE TABLE t (id BIGINT NOT NULL, a BIGINT NOT NULL, o BIGINT NOT NULL)
AS VALUES (1, 10, 1), (1, 11, 2), (2, 20, 1);

WITH d AS (SELECT DISTINCT ON (id) id, a FROM t ORDER BY id, o)
SELECT a FROM d UNION SELECT a FROM d;

On main (238ae32) it fails with Internal error: Physical input schema should be the same as the one converted from logical input schema. Differences: field nullability at index 0 [a]: (physical) true vs (logical) false.; on this branch it returns {10, 20}. Might be worth a sqllogictest case, since it's nullability rather than metadata.

A width check on the first child was hiding real try_new failures. Also cover the DISTINCT ON nullability repro.

CI checks the pull request merged with main, which now requires the short SPDX source header.

Amogh-2404 · 2026-11-05T03:22:23Z

Thanks, that repro is clearer than the one I started from. It is in distinct_on.slt. The query returns 10 and 20, and a comes out nullable.

kosiew

@Amogh-2404, thank you for addressing the earlier review comments. The revised fallback now distinguishes legitimate loose-type UNIONs from structural errors, and the additional regression tests cover metadata propagation, nullability changes, and schema recomputation. The approach looks sound.

Refreshing the schema at the logical UNION boundary is the right architectural choice. It restores consistency with rewritten children without changing CAST metadata semantics or relying on physical planning to compensate for a stale logical contract. Preserving the stored schema temporarily for inputs awaiting type coercion is also justified.

I have reviewed the follow-up changes and consider my previous blocking concern resolved. The additional tests address the suggested coverage, and I have no further technical changes to request.

requested changes have been addressed

fix: preserve schema consistency across union and cast

ced67e8

github-actions Bot added logical-expr Logical plan and expressions core Core DataFusion crate labels Oct 11, 2026

kosiew previously requested changes Nov 3, 2026

View reviewed changes

fix: keep a stale union schema only while types still disagree

50c0f15

A width check on the first child was hiding real try_new failures. Also cover the DISTINCT ON nullability repro.

github-actions Bot added the sqllogictest SQL Logic Tests (.slt) label Nov 5, 2026

chore: use the SPDX header on the union metadata test

7e1cfa0

CI checks the pull request merged with main, which now requires the short SPDX source header.

kosiew reviewed Nov 5, 2026

View reviewed changes

Navigation Menu

Repository navigation

fix: preserve schema consistency across union and cast - #26201

fix: preserve schema consistency across union and cast#26201
Amogh-2404 wants to merge 3 commits into
apache:mainfrom
Amogh-2404:fix/union-recompute-schema-field-metadata

Amogh-2404 commented Oct 11, 2026

Uh oh!

kosiew left a comment

Uh oh!

kosiew Nov 3, 2026

Uh oh!

Amogh-2404 Nov 5, 2026

Uh oh!

kosiew Nov 3, 2026

Uh oh!

Amogh-2404 Nov 5, 2026

Uh oh!

kosiew Nov 3, 2026

Uh oh!

Amogh-2404 Nov 5, 2026

Uh oh!

kosiew Nov 3, 2026

Uh oh!

Amogh-2404 Nov 5, 2026

Uh oh!

akupchinskiy commented Nov 4, 2026

Uh oh!

Amogh-2404 commented Nov 5, 2026

Uh oh!

kosiew left a comment •

edited

Loading

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Amogh-2404 commented Oct 11, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

What is the testing strategy for this PR?

Are there any user-facing changes?

Uh oh!

kosiew left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

akupchinskiy commented Nov 4, 2026

Uh oh!

Amogh-2404 commented Nov 5, 2026

Uh oh!

kosiew left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kosiew left a comment •

edited

Loading