Order-preserving RepartitionExec fails with ResourcesExhausted on non-spillable merge reservation for Utf8View data #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: 34eb6dc (main).

To Reproduce

Run with datafusion-cli -m 200M -f repro.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:

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 (

   datafusion/datafusion/physical-plan/src/repartition/mod.rs

   Line 1376 in 34eb6dc

   let columns = take_arrays(batch.columns(), &indices_array, None)?;

   ). 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 (

   datafusion/datafusion/physical-plan/src/sorts/builder.rs

   Lines 89 to 95 in 34eb6dc

   	pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<usize> {
   	let size = get_record_batch_memory_size(&batch);
   	self.batches_mem_used += size;
   	// Only request additional memory from the pool when actual batch
   	// usage exceeds the current reservation (which may include
   	// pre-reserved bytes from sort_spill_reservation_bytes).
   	try_grow_reservation_to_at_least(&mut self.reservation, self.batches_mem_used)?;

   ,

   datafusion/datafusion/physical-plan/src/sorts/builder.rs

   Lines 163 to 170 in 34eb6dc

   	// Remove consumed indices, keeping any remaining for the next call.
   	self.indices.drain(..rows_to_emit);
   	if self.indices.is_empty() {
   	self.retain_cursor_batches();
   	} else {
   	self.retain_live_batches();
   	}

   ). 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

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.