This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new e948f17dbf fix: UnionExec now conforms each batch to the union's
declared schema (#23861)
e948f17dbf is described below
commit e948f17dbf7c67f753a41639a143f993dc945fe6
Author: dario curreri <[email protected]>
AuthorDate: Tue Aug 4 22:18:57 2026 +0200
fix: UnionExec now conforms each batch to the union's declared schema
(#23861)
## Which issue does this PR close?
- Closes #23862.
- Related to #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 `#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.
---
.../tests/memory_limit/union_nullable_spill.rs | 13 +-
datafusion/core/tests/sql/mod.rs | 1 +
datafusion/core/tests/sql/union_nullable.rs | 204 +++++++++++++++++++++
datafusion/physical-plan/src/union.rs | 130 ++++++++++++-
4 files changed, 342 insertions(+), 6 deletions(-)
diff --git a/datafusion/core/tests/memory_limit/union_nullable_spill.rs
b/datafusion/core/tests/memory_limit/union_nullable_spill.rs
index c5ef2387d3..d04273bc7f 100644
--- a/datafusion/core/tests/memory_limit/union_nullable_spill.rs
+++ b/datafusion/core/tests/memory_limit/union_nullable_spill.rs
@@ -103,10 +103,15 @@ fn build_task_ctx(pool_size: usize) ->
Arc<datafusion_execution::TaskContext> {
/// have mismatched nullability (one child's `val` is non-nullable, the other's
/// is nullable with NULLs). A tiny FairSpillPool forces all batches to spill.
///
-/// UnionExec returns child streams without schema coercion, so batches from
-/// different children carry different per-field nullability into the shared
-/// SpillPool. The IPC writer must use the SpillManager's canonical (nullable)
-/// schema — not the first batch's schema — so readback batches are valid.
+/// `UnionExec` now re-stamps every child batch with its own declared
(nullable)
+/// schema before they reach `RepartitionExec` (see
+/// <https://github.com/apache/datafusion/issues/15394>), so this no longer
+/// exercises mismatched-nullability batches arriving at the SpillManager via
+/// `UnionExec` specifically. It's kept as a regression test for the
+/// SpillManager fix itself: the IPC writer must use the SpillManager's
+/// canonical schema -- not the first batch's schema -- so readback batches
+/// stay valid for any caller that does hand it batches with differing
+/// nullability. See <https://github.com/apache/datafusion/issues/21292>.
///
/// Otherwise, sort_batch will panic with
/// `Column 'val' is declared as non-nullable but contains null values`
diff --git a/datafusion/core/tests/sql/mod.rs b/datafusion/core/tests/sql/mod.rs
index 33f9d3c02c..afed2f82d5 100644
--- a/datafusion/core/tests/sql/mod.rs
+++ b/datafusion/core/tests/sql/mod.rs
@@ -71,6 +71,7 @@ mod runtime_config;
pub mod select;
mod sql_api;
mod union_comparison;
+mod union_nullable;
mod unparser;
async fn register_aggregate_csv_by_sql(ctx: &SessionContext) {
diff --git a/datafusion/core/tests/sql/union_nullable.rs
b/datafusion/core/tests/sql/union_nullable.rs
new file mode 100644
index 0000000000..d2dc663366
--- /dev/null
+++ b/datafusion/core/tests/sql/union_nullable.rs
@@ -0,0 +1,204 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Regression tests asserting that every batch yielded by a `UNION ALL`
+//! reports the union's own declared schema, even when the same column is
+//! `NOT NULL` on one leg and nullable on another. See
+//! <https://github.com/apache/datafusion/issues/15394>.
+
+use std::sync::Arc;
+
+use arrow::array::{Int64Array, RecordBatch, StringArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::prelude::*;
+use datafusion_common::Result;
+
+/// Builds two single-partition tables that agree on `id`/`status` types but
+/// disagree on whether `status` is nullable, then runs `UNION ALL` over them.
+async fn union_all_mismatched_nullable(
+ left_nullable: bool,
+ right_nullable: bool,
+) -> Result<DataFrame> {
+ let ctx = SessionContext::new();
+
+ let schema_a = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int64, false),
+ Field::new("status", DataType::Utf8, left_nullable),
+ ]));
+ let batch_a = RecordBatch::try_new(
+ Arc::clone(&schema_a),
+ vec![
+ Arc::new(Int64Array::from(vec![1, 2])),
+ Arc::new(StringArray::from(vec!["ok", "ok"])),
+ ],
+ )?;
+ ctx.register_batch("table_a", batch_a)?;
+
+ let schema_b = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int64, false),
+ Field::new("status", DataType::Utf8, right_nullable),
+ ]));
+ let status_values: Vec<Option<&str>> = if right_nullable {
+ vec![Some("done"), None]
+ } else {
+ vec![Some("done"), Some("also-done")]
+ };
+ let batch_b = RecordBatch::try_new(
+ Arc::clone(&schema_b),
+ vec![
+ Arc::new(Int64Array::from(vec![3, 4])),
+ Arc::new(StringArray::from(status_values)),
+ ],
+ )?;
+ ctx.register_batch("table_b", batch_b)?;
+
+ ctx.sql(
+ "SELECT id, status FROM table_a \
+ UNION ALL \
+ SELECT id, status FROM table_b",
+ )
+ .await
+}
+
+/// The schema DataFusion actually commits to for a query: the logical plan
+/// after the `Analyzer` (which includes the `UNION` nullability/type
+/// coercion this test targets) and `Optimizer` have run. `DataFrame::schema`
+/// alone is not enough here -- it reflects the raw, pre-`Analyzer` plan (see
+/// `SessionState::create_logical_plan`), which for a `UNION` still has the
+/// first leg's un-coerced type.
+fn analyzed_schema(df: &DataFrame) -> Result<Schema> {
+ Ok(df
+ .clone()
+ .into_optimized_plan()?
+ .schema()
+ .as_arrow()
+ .clone())
+}
+
+/// Every `RecordBatch` actually produced by a `UNION ALL` must match the
+/// query's analyzed output schema field-for-field -- including
+/// nullability -- no matter which leg it came from.
+async fn assert_every_batch_matches_declared_schema(df: DataFrame) ->
Result<()> {
+ let declared_schema = analyzed_schema(&df)?;
+
+ let batches = df.collect().await?;
+ assert!(!batches.is_empty());
+ for batch in &batches {
+ assert_eq!(
+ batch.schema().as_ref(),
+ &declared_schema,
+ "a UNION ALL leg produced a RecordBatch whose schema disagrees \
+ with the union's declared output schema (commonly a dropped \
+ nullable flag) -- this is what downstream consumers that check \
+ schema equality across batches (e.g. pyarrow) reject with \
+ `ArrowInvalid: Schema at index N was different`"
+ );
+ }
+ Ok(())
+}
+
+#[tokio::test]
+async fn union_all_same_type_left_not_null_right_nullable() -> Result<()> {
+ let df = union_all_mismatched_nullable(false, true).await?;
+ assert!(
+ analyzed_schema(&df)?
+ .field_with_name("status")?
+ .is_nullable()
+ );
+ assert_every_batch_matches_declared_schema(df).await
+}
+
+#[tokio::test]
+async fn union_all_same_type_left_nullable_right_not_null() -> Result<()> {
+ let df = union_all_mismatched_nullable(true, false).await?;
+ assert!(
+ analyzed_schema(&df)?
+ .field_with_name("status")?
+ .is_nullable()
+ );
+ assert_every_batch_matches_declared_schema(df).await
+}
+
+#[tokio::test]
+async fn union_all_same_type_both_not_null_stays_not_null() -> Result<()> {
+ let df = union_all_mismatched_nullable(false, false).await?;
+ let declared_schema = analyzed_schema(&df)?;
+ assert!(
+ !declared_schema.field_with_name("status")?.is_nullable(),
+ "status should remain NOT NULL when neither leg is nullable"
+ );
+ assert_every_batch_matches_declared_schema(df).await
+}
+
+/// Same bug, but the coercion also has to widen the *type* (Int32 -> Int64)
+/// on one leg. The leg that already matched the target type still needed
+/// its nullability reconciled at execution time, independent of whichever
+/// legs needed a `CAST`.
+#[tokio::test]
+async fn union_all_widening_cast_also_fixes_nullable() -> Result<()> {
+ use arrow::array::Int32Array;
+
+ let ctx = SessionContext::new();
+
+ let schema_a = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int64, false),
+ Field::new("val", DataType::Int32, false),
+ ]));
+ let batch_a = RecordBatch::try_new(
+ Arc::clone(&schema_a),
+ vec![
+ Arc::new(Int64Array::from(vec![1, 2])),
+ Arc::new(Int32Array::from(vec![10, 20])),
+ ],
+ )?;
+ ctx.register_batch("table_a", batch_a)?;
+
+ let schema_b = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int64, false),
+ Field::new("val", DataType::Int64, true),
+ ]));
+ let batch_b = RecordBatch::try_new(
+ Arc::clone(&schema_b),
+ vec![
+ Arc::new(Int64Array::from(vec![3, 4])),
+ Arc::new(Int64Array::from(vec![Some(30), None])),
+ ],
+ )?;
+ ctx.register_batch("table_b", batch_b)?;
+
+ let df = ctx
+ .sql(
+ "SELECT id, val FROM table_a \
+ UNION ALL \
+ SELECT id, val FROM table_b",
+ )
+ .await?;
+
+ let declared_schema = analyzed_schema(&df)?;
+ assert_eq!(
+ declared_schema.field_with_name("val")?.data_type(),
+ &DataType::Int64
+ );
+ assert!(declared_schema.field_with_name("val")?.is_nullable());
+
+ let batches = df.collect().await?;
+ assert!(!batches.is_empty());
+ for batch in &batches {
+ assert_eq!(batch.schema().as_ref(), &declared_schema);
+ }
+ Ok(())
+}
diff --git a/datafusion/physical-plan/src/union.rs
b/datafusion/physical-plan/src/union.rs
index 4722329ea5..8d77556509 100644
--- a/datafusion/physical-plan/src/union.rs
+++ b/datafusion/physical-plan/src/union.rs
@@ -46,6 +46,7 @@ use crate::projection::{ProjectionExec, make_with_child};
use crate::statistics::{ChildStats, StatisticsArgs};
use crate::stream::ObservedStream;
+use arrow::array::RecordBatchOptions;
use arrow::datatypes::{Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use datafusion_common::config::ConfigOptions;
@@ -58,11 +59,88 @@ use datafusion_physical_expr::{
EquivalenceProperties, PhysicalExpr, calculate_union, conjunction,
};
-use futures::Stream;
+use futures::{Stream, StreamExt};
use itertools::Itertools;
use log::{debug, trace, warn};
use tokio::macros::support::thread_rng_n;
+/// Wraps a child stream so that every batch it yields is re-stamped with
+/// `schema` instead of the child's own schema.
+///
+/// This is used by both [`UnionExec`] and [`InterleaveExec`] when a child's
+/// output schema disagrees with the operator's declared output schema --
+/// in practice this only happens for nullability (the declared schema is
+/// nullable wherever *any* input's field is, but casts are only inserted
+/// between inputs when the *type* differs, not when only nullability
+/// does). For [`UnionExec`], [`UnionExec::try_new`] guarantees this: it
+/// calls `calculate_union`, which rejects any input whose field data types
+/// don't match the computed union schema. [`InterleaveExec::try_new`] does
+/// not repeat that check -- its inputs are only ever produced by the
+/// optimizer rewriting an already-validated `UnionExec`, whose children's
+/// types are therefore already known to agree -- but if this wrapper ever
+/// did see a genuine data type mismatch (e.g. from a hand-built
+/// `InterleaveExec`), `RecordBatch::try_new_with_options` below reports it
+/// as an error rather than silently yielding a corrupt batch.
+struct SchemaConformingStream {
+ schema: SchemaRef,
+ inner: SendableRecordBatchStream,
+}
+
+impl SchemaConformingStream {
+ fn new(schema: SchemaRef, inner: SendableRecordBatchStream) -> Self {
+ Self { schema, inner }
+ }
+}
+
+impl RecordBatchStream for SchemaConformingStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl Stream for SchemaConformingStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx).map(|opt| {
+ opt.map(|batch_result| {
+ batch_result.and_then(|batch| {
+ let options =
+
RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+ RecordBatch::try_new_with_options(
+ Arc::clone(&self.schema),
+ batch.columns().to_vec(),
+ &options,
+ )
+ .map_err(Into::into)
+ })
+ })
+ })
+ }
+
+ fn size_hint(&self) -> (usize, Option<usize>) {
+ self.inner.size_hint()
+ }
+}
+
+/// Wraps `stream` in a [`SchemaConformingStream`] if its schema disagrees
+/// with `schema`, otherwise returns it unchanged. See
+/// [`SchemaConformingStream`] and
+/// <https://github.com/apache/datafusion/issues/15394>.
+fn conform_stream_schema(
+ schema: SchemaRef,
+ stream: SendableRecordBatchStream,
+) -> SendableRecordBatchStream {
+ if stream.schema() == schema {
+ stream
+ } else {
+ Box::pin(SchemaConformingStream::new(schema, stream))
+ }
+}
+
/// `UnionExec`: `UNION ALL` execution plan.
///
/// `UnionExec` combines multiple inputs with the same schema by
@@ -294,6 +372,7 @@ impl ExecutionPlan for UnionExec {
if partition < input.output_partitioning().partition_count() {
let stream = input.execute(partition, context)?;
debug!("Found a Union partition to execute");
+ let stream = conform_stream_schema(self.schema(), stream);
return Ok(Box::pin(ObservedStream::new(
stream,
baseline_metrics,
@@ -668,7 +747,8 @@ impl ExecutionPlan for InterleaveExec {
let mut input_stream_vec = vec![];
for input in self.inputs.iter() {
if partition < input.output_partitioning().partition_count() {
- input_stream_vec.push(input.execute(partition,
Arc::clone(&context))?);
+ let stream = input.execute(partition, Arc::clone(&context))?;
+ input_stream_vec.push(conform_stream_schema(self.schema(),
stream));
} else {
// Do not find a partition to execute
break;
@@ -982,6 +1062,52 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_interleave_conforms_batch_schema() -> Result<()> {
+ // Two inputs agree on the column's type but disagree on nullability;
+ // InterleaveExec's declared schema ORs nullability across inputs, so
+ // every yielded batch must be re-stamped with that schema. See
+ // <https://github.com/apache/datafusion/issues/15394>.
+ let task_ctx = Arc::new(TaskContext::default());
+
+ let schema_not_null =
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32,
false)]));
+ let batch_not_null = RecordBatch::try_new(
+ Arc::clone(&schema_not_null),
+ vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2]))],
+ )?;
+
+ let schema_nullable =
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32,
true)]));
+ let batch_nullable = RecordBatch::try_new(
+ Arc::clone(&schema_nullable),
+ vec![Arc::new(arrow::array::Int32Array::from(vec![3, 4]))],
+ )?;
+
+ let hash_expr = vec![col("a", schema_not_null.as_ref())?];
+ let left: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new(
+ TestMemoryExec::try_new_exec(&[vec![batch_not_null]],
schema_not_null, None)?,
+ Partitioning::Hash(hash_expr.clone(), 1),
+ )?);
+ let right: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new(
+ TestMemoryExec::try_new_exec(&[vec![batch_nullable]],
schema_nullable, None)?,
+ Partitioning::Hash(hash_expr, 1),
+ )?);
+
+ let interleave: Arc<dyn ExecutionPlan> =
+ Arc::new(InterleaveExec::try_new(vec![left, right])?);
+ let interleave_schema = interleave.schema();
+ assert!(interleave_schema.field(0).is_nullable());
+
+ let batches = collect(interleave, task_ctx).await?;
+ assert!(!batches.is_empty());
+ for batch in &batches {
+ assert_eq!(batch.schema(), interleave_schema);
+ }
+
+ Ok(())
+ }
+
fn stats_merge_inputs() -> (SchemaRef, Statistics, Statistics, Statistics)
{
let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::UInt32, true)]));
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
