diff --git a/datafusion/common/src/nested_struct.rs b/datafusion/common/src/nested_struct.rs index e915b91b911cc..24c96ee1f0ceb 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_empty_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,96 @@ 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. +/// - Validates that all source type IDs exist in the target union. +/// - Recursively adapts each matching child array using `cast_column`. +/// - Fills target children not present in the source with empty arrays (dense) or typed nulls (sparse). +/// - 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 { + if source_mode != target_mode { + return _plan_err!( + "Cannot cast Union with mode {source_mode:?} to Union with mode {target_mode:?}" + ); + } + + if let Some(source_union) = source_col.as_any().downcast_ref::() { + // Validate all source type IDs exist in target + for (source_type_id, source_field) in source_fields.iter() { + if target_fields + .iter() + .all(|(target_type_id, _)| target_type_id != source_type_id) + { + return _plan_err!( + "Cannot cast union: source type ID {source_type_id} ('{}') is missing from target union", + source_field.name() + ); + } + } + + let mut adapted_children: Vec = Vec::with_capacity(target_fields.len()); + + for (target_type_id, target_child_field) in target_fields.iter() { + let has_source_child = + source_fields.iter().any(|(s_id, _)| s_id == target_type_id); + + if has_source_child { + let source_child = source_union.child(target_type_id); + let adapted_child = cast_column( + source_child, + target_child_field.data_type(), + cast_options, + ) + .map_err(|e| { + e.context(format!( + "While casting union child with type ID {target_type_id} ('{}')", + target_child_field.name() + )) + })?; + adapted_children.push(adapted_child); + } else { + match target_mode { + UnionMode::Dense => { + adapted_children + .push(new_empty_array(target_child_field.data_type())); + } + UnionMode::Sparse => { + adapted_children.push(new_null_array( + target_child_field.data_type(), + source_col.len(), + )); + } + } + } + } + + let type_ids = source_union.type_ids().clone(); + let offsets = source_union.offsets().cloned(); + + let union_array = UnionArray::try_new( + target_fields.clone(), + type_ids, + offsets, + adapted_children, + )?; + Ok(Arc::new(union_array)) + } else { + _plan_err!( + "Cannot cast column of type {} to union type. Source must be a UnionArray.", + source_col.data_type() + ) + } +} + /// 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 +308,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)?), } } @@ -260,8 +364,8 @@ fn cast_list_view_column( source_list.sizes().clone(), cast_values, source_list.nulls().cloned(), - )?; - Ok(Arc::new(result)) + ); + Ok(Arc::new(result?)) } fn cast_fixed_size_list_column( @@ -524,6 +628,28 @@ 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), + ) => { + if source_mode != target_mode { + return _plan_err!( + "Cannot cast Union with mode {source_mode:?} to Union with mode {target_mode:?}" + ); + } + for (source_type_id, source_field) in source_fields.iter() { + let Some((_, target_field)) = target_fields + .iter() + .find(|(target_type_id, _)| *target_type_id == source_type_id) + else { + return _plan_err!( + "Cannot cast union: source type ID {source_type_id} ('{}') is missing from target union", + source_field.name() + ); + }; + validate_field_compatibility(source_field, target_field)?; + } + } _ => { if !can_cast_types(source_type, target_type) { return _plan_err!( @@ -543,7 +669,8 @@ pub fn validate_data_type_compatibility( /// /// This is the case when both types are struct types, or both are the same /// container type (List, LargeList, equal-width FixedSizeList, ListView, -/// LargeListView, Dictionary) wrapping types that recursively contain structs. +/// LargeListView, Dictionary) wrapping types that recursively contain structs, +/// or both are Union types with matching modes. /// /// Use this predicate at both planning time (to decide whether to apply struct /// compatibility validation) and execution time (to decide whether to route @@ -569,6 +696,9 @@ pub fn requires_nested_struct_cast( (DataType::Dictionary(_, s_val), DataType::Dictionary(_, t_val)) => { requires_nested_struct_cast(s_val, t_val) } + (DataType::Union(_, source_mode), DataType::Union(_, target_mode)) => { + source_mode == target_mode + } _ => false, } } @@ -1703,3 +1833,525 @@ 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}" + ); + } + + #[test] + fn test_adapt_batch_to_schema_stricter_sparse_union() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields.clone(), UnionMode::Sparse), + false, + )])); + + let source_union_fields = UnionFields::try_new( + vec![0, 1], + 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::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( + source_union_fields, + 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.type_ids(), &type_ids); + assert_eq!(adapted_union.child(0).data_type(), &DataType::Int32); + assert_eq!(adapted_union.child(1).data_type(), &DataType::Utf8); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_stricter_dense_union() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, true), + Field::new("str", DataType::Utf8, 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![0, 1], + 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 = [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.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.type_ids(), &type_ids); + assert_eq!(adapted_union.offsets(), Some(&offsets)); + 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::{UnionFields, UnionMode}; + + let target_union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("value", DataType::Int32, true)], + ) + .unwrap(); + // Target is Dense + 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![0], + vec![Field::new("value", DataType::Int32, false)], + ) + .unwrap(); + // Source is Sparse + 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 type_ids = [0, 0].into_iter().collect::>(); + let source_union = + UnionArray::try_new(source_union_fields, type_ids, None, vec![int_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_missing_type_id_rejected() { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + 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 ID 1 which is not in target (only has 0) + let source_union_fields = UnionFields::try_new( + vec![1], + vec![Field::new("value", DataType::Int32, 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 type_ids = [1, 1].into_iter().collect::>(); + let source_union = + UnionArray::try_new(source_union_fields, type_ids, None, vec![int_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_validate_data_type_compatibility_union() { + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ], + ) + .unwrap(); + let target_type = DataType::Union(target_fields, UnionMode::Dense); + + // Compatible: subset of type IDs with stricter nullability + let source_fields = + UnionFields::try_new(vec![0], vec![Field::new("a", DataType::Int32, false)]) + .unwrap(); + let source_type = DataType::Union(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( + UnionFields::try_new(vec![0], vec![Field::new("a", DataType::Int32, false)]) + .unwrap(), + UnionMode::Sparse, + ); + assert!( + validate_data_type_compatibility("u", &sparse_source_type, &target_type) + .is_err() + ); + + // Incompatible: source has extra type ID not in target + let extra_id_source = DataType::Union( + UnionFields::try_new(vec![2], vec![Field::new("c", DataType::Int32, false)]) + .unwrap(), + UnionMode::Dense, + ); + assert!( + validate_data_type_compatibility("u", &extra_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..fa5fb35453e5e --- /dev/null +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -0,0 +1,443 @@ +// 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(()) +} + +/// Direct unit test for `AggregateExec` boundary adaptation with Dense Union input: +/// Feeds `AggregateExec` with batches containing a Union with stricter child nullability +/// than the plan schema. +#[tokio::test] +async fn test_aggregate_exec_direct_input_adaptation_dense_union() -> Result<()> { + use arrow::array::{ArrayRef, Int32Array, StringArray, UnionArray}; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + ), + ])); + + let source_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + ], + )?; + let source_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + ), + ])); + + let num_rows = 100_u32; + let int_array: ArrayRef = + Arc::new(Int32Array::from_iter_values((0..50_i32).map(|i| i * 10))); + let str_array: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..50).map(|i| format!("val_{i}")), + )); + let type_ids = (0..num_rows) + .map(|i| (i % 2) as i8) + .collect::>(); + let offsets = (0..num_rows) + .map(|i| (i / 2) as i32) + .collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + Some(offsets), + vec![int_array, str_array], + )?; + + let stricter_batch = RecordBatch::try_new( + source_schema, + vec![ + Arc::new(UInt32Array::from_iter_values(0..num_rows)), + Arc::new(source_union), + ], + )?; + + 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(()) +} + +/// Direct unit test for `AggregateExec` boundary adaptation with Sparse Union input: +/// Feeds `AggregateExec` with batches containing a Sparse Union with stricter child nullability. +#[tokio::test] +async fn test_aggregate_exec_direct_input_adaptation_sparse_union() -> Result<()> { + use arrow::array::{ArrayRef, Int32Array, StringArray, UnionArray}; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Union(target_union_fields, UnionMode::Sparse), + false, + ), + ])); + + let source_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + ], + )?; + let source_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Union(source_union_fields.clone(), UnionMode::Sparse), + false, + ), + ])); + + let num_rows = 100_u32; + let int_array: ArrayRef = Arc::new(Int32Array::from_iter_values( + (0..num_rows as i32).map(|i| i * 10), + )); + let str_array: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..num_rows).map(|i| format!("val_{i}")), + )); + let type_ids = (0..num_rows) + .map(|i| (i % 2) as i8) + .collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + None, + vec![int_array, str_array], + )?; + + let stricter_batch = RecordBatch::try_new( + source_schema, + vec![ + Arc::new(UInt32Array::from_iter_values(0..num_rows)), + Arc::new(source_union), + ], + )?; + + 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/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index ac7727b459300..01bb7ef764c07 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -290,7 +290,7 @@ impl AggregateStream { let agg_filter_expr = Arc::clone(&agg.filter_expr); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?; let filter_expressions = match agg.mode.input_mode() { diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index c0253093c8a7b..a571cf5d9caf9 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -395,7 +395,7 @@ impl GroupedHashAggregateStream { let agg_filter_expr = Arc::clone(&agg.filter_expr); let batch_size = context.session_config().batch_size(); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); let aggregate_argument_metrics = AggregateArgumentMetrics::new( diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 193fdba4b0198..6ab122f6efa3c 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -68,7 +68,7 @@ impl GroupedTopKAggregateStream { ) -> Result { let agg_schema = Arc::clone(&aggr.schema); let group_by = Arc::clone(&aggr.group_by); - let input = aggr.input.execute(partition, Arc::clone(context))?; + let input = aggr.execute_input(partition, context)?; let baseline_metrics = BaselineMetrics::new(&aggr.metrics, partition); let group_by_metrics = GroupByMetrics::new(&aggr.metrics, partition); let aggregate_arguments = diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 2df5960188a2b..7ec751f02189e 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -434,7 +434,7 @@ impl PartialHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); @@ -1020,7 +1020,7 @@ impl FinalHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0f1b718a3c8e3..8dd62b9edd816 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -170,7 +170,8 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, - InputOrderMode, Partitioning, SendableRecordBatchStream, Statistics, + InputOrderMode, Partitioning, RecordBatchStream, SendableRecordBatchStream, + Statistics, }; use datafusion_common::config::ConfigOptions; use parking_lot::Mutex; @@ -876,6 +877,45 @@ pub struct AggregateExec { dynamic_filter: Option>, } +/// A stream wrapper that ensures every yielded batch matches the declared input schema. +pub(crate) struct AdaptedInputRecordBatchStream { + inner: SendableRecordBatchStream, + schema: SchemaRef, +} + +impl AdaptedInputRecordBatchStream { + pub(crate) fn new(inner: SendableRecordBatchStream, schema: SchemaRef) -> Self { + Self { inner, schema } + } +} + +impl futures::Stream for AdaptedInputRecordBatchStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + use futures::StreamExt; + match self.inner.poll_next_unpin(cx) { + std::task::Poll::Ready(Some(Ok(batch))) => { + let adapted = datafusion_common::nested_struct::adapt_batch_to_schema( + batch, + &self.schema, + ); + std::task::Poll::Ready(Some(adapted)) + } + other => other, + } + } +} + +impl RecordBatchStream for AdaptedInputRecordBatchStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + impl AggregateExec { /// Function used in `OptimizeAggregateOrder` optimizer rule, /// where we need parts of the new value, others cloned from the old one @@ -1164,6 +1204,18 @@ impl AggregateExec { /// Aggregation has multiple specialized implementations optimized for /// different workloads. This function picks the best available path. + pub(crate) fn execute_input( + &self, + partition: usize, + context: &Arc, + ) -> Result { + let input = self.input.execute(partition, Arc::clone(context))?; + Ok(Box::pin(AdaptedInputRecordBatchStream::new( + input, + self.input.schema(), + ))) + } + fn execute_typed( &self, partition: usize, diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 2c26b74da7748..c6726d6c6ce1e 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -271,7 +271,7 @@ impl OrderedFinalAggregateStream { )); debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; Self::new_with_input(agg, context, partition, input, &agg.input_order_mode) } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 9e93a111a6466..afc2fa5b30ad1 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -129,7 +129,7 @@ impl OrderedPartialAggregateStream { debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index da00b42e5c3ed..8025146a10db4 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -332,7 +332,7 @@ impl OrderedSingleAggregateStream { debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 2f4535e66f4ef..06e8da1c7d76e 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -157,7 +157,7 @@ impl PartialReduceHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 3e306d72a7e82..10d3519c43e30 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -323,7 +323,7 @@ impl SingleHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);