fix: frame IPC messages when reading spill files so batches pin only their own bytes - #24592
Open
jayzhan211 wants to merge 1 commit into
Open
fix: frame IPC messages when reading spill files so batches pin only their own bytes#24592jayzhan211 wants to merge 1 commit into
jayzhan211 wants to merge 1 commit into
Conversation
…their own bytes The spill reader fed raw 128 KB read chunks to arrow's zero-copy StreamDecoder. A batch whose message fit inside a chunk kept the whole chunk alive, and a message spanning chunks was gathered into a Vec grown by doubling, so read-back batches retained, and were accounted for, up to 27x the memory recorded for them at spill time. Reassemble each IPC message into an exactly sized allocation before decoding, using the metadata's bodyLength to size the body buffer. Closes apache#17340
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24592 +/- ##
========================================
Coverage 81.38% 81.38%
========================================
Files 1116 1116
Lines 397960 398181 +221
Branches 397960 398181 +221
========================================
+ Hits 323880 324071 +191
- Misses 55120 55139 +19
- Partials 18960 18971 +11 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
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 limitwarning in #17340 reports. Under a memory limit this is not just a noisy log: the multi-level merge inSortExecand the spilling aggregate size their merge fan-in frommax_record_batch_memoryrecorded 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'sStreamDecoder, 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:
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
Vecgrown by doubling and the batch keeps the spare capacity:In the
spill_iobench 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?
SpillReaderStreamnow reassembles each IPC message into allocations sized from the message's own headers before decoding, via a smallMessageFramerstate machine:meta_len;Vecof exactlyprefix + meta_lenbytes;bodyLengthfrom the flatbuffer metadata (arrow_ipc::root_as_message);Vec::with_capacity(body_len)from however many chunks it spans;[head, body]to the sameStreamDecoder.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:A 5 KB batch retains 5 KB and a 256 KB batch retains 256 KB, so
max_record_batch_memoryrecorded 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_iobench vsmain(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):mainmemory_limit::test_stringview_external_sort(the reproducer in #17340)memory_limitintegration suitespilling_fuzz_in_memory_constrained_env+sort_fuzz+aggregate_fuzzAre 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'sget_record_batch_memory_sizeis within the margin of the written maximum. Fails onmainwithread-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.memory_limitintegration 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.