From d9f53b3ab677180d9a74a9273e61ce242d8c49a2 Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Mon, 24 Aug 2026 15:40:30 +0000 Subject: [PATCH 1/2] fix: adapt input batches with stricter nested nullability to planned schema in aggregation (#24394) Closes #24069. In DataFusion, in-memory table sources such as `MemTable::try_new` accept `RecordBatch`es whose schemas are stricter than the table's declared schema via `Schema::contains(&batches_schema)` (e.g. nullable nested fields declared on the table vs non-nullable nested fields in the input batches). However, `MemoryStream` previously advertised the declared table schema while emitting the underlying stricter `RecordBatch`es without adapting them. When downstream operators (such as `AggregateExec` with `array_agg` or distinct aggregation) received batches with stricter nested schemas, runtime type mismatch errors occurred (e.g. #24069). 1. **`datafusion_common::nested_struct::adapt_batch_to_schema`**: - Adapts `RecordBatch`es whose nested schemas are stricter than a target schema. - Recursively reconstructs compatible nested Struct/List types. - Explicitly handles Dense and Sparse Union schema conformance while preserving type IDs and dense offsets without copying buffer data. - Requires exact Union type-ID sets and matching modes. - Does not broaden generic SQL CAST semantics (`requires_nested_struct_cast` remains untouched). 2. **`MemoryStream` Producer Boundary Normalization (`datafusion-physical-plan/src/memory.rs`)**: - Fixes the producer-side invariant exposed by #24069. - When batches have stricter schemas accepted by `MemTable::try_new`, `MemoryStream::poll_next` normalizes emitted batches using `adapt_batch_to_schema(batch, &self.schema)` whenever the runtime batch schema differs from `self.schema` and `self.schema.contains(batch.schema())`. - Ensures all `RecordBatch`es emitted by `MemoryStream` conform exactly to `stream.schema()`. 3. **Regression Coverage**: - Direct `MemoryStream` regressions in `memory.rs` verifying emitted batches match the advertised schema, including projection handling. - Unit tests in `nested_struct.rs` covering nested Struct and Dense/Sparse Union adaptation, unpacked scalar values, type IDs, offsets, reordered IDs, and incompatible Union layouts. - End-to-end SQL aggregation integration tests in `nested_nullability.rs` covering standard, DISTINCT, and spilling aggregations for #24069. Yes: - `datafusion-common` unit tests for `adapt_batch_to_schema` and Union adaptation (`test_adapt_batch_to_schema_*`). - `datafusion-physical-plan` unit tests for `MemoryStream` emitted batch schema conformance and projection (`test_memory_stream_emitted_batch_matches_declared_schema*`). - `datafusion` core SQL integration tests in `datafusion/core/tests/sql/aggregates/nested_nullability.rs`. No. Queries aggregating in-memory tables whose batches have stricter nested nullability than the table schema now succeed as expected. (cherry picked from commit 2326917aa8d691aa5123518ca602ce63ab18ffc6) --- datafusion/common/src/nested_struct.rs | 713 +++++++++++++++++- datafusion/core/tests/sql/aggregates/mod.rs | 1 + .../sql/aggregates/nested_nullability.rs | 246 ++++++ datafusion/physical-plan/src/memory.rs | 135 ++++ 4 files changed, 1093 insertions(+), 2 deletions(-) create mode 100644 datafusion/core/tests/sql/aggregates/nested_nullability.rs diff --git a/datafusion/common/src/nested_struct.rs b/datafusion/common/src/nested_struct.rs index e915b91b911cc..a6eb7df22a2cf 100644 --- a/datafusion/common/src/nested_struct.rs +++ b/datafusion/common/src/nested_struct.rs @@ -19,11 +19,14 @@ use crate::error::{_plan_err, Result}; use arrow::{ array::{ Array, ArrayRef, AsArray, DictionaryArray, FixedSizeListArray, GenericListArray, - GenericListViewArray, StructArray, downcast_integer, make_array, new_null_array, + GenericListViewArray, RecordBatch, StructArray, UnionArray, downcast_integer, + make_array, new_null_array, }, buffer::NullBuffer, compute::{CastOptions, can_cast_types, cast_with_options}, - datatypes::{DataType, DataType::Struct, Field, FieldRef}, + datatypes::{ + DataType, DataType::Struct, Field, FieldRef, SchemaRef, UnionFields, UnionMode, + }, }; use std::{collections::HashSet, sync::Arc}; @@ -121,6 +124,63 @@ fn cast_struct_column( } } +/// Cast a union column to match target union fields, handling child fields recursively. +/// +/// ## Casting Behavior +/// - Preserves union mode (sparse or dense). Incompatible modes are rejected. +/// - Requires exact matching union type ID sets (order may differ). +/// - Recursively adapts each matching child array using `cast_column`. +/// - Preserves row-level `type_ids` and dense `offsets` buffers without copying primitive data. +fn cast_union_column( + source_col: &ArrayRef, + source_fields: &UnionFields, + source_mode: &UnionMode, + target_fields: &UnionFields, + target_mode: &UnionMode, + cast_options: &CastOptions, +) -> Result { + validate_union_schema_compatibility( + source_fields, + source_mode, + target_fields, + target_mode, + )?; + + let source_union = source_col + .as_any() + .downcast_ref::() + .ok_or_else(|| { + crate::error::DataFusionError::Plan(format!( + "Expected UnionArray for Union data type, got {}", + source_col.data_type() + )) + })?; + + let mut children = Vec::with_capacity(target_fields.len()); + + for (target_type_id, target_field) in target_fields.iter() { + let source_child = source_union.child(target_type_id); + + children.push( + cast_column(source_child, target_field.data_type(), cast_options).map_err( + |e| { + e.context(format!( + "While adapting Union child type ID {target_type_id} ('{}')", + target_field.name() + )) + }, + )?, + ); + } + + Ok(Arc::new(UnionArray::try_new( + target_fields.clone(), + source_union.type_ids().clone(), + source_union.offsets().cloned(), + children, + )?)) +} + /// Cast a column to match the target field type, with special handling for nested structs. /// /// This function serves as the main entry point for column casting operations. For struct @@ -215,6 +275,17 @@ pub fn cast_column( target_value_type, cast_options, ), + ( + DataType::Union(source_fields, source_mode), + DataType::Union(target_fields, target_mode), + ) => cast_union_column( + source_col, + source_fields, + source_mode, + target_fields, + target_mode, + cast_options, + ), _ => Ok(cast_with_options(source_col, target_type, cast_options)?), } } @@ -490,6 +561,48 @@ fn validate_field_compatibility( ) } +fn validate_union_schema_compatibility( + source_fields: &UnionFields, + source_mode: &UnionMode, + target_fields: &UnionFields, + target_mode: &UnionMode, +) -> Result<()> { + if source_mode != target_mode { + return _plan_err!( + "Cannot adapt Union from mode {source_mode:?} to {target_mode:?}" + ); + } + + // This adapter is for schema conformance, not general Union variant-set evolution. + if source_fields.len() != target_fields.len() { + return _plan_err!( + "Cannot adapt Union schema with different field sets: source has {} fields, target has {}", + source_fields.len(), + target_fields.len() + ); + } + + for (target_type_id, target_field) in target_fields.iter() { + let Some((_, source_field)) = source_fields + .iter() + .find(|(source_type_id, _)| *source_type_id == target_type_id) + else { + return _plan_err!( + "Cannot adapt Union schema: target type ID {target_type_id} ('{}') is missing from source", + target_field.name() + ); + }; + + if !target_field.contains(source_field) { + return _plan_err!( + "Cannot adapt Union child with type ID {target_type_id}: source field {source_field} is not contained by target field {target_field}" + ); + } + } + + Ok(()) +} + /// Validates that `source_type` can be cast to `target_type`, recursively /// handling container types that wrap structs. pub fn validate_data_type_compatibility( @@ -524,6 +637,17 @@ pub fn validate_data_type_compatibility( } validate_data_type_compatibility(field_name, s_val, t_val)?; } + ( + DataType::Union(source_fields, source_mode), + DataType::Union(target_fields, target_mode), + ) => { + validate_union_schema_compatibility( + source_fields, + source_mode, + target_fields, + target_mode, + )?; + } _ => { if !can_cast_types(source_type, target_type) { return _plan_err!( @@ -1703,3 +1827,588 @@ mod tests { )); } } + +/// Adapts a `RecordBatch` to conform to `target_schema`, verifying that each target field +/// type contains the incoming column data type (as verified by [`arrow::datatypes::DataType::contains`]) +/// and transforms the metadata/types of differing columns to match `target_schema` +/// without copying primitive buffer data. +/// +/// If `batch` has an incompatible column count or incompatible column data types, +/// an error is returned. +pub fn adapt_batch_to_schema( + batch: RecordBatch, + target_schema: &SchemaRef, +) -> Result { + if Arc::ptr_eq(batch.schema_ref(), target_schema) + || batch.schema().as_ref() == target_schema.as_ref() + { + return Ok(batch); + } + + if batch.num_columns() != target_schema.fields().len() { + return _plan_err!( + "Batch schema does not conform to expected schema (column count mismatch). Expected: {target_schema}, got: {}", + batch.schema() + ); + } + + let mut columns = Vec::with_capacity(batch.num_columns()); + let mut needs_column_adaptation = false; + let cast_options = CastOptions::default(); + + for (target_field, col) in target_schema.fields().iter().zip(batch.columns()) { + if target_field.data_type() != col.data_type() { + // If data types differ, verify that target_field's data type contains + // the column's data type (e.g. stricter nested struct / list field nullability). + if !target_field.data_type().contains(col.data_type()) { + return _plan_err!( + "Batch column '{}' with type {} cannot be adapted to expected type {}", + target_field.name(), + col.data_type(), + target_field.data_type() + ); + } + needs_column_adaptation = true; + let adapted_col = cast_column(col, target_field.data_type(), &cast_options)?; + columns.push(adapted_col); + } else { + columns.push(Arc::clone(col)); + } + } + + if needs_column_adaptation { + Ok(RecordBatch::try_new(Arc::clone(target_schema), columns)?) + } else { + // Schema differs only in top-level metadata or field nullability, while + // column data types match exactly. Replace the schema on the batch. + Ok(RecordBatch::try_new( + Arc::clone(target_schema), + batch.columns().to_vec(), + )?) + } +} + +#[cfg(test)] +mod adapt_schema_tests { + use super::*; + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{Field, Fields, Schema}; + + #[test] + fn test_adapt_batch_to_schema_identical() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ])); + + let a = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let b = Arc::new(StringArray::from(vec![Some("x"), None, Some("z")])) as ArrayRef; + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a, b])?; + + let adapted = adapt_batch_to_schema(batch.clone(), &schema)?; + assert_eq!(adapted, batch); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_stricter_nested_struct() -> Result<()> { + // Declared table schema: {a: Struct({x: Int32 (nullable), y: Utf8 (nullable)})} + let declared_inner_fields = Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Utf8, true), + ]); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "a", + Struct(declared_inner_fields), + false, + )])); + + // Runtime batch schema: {a: Struct({x: Int32 (NON-nullable), y: Utf8 (NON-nullable)})} + let runtime_inner_fields = Fields::from(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Utf8, false), + ]); + let runtime_schema = Arc::new(Schema::new(vec![Field::new( + "a", + Struct(runtime_inner_fields.clone()), + false, + )])); + + let x = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let y = Arc::new(StringArray::from(vec!["x", "y", "z"])) as ArrayRef; + let struct_array = + Arc::new(StructArray::new(runtime_inner_fields, vec![x, y], None)) + as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![struct_array])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + assert_eq!(adapted.num_rows(), 3); + + // Verify nested fields now have the declared nullability + let Struct(fields) = adapted.column(0).data_type() else { + panic!("expected struct"); + }; + assert!(fields[0].is_nullable()); + assert!(fields[1].is_nullable()); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_top_level_nullability_only() -> Result<()> { + // Declared schema has nullable column 'a', runtime batch has non-nullable 'a' + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let runtime_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + let a = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![a])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + assert!(adapted.schema().field(0).is_nullable()); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_null_into_non_nullable_rejected() { + // Declared schema is non-nullable, but runtime batch is nullable + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let runtime_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let a = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![a]).unwrap(); + + // Must reject because nullable is not contained by non-nullable + let result = adapt_batch_to_schema(batch, &declared_schema); + assert!(result.is_err()); + } + + #[test] + fn test_adapt_batch_to_schema_incompatible_type_rejected() { + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let runtime_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + + let a = Arc::new(StringArray::from(vec!["1", "2"])) as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![a]).unwrap(); + + let result = adapt_batch_to_schema(batch, &declared_schema); + assert!(result.is_err()); + } + + fn test_two_field_union(nullable: bool) -> UnionFields { + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, nullable), + Field::new("str", DataType::Utf8, nullable), + ], + ) + .unwrap() + } + + #[test] + fn test_adapt_batch_to_schema_stricter_sparse_union() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let target_union_fields = test_two_field_union(true); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Sparse), + false, + )])); + + let source_union_fields = test_two_field_union(false); + let runtime_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Sparse), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); + let str_array: ArrayRef = + Arc::new(StringArray::from(vec!["hello", "world", "!"])); + let type_ids = [0, 0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + None, + vec![int_array, str_array], + )?; + let batch = RecordBatch::try_new(runtime_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let DataType::Union(fields, mode) = adapted_union.data_type() else { + panic!("expected union"); + }; + assert_eq!(*mode, UnionMode::Sparse); + assert!(fields.iter().all(|(_, f)| f.is_nullable())); + + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_stricter_dense_union() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let target_union_fields = test_two_field_union(true); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + )])); + + let source_union_fields = test_two_field_union(false); + let runtime_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 30])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["hello"])); + let type_ids = [0, 1, 0].into_iter().collect::>(); + let offsets = [0, 0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + Some(offsets), + vec![int_array, str_array], + )?; + let batch = RecordBatch::try_new(runtime_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let DataType::Union(fields, mode) = adapted_union.data_type() else { + panic!("expected union"); + }; + assert_eq!(*mode, UnionMode::Dense); + assert!(fields.iter().all(|(_, f)| f.is_nullable())); + + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_union_reordered_and_non_contiguous_type_ids() + -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let target_union_fields = UnionFields::try_new( + vec![3, 1], + vec![ + Field::new("str", DataType::Utf8, true), + Field::new("int", DataType::Int32, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + )])); + + let source_union_fields = UnionFields::try_new( + vec![1, 3], + vec![ + Field::new("int", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + ], + )?; + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 30])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["b"])); + let type_ids = [1, 3, 1].into_iter().collect::>(); + let offsets = [0, 0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids.clone(), + Some(offsets.clone()), + vec![int_array, str_array], + )?; + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(source_batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + adapted_union.data_type(), + declared_schema.field(0).data_type() + ); + assert_eq!(adapted_union.type_ids(), &type_ids); + assert_eq!(adapted_union.offsets(), Some(&offsets)); + + // Child 1 is int, Child 3 is str (accessed by type ID) + let int_child = adapted_union + .child(1) + .as_any() + .downcast_ref::() + .unwrap(); + let str_child = adapted_union + .child(3) + .as_any() + .downcast_ref::() + .unwrap(); + + // Row 0: type_id 1 -> int value 10 + assert_eq!(adapted_union.type_id(0), 1); + assert_eq!(int_child.value(adapted_union.value_offset(0)), 10); + + // Row 1: type_id 3 -> str value "b" + assert_eq!(adapted_union.type_id(1), 3); + assert_eq!(str_child.value(adapted_union.value_offset(1)), "b"); + + // Row 2: type_id 1 -> int value 30 + assert_eq!(adapted_union.type_id(2), 1); + assert_eq!(int_child.value(adapted_union.value_offset(2)), 30); + + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_union_nested_struct() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_struct_fields = vec![Field::new("x", DataType::Int32, true)]; + let target_union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("s", Struct(target_struct_fields.into()), true)], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + )])); + + let source_struct_fields = vec![Field::new("x", DataType::Int32, false)]; + let source_union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("s", Struct(source_struct_fields.into()), false)], + )?; + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let struct_child: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("x", DataType::Int32, false)].into(), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + None, + )); + let type_ids = [0, 0].into_iter().collect::>(); + let offsets = [0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids.clone(), + Some(offsets.clone()), + vec![struct_child], + )?; + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(source_batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let adapted_child = adapted_union.child(0); + let struct_arr = adapted_child + .as_any() + .downcast_ref::() + .unwrap(); + assert!(struct_arr.fields()[0].is_nullable()); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_union_incompatible_mode_rejected() { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(test_two_field_union(true), UnionMode::Dense), + false, + )])); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(test_two_field_union(false), UnionMode::Sparse), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let type_ids = [0, 0].into_iter().collect::>(); + let source_union = UnionArray::try_new( + test_two_field_union(false), + type_ids, + None, + vec![int_array, str_array], + ) + .unwrap(); + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)]).unwrap(); + + let res = adapt_batch_to_schema(source_batch, &declared_schema); + assert!(res.is_err()); + } + + #[test] + fn test_adapt_batch_to_schema_union_field_set_mismatch_rejected() { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + // Target has type ID [0] + let target_union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("value", DataType::Int32, true)], + ) + .unwrap(); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Sparse), + false, + )])); + + // Source has type IDs [0, 1] (where ID 0 is compatible) + let source_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, false), + Field::new("extra", DataType::Utf8, false), + ], + ) + .unwrap(); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Sparse), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let type_ids = [0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + None, + vec![int_array, str_array], + ) + .unwrap(); + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)]).unwrap(); + + let res = adapt_batch_to_schema(source_batch, &declared_schema); + assert!(res.is_err()); + let err = res.unwrap_err().to_string(); + assert!( + err.contains("different field sets") + || err.contains("cannot be adapted to expected type"), + "unexpected error message: {err}" + ); + } + + #[test] + fn test_validate_data_type_compatibility_union() { + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_type = DataType::Union(test_two_field_union(true), UnionMode::Dense); + + // Compatible: exact same type IDs in different order with stricter nullability + let reordered_source_fields = UnionFields::try_new( + vec![1, 0], + vec![ + Field::new("str", DataType::Utf8, false), + Field::new("value", DataType::Int32, false), + ], + ) + .unwrap(); + let source_type = DataType::Union(reordered_source_fields, UnionMode::Dense); + assert!( + validate_data_type_compatibility("u", &source_type, &target_type).is_ok() + ); + + // Incompatible: mismatched mode + let sparse_source_type = + DataType::Union(test_two_field_union(false), UnionMode::Sparse); + assert!( + validate_data_type_compatibility("u", &sparse_source_type, &target_type) + .is_err() + ); + + // Incompatible: field-set mismatch (extra source ID 2) + let extra_id_source = DataType::Union( + UnionFields::try_new( + vec![0, 1, 2], + vec![ + Field::new("value", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + Field::new("extra", DataType::Int32, false), + ], + ) + .unwrap(), + UnionMode::Dense, + ); + assert!( + validate_data_type_compatibility("u", &extra_id_source, &target_type) + .is_err() + ); + + // Incompatible: field-set mismatch (missing source ID 1) + let missing_id_source = DataType::Union( + UnionFields::try_new( + vec![0], + vec![Field::new("value", DataType::Int32, false)], + ) + .unwrap(), + UnionMode::Dense, + ); + assert!( + validate_data_type_compatibility("u", &missing_id_source, &target_type) + .is_err() + ); + } +} diff --git a/datafusion/core/tests/sql/aggregates/mod.rs b/datafusion/core/tests/sql/aggregates/mod.rs index b209e91cc81e7..186297b639cbd 100644 --- a/datafusion/core/tests/sql/aggregates/mod.rs +++ b/datafusion/core/tests/sql/aggregates/mod.rs @@ -1021,3 +1021,4 @@ pub fn split_fuzz_timestamp_data_into_batches( pub mod basic; pub mod dict_nulls; +mod nested_nullability; diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs new file mode 100644 index 0000000000000..448759ad74c54 --- /dev/null +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -0,0 +1,246 @@ +// 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 for aggregating batches whose data types are *stricter* +//! than the table's declared schema. See +//! . +//! +//! Builds on the end-to-end reproducer from #24278 by @alamb. +//! +//! A `RecordBatch` is a valid instance of a schema that is a superset of its +//! own (see [`Schema::contains`] / `Field::contains`): most commonly the +//! schema declares a (possibly nested) field as nullable while the batch's +//! arrays mark it non-nullable. `MemTable::try_new` accepts such batches via +//! exactly that check, and engines embedding DataFusion (e.g. Comet) feed +//! such batches over FFI. Aggregations must therefore not fail when the +//! runtime arrays are stricter than the planned schema. +//! +//! [`Schema::contains`]: arrow::datatypes::Schema::contains + +use std::sync::Arc; + +use arrow::array::{BooleanArray, RecordBatch, StructArray, UInt32Array}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion::physical_plan::collect; +use datafusion::physical_plan::expressions::col; +use datafusion::prelude::*; +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::FairSpillPool; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_functions_aggregate::array_agg::array_agg_udaf; + +/// Returns the fields of the struct column `b`: a single `colA Boolean`. +/// +/// `col_a_nullable` controls whether `colA` is declared nullable — the only +/// difference between the table's declared schema (`true`) and the actual +/// batches (`false`). +fn make_struct_fields(col_a_nullable: bool) -> Fields { + Fields::from(vec![Field::new("colA", DataType::Boolean, col_a_nullable)]) +} + +/// Returns the schema `(a UInt32 NOT NULL, b Struct("colA" Boolean) NOT NULL)` +/// with the nested field `b.colA` nullable per `col_a_nullable`. +/// +/// See [`make_struct_fields`]. +fn make_schema(col_a_nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Struct(make_struct_fields(col_a_nullable)), + false, + ), + ])) +} + +/// Runs a SQL aggregation over a table whose batches are stricter than its +/// declared schema. +/// +/// [`Self::run`] registers table `t(a UInt32, b Struct("colA" Boolean))` +/// where the declared schema marks the nested field `colA` as nullable, but +/// the batches carry a stricter, non-nullable `colA`, then runs the query +/// and returns the collected result. +struct AggregateBatchesTest { + /// Number of rows in the table. `a` is `0..num_rows` (so also the number + /// of groups for `GROUP BY a`) and `b.colA` alternates `true` / `false`. + num_rows: u32, + /// If set, the context uses a [`FairSpillPool`] of this size (and a small + /// batch size) so the aggregation is forced to spill. + memory_limit: Option, +} + +impl AggregateBatchesTest { + fn new() -> Self { + Self { + num_rows: 100, + memory_limit: None, + } + } + + fn with_num_rows(mut self, num_rows: u32) -> Self { + self.num_rows = num_rows; + self + } + + fn with_memory_limit(mut self, memory_limit: usize) -> Self { + self.memory_limit = Some(memory_limit); + self + } + + /// Runs `sql` against the table described above and asserts the result + /// has one output row per group (i.e. [`Self::num_rows`] rows in total). + async fn run(self, sql: &str) -> Result<()> { + // The table's declared schema: the nested field `b.colA` is + // nullable ... + let declared_schema = make_schema(true); + + // ... while the batches are stricter: `b.colA` is non-nullable. + // `MemTable::try_new` accepts this combination via + // `Schema::contains`. + let batch_struct_fields = make_struct_fields(false); + let batch = RecordBatch::try_new( + make_schema(false), + vec![ + Arc::new(UInt32Array::from_iter_values(0..self.num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..self.num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let table = MemTable::try_new(declared_schema, vec![vec![batch]])?; + + let ctx = match self.memory_limit { + Some(limit) => { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(limit))) + .build_arc()?; + SessionContext::new_with_config_rt( + SessionConfig::new().with_batch_size(100), + runtime, + ) + } + None => SessionContext::new(), + }; + ctx.register_table("t", Arc::new(table))?; + + let result = ctx.sql(sql).await?.collect().await?; + + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, self.num_rows as usize); + Ok(()) + } +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +} + +/// Direct unit test for `AggregateExec` boundary adaptation: +/// Feeds `AggregateExec` directly from a `MemorySourceConfig` whose batches carry +/// a stricter nested struct nullability than the plan schema without going +/// through `MemTable`. +#[tokio::test] +async fn test_aggregate_exec_direct_input_adaptation() -> Result<()> { + let declared_schema = make_schema(true); + let batch_struct_fields = make_struct_fields(false); + let num_rows = 100_u32; + let stricter_batch = RecordBatch::try_new( + make_schema(false), + vec![ + Arc::new(UInt32Array::from_iter_values(0..num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let input_plan: Arc = MemorySourceConfig::try_new_exec( + &[vec![stricter_batch]], + Arc::clone(&declared_schema), + None, + )?; + + let grouping_set = + PhysicalGroupBy::new_single(vec![(col("a", &declared_schema)?, "a".to_string())]); + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("b", &declared_schema)?]) + .schema(Arc::clone(&declared_schema)) + .alias("array_agg(b)") + .build()?, + )]; + + let agg_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + grouping_set, + aggregates, + vec![None], + input_plan, + Arc::clone(&declared_schema), + )?); + + let task_ctx = Arc::new(TaskContext::default()); + let results = collect(agg_exec, task_ctx).await?; + + let total_rows: usize = results.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, num_rows as usize); + Ok(()) +} diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index efe42c7ebc5f0..0b6bdf4490d8b 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -106,6 +106,17 @@ impl Stream for MemoryStream { None => batch.clone(), }; + // MemoryStream advertises `self.schema`, therefore emitted RecordBatches + // must conform to it when batches were provided with stricter nested types + // (e.g. MemTable accepts stricter batches via Schema::contains). + let batch = if batch.schema().as_ref() != self.schema.as_ref() + && self.schema.contains(batch.schema().as_ref()) + { + datafusion_common::nested_struct::adapt_batch_to_schema(batch, &self.schema)? + } else { + batch + }; + let Some(&fetch) = self.fetch.as_ref() else { return Poll::Ready(Some(Ok(batch))); }; @@ -673,4 +684,128 @@ mod lazy_memory_tests { Ok(()) } + + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema() -> Result<()> { + use arrow::array::{ArrayRef, BooleanArray, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + // Declared schema expects nullable struct field colA + let declared_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, true)]); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "b", + DataType::Struct(declared_fields), + false, + )])); + + // Runtime batch has stricter non-nullable struct field colA + let source_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, false)]); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "b", + DataType::Struct(source_fields.clone()), + false, + )])); + + let struct_array: ArrayRef = Arc::new(StructArray::new( + source_fields, + vec![Arc::new(BooleanArray::from(vec![true, false]))], + None, + )); + let stricter_batch = RecordBatch::try_new(source_schema, vec![struct_array])?; + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), declared_schema); + + let struct_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(struct_col.fields()[0].is_nullable()); + let bool_child = struct_col + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(bool_child.value(0)); + assert!(!bool_child.value(1)); + + Ok(()) + } + + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_with_projection() + -> Result<()> { + use arrow::array::{ArrayRef, BooleanArray, Int32Array, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + // Declared full schema: col a (Int32), col b (Struct) + let declared_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, true)]); + let full_declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Struct(declared_fields), false), + ])); + + // Projected schema for column "b" (projection = [1]) + let projected_schema = Arc::new(full_declared_schema.project(&[1])?); + + // Runtime batch has stricter struct + let source_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, false)]); + let source_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Struct(source_fields.clone()), false), + ])); + + let struct_array: ArrayRef = Arc::new(StructArray::new( + source_fields, + vec![Arc::new(BooleanArray::from(vec![true, false]))], + None, + )); + let stricter_batch = RecordBatch::try_new( + source_schema, + vec![Arc::new(Int32Array::from(vec![10, 20])), struct_array], + )?; + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&projected_schema), + Some(vec![1]), + )?; + + assert_eq!(stream.schema(), projected_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), projected_schema); + assert_eq!(emitted_batch.num_columns(), 1); + + let struct_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(struct_col.fields()[0].is_nullable()); + let bool_child = struct_col + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(bool_child.value(0)); + assert!(!bool_child.value(1)); + + Ok(()) + } } From f23582e370f81fa4b0be72020f427bc80a4dcb87 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 28 Aug 2026 10:22:01 -0400 Subject: [PATCH 2/2] test: add MemoryStream regressions for Union and Map<.., Struct> adaptation Adds two `MemoryStream` regression tests requested during review of the #24394 backport: - Dense Union with nullable declared children and non-nullable runtime children, exercising the Union reconstruction path at the producer boundary and asserting preserved type IDs and child values. - `Map>` where the runtime nested field is non-nullable and the declared nested field is nullable. This passes on branch-55 without a `DataType::Map` arm in `cast_column`: `Schema::contains` accepts the stricter shape and Arrow's generic Map cast recursively casts the value Struct and rebuilds it with the target fields. No backport of the Map prerequisite from #23914 is required. Both assert `emitted_batch.schema() == stream.schema()`. Co-Authored-By: Claude Opus 5 (1M context) --- datafusion/physical-plan/src/memory.rs | 182 +++++++++++++++++++++++++ 1 file changed, 182 insertions(+) diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index 0b6bdf4490d8b..0c77d7e773265 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -808,4 +808,186 @@ mod lazy_memory_tests { Ok(()) } + + /// Regression for the Union reconstruction path at the `MemoryStream` + /// producer boundary: a declared nullable Union child vs a stricter + /// non-nullable runtime child. + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_union() -> Result<()> + { + use arrow::array::{Array, ArrayRef, Float64Array, Int32Array, UnionArray}; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{DataType, Field, Schema, UnionFields, UnionMode}; + use futures::StreamExt; + + let declared_union_fields = UnionFields::try_new( + vec![0_i8, 1], + vec![ + Field::new("i", DataType::Int32, true), + Field::new("f", DataType::Float64, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(declared_union_fields, UnionMode::Dense), + false, + )])); + + let source_union_fields = UnionFields::try_new( + vec![0_i8, 1], + vec![ + Field::new("i", DataType::Int32, false), + Field::new("f", DataType::Float64, false), + ], + )?; + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let type_ids = ScalarBuffer::from(vec![0_i8, 1, 0]); + let offsets = ScalarBuffer::from(vec![0_i32, 0, 1]); + let union_array: ArrayRef = Arc::new(UnionArray::try_new( + source_union_fields, + type_ids, + Some(offsets), + vec![ + Arc::new(Int32Array::from(vec![10, 20])), + Arc::new(Float64Array::from(vec![1.5])), + ], + )?); + let stricter_batch = RecordBatch::try_new(source_schema, vec![union_array])?; + + assert!(declared_schema.contains(stricter_batch.schema().as_ref())); + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), stream.schema()); + assert_eq!(emitted_batch.schema(), declared_schema); + + let union_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(union_col.len(), 3); + assert_eq!(union_col.type_id(0), 0); + assert_eq!(union_col.type_id(1), 1); + assert_eq!(union_col.type_id(2), 0); + let i_child = union_col + .child(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(i_child.values(), &[10, 20]); + + Ok(()) + } + + /// Regression for a contained `Map<.., Struct>` whose runtime nested field + /// is non-nullable while the declared nested field is nullable. + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_map_of_struct() + -> Result<()> { + use arrow::array::{ + Array, ArrayRef, Int32Array, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + fn map_field(value_child_nullable: bool) -> Field { + let value_struct = DataType::Struct(Fields::from(vec![Field::new( + "v", + DataType::Int32, + value_child_nullable, + )])); + let entries = Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", value_struct, true), + ])), + false, + ); + Field::new("m", DataType::Map(Arc::new(entries), false), true) + } + + let declared_schema = Arc::new(Schema::new(vec![map_field(true)])); + let source_schema = Arc::new(Schema::new(vec![map_field(false)])); + + let value_fields = Fields::from(vec![Field::new("v", DataType::Int32, false)]); + let values_struct = StructArray::new( + value_fields, + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + None, + ); + let entries = StructArray::new( + Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", values_struct.data_type().clone(), true), + ]), + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef, + Arc::new(values_struct) as ArrayRef, + ], + None, + ); + let DataType::Map(source_entries_field, _) = source_schema.field(0).data_type() + else { + unreachable!("map field") + }; + let map_array: ArrayRef = Arc::new(MapArray::try_new( + Arc::clone(source_entries_field), + OffsetBuffer::new(vec![0, 2, 3].into()), + entries, + None, + false, + )?); + let stricter_batch = RecordBatch::try_new(source_schema, vec![map_array])?; + + // The stricter batch is accepted by `MemTable::try_new`-style checks. + assert!(declared_schema.contains(stricter_batch.schema().as_ref())); + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), stream.schema()); + assert_eq!(emitted_batch.schema(), declared_schema); + + let map_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(map_col.len(), 2); + let values = map_col + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(values.fields()[0].is_nullable()); + let ints = values + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ints.values(), &[1, 2, 3]); + + Ok(()) + } }