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

[patrickswedish]

Which issue does this PR close?

Closes #24069.

Rationale for this change

In DataFusion, in-memory table sources such as MemTable::try_new accept RecordBatches 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 RecordBatches 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?

1. datafusion_common::nested_struct::adapt_batch_to_schema:

   • Adapts RecordBatches 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 GroupedHashAggregateStream::emit throws ArrowError: column types must match schema types #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 RecordBatches 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 GroupedHashAggregateStream::emit throws ArrowError: column types must match schema types #24069.

Are these changes tested?

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?

No. Queries aggregating in-memory tables whose batches have stricter nested nullability than the table schema now succeed as expected.

[kosiew]

hi @patrickswedish

Can you fix the CI tests?

[patrickswedish]

Hi @kosiew, thank you very much for reviewing and pointing that out!

I have resolved the missing \RecordBatchStream\ import, addressed the qualification and integration test requirements, and fixed \AggregateExec::execute_input\ to adapt batches directly to \self.input.schema().

The entire local validation battery (\cargo fmt --all --check, \cargo test -p datafusion-common --lib nested_struct::adapt_schema_tests, \cargo test -p datafusion-physical-plan --lib aggregates, and \cargo test -p datafusion --test core_integration nested_nullability, plus \cargo clippy --all-targets -- -D warnings) is now completely green with zero errors.

The branch has been squashed into 1 clean atomic commit (\47de466) and pushed. Whenever you have a chance, please approve the workflow runs on GitHub Actions. Thank you!

[kosiew]

@patrickswedish
Thanks for the iteration.
Can you also resolve the merge conflicts?

[patrickswedish]

@kosiew Resolved the conflicts against current \main, including the overlapping aggregate metrics changes, and fixed the schema-adaptation regression exposed by \ est_no_pushdown_through_global_aggregate_with_name_collision.

Re-ran the aggregate test suites and filter-pushdown integration tests successfully, and the branch is clean and mergeable again. Thanks!

[codecov-commenter]

Codecov Report

❌ Patch coverage is 86.70520% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.24%. Comparing base (c429919) to head (8390472).
⚠️ Report is 15 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/common/src/nested_struct.rs	91.04%	6 Missing and 6 partials ⚠️
...fusion/physical-plan/src/aggregates/hash_stream.rs	0.00%	0 Missing and 2 partials ⚠️
...n/physical-plan/src/aggregates/aggregate_stream.rs	0.00%	0 Missing and 1 partial ⚠️
...hysical-plan/src/aggregates/grouped_hash_stream.rs	0.00%	0 Missing and 1 partial ⚠️
...hysical-plan/src/aggregates/grouped_topk_stream.rs	0.00%	0 Missing and 1 partial ⚠️
datafusion/physical-plan/src/aggregates/mod.rs	96.55%	0 Missing and 1 partial ⚠️
...ysical-plan/src/aggregates/ordered_final_stream.rs	0.00%	0 Missing and 1 partial ⚠️
...ical-plan/src/aggregates/ordered_partial_stream.rs	0.00%	0 Missing and 1 partial ⚠️
...sical-plan/src/aggregates/ordered_single_stream.rs	0.00%	0 Missing and 1 partial ⚠️
...sical-plan/src/aggregates/partial_reduce_stream.rs	0.00%	0 Missing and 1 partial ⚠️
... and 1 more

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #24394    +/-   ##
========================================
  Coverage   81.24%   81.24%            
========================================
  Files        1113     1113            
  Lines      392744   392907   +163     
  Branches   392744   392907   +163     
========================================
+ Hits       319090   319232   +142     
- Misses      54900    54912    +12     
- Partials    18754    18763     +9     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[kosiew]

@patrickswedish,

Thanks for the follow-up work here. The earlier items around the RecordBatchStream import, integration coverage, adapting against self.input.schema(), and the conflict resolution all look addressed.

I found one remaining issue with Union schemas that I think needs to be fixed before this lands. DataType::contains can accept a stricter incoming Union schema, but the adapter then delegates to Arrow's cast implementation, which does not support Union-to-Union casts. As a result, an input schema that should be compatible can still fail at runtime when it reaches the aggregate wrapper.

I reproduced this directly with adapt_batch_to_schema. The existing adapt_schema_tests also pass, so this looks isolated to Union adaptation rather than the broader schema adaptation work.

Once Union fields are adapted recursively, or unsupported Union shapes are rejected before reaching the cast path, I think this should be in good shape. It would also be useful to add both a direct adapter regression and an aggregate execution regression covering stricter Union child nullability.

[kosiew]

I think there is still an issue here for Union schemas.

DataType::contains accepts a target Union when its child field is nullable and the corresponding incoming Union child is non-nullable. That means the incoming schema is considered a valid stricter version of the target schema.

We then reach cast_column, but Arrow does not support casting a Union to another Union. Its Union cast support is for extracting a child into a non-Union target. In practice, this means a compatible stricter Union batch now fails with cannot cast Union ... to Union ... instead of being adapted to the declared input schema.

Could we handle Union arrays explicitly here by recursively adapting each child to the corresponding target field, then rebuilding the UnionArray with the target UnionFields while preserving the type IDs, dense offsets, and mode?

Another option would be to reject Union shapes that the adapter cannot actually construct, even if DataType::contains currently considers them compatible.

It would be good to add a direct adapt_batch_to_schema regression for stricter Union child nullability, plus an aggregate execution regression that exercises the same case through AggregateExec.

[patrickswedish]

Hi @kosiew,

Thank you for the guidance! We have addressed this in the latest update:

1. Narrowed Schema Conformance for Unions (cast_union_column & �alidate_union_schema_compatibility):

   • Explicitly handles Union arrays (both Sparse and Dense modes) in
     ested_struct::cast_column.
   • Requires matching union modes, exact type ID set equality, and recursive containment ( arget_child.contains(source_child)).
   • Recursively adapts matching children using cast_column and reconstructs the UnionArray with arget_fields, preserving ype_ids and dense offsets buffers without copying buffer data.
   • Rejects unsupported field set evolution (extra/missing type IDs).
   • Removed Union from
     equires_nested_struct_cast so that generic SQL CAST semantics across DataFusion are untouched.

2. Unit and Integration Regressions:

   • Added unit tests verifying exact unpacked scalar values (10, "b", 30), row-level ype_ids, offsets, and target schema containment for Dense and Sparse unions.
   • Added non-contiguous/reordered type-ID test ([(1, int), (3, str)] -> [(3, str), (1, int)]).
   • Added negative tests verifying rejection of field-set mismatches and mode mismatches.
   • Added direct AggregateExec input boundary integration tests in
     ested_nullability.rs asserting unpacked �rray_agg(b) values and output schema nullability.

[patrickswedish]

Hi @kosiew,

Thank you very much for catching this edge case and for the detailed review!

We have implemented recursive adaptation for Union arrays and addressed all your points:

1. Recursive Union Adaptation (\cast_union_column):

   • In \datafusion_common::nested_struct, implemented \cast_union_column\ supporting both \Sparse\ and \Dense\ union modes.
   • Recursively casts each child array to the target child field's data type using \cast_column.
   • Preserves row-level \ ype_ids\ buffer and dense \offsets\ buffer without copying unnecessary primitive data.
   • For extra fields present only in the target union, populates them with empty arrays (Dense) or typed null arrays (Sparse).
   • Validates that union modes match (\source_mode == target_mode) and that all source type IDs exist in the target union, cleanly returning a deterministic planning error if mismatched.

2. Integration with Cast and Validation:

   • Updated \cast_column\ to route (DataType::Union, DataType::Union)\ through \cast_union_column.
   • Updated \�alidate_data_type_compatibility\ and
     equires_nested_struct_cast\ to validate matching union modes and recursively validate child field compatibility.

3. Comprehensive Regression Tests:

   • Direct unit tests in \datafusion_common::nested_struct::adapt_schema_tests:
     • \ est_adapt_batch_to_schema_stricter_sparse_union\
     • \ est_adapt_batch_to_schema_stricter_dense_union\
     • \ est_adapt_batch_to_schema_union_nested_struct\
     • \ est_adapt_batch_to_schema_union_incompatible_mode_rejected\
     • \ est_adapt_batch_to_schema_union_missing_type_id_rejected\
     • \ est_validate_data_type_compatibility_union\
   • Direct \AggregateExec\ input boundary integration tests in \datafusion/core/tests/sql/aggregates/nested_nullability.rs:
     • \ est_aggregate_exec_direct_input_adaptation_dense_union\
     • \ est_aggregate_exec_direct_input_adaptation_sparse_union\

The branch has been rebased onto the latest \main, formatted with \cargo fmt, and validated with \cargo clippy --all-targets --all-features -- -D warnings\ and the full test suite.

[alamb]

Thank you @patrickswedish and @kosiew -- I left some comments. Let me know what you think

[alamb]

There seems to be a lot of repetition in these tests -- perhaps some of the hints here could be applied: https://datafusion.apache.org/contributor-guide/pr_review.html#review-the-test-coverage

[patrickswedish]

Hi @alamb,

Thank you for pointing us to the PR review test coverage guide!

We refactored the test suite in \datafusion/common/src/nested_struct.rs\ and \datafusion/core/tests/sql/aggregates/nested_nullability.rs:

