patrickswedish opened a new pull request, #24394: URL: https://github.com/apache/datafusion/pull/24394
## Which issue does this PR close? - Closes #24069 - Fixes https://github.com/apache/datafusion-comet/issues/5239 - Supersedes test reproducer PR #24278 ## Rationale for this change In Apache Arrow and DataFusion, schema containment (`Schema::contains` / `Field::contains`) allows input data sources (such as `MemTable`, Comet over FFI, or external partitions) to supply `RecordBatch`es 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 `RecordBatch`es with `ArrowError: column types must match schema types`. ## What changes are included in this PR? Instead of scattering ad-hoc casting logic inside individual accumulators or emit sites, this PR enforces schema conformance at the stream boundary: 1. **`datafusion_common::nested_struct::adapt_batch_to_schema`**: - Fast path: Pointer equality (`Arc::ptr_eq`) returns immediately with $\sim 1\text{ CPU cycle}$ overhead. - Adaptation path: When column data types differ and `target_schema.contains(&batch.schema())`, adapts columns via metadata transformation without copying primitive buffer data. 2. **`MemoryStream::poll_next`**: Adapts projected batches to `self.schema`, ensuring `MemTable` and in-memory scans always fulfill the `RecordBatchStream::schema()` contract. 3. **`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`. 4. **End-to-End Regression Tests (`nested_nullability.rs`)**: - `array_agg_struct_from_stricter_batches`: Tests non-spill aggregation over stricter nested struct fields. - `array_agg_distinct_struct_from_stricter_batches`: Tests single distinct aggregation (`RowConverter` group values) over stricter struct fields. - `array_agg_struct_from_stricter_batches_with_spilling`: Tests spill / emit path under memory pressure. - `array_agg_distinct_struct_from_stricter_batches_with_spilling`: Tests distinct spill / emit path under memory pressure. ## Are these changes tested? Yes, tested via the new `datafusion/core/tests/sql/aggregates/nested_nullability.rs` test suite. ## Are there any user-facing changes? No API changes. Queries aggregating batches with stricter nested schemas that previously panicked or errored now execute successfully. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
