diff --git a/datafusion/core/tests/memory_limit/union_nullable_spill.rs b/datafusion/core/tests/memory_limit/union_nullable_spill.rs index c5ef2387d3cdc..d04273bc7fdb1 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 { /// 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 +/// ), 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 . /// /// 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 33f9d3c02ce87..afed2f82d57a8 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 0000000000000..d2dc66336621a --- /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 +//! . + +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 { + 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> = 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 { + 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 4722329ea55a4..8d77556509b9e 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; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + 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) { + self.inner.size_hint() + } +} + +/// Wraps `stream` in a [`SchemaConformingStream`] if its schema disagrees +/// with `schema`, otherwise returns it unchanged. See +/// [`SchemaConformingStream`] and +/// . +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 + // . + 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 = 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 = Arc::new(RepartitionExec::try_new( + TestMemoryExec::try_new_exec(&[vec![batch_nullable]], schema_nullable, None)?, + Partitioning::Hash(hash_expr, 1), + )?); + + let interleave: Arc = + 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)]));