Skip to content

fix: union all by name - #15603

Closed
chenkovsky wants to merge 3 commits into
apache:mainfrom
chenkovsky:fix/physical-expr
Closed

fix: union all by name#15603
chenkovsky wants to merge 3 commits into
apache:mainfrom
chenkovsky:fix/physical-expr

Conversation

@chenkovsky

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

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

@chenkovsky
chenkovsky marked this pull request as ready for review April 6, 2025 10:55
@github-actionsgithub-actionsBot added the sqllogictest SQL Logic Tests (.slt) label Apr 6, 2025

@alambalamb left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for this @chenkovsky (and all the other PRs recently -- very much appreciated)

#[pin]
stream: S,

transform_schema: bool,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 🤔

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

LogicalPlan::Projection(Projection{ input, expr, .. }) => self

in physical plan, it will recompute nullaibility from bottom to top.

e.nullable(&input_schema)?,

but in this scenario, it seems that we need to pass nullability from top to bottom.

I need more suggestions.

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I want to learn some experience from spark.

for logical plan, I haven't found any logic to handle this problem.

https://github.com/apache/spark/blob/75d80c7795ca71d24229010ab04ae740473126aa/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala#L475

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.

https://github.com/apache/spark/blob/75d80c7795ca71d24229010ab04ae740473126aa/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala#L688

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| {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

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.

@github-actions

Copy link
Copy Markdown

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.

@github-actionsgithub-actionsBot added the Stale PR has not had any activity for some time label Jun 11, 2025
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.
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

sqllogictestSQL Logic Tests (.slt)StalePR has not had any activity for some time

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Each Union branch should produce fields with nullable: true if any of the branches produces fields with nullable: true on the same position

4 participants

@chenkovsky@Omega359@alamb@comphead