ExternalSorter over-reserves memory for Utf8View batches that share or repeat data buffers, so a sort after WindowAggExec fails with ResourcesExhausted

[adriangb]

Describe the bug

A sort that follows WindowAggExec and a hash RepartitionExec can fail with ResourcesExhausted on its first input batch when the data has Utf8View columns. The batch holds about 5 MB of string data, but ExternalSorter tries to reserve about 700 MB for it. The same query with Utf8 columns runs in a 2 GB pool. With Utf8View it needs a pool of more than 8 GB, even though its peak RSS without a limit is lower (721 MB with views, 1.09 GB with Utf8).

To Reproduce

With datafusion-cli 55.1.0 (also reproduced on main at 33574a1):

-- gen.sql
COPY (
  SELECT
    i % 100000 AS k1,
    (i * 7919) % 100003 AS k2,
    arrow_cast(concat(md5(CAST(i AS VARCHAR)), md5(CAST(i + 1 AS VARCHAR)), md5(CAST(i + 2 AS VARCHAR)), md5(CAST(i + 3 AS VARCHAR))), 'Utf8') AS payload
  FROM generate_series(1, 2000000) AS t(i)
) TO 'data.parquet' STORED AS PARQUET;

-- query.sql
SET datafusion.execution.target_partitions = 8;
SET datafusion.execution.parquet.schema_force_view_types = true;
SELECT count(*), sum(length(payload)), max(m1), max(m2)
FROM (
  SELECT k2, payload, m1, MAX(k1) OVER (PARTITION BY k2) AS m2
  FROM (
    SELECT k1, k2, payload, MAX(k2) OVER (PARTITION BY k1) AS m1
    FROM 'data.parquet'
  )
);

datafusion-cli -q -f gen.sql
datafusion-cli -q -m 4g --mem-pool-type greedy -f query.sql

Actual:

Resources exhausted: Failed to allocate additional 674.5 MB for ExternalSorter[5] with 0.0 B already allocated for this reservation - 663.6 MB remain available for the total memory pool: greedy(used: 3.4 GB, pool_size: 4.0 GB)

Pool (greedy)	schema_force_view_types = true	= false
2g	fails (715 MB request)	ok
4g	fails (675 MB request)	ok
8g	fails (660 MB request)	ok
16g	ok	ok

Expected behavior

The Utf8View query should run with a memory limit close to the one the Utf8 query needs. A sort should not need hundreds of MB to accept one input batch that holds about 5 MB of data.

Analysis

The plan is DataSourceExec → RepartitionExec(Hash k1) → SortExec → WindowAggExec → RepartitionExec(Hash k2) → SortExec → WindowAggExec. I measured the batches between the operators with a small harness. "Accounted" is get_record_batch_memory_size(batch) + batch.get_sliced_size(), which ExternalSorter::insert_batch tries to reserve. "Referenced" is the view bytes plus the payload bytes the views use.

Output of	Rows per batch	Accounted	Referenced	Buffer entries / distinct buffer bytes
first SortExec	8192	52 MB	1.2 MB	26 / 26 MB
first WindowAggExec	250k (one batch per partition)	714 MB	40 MB	668 / 33 MB
second RepartitionExec (input of second sort)	31k	727 MB	5.3 MB	662 / 33 MB
second SortExec	8192	4.3 GB	1.3 MB	about 3,900 / 267 MB

Three things combine:

1. WindowAggExec::compute_aggregates calls concat_batches over all buffered input. For view arrays, concat uses GenericByteViewBuilder::append_array, which appends each input array's full data-buffer list. The sorted input batches share the same buffers, so the output repeats each buffer about 25 times (668 entries, 33 MB distinct). Arrow chose not to deduplicate here (InProgressByteViewArray::append_views_and_update_buffer_index/GenericByteViewBuilder::append_array should deduplicate buffers arrow-rs#10692).
2. The hash RepartitionExec take keeps that buffer list. Its LimitedBatchCoalescer sets biggest_coalesce_batch_size = batch_size / 2, so outputs of at least 4096 rows pass through without the byte-view gc the coalescer does for sparse sources.
3. get_reserved_bytes_for_record_batch adds get_sliced_size(). For view arrays that sum includes the capacity of every data-buffer entry, and it does not deduplicate. On main this comes from ArrayData::get_slice_memory_size (fix(arrow-data): account for view payload buffers in slice memory size arrow-rs#10519). On 55.1.0 it comes from the byte_view_data_buffer_size workaround. In the table, get_record_batch_memory_size (which counts each buffer once) is 43 MB, but get_sliced_size is 684 MB.

ExternalSorter::reserve_memory_for_batch_and_maybe_spill returns the error when the reservation fails and in_mem_batches is empty. Splitting the batch would not help, because each slice keeps the same buffers.

Possible fixes

• Count each data buffer once in the sort reservation, for example with RecordBatchMemoryCounter or by deduplicating by pointer in get_sliced_size. Here this reduces the request from 727 MB to about 78 MB.
• Compact view arrays whose buffer capacity is much larger than the bytes they use, with the same 2x rule as BatchCoalescer, before ExternalSorter reserves memory for them. As a minimum, do this before returning ResourcesExhausted when there is nothing to spill. Here this reduces the request to about 11 MB and frees the pinned upstream buffers.

Related: #22862 (the same fix for hash join build batches), #25712, #25271, #19679, #14748.

[tohuya6]

hey @adriangb, could i help out with this? similar to #22682 right?

[adriangb]

Yes, please go ahead, thanks! I think you mean #22862 (the RecordBatchMemoryCounter fix for hash join build batches), and yes, this is the same problem in ExternalSorter.

Two things I found while prototyping locally, in case they help:

1. Deduplicating the first-batch reservation is not enough on its own. sort_batch_stream charges each sorted output chunk with get_record_batch_memory_size, and every chunk still references the input's view buffers, so they are charged once per chunk. With only the first-batch fix, the repro above still fails with "Failed to allocate additional 109.7 MB for ExternalSorter[0] with 23.8 MB already allocated". Charging the sorted output through RecordBatchMemoryCounter fixes it.
2. That second change makes test_single_batch_spill_returns_live_workspace_loan_on_drop (from fix: preserve external sort workspace across spilling #24740) fail. The workspace loan there only happens when sorted output costs more than the input estimate, and with deduplication I could not find a normal input that does that. It may be worth checking with the fix: preserve external sort workspace across spilling #24740 author whether the loan path is still needed.

Separately, compacting sparse view arrays (the BatchCoalescer 2x rule) when the sort's reservation fails with nothing to spill also rescues the first-batch case. That could be a follow-up if you want to keep the first PR small.

[tohuya6]

Thanks for your prototyping notes, the PR addresses both points and is up for review. I also asked the 24740 author what he thinks, appreciate the suggestion!