Uh oh!
There was an error while loading. Please reload this page.
Fix spurious failure in convert_batches test helper - #16627
Merged
findepi merged 1 commit intoJul 1, 2025
Merged
Conversation
When query involves e.g. UNION ALL, it may produces record batches with incompatible schema. For example, one union branch may produce a nullable field while the order may produce a non-null field. Before the change, `convert_batches` success or failure depended on ordering of record batches returned by the query, and thus could lead (and did lead) to test flakiness. This changes schema source used by the `convert_batches`. Instead of peaking the first schema, it used the declared schema from the data frame, which should be correct with respect to nullability.
findepiforce-pushed
the
findepi/fix-spurious-failure-in-convert-batches-test-helper-aff883
branch
from
June 30, 2025 14:41
8fa74c0 to
36f21cfComparealamb
commented
Jun 30, 2025
Contributor
THANK YOU |
Uh oh!
There was an error while loading. Please reload this page.
findepi
deleted the
findepi/fix-spurious-failure-in-convert-batches-test-helper-aff883
branch
July 1, 2025 07:44
avantgardnerio added a commit
to coralogix/arrow-datafusion
that referenced
this pull request
Mar 18, 2026
avantgardnerio added a commit
to coralogix/arrow-datafusion
that referenced
this pull request
Mar 18, 2026
yliang412 pushed a commit
to yliang412/datafusion
that referenced
this pull request
Aug 4, 2026
…apache#23861) ## Which issue does this PR close? - Closesapache#23862. - Related to apache#15394. ## 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 `RecordBatch`es 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) ```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 `apache#16627` had to make the sqllogictest `convert_batches` helper tolerant of this exact mismatch instead of failing, and why `apache#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 apache#23862 / apache#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.
kosiew pushed a commit
to kosiew/datafusion
that referenced
this pull request
Aug 12, 2026
…apache#23861) ## Which issue does this PR close? - Closesapache#23862. - Related to apache#15394. ## 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 `RecordBatch`es 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) ```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 `apache#16627` had to make the sqllogictest `convert_batches` helper tolerant of this exact mismatch instead of failing, and why `apache#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 apache#23862 / apache#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.
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 freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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 #15394Projection#15242 (comment)Rationale for this change
When query involves e.g. UNION ALL, it may produces record batches with incompatible schema. For example, one union branch may produce a nullable field while the order may produce a non-null field. Before the change,
convert_batchessuccess or failure depended on ordering of record batches returned by the query, and thus could lead (and did lead) to test flakiness.What changes are included in this PR?
This changes schema source used by the
convert_batches. Instead of peaking the first schema, it used the declared schema from the data frame, which should be correct with respect to nullability.Are these changes tested?
yes
Are there any user-facing changes?
no