perf: keep one shuffle block read buffer per reduce task - #5913
Draft
andygrove wants to merge 1 commit into
Draft
Conversation
NativeBatchDecoderIterator read each compressed block into a thread-local direct buffer and reset it to 128 KB on close. CometBlockStoreShuffleReader closes an iterator per fetched map output, so every map output larger than 128 KB compressed cost two direct allocations, each zero-filled and each passing through the JDK's direct-memory reservation. Replace the thread-local with a ShuffleBlockBuffer owned by the reader for the whole task and shared by every iterator it creates. The buffer grows to twice the largest block and is released with the task, matching the direct-read path's CometShuffleBlockIterator.
andygrove
marked this pull request as draft
September 13, 2026 23:03
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?
Part of #5905 (finding R5). Does not close it.
Rationale for this change
CometBlockStoreShuffleReadercreates oneNativeBatchDecoderIteratorper fetched map output and closes the previous one as it moves on. The iterator read each compressed block into a thread-local directByteBuffer, andclose()reset that buffer to its 128 KB initial size whenever it had grown. With 8192-row batches almost every block is larger than 128 KB compressed, so for every map output a reducer paid twoByteBuffer.allocateDirectcalls: one to shrink on close, one to regrow on the next block. Each is zero-filled and goes through the JDK's direct-memory reservation, which can triggerSystem.gc()when close toMaxDirectMemorySize. A reducer over thousands of map outputs did this thousands of times per task.The reset was also running on whichever thread called
close(). When task completion closed the iterator from another thread, it replaced that thread's thread-local buffer rather than the reader's.The direct-read path (
CometShuffleBlockIterator) already keeps one buffer per task that grows and stays; this brings the JVM-consumer path in line with it.What changes are included in this PR?
ShuffleBlockBuffer, a task-scoped growable direct buffer:acquire(n)returns the buffer positioned at zero with the limit set, allocating only when the block does not fit (doubling, starting at 128 KB). It never shrinks; it is released with the task.NativeBatchDecoderIteratortakes adataBuffer: ShuffleBlockBuffer(defaulting to a fresh one) instead of using a thread-local, andclose()no longer touches the buffer. The thread-local and its reset are removed.CometBlockStoreShuffleReader.read()creates oneShuffleBlockBufferand passes it to every iterator it creates for the task.CometCelebornShuffleReadercreates one iterator per task already, so its behaviour is unchanged.Memory: at most one direct buffer of twice the largest compressed block per running task, held for the task's duration, versus the old steady state of the same buffer plus a 128 KB one being churned per map output.
How are these changes tested?
reusesTaskScopedBufferAcrossIterators(run fromCometCelebornShuffleReaderSuitelike the other decoder lifecycle checks): the buffer allocates once at the initial size, hands the same instance back while blocks fit, grows to twice the block when one does not, and is shared by three successive iterators whoseclose()calls do not reallocate it.CometCelebornShuffleReaderSuite,CometNativeShuffleSuiteandCometShuffleSuitepass (171 tests), covering the block-store reader path end to end.