TinyMurky opened a new pull request, #25497:
URL: https://github.com/apache/datafusion/pull/25497
## Which issue does this PR close?
- Closes #25199
## Rationale for this change
#24392 changed OrderSensitiveArrayAggAccumulator to retain payloads as Arrow
arrays instead of converting each value to a ScalarValue. This significantly
reduced memory usage for reasonably sized input batches.
However, the accumulator currently stores one ArrayRef per update_batch()
call. Each array has a fixed cost from its ArrayData, buffer allocation, and
Arc, regardless of how many rows it contains.
Ordered `ARRAY_AGG` does not support a native `GroupsAccumulator` because
`groups_accumulator_supported` requires `order_bys` to be empty. Grouped
execution therefore falls back to `GroupsAccumulatorAdapter`, which calls
`update_batch()` once per group per input batch.
With a high-cardinality `GROUP BY`, these calls often contain only one or
two rows. Without this change, each call retains a separate `ArrayRef`, so the
accumulator pays the fixed per-array allocation cost for many small Arrow
arrays, increasing the per-row memory footprint.
## What changes are included in this PR?
- Coalesce consecutive small ordered ARRAY_AGG payload batches up to 64
rows.
- Add tests covering:
- coalescing exactly up to the threshold;
- creating a new batch after the threshold is reached;
- entry indices across coalesced and newly created batches;
- ordering across multiple coalesced batches;
- coalescing payloads received through partial-state merge_batch().
- Add ordered ARRAY_AGG benchmarks covering:
- 1, 8, 64, and 2,048 rows per update_batch();
- random input;
- preordered input.
### Retained memory
Retained memory was measured using `Accumulator::size()` after inserting
2,048 Int64 payloads with an Int64 ordering key and before calling evaluate().
The same measurement code was used for the baseline after
[#24392](https://github.com/apache/datafusion/pull/24392) and for this change.
| Rows/update | Batches | Total retained (with this change) | Bytes/row
(with this change) | Total retained (baseline
[#24392](https://github.com/apache/datafusion/pull/24392)| Bytes/row (baseline
[#24392](https://github.com/apache/datafusion/pull/24392)) |
| ----------: | ------: | --------------------------------: |
---------------------------: |
-----------------------------------------------------------------------------------
|
------------------------------------------------------------------------------ |
| 1 | 32 | 104,701 B |
51.12 B/row | 445,181 B
| 217.37 B/row
|
| 8 | 32 | 90,365 B |
44.12 B/row | 115,453 B
| 56.37 B/row
|
| 64 | 32 | 88,573 B |
43.25 B/row | 88,573 B
| 43.25 B/row
|
| 2048 | 1 | 84,901 B |
41.46 B/row | 84,901 B
| 41.46 B/row
|
The worst-case one-row update footprint decreases from 217.37 B/row to 51.12
B/row. Inputs already at or above the 64-row coalescing threshold retain the
existing memory footprint.
## Are these changes tested?
Added unit tests covering:
- coalescing small batches exactly up to the threshold;
- creating a new batch when the threshold would be exceeded;
- preserving the correct batch_idx and row_idx for entries;
- sorting values across multiple coalesced batches;
- coalescing small partial-state payloads passed through merge_batch().
Added Criterion benchmarks for ordered ARRAY_AGG using random and preordered
input with 1, 8, 64, and 2,048 rows per update_batch().
The following commands have been executed and passed:
- `cargo test -p datafusion-functions-aggregate --lib array_agg::tests`
- `cargo bench -p datafusion-functions-aggregate --bench array_agg --no-run`
- `cargo bench -p datafusion-functions-aggregate --bench array_agg
--ordered_array_agg`
- `cargo test --profile=ci --test sqllogictests`
- `cargo test -p datafusion`
- `cargo test -p datafusion-cli`
## Are there any user-facing changes?
No
--
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]
