Skip to content

perf: keep one shuffle block read buffer per reduce task - #5913

Draft
andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:perf/shuffle-decoder-buffer-reuse
Draft

perf: keep one shuffle block read buffer per reduce task#5913
andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:perf/shuffle-decoder-buffer-reuse

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #5905 (finding R5). Does not close it.

Rationale for this change

CometBlockStoreShuffleReader creates one NativeBatchDecoderIterator per fetched map output and closes the previous one as it moves on. The iterator read each compressed block into a thread-local direct ByteBuffer, and close() 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 two ByteBuffer.allocateDirect calls: 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 trigger System.gc() when close to MaxDirectMemorySize. 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?

  • New 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.
  • NativeBatchDecoderIterator takes a dataBuffer: ShuffleBlockBuffer (defaulting to a fresh one) instead of using a thread-local, and close() no longer touches the buffer. The thread-local and its reset are removed.
  • CometBlockStoreShuffleReader.read() creates one ShuffleBlockBuffer and passes it to every iterator it creates for the task. CometCelebornShuffleReader creates 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?

  • New lifecycle check reusesTaskScopedBufferAcrossIterators (run from CometCelebornShuffleReaderSuite like 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 whose close() calls do not reallocate it.
  • CometCelebornShuffleReaderSuite, CometNativeShuffleSuite and CometShuffleSuite pass (171 tests), covering the block-store reader path end to end.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:shuffle Shuffle (JVM and native) enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant