alamb opened a new pull request, #24237:
URL: https://github.com/apache/datafusion/pull/24237
## Which issue does this PR close?
- Closes #24069.
## Rationale for this change
As reported in #24069 (found via DataFusion Comet,
apache/datafusion-comet#5239), `collect_list` / `collect_set` (`array_agg` /
`array_agg(DISTINCT)`) over a struct column fails during aggregation with
```
ArrowError: column types must match schema types,
expected List(Struct("colA": Boolean, ...))
but found List(Struct("colA": non-null Boolean, ...)) at column index 1
```
via `GroupedHashAggregateStream::emit` ← `spill` ←
`try_update_memory_reservation`, and the non-distinct `array_agg` groups
accumulator can even panic inside `ListArray::new` with the same type mismatch.
The trigger is input batches whose data types are *stricter* than the plan
schema declares — most commonly a nested struct field that is non-nullable even
though the schema declares it nullable. This is explicitly allowed by
DataFusion's schema contract (`Field::contains` / `Schema::contains` treat a
non-nullable field as a valid instance of a nullable one) and happens in
practice when batches come from FFI / engines like Comet, whose runtime arrays
can be stricter than the declared plan schema.
Aggregate group and accumulator output columns are derived from the actual
input data, so the stricter type propagates into the emitted columns. The
output and spill schemas, however, are computed statically from the declared
input fields, and `RecordBatch::try_new` requires *exact* data type equality
(including nested field nullability), so assembling the output/spill batch
fails. The same mismatch affects both the legacy `GroupedHashAggregateStream`
and the new migrated aggregation streams, on both the spill and the regular
output paths.
## What changes are included in this PR?
- Adds `new_batch_conforming_schema` to `physical-plan::aggregates`: builds
a `RecordBatch` from the declared schema and the output columns, first casting
any column whose data type is contained by (but not equal to) the declared
field type. Such casts only rewrite nested type nullability metadata and do not
copy data buffers. Genuinely incompatible types still surface the existing
`RecordBatch::try_new` error.
- Uses this helper at the aggregation batch-assembly sites:
`GroupedHashAggregateStream::emit` / `transform_to_states`,
`AggregateHashTable::{next_output_batch, take_state_batch}` (unordered and
ordered variants), the partial-skip `convert_to_state` path, and the
no-grouping `AggregateStream`.
- `ArrayAggGroupsAccumulator::{evaluate, convert_to_state}` now build the
`ListArray` item field from the actual values' data type instead of the
planner-declared type, fixing the panic in `ListArray::new` for the same
divergence.
## Are these changes tested?
Yes. A new regression test,
`test_spill_array_agg_of_struct_with_non_nullable_field`, aggregates a struct
column whose batches carry a non-nullable nested field while the plan schema
declares it nullable, forces spilling with a small memory pool, and runs the
full matrix of `array_agg` / `array_agg(DISTINCT)` ×
`enable_migration_aggregate` on/off. On `main`, all four combinations fail (the
reported `ArrowError` in the distinct cases, the `ListArray::new` panic in the
non-distinct cases); with this change all pass.
Also ran the `datafusion-physical-plan` and `datafusion-functions-aggregate`
test suites, the `aggregate*` sqllogictests, the core `aggregate` integration
tests, and the aggregate/spilling fuzz tests (`--features extended_tests`) —
all pass.
## Are there any user-facing changes?
Queries that previously failed (or panicked) with the above error now work.
No API changes.
--
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]
