fix: UnionExec now conforms each batch to the union's declared schema#23861
Open
dariocurr wants to merge 1 commit into
Open
fix: UnionExec now conforms each batch to the union's declared schema#23861dariocurr wants to merge 1 commit into
dariocurr wants to merge 1 commit into
Conversation
Codecov Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
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
force-pushed
the
fix/union-nullable-schema-mismatch
branch
from
July 24, 2026 13:30
c122156 to
bed030d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
nullable: trueif any of the branches produces fields withnullable: trueon the same position #15394.Rationale for this change
UNION ALLbetween an input whose column isNOT NULLand an input where the same column is nullable produces a valid, correctly-typed logical plan — the analyzer already OR's nullability across legs incoerce_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'sRecordBatches completely unchanged. A leg whose column was alreadyNOT NULL(and therefore needed noCASTfrom the analyzer) keeps emitting batches with aNOT NULLfield, 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 notablypyarrow.Table.from_batchesvia the Arrow C Stream FFI used by thedatafusionPython bindings — rejects the stream withArrowInvalid: Schema at index N was different, even though every individualSELECTruns fine on its own.Minimal reproducible example (Python)
The same root cause is why
#16627had to make the sqllogictestconvert_batcheshelper 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 againstUnionExec's own declared schema, and if they disagree, wraps the child stream in a small newSchemaConformingStreamthat 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 existingcoerce_union_schemadocs), and only theField::nullablemetadata 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 stayNOT NULL), and a case where one leg also needs a realCAST(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 plainUnionExec, 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.rswith 4 new tests. I verified each one fails with a clear schema-mismatch assertion onmain(i.e. before this fix) and passes with it applied. Also ran the fulldatafusion-physical-plananddatafusion-optimizerunit suites,union.slt/union_by_name.sltsqllogictests, andcargo fmt/clippy(--no-deps, since an unrelated pre-existing dead-code lint indatafusion-physical-exprfails-D warningsonmaineven without this change).Are there any user-facing changes?
UNION ALLresults now consistently report the analyzer's declared nullability on every batch, regardless of which leg produced it. No public API changes.