fix: adapt input batches with stricter nested nullability to planned schema in aggregation - #24394
Open
patrickswedish wants to merge 3 commits into
Open
Conversation
…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
…tative boundary in aggregation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
GroupedHashAggregateStream::emitthrows ArrowError: column types must match schema types #24069collect_set(struct)aggregate crashes on spill emit datafusion-comet#5239Rationale for this change
In Apache Arrow and DataFusion, schema containment (
arrow::datatypes::Schema::contains/Field::contains) allows input data sources (such asMemTable, Comet over FFI, or external partitions) to supplyRecordBatches whose data types are stricter than the declared plan schema (e.g. a nested struct or list child field is markednon-nullablein runtime batches while the catalog/planner schema declared itnullable).However, Arrow-rs operators (
RecordBatch::try_new,ListArray::new, andRowConverter::convert_columns) enforce strictDataTypeequality. When un-adapted stricter batches enterAggregateExec:ArrayAggGroupsAccumulatorconstructsListArrayusingself.datatype(declared type) andflat_values(runtime stricter type), panicking inListArray::new.GroupValuesByRow(RowConverter) is initialized from the planner schema and fails withRowConverter column schema mismatchwhenarray_agg(DISTINCT struct)is planned.GroupedHashAggregateStream::spill()andemit()fail when buildingRecordBatches withArrowError: 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:
datafusion_common::nested_struct::adapt_batch_to_schema:Arc::ptr_eqor structural equality).target_schema.contains(batch.schema().as_ref()). Incompatible schemas return an explicit error and are never cast.adapted_batch.schema() == target_schemaeven if only top-level field nullability differed.AggregateExec::execute_input:AdaptedInputRecordBatchStreamso all aggregate stream variants (AggregateStream,GroupedHashAggregateStream,GroupedTopKAggregateStream,OrderedPartialAggregateStream,PartialHashAggregateStream,FinalHashAggregateStream, etc.) seamlessly receive batches conforming toinput_schema.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 withRowConverter.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: DirectAggregateExecboundary adaptation over customTestMemoryExec.Are these changes tested?
Yes, tested via unit tests in
datafusion_commonand the integration suite indatafusion/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.