adriangb opened a new issue, #27323:
URL: https://github.com/apache/datafusion/issues/27323
### Describe the bug
An order-preserving hash `RepartitionExec` fails with `ResourcesExhausted`
on `RepartitionExec[Merge N]` reservations, which report `can spill: false`.
The query returns only 7000 rows. The merge reservations grow until they fill
the pool, whatever the pool size.
Commit: 34eb6dc5bf057a35c471a4fbfd2eece2a7937afe (main).
### To Reproduce
Run with `datafusion-cli -m 200M -f repro.sql`:
```sql
COPY (SELECT value % 1000 AS k, repeat('x', 100) || (value % 7)::varchar AS
s, value AS v
FROM generate_series(1, 2000000) ORDER BY k)
TO 'data.parquet' STORED AS PARQUET;
SET datafusion.execution.target_partitions = 8;
CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'data.parquet' WITH ORDER
(k ASC);
SELECT count(*) FROM (SELECT k, s, count(*) FROM t GROUP BY k, s ORDER BY k);
```
The file is 5.6 MB. Error:
```
Resources exhausted: Additional allocation failed for RepartitionExec[Merge
6] with top memory consumers (across reservations) as:
RepartitionExec[Merge 3]#44(can spill: false) consumed 25.5 MB, peak 25.5
MB,
RepartitionExec[Merge 5]#43(can spill: false) consumed 25.3 MB, peak 25.3
MB,
RepartitionExec[Merge 0]#35(can spill: false) consumed 25.2 MB, peak 25.2
MB.
Error: Failed to allocate additional 640.0 B for RepartitionExec[Merge 6]
with 24.8 MB already allocated for this reservation - 6.0 B remain available
for the total memory pool: greedy(used: 200.0 MB, pool_size: 200.0 MB)
```
With `-m 50M` each merge holds about 6 MB. With `-m 20M` each holds about
2.5 MB. The query passes at `-m 400M`.
Plan (relevant part):
```
AggregateExec: mode=FinalPartitioned, gby=[k, s]
RepartitionExec: partitioning=Hash([k@0, s@1], 8), input_partitions=8,
preserve_order=true, sort_exprs=k@0 ASC NULLS LAST
AggregateExec: mode=Partial, gby=[k, s],
group_clustering_mode=Partial([0])
DataSourceExec: file_groups=8, format=parquet
```
`EXPLAIN ANALYZE` metrics without a limit:
| operator | output_rows | output_batches | output_bytes |
|---|---|---|---|
| DataSourceExec | 2.00 M | 245 | 46.1 MB |
| AggregateExec Partial | 7.01 K | 246 | 42.7 MB |
| RepartitionExec | 7.01 K | 8 | 330.1 MB |
If `s` is `Utf8` (`arrow_cast(s, 'Utf8')`), the partial aggregate reports
1.7 MB, the repartition reports 828 KB, and the query passes at `-m 10M`.
### Expected behavior
The query succeeds with a small limit. The data that the merge must hold is
a few KB per output partition.
### Additional context
Three effects combine.
1. The partial aggregate emits `Utf8View` batches of about 28 rows that
report about 174 KB each. The view data buffers dominate that size.
2. The hash partitioner uses `take_arrays`
(https://github.com/apache/datafusion/blob/34eb6dc5bf057a35c471a4fbfd2eece2a7937afe/datafusion/physical-plan/src/repartition/mod.rs#L1376).
For `Utf8View` columns, `take` copies only the views and shares the data
buffers. Each of the 8 output batches reports the full buffer size. The sum
over all merges is about 8 times the real memory.
3. `BatchBuilder` keeps every input batch that has a row in `indices` until
it builds an output batch
(https://github.com/apache/datafusion/blob/34eb6dc5bf057a35c471a4fbfd2eece2a7937afe/datafusion/physical-plan/src/sorts/builder.rs#L89-L95,
https://github.com/apache/datafusion/blob/34eb6dc5bf057a35c471a4fbfd2eece2a7937afe/datafusion/physical-plan/src/sorts/builder.rs#L163-L170).
Each output partition has fewer than `batch_size` rows, so the merge never
emits before the end and keeps all of its input batches. The reservation grows
with the number of input batches. It is not bounded by `N inputs x batch size`.
The merge reservation is created at
https://github.com/apache/datafusion/blob/34eb6dc5bf057a35c471a4fbfd2eece2a7937afe/datafusion/physical-plan/src/repartition/mod.rs#L1857-L1869
without `with_can_spill(true)`. A refused `try_grow` in
`BatchBuilder::push_batch` fails the query. The repartition channel spill path
(https://github.com/apache/datafusion/blob/34eb6dc5bf057a35c471a4fbfd2eece2a7937afe/datafusion/physical-plan/src/repartition/mod.rs#L218-L230)
does not help, because the memory is held after the channel.
`SortPreservingMergeExec` uses the same `StreamingMergeBuilder` and
`BatchBuilder`, so it has the same retention behavior.
Possible fixes:
- In `BatchBuilder`, when `push_batch` cannot grow the reservation, emit the
rows in `indices` as a short batch, then release the dead batches and retry.
This bounds the retained batches to about one per input stream.
- Avoid the shared-buffer over-count: compact (`gc`) view arrays after
`take` in the hash partitioner, or account shared buffers once.
--
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]
