jayzhan211 opened a new pull request, #24592:
URL: https://github.com/apache/datafusion/pull/24592
## Which issue does this PR close?
- Closes #17340.
## Rationale for this change
Record batches read back from a spill file retain far more memory than was
recorded for them when they were written, which is what the `Record batch
memory usage (...) exceeds the expected limit` warning in #17340 reports. Under
a memory limit this is not just a noisy log: the multi-level merge in
`SortExec` and the spilling aggregate size their merge fan-in from
`max_record_batch_memory` recorded at spill time, so inflated read-back batches
consume memory the operator never budgeted for.
The cause is in `SpillReaderStream`. It reads the file in 128 KB chunks and
hands them straight to arrow's `StreamDecoder`, which builds arrays on *slices*
of the buffer it is given — and a slice keeps its whole backing allocation
alive. That goes wrong in two ways.
**Small batches pin the whole chunk.** With ~5 KB batches one chunk holds
~27 messages. Each decoded batch's buffers are slices of that chunk, so each
batch retains, and is accounted for, 128 KB:
```
chunk (128 KB allocation)
┌──────┬──────┬──────┬─────┬───────┐
│ msg1 │ msg2 │ msg3 │ ... │ msg27 │
└──────┴──────┴──────┴─────┴───────┘
▲
batch1's buffers slice here, yet keep all 128 KB alive
```
Traced in the sort tests: a 100-row Utf8 batch read back with `caps=[(404,
131072), (4316, 131072)]` — 4.7 KB of data, 128 KB retained, 27× what was
recorded at spill time. And it is real retention, not just accounting: while
the merge holds that batch, the other 26 in the chunk stay alive too.
**Straddling batches double.** A message spanning two chunks cannot be
sliced, so the decoder gathers it into a `Vec` grown by doubling and the batch
keeps the spare capacity:
```
chunk N chunk N+1
┌──────┬────────────────────┬──────────────┬────────┐
│ ... │ msgK (first part) │ msgK (rest) │ msgK+1 │
└──────┴────────────────────┴──────────────┴────────┘
```
In the `spill_io` bench about half of all 256 KB batches came back in a 512
KB allocation (`retained=523264 data=262144`). How much spare capacity a
straddling batch ends up with depends on where the chunk boundary fell, which
is why the reports on #17340 range from ~10% over (`967744 vs 877568`) up to 2×.
## What changes are included in this PR?
`SpillReaderStream` now reassembles each IPC message into allocations sized
from the message's own headers before decoding, via a small `MessageFramer`
state machine:
1. read the 4-byte length prefix (skipping the continuation marker) →
`meta_len`;
2. fill a head `Vec` of exactly `prefix + meta_len` bytes;
3. take `bodyLength` from the flatbuffer metadata
(`arrow_ipc::root_as_message`);
4. fill a body `Vec::with_capacity(body_len)` from however many chunks it
spans;
5. hand `[head, body]` to the same `StreamDecoder`.
The whole body is now inside one buffer whose allocation is exactly
`body_len`, so the decoder takes its zero-copy path and the batch pins exactly
its own message:
```
body for msg1 (5 KB) body for msgK (256 KB)
┌──────┐ ┌────────────────────┐
│ msg1 │ ◀── batch1 │ msgK │ ◀── batchK
└──────┘ └────────────────────┘
```
A 5 KB batch retains 5 KB and a 256 KB batch retains 256 KB, so
`max_record_batch_memory` recorded at write time matches what the merge
actually gets back.
This costs one memcpy per message — the one the decoder already paid for
straddling messages — minus the doubling reallocation, so it is not slower.
`spill_io` bench vs `main` (two runs on a quiet machine):
`StreamReader/read_100` −7.7%, `q2/lz4_frame` −5.9%, all other cases within
noise.
Warnings from the #17340 check
(`RUST_LOG=datafusion_physical_plan::spill=debug`):
| | `main` | this PR |
|---|---|---|
| `memory_limit::test_stringview_external_sort` (the reproducer in #17340) |
26 | **0** |
| whole `memory_limit` integration suite | ~5000 | **0** |
| `spilling_fuzz_in_memory_constrained_env` + `sort_fuzz` + `aggregate_fuzz`
| 4170 | **0** |
## Are these changes tested?
Yes:
- `test_read_back_does_not_inflate_batch_memory`: spills 50 small batches
(Int32, Utf8, Utf8View, List) and asserts every read-back batch's
`get_record_batch_memory_size` is within the margin of the written maximum.
Fails on `main` with `read-back batch retains 131072 bytes, written max was
24196`.
- `test_message_framer_across_chunk_boundaries`: frames and decodes an IPC
stream delivered in chunks of 1, 3, 7, 64, 1000 bytes and as a whole, checking
the batches are intact and each retains no more than its own message body.
- Existing spill, sort, aggregate, repartition unit tests, the
`memory_limit` integration tests and the extended spilling fuzz suites pass.
## Are there any user-facing changes?
No API changes. Queries that spill use less memory when reading spills back,
and the spurious accounting warning from #17340 no longer fires.
--
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]