• Extracted shared test field builders (\ est_two_field_union, etc.) to eliminate repeated boilerplate across test cases.
• Streamlined unit tests to focus on distinct semantic cases: Sparse adaptation, Dense adaptation, non-contiguous/reordered type-ID mappings, nested Structs, mode mismatches, and field-set mismatches.
• Retained exact row-level unpacked value assertions (\10, \b\, \30), active type IDs, and dense offsets without test repetition.
• Removed oversized duplicate tests from
  ested_nullability.rs\ to keep integration tests concise and focused on the end-to-end bug report reproducer.

[alamb]

I feel like this is a solution to a symptom (nullability mismatch) rather than the underlying problem (an operator is not declaring its output schema correctly and producing record batches with the wrong shape)

[patrickswedish]

Hi @alamb,

Thank you for this key architectural insight! You are completely right: fixing this at the \AggregateExec\ boundary was treating a symptom rather than addressing the producer contract.

Root Cause Analysis

In DataFusion, in-memory table sources like \MemTable::try_new\ accept batches whose schemas are stricter than the declared table schema using \Schema::contains(&batches_schema)\ (e.g. nullable nested fields in the declared table schema vs non-nullable nested fields in the input batches).

When \MemTable::scan\ creates \MemorySourceConfig\ / \MemoryExec, \MemoryStream\ was constructed with the declared schema, but its \poll_next\ emitted the underlying stricter \RecordBatch\�s without adapting them. Downstream operators (like \AggregateExec) received batches that did not conform to the stream's advertised output schema.

Architectural Solution

1. Reverted AggregateExec Changes: Completely removed \AdaptedInputRecordBatchStream\ and \AggregateExec::execute_input, restoring all aggregate physical plans to clean upstream state.
2. Fixed Producer Invariant in \MemoryStream\ (\physical-plan/src/memory.rs):
   • \MemoryStream::poll_next\ now normalizes emitted batches using \�dapt_batch_to_schema(batch, &self.schema)\ whenever a batch differs from \self.schema\ and \self.schema.contains(batch.schema()).
   • Every \RecordBatch\ emitted by \MemoryStream\ is guaranteed to conform to \stream.schema().
3. Retained Narrow Schema Conformance in
   ested_struct\ (\common/src/nested_struct.rs):
   • \�dapt_batch_to_schema\ supports Structs, Lists, and Unions (Dense/Sparse) without changing general DataFusion CAST behavior (
     equires_nested_struct_cast\ remains unchanged).
4. Unit and Integration Regressions:
   • Added unit tests in \physical-plan/src/memory.rs\ directly verifying that \MemoryStream\ emits batches matching \self.schema\ (with and without projection).
   • Retained end-to-end SQL aggregation regressions in
     ested_nullability.rs\ covering standard, distinct, and spilling aggregations.

[patrickswedish]

Update Summary

Thank you @alamb and @kosiew for the thorough review and guidance!

Following the feedback, we have restructured the fix around the producer boundary rather than wrapping individual downstream consumer operators:

1. Producer Normalization in MemoryStream (datafusion/physical-plan/src/memory.rs):

   • When MemTable accepts batches whose nested schemas are stricter than the table's declared schema (via Schema::contains), MemoryStream::poll_next now normalizes emitted batches using datafusion_common::nested_struct::adapt_batch_to_schema whenever the batch schema differs from the stream's advertised schema and self.schema.contains(batch.schema()).
   • Ensures every RecordBatch emitted by MemoryStream conforms to stream.schema() before downstream operators consume it.

2. Reverted AggregateExec Wrapper:

   • Removed AdaptedInputRecordBatchStream and restored datafusion/physical-plan/src/aggregates/ to clean upstream state.

3. Narrow Schema-Conformance Support for Unions (datafusion/common/src/nested_struct.rs):

   • Retained recursive child adaptation and exact type-ID set matching for Arrow UnionArray (sparse and dense) in adapt_batch_to_schema without expanding general DataFusion SQL CAST behavior.

4. Test Suite Refactoring:

   • Streamlined unit tests in nested_struct.rs using helper builders to remove repeated boilerplate while asserting exact values (10, "b", 30), type IDs, and dense offsets.
   • Added invariant unit tests in physical-plan/src/memory.rs validating emitted batch schema conformance and projection handling.
   • Cleaned up nested_nullability.rs to focus on end-to-end bug report GroupedHashAggregateStream::emit throws ArrowError: column types must match schema types #24069 reproducer regressions (standard, distinct, and spilling).

5. Validation Battery:

   • cargo fmt --check passed across all workspace packages.
   • cargo clippy --workspace ... -- -D warnings passed with 0 warnings.
   • cargo test -p datafusion-common --lib passed (562 passed; 0 failed).
   • cargo test -p datafusion-physical-plan --lib passed (1770 passed; 0 failed).
   • cargo test -p datafusion --test core_integration -- nested_nullability passed (5 passed; 0 failed).
   • Rebased cleanly onto the latest upstream/main as 1 atomic commit.