From a37166dd1514a0e4ef7340e4f707f2a2fc6d7a3e Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Tue, 18 Aug 2026 22:22:56 +0200 Subject: [PATCH] fix: normalize emitted MemoryStream batches to declared table schema (#24069) Ensures batches emitted by MemoryStream conform to its advertised schema by adapting batches when runtime nested data types have stricter nullability than the table's declared schema (e.g. accepted by MemTable via Schema::contains). Also extends datafusion_common::nested_struct::adapt_batch_to_schema with narrow schema-conformance support for Arrow UnionArray (sparse and dense) without changing general SQL CAST behavior. Fixes #24069 Closes #24394 --- datafusion/common/src/nested_struct.rs | 802 +++++++++++++++++- datafusion/core/tests/sql/aggregates/mod.rs | 1 + .../sql/aggregates/nested_nullability.rs | 246 ++++++ datafusion/physical-plan/src/memory.rs | 135 +++ 4 files changed, 1182 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..fb8960387a807 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,51 @@ 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 +640,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 +1830,674 @@ mod tests { )); } } + +/// Adapts a [`RecordBatch`] to a target [`SchemaRef`]. +/// +/// If `batch` already has the target schema, it is returned immediately. +/// +/// If `batch` has columns whose data types differ from `target_schema` (e.g. stricter +/// nested struct or list nullabilities), this function verifies that each target data +/// 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::{BooleanArray, 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 batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec![Some("x"), None, Some("z")])), + ], + )?; + + let adapted = adapt_batch_to_schema(batch.clone(), &schema)?; + assert!(Arc::ptr_eq(batch.schema_ref(), adapted.schema_ref())); + assert_eq!(batch, adapted); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_stricter_nested_struct() -> Result<()> { + let declared_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "nested", + Struct(Fields::from(vec![Field::new( + "val", + DataType::Boolean, + true, + )])), + false, + ), + ])); + + let stricter_batch_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "nested", + Struct(Fields::from(vec![Field::new( + "val", + DataType::Boolean, + false, + )])), + false, + ), + ])); + + let struct_col = Arc::new(StructArray::new( + Fields::from(vec![Field::new("val", DataType::Boolean, false)]), + vec![Arc::new(BooleanArray::from(vec![true, false, true]))], + None, + )); + + let batch = RecordBatch::try_new( + stricter_batch_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3])), struct_col], + )?; + + 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); + assert_eq!(adapted.num_columns(), 2); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_top_level_nullability_only() -> Result<()> { + // Target is nullable, batch is non-nullable + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let stricter_batch_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + let batch = RecordBatch::try_new( + stricter_batch_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + assert_eq!(adapted.column(0).len(), 3); + + // Target is non-nullable, batch is nullable (with no nulls) + let non_null_target = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let nullable_batch_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch2 = RecordBatch::try_new( + nullable_batch_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + let adapted2 = adapt_batch_to_schema(batch2, &non_null_target)?; + assert_eq!(adapted2.schema().as_ref(), non_null_target.as_ref()); + assert_eq!(adapted2.column(0).len(), 3); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_null_into_non_nullable_rejected() { + let declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), // target is non-nullable + ])); + + let batch_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), // batch is nullable and contains nulls + ])); + + let batch = RecordBatch::try_new( + batch_schema, + vec![Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]))], + ) + .unwrap(); + + let res = adapt_batch_to_schema(batch, &declared_schema); + assert!(res.is_err()); + } + + #[test] + fn test_adapt_batch_to_schema_incompatible_type_rejected() { + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + + let batch_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let batch = RecordBatch::try_new( + batch_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + let res = adapt_batch_to_schema(batch, &declared_schema); + assert!(res.is_err()); + let err_msg = res.unwrap_err().to_string(); + assert!( + err_msg.contains("cannot be adapted to expected type"), + "unexpected error message: {err_msg}" + ); + } + + 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 declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(test_two_field_union(true), UnionMode::Sparse), + 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, 30])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let type_ids = [0, 1, 0].into_iter().collect::>(); + let source_union = UnionArray::try_new( + test_two_field_union(false), + type_ids.clone(), + None, + 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); + + // Verify exact active values + let int_child = adapted_union + .child(0) + .as_any() + .downcast_ref::() + .unwrap(); + let str_child = adapted_union + .child(1) + .as_any() + .downcast_ref::() + .unwrap(); + + // Row 0: type_id 0 -> 10 + assert_eq!(adapted_union.type_id(0), 0); + assert_eq!(int_child.value(adapted_union.value_offset(0)), 10); + + // Row 1: type_id 1 -> "b" + assert_eq!(adapted_union.type_id(1), 1); + assert_eq!(str_child.value(adapted_union.value_offset(1)), "b"); + + // Row 2: type_id 0 -> 30 + assert_eq!(adapted_union.type_id(2), 0); + assert_eq!(int_child.value(adapted_union.value_offset(2)), 30); + + 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 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::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 = [0, 1, 0].into_iter().collect::>(); + let offsets = [0, 0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + test_two_field_union(false), + 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)); + + // Verify exact active values + let int_child = adapted_union + .child(0) + .as_any() + .downcast_ref::() + .unwrap(); + let str_child = adapted_union + .child(1) + .as_any() + .downcast_ref::() + .unwrap(); + + // Row 0: type_id 0 -> offset 0 -> 10 + assert_eq!(adapted_union.type_id(0), 0); + assert_eq!(int_child.value(adapted_union.value_offset(0)), 10); + + // Row 1: type_id 1 -> offset 0 -> "b" + assert_eq!(adapted_union.type_id(1), 1); + assert_eq!(str_child.value(adapted_union.value_offset(1)), "b"); + + // Row 2: type_id 0 -> offset 1 -> 30 + assert_eq!(adapted_union.type_id(2), 0); + assert_eq!(int_child.value(adapted_union.value_offset(2)), 30); + + 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::{UnionFields, UnionMode}; + + // Target: fields ordered [(3, str), (1, int)] + let target_union_fields = UnionFields::try_new( + vec![3, 1], + vec![ + Field::new("str", DataType::Utf8, true), + Field::new("value", DataType::Int32, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + )])); + + // Source: fields ordered [(1, int), (3, str)] with stricter nullability + let source_union_fields = UnionFields::try_new( + vec![1, 3], + vec![ + Field::new("value", 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(()) + } }