You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/24 15:40:55 UTC
(datafusion) 01/01: fix: adapt input batches with stricter nested nullability to planned schema in aggregation (#24394)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24394-482d1fd1fe264672090d8e2f5d98e35f73c87614
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 2326917aa8d691aa5123518ca602ce63ab18ffc6
Author: Patrick Ribbsaeter <pa...@gmail.com>
AuthorDate: Mon Aug 24 15:40:30 2026 +0000
fix: adapt input batches with stricter nested nullability to planned schema in aggregation (#24394)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
Closes #24069.
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
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).
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
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.
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
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`.
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
No. Queries aggregating in-memory tables whose batches have stricter
nested nullability than the table schema now succeed as expected.
---
datafusion/common/src/nested_struct.rs | 714 ++++++++++++++++++++-
datafusion/core/tests/sql/aggregates/mod.rs | 1 +
.../tests/sql/aggregates/nested_nullability.rs | 246 +++++++
datafusion/physical-plan/src/memory.rs | 135 ++++
4 files changed, 1093 insertions(+), 3 deletions(-)
diff --git a/datafusion/common/src/nested_struct.rs b/datafusion/common/src/nested_struct.rs
index 3f5e4ab4b7..32736b069f 100644
--- a/datafusion/common/src/nested_struct.rs
+++ b/datafusion/common/src/nested_struct.rs
@@ -19,12 +19,14 @@ use crate::error::{_plan_err, Result};
use arrow::{
array::{
Array, ArrayRef, AsArray, DictionaryArray, FixedSizeListArray, GenericListArray,
- GenericListViewArray, MapArray, StructArray, UInt64Array, downcast_integer,
- make_array, new_null_array,
+ GenericListViewArray, MapArray, RecordBatch, StructArray, UInt64Array,
+ UnionArray, downcast_integer, make_array, new_null_array,
},
buffer::NullBuffer,
compute::{CastOptions, can_cast_types, cast_with_options, take},
- datatypes::{DataType, DataType::Struct, Field, FieldRef},
+ datatypes::{
+ DataType, DataType::Struct, Field, FieldRef, SchemaRef, UnionFields, UnionMode,
+ },
};
use std::{
collections::{HashMap, HashSet},
@@ -125,6 +127,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<ArrayRef> {
+ validate_union_schema_compatibility(
+ source_fields,
+ source_mode,
+ target_fields,
+ target_mode,
+ )?;
+
+ let source_union = source_col
+ .as_any()
+ .downcast_ref::<UnionArray>()
+ .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
@@ -231,6 +290,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)?),
}
}
@@ -751,6 +821,48 @@ fn validate_map_entries_field(entries: &Field) -> Result<(&FieldRef, &FieldRef)>
/// not remove source fields. Sorted Maps require an unchanged key type. Map
/// entries and keys must remain non-nullable. Source and target sorted flags
/// must match. Maps without semantic Struct children use Arrow's normal cast.
+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(())
+}
+
pub fn validate_data_type_compatibility(
field_name: &str,
source_type: &DataType,
@@ -801,6 +913,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!(
@@ -2805,3 +2928,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<RecordBatch> {
+ 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::<ScalarBuffer<i8>>();
+ 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::<UnionArray>()
+ .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::<ScalarBuffer<i8>>();
+ let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();
+ 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::<UnionArray>()
+ .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::<ScalarBuffer<i8>>();
+ let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();
+ 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::<UnionArray>()
+ .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::<Int32Array>()
+ .unwrap();
+ let str_child = adapted_union
+ .child(3)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .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::<ScalarBuffer<i8>>();
+ let offsets = [0, 1].into_iter().collect::<ScalarBuffer<i32>>();
+ 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::<UnionArray>()
+ .unwrap();
+ let adapted_child = adapted_union.child(0);
+ let struct_arr = adapted_child
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .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::<ScalarBuffer<i8>>();
+ 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::<ScalarBuffer<i8>>();
+ 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 b209e91cc8..186297b639 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 0000000000..448759ad74
--- /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
+//! <https://github.com/apache/datafusion/issues/24069>.
+//!
+//! 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<usize>,
+}
+
+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<dyn ExecutionPlan> = 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 efe42c7ebc..0b6bdf4490 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::<StructArray>()
+ .unwrap();
+ assert!(struct_col.fields()[0].is_nullable());
+ let bool_child = struct_col
+ .column(0)
+ .as_any()
+ .downcast_ref::<BooleanArray>()
+ .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<colA: nullable Boolean>)
+ 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::<StructArray>()
+ .unwrap();
+ assert!(struct_col.fields()[0].is_nullable());
+ let bool_child = struct_col
+ .column(0)
+ .as_any()
+ .downcast_ref::<BooleanArray>()
+ .unwrap();
+ assert!(bool_child.value(0));
+ assert!(!bool_child.value(1));
+
+ Ok(())
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org