fix: UnionExec now conforms each batch to the union's declared schema

[dariocurr]

Which issue does this PR close?

• Closes UNION ALL between NOT NULL and nullable columns produces batches with inconsistent nullability #23862.
• Related to Each Union branch should produce fields with nullable: true if any of the branches produces fields with nullable: true on the same position #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 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.

Minimal reproducible example (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 #16627 had to make the sqllogictest convert_batches helper 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 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 #23862 / #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.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 82.19178% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.87%. Comparing base (3f81613) to head (4af8735).
⚠️ Report is 107 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/union.rs	82.19%	7 Missing and 6 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23861      +/-   ##
==========================================
+ Coverage   80.71%   80.87%   +0.16%     
==========================================
  Files        1090     1101      +11     
  Lines      370339   375837    +5498     
  Branches   370339   375837    +5498     
==========================================
+ Hits       298926   303967    +5041     
- Misses      53603    53757     +154     
- Partials    17810    18113     +303     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[dariocurr]

Is there any news on that?

[kosiew]

@dariocurr
Thanks for working on this schema consistency issue.

The UnionExec fix looks useful, but the same mismatch can still occur after the optimizer replaces it with InterleaveExec. I have left one blocking comment for that path and one small test documentation suggestion.

[kosiew]

I think the same schema mismatch can still surface when the optimizer replaces a UnionExec with an InterleaveExec.

InterleaveExec constructs its declared schema using the same union_schema helper, but execute() passes the child streams directly to CombinedRecordBatchStream. Its poll_next implementation then returns each child RecordBatch unchanged.

Since ensure_distribution can rewrite a UnionExec to an InterleaveExec when the children are interleavable (datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:1493), a SQL-visible UNION ALL plan may still emit batches whose schemas do not match the declared interleave or union schema.

Could you apply the same schema-conforming wrapper to each InterleaveExec child stream, or update CombinedRecordBatchStream so yielded batches are re-stamped with its declared schema?

[dariocurr]

Good catch, thanks. Pushed 131796d: InterleaveExec::execute() now wraps each child stream through the same conform_stream_schema/SchemaConformingStream helper (extracted from the UnionExec path) before handing them to CombinedRecordBatchStream, so batches are re-stamped with the interleave's declared schema whenever nullability disagrees. Added test_interleave_conforms_batch_schema (in union.rs) covering this directly.

[kosiew]

The regression coverage is helpful. The module-level comment currently repeats much of the PR and issue description, though.

Could we trim it down to the invariant being tested and the issue link? The assertions already make the downstream PyArrow failure mode clear.

[dariocurr]

Trimmed in 131796d — the module doc-comment is now just the invariant being tested plus the issue link.

[kosiew]

@dariocurr
I would love to review this again, after you fix the clippy errors.

[dariocurr]

Follow-up in 4af8735, fixing a few loose ends from a closer pass:

• SchemaConformingStream's doc comment overstated the guarantee for InterleaveExec -- unlike UnionExec::try_new, InterleaveExec::compute_properties doesn't re-validate that child field types match. Comment now states the real invariant and notes a mismatch would surface as an error, not a silent bad batch.
• Two tests iterated batches without asserting the list was non-empty, so they could pass vacuously.
• union_nullable_spill.rs's comment claimed UnionExec skips schema coercion, which this PR makes false. Updated to describe what it now guards.

Tests, fmt, and clippy still clean.

[kosiew]

@dariocurr
Thanks for the follow-up.

The earlier feedback has been addressed:

• InterleaveExec now conforms each child stream before passing it to CombinedRecordBatchStream, and the direct regression test confirms that every emitted batch uses the interleave schema.
• The SQL test module comment now focuses on the invariant and the related issue link.
• The documentation guarantee now matches the actual constructor and runtime behavior.
• Both batch loops that were previously vacuous now verify that output is non-empty.
• The spill test comment is up to date.

I did not find any new defects or have any inline comments.

Looks good to me. Approving.

[alamb]

This is a nice change -- thank you @dariocurr and @kosiew -- ideally we could coerce the schema at plan time but I don't know how to coerce nullability

[dariocurr]

Following up on the schema-coercion-at-plan-time idea from this thread: opened #24094, which moves the re-stamping from inline logic in execute() into an explicit CoerceSchemaExec node inserted by UnionExec::try_new/InterleaveExec::try_new. It shows up in EXPLAIN now. Benchmarked it against this PR's approach and didn't find a measurable performance difference either way.