Skip to content

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

Description

@dariocurr

Describe the bug

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.

This is the same underlying nullability-merging gap as #15394 (which focuses on UNION ... BY NAME and literal expressions specifically) — I'm filing this separately because it reproduces with a plain UNION ALL, no BY NAME or literals involved, and includes a concrete failure at the DataFusion → pyarrow boundary along with a fix. Happy to have this closed as a duplicate if the maintainers prefer to track it under #15394 instead.

To Reproduce

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

Output:

pyarrow.lib.ArrowInvalid: Schema at index 1 was different:
id: int64 not null
status: string not null
vs
id: int64 not null
status: string

The same query runs fine entirely within Rust/DataFusion (ctx.sql(...).collect().await) — the mismatch only surfaces once something (pyarrow, in this case) enforces schema equality across the batches in the stream.

Expected behavior

Every RecordBatch produced by a UNION ALL should report the same schema — the union's own declared output schema — regardless of which leg produced it.

Additional context

I've opened #23861 with a fix: UnionExec::execute() now compares each child stream's schema against UnionExec's own declared schema and, if they disagree, re-stamps every batch with the union's schema before yielding it. This is safe because the union's schema can only be more permissive than any single input's (nullability is combined with logical OR, never narrowed) — only the Field::nullable metadata changes, not the underlying data.

This is also the same root cause #16627 had to work around at the sqllogictest tooling level (making convert_batches tolerant of the mismatch instead of failing), and what #15603 (stale, closed for inactivity) attempted to fix at the physical-execution layer.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions