Skip to content

fix: adapt input batches with stricter nested nullability to planned schema in aggregation - #24394

Open
patrickswedish wants to merge 3 commits into
apache:mainfrom
patrickswedish:fix/agg-stricter-nested-nullability-24069
Open

fix: adapt input batches with stricter nested nullability to planned schema in aggregation#24394
patrickswedish wants to merge 3 commits into
apache:mainfrom
patrickswedish:fix/agg-stricter-nested-nullability-24069

Conversation

@patrickswedish

@patrickswedish patrickswedish commented Aug 15, 2026

Copy link
Copy Markdown

Which issue does this PR close?

Rationale for this change

In Apache Arrow and DataFusion, schema containment (arrow::datatypes::Schema::contains / Field::contains) allows input data sources (such as MemTable, Comet over FFI, or external partitions) to supply RecordBatches whose data types are stricter than the declared plan schema (e.g. a nested struct or list child field is marked non-nullable in runtime batches while the catalog/planner schema declared it nullable).

However, Arrow-rs operators (RecordBatch::try_new, ListArray::new, and RowConverter::convert_columns) enforce strict DataType equality. When un-adapted stricter batches enter AggregateExec:

  1. ArrayAggGroupsAccumulator constructs ListArray using self.datatype (declared type) and flat_values (runtime stricter type), panicking in ListArray::new.
  2. GroupValuesByRow (RowConverter) is initialized from the planner schema and fails with RowConverter column schema mismatch when array_agg(DISTINCT struct) is planned.
  3. GroupedHashAggregateStream::spill() and emit() fail when building RecordBatches with ArrowError: column types must match schema types.

What changes are included in this PR?

This PR enforces schema conformance at the single authoritative boundary where planned schemas meet runtime data streams in aggregation:

  1. datafusion_common::nested_struct::adapt_batch_to_schema:
    • Fast path: Returns immediately if the batch already uses the target schema (Arc::ptr_eq or structural equality).
    • Strict containment check: Verifies target_schema.contains(batch.schema().as_ref()). Incompatible schemas return an explicit error and are never cast.
    • Metadata transformation: Adapts nested array/schema metadata to match the target schema while reusing existing array data where supported by the existing nested cast helpers.
    • Schema normalization: Guarantees the returned batch strictly has adapted_batch.schema() == target_schema even if only top-level field nullability differed.
  2. AggregateExec::execute_input:
    • Wraps the input stream in AdaptedInputRecordBatchStream so all aggregate stream variants (AggregateStream, GroupedHashAggregateStream, GroupedTopKAggregateStream, OrderedPartialAggregateStream, PartialHashAggregateStream, FinalHashAggregateStream, etc.) seamlessly receive batches conforming to input_schema.
  3. Unit and End-to-End Regression Tests:
    • test_adapt_batch_to_schema_identical: Identical schema passthrough.
    • test_adapt_batch_to_schema_stricter_nested_struct: Stricter nested struct adaptation.
    • test_adapt_batch_to_schema_top_level_nullability_only: Top-level nullability replacement.
    • test_adapt_batch_to_schema_incompatible_rejected: Incompatible schema rejection.
    • test_adapt_batch_to_schema_incompatible_type_rejected: Incompatible type rejection.
    • array_agg_struct_from_stricter_batches: End-to-end non-spill aggregation.
    • array_agg_distinct_struct_from_stricter_batches: End-to-end distinct aggregation with RowConverter.
    • array_agg_struct_from_stricter_batches_with_spilling: End-to-end spill / emit under memory pressure.
    • array_agg_distinct_struct_from_stricter_batches_with_spilling: End-to-end distinct spill / emit under memory pressure.
    • test_aggregate_exec_direct_input_adaptation: Direct AggregateExec boundary adaptation over custom TestMemoryExec.

Are these changes tested?

Yes, tested via unit tests in datafusion_common and the integration suite in datafusion/core/tests/sql/aggregates/nested_nullability.rs.

Are there any user-facing changes?

No API changes. Queries aggregating batches with stricter nested schemas that previously panicked or errored now execute successfully.

…an schema

In DataFusion, input batches may carry data types with stricter nested nullability
than declared by the plan schema (e.g. non-nullable child struct/list fields where
the plan allows nulls, as permitted by `Field::contains` / `Schema::contains`).

Previously, feeding such stricter batches into `AggregateExec` caused runtime failures
across several components (panics in `ArrayAggGroupsAccumulator`, `RowConverter` schema
mismatch in distinct aggregation, and `ArrowError` on spill/emit).

This commit:
1. Adds `datafusion_common::nested_struct::adapt_batch_to_schema` to zero-copy adapt
   stricter batches to a target schema via metadata transformation.
2. Adapts outgoing batches in `MemoryStream::poll_next` to match `MemoryStream::schema()`.
3. Wraps `AggregateExec` input stream execution to ensure all aggregate stream variants
   receive batches conforming to `input_schema`.
4. Adds end-to-end regression tests in `nested_nullability.rs` covering regular aggregation,
   distinct aggregation, and memory spilling.

Closes apache#24069
@github-actions github-actions Bot added core Core DataFusion crate common Related to common crate physical-plan Changes to the physical-plan crate labels Aug 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

common Related to common crate core Core DataFusion crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

GroupedHashAggregateStream::emit throws ArrowError: column types must match schema types Native collect_set(struct) aggregate crashes on spill emit

1 participant