Uh oh!
There was an error while loading. Please reload this page.
fix: union all by name - #15603
Conversation
alamb
left a comment
There was a problem hiding this comment.
Thank you for this @chenkovsky (and all the other PRs recently -- very much appreciated)
| #[pin] | ||
| stream: S, | ||
| transform_schema: bool, |
There was a problem hiding this comment.
this seems like this is fixing the symptom rather than the root cause
I think it would be better to have the correct schema reflected in the plan in the first place 🤔
There was a problem hiding this comment.
yes, correct nullability in schema is better. I tried to fix logical plan before.
But nullability in logical plan won't affect physical plan. it's ignored.
in physical plan, it will recompute nullaibility from bottom to top.
but in this scenario, it seems that we need to pass nullability from top to bottom.
I need more suggestions.
There was a problem hiding this comment.
I want to learn some experience from spark.
for logical plan, I haven't found any logic to handle this problem.
for physical plan, spark is much easier, its InternalRow is schemaless. so it will use the schema of physical plan by default. but recordbatch contains schema.
I'm not 100% sure, I think current logical plan and physical plan schema is correct. the root cause is that recordbatch's schema doesn't match physical plan's. so adding an adapter is a proper way.
| let ret = this.stream.poll_next(cx); | ||
| if transform_schema { | ||
| if let Poll::Ready(Some(Ok(batch))) = ret { | ||
| return Poll::Ready(Some(batch.with_schema(schema).map_err(|e| { |
There was a problem hiding this comment.
I think this is one of the notorious problems when logical schema doesn't match the physical one on nullability/metadata. But this change might bring a performance impact, although the schema change is just reassigning the value but it also calls schema_contains which may be expensive
There was a problem hiding this comment.
yes, there's performance concern now. if this approach is feasible, I can try to optimize it, maybe use RecordBatch::try_new_with_options
Omega359
commented
Apr 9, 2025
Thanks for looking into the nullable issue, it's been on my plate for a bit to look into some more. It's really the last blocker I know of for union by name to work correctly. |
Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days. |
…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.
…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.
Which issue does this PR close?
nullable: trueif any of the branches produces fields withnullable: trueon the same position #15394.Rationale for this change
schema from inner physical plan is returned.
What changes are included in this PR?
update UnionExec and RecordBatchStreamAdapter to transform schema.
I found that logical plan's nullability for Projection is also not correct after optimization,
But this won't make the test fail. So I haven't included this part in this PR. Do we need to correct logical plan?
Are these changes tested?
UT
Are there any user-facing changes?
No