Skip to content

fix: UnionExec now conforms each batch to the union's declared schema#23861

Open
dariocurr wants to merge 1 commit into
apache:mainfrom
dariocurr:fix/union-nullable-schema-mismatch
Open

fix: UnionExec now conforms each batch to the union's declared schema#23861
dariocurr wants to merge 1 commit into
apache:mainfrom
dariocurr:fix/union-nullable-schema-mismatch

Conversation

@dariocurr

@dariocurr dariocurr commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

UNION ALL between an input whose column is NOT NULL and an input where the same column is nullable produces a valid, correctly-typed logical plan — the analyzer already OR's nullability across legs in coerce_union_schema (datafusion/optimizer/src/analyzer/type_coercion.rs), so the union's declared schema correctly reports the field as nullable.

The bug is at execution time: UnionExec::execute() hands out each child's RecordBatches completely unchanged. A leg whose column was already NOT NULL (and therefore needed no CAST from the analyzer) keeps emitting batches with a NOT NULL field, contradicting the union's own declared (nullable) schema. DataFusion's own execution tolerates this silently, but any consumer that checks schema equality across batches from the same stream — most notably pyarrow.Table.from_batches via the Arrow C Stream FFI used by the datafusion Python bindings — rejects the stream with ArrowInvalid: Schema at index N was different, even though every individual SELECT runs fine on its own.

Minimal reproducible example (Python)

import pyarrow as pa
from datafusion import SessionContext

ctx = SessionContext()
ctx.register_record_batch(
    "table_a",
    pa.record_batch(
        {"id": [1, 2], "status": ["ok", "ok"]},
        schema=pa.schema([("id", pa.int64()), ("status", pa.string())]),  # NOT NULL
    ),
)
ctx.register_record_batch(
    "table_b",
    pa.record_batch(
        {"id": [3, 4], "status": ["done", None]},
        schema=pa.schema([("id", pa.int64()), ("status", pa.string(), True)]),  # nullable
    ),
)

df = ctx.sql("SELECT id, status FROM table_a UNION ALL SELECT id, status FROM table_b")
print(df.schema())  # status: string, nullable -- correct
df.to_pandas()      # raises pyarrow.lib.ArrowInvalid: Schema at index 1 was different

The same root cause is why #16627 had to make the sqllogictest convert_batches helper tolerant of this exact mismatch instead of failing, and why #15603 (stale, closed for inactivity) attempted a similar fix at the physical-execution layer but didn't land.

What changes are included in this PR?

  • datafusion/physical-plan/src/union.rs: UnionExec::execute() now compares each child stream's schema against UnionExec's own declared schema, and if they disagree, wraps the child stream in a small new SchemaConformingStream that re-stamps every batch with the union's schema before yielding it. This is always safe: the union's schema can only be more permissive than any single input's (nullability is combined with logical OR, never narrowed — see the existing coerce_union_schema docs), and only the Field::nullable metadata changes; the underlying array data and data type are untouched.
  • datafusion/core/tests/sql/union_nullable.rs (new): regression tests covering same-type nullable/non-nullable mismatches in both leg orders, the "both legs NOT NULL" case (schema should stay NOT NULL), and a case where one leg also needs a real CAST (Int32 -> Int64) in addition to the nullability fix.

InterleaveExec (used for sorted unions) may have an analogous issue, but I kept this PR scoped to plain UnionExec, which is what's reported in #23862 / #15394 and reproduces the Python-binding failure above.

Are these changes tested?

Yes — added datafusion/core/tests/sql/union_nullable.rs with 4 new tests. I verified each one fails with a clear schema-mismatch assertion on main (i.e. before this fix) and passes with it applied. Also ran the full datafusion-physical-plan and datafusion-optimizer unit suites, union.slt/union_by_name.slt sqllogictests, and cargo fmt/clippy (--no-deps, since an unrelated pre-existing dead-code lint in datafusion-physical-expr fails -D warnings on main even without this change).

Are there any user-facing changes?

UNION ALL results now consistently report the analyzer's declared nullability on every batch, regardless of which leg produced it. No public API changes.

@codecov-commenter

codecov-commenter commented Jul 24, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.00000% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.71%. Comparing base (3f81613) to head (bed030d).

Files with missing lines Patch % Lines
datafusion/physical-plan/src/union.rs 90.00% 3 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #23861      +/-   ##
==========================================
- Coverage   80.71%   80.71%   -0.01%     
==========================================
  Files        1090     1090              
  Lines      370339   370369      +30     
  Branches   370339   370369      +30     
==========================================
+ Hits       298926   298946      +20     
- Misses      53603    53613      +10     
  Partials    17810    17810              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

UNION ALL between an input with a NOT NULL column and one where the same
column is nullable produced a valid, correctly-typed logical plan (the
analyzer already OR's nullability across legs in coerce_union_schema), but
UnionExec::execute() handed out each child's RecordBatches completely
unchanged. The leg whose column was already NOT NULL (and so needed no
CAST) kept emitting batches with a NOT NULL field, contradicting the
union's own declared (nullable) schema.

Any consumer that checks schema equality across batches from the same
stream -- e.g. pyarrow.Table.from_batches via the Arrow C Stream FFI used
by the datafusion Python bindings -- then rejects the stream with
`ArrowInvalid: Schema at index N was different`, even though every
individual SELECT runs fine on its own and DataFusion's own execution
never errors.

UnionExec::execute() now re-stamps each child's batches with the union's
own schema whenever they disagree. This is always safe: the union's schema
can only be more permissive than any single input's (nullability is OR'd,
never narrowed), and only the Field::nullable flag changes -- the
underlying array data and data type are untouched.

Closes apache#15394.
@dariocurr
dariocurr force-pushed the fix/union-nullable-schema-mismatch branch from c122156 to bed030d Compare July 24, 2026 13:30
@github-actions github-actions Bot added core Core DataFusion crate physical-plan Changes to the physical-plan crate labels Jul 24, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

UNION ALL between NOT NULL and nullable columns produces batches with inconsistent nullability

2 participants