From bc3e57c9058a4411d6d385238c2b39746a6a0287 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 13 Sep 2026 15:20:12 -0600 Subject: [PATCH] perf: keep one shuffle block read buffer per reduce task 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. --- .../CometBlockStoreShuffleReader.scala | 6 +- .../shuffle/NativeBatchDecoderIterator.scala | 62 ++++++++++++------- .../CometCelebornShuffleReaderSuite.scala | 4 ++ ...eBatchDecoderIteratorLifecycleChecks.scala | 48 ++++++++++++++ 4 files changed, 97 insertions(+), 23 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala index 3048456ea78..f9983b651e4 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala @@ -91,6 +91,9 @@ class CometBlockStoreShuffleReader[K, C]( val nativeLib = new Native() val nativeUtil = new NativeUtil() val tracingEnabled = CometConf.COMET_TRACING_ENABLED.get() + // One read buffer for the whole task. A decoder iterator is created per fetched map output, + // and each would otherwise allocate its own direct buffer. + val dataBuffer = new ShuffleBlockBuffer() // Closes last read iterator and shared resources after the task is finished. // We need to close read iterator during iterating input streams, @@ -113,7 +116,8 @@ class CometBlockStoreShuffleReader[K, C]( dep.decodeTime, nativeLib, nativeUtil, - tracingEnabled) + tracingEnabled, + dataBuffer = dataBuffer) currentReadIterator }) .map(b => (0, b)) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala index 6227da4bf4f..9db15f0e593 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala @@ -42,7 +42,8 @@ case class NativeBatchDecoderIterator( nativeLib: Native, nativeUtil: NativeUtil, tracingEnabled: Boolean, - expectedSchema: Option[Array[Byte]] = None) + expectedSchema: Option[Array[Byte]] = None, + dataBuffer: ShuffleBlockBuffer = new ShuffleBlockBuffer()) extends Iterator[ColumnarBatch] { // One consumer reads this iterator, while task completion may close it from another thread. @@ -58,8 +59,6 @@ case class NativeBatchDecoderIterator( !validateRemoteFrames || expectedSchema.exists(_ != null), "Remote shuffle decoding requires the expected Spark schema") - import NativeBatchDecoderIterator._ - private val channel: ReadableByteChannel = if (in != null) { Channels.newChannel(in) } else { @@ -204,16 +203,7 @@ case class NativeBatchDecoderIterator( s"Native shuffle block size of $bytesToRead exceeds " + s"maximum of ${Integer.MAX_VALUE}. Try reducing shuffle batch size.") } - var dataBuf = threadLocalDataBuf.get() - if (dataBuf.capacity() < bytesToRead) { - // it is unlikely that we would overflow here since it would - // require a 1GB compressed shuffle block but we check anyway - val newCapacity = (bytesToRead * 2L).min(Integer.MAX_VALUE).toInt - dataBuf = ByteBuffer.allocateDirect(newCapacity) - threadLocalDataBuf.set(dataBuf) - } - dataBuf.clear() - dataBuf.limit(bytesToRead.toInt) + val dataBuf = dataBuffer.acquire(bytesToRead.toInt) while (dataBuf.hasRemaining && channel.read(dataBuf) >= 0) {} if (dataBuf.hasRemaining) { throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch") @@ -250,24 +240,52 @@ case class NativeBatchDecoderIterator( prefetched.filterNot(_ eq previous).foreach(pending => release(pending.close())) if (decoderHandle != 0L) release(nativeLib.releaseRemoteShuffleDecoder(decoderHandle)) if (in != null) release(in.close()) - release(resetDataBuf()) if (failure != null) throw failure } } } } -object NativeBatchDecoderIterator { +/** + * The direct buffer a compressed shuffle block is read into before it is handed to native code. + * + * One instance is meant to live for a whole reduce task and be shared by every + * [[NativeBatchDecoderIterator]] the task creates: the block-store reader creates one iterator + * per fetched map output, and a reducer over thousands of map outputs would otherwise allocate a + * fresh direct buffer per map output. The buffer only grows, to twice the largest block seen, and + * is released with the task. This mirrors the direct-read path's `CometShuffleBlockIterator`, + * which keeps one buffer per task for the same reason. + * + * Not thread-safe: a block is read into the buffer and decoded from it by the same thread before + * the next block is read. + */ +final class ShuffleBlockBuffer { + import ShuffleBlockBuffer._ - private val INITIAL_BUFFER_SIZE = 128 * 1024 + private var buffer: ByteBuffer = _ - private val threadLocalDataBuf: ThreadLocal[ByteBuffer] = ThreadLocal.withInitial(() => { - ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE) - }) + /** Number of direct buffers allocated so far; exposed for tests. */ + private[shuffle] var allocations = 0 - private def resetDataBuf(): Unit = { - if (threadLocalDataBuf.get().capacity() > INITIAL_BUFFER_SIZE) { - threadLocalDataBuf.set(ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE)) + /** + * Returns a buffer positioned at zero with its limit set to `bytesToRead`, growing the + * underlying allocation when the block does not fit. + */ + def acquire(bytesToRead: Int): ByteBuffer = { + if (buffer == null || buffer.capacity() < bytesToRead) { + // Doubling keeps the number of reallocations logarithmic in the largest block. Clamp so a + // block near the 2 GB direct-buffer limit still gets a buffer it fits in. + val newCapacity = + (bytesToRead * 2L).max(INITIAL_BUFFER_SIZE).min(Integer.MAX_VALUE).toInt + buffer = ByteBuffer.allocateDirect(newCapacity) + allocations += 1 } + buffer.clear() + buffer.limit(bytesToRead) + buffer } } + +object ShuffleBlockBuffer { + private val INITIAL_BUFFER_SIZE = 128 * 1024 +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index 54448ec2039..8d7608eb73b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -1627,6 +1627,10 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { context.markTaskCompleted(None) } + test("decoder iterators share one task-scoped read buffer without reallocating") { + NativeBatchDecoderIteratorLifecycleChecks.reusesTaskScopedBufferAcrossIterators() + } + test("decoder cleanup releases a prefetched batch that was not consumed") { NativeBatchDecoderIteratorLifecycleChecks.closesPrefetchedBatch() } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala index 52d11d94e30..ac8ff230d8f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala @@ -117,6 +117,54 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } } + /** + * A task-scoped buffer must be handed from one decoder iterator to the next without a new + * direct allocation, grow only when a block does not fit, and never shrink back when an + * iterator is closed; the old per-iterator reset reallocated on every fetched map output. + */ + def reusesTaskScopedBufferAcrossIterators(): Unit = { + val buffer = new ShuffleBlockBuffer() + val small = buffer.acquire(12) + assert(buffer.allocations == 1) + assert(small.isDirect && small.position() == 0 && small.limit() == 12) + assert(small.capacity() >= 128 * 1024, "first allocation starts at the initial size") + + // Fits in the initial allocation: same buffer, no new allocation. + assert(buffer.acquire(64 * 1024) eq small) + assert(buffer.allocations == 1) + + // Larger than the allocation: grows to twice the block. + val large = buffer.acquire(300 * 1024) + assert(buffer.allocations == 2) + assert(large.capacity() == 600 * 1024 && large.limit() == 300 * 1024) + + // Every iterator that shares the buffer sees the grown allocation, and closing one does not + // give it back. + val batch = new TrackingBatch() + val util = new NativeUtil { + override def getNextBatch( + numOutputCols: Int, + decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = Some(batch) + } + try { + (0 until 3).foreach { _ => + val decoder = NativeBatchDecoderIterator( + new ByteArrayInputStream(frame), + new SQLMetric("nsTiming", 0L), + null, + util, + tracingEnabled = false, + dataBuffer = buffer) + assert(decoder.hasNext) + decoder.close() + } + } finally { + util.close() + } + assert(buffer.allocations == 2, "closing iterators must not reallocate the shared buffer") + assert(buffer.acquire(200 * 1024) eq large) + } + def closesPrefetchedBatch(): Unit = { val batch = new TrackingBatch() withDecoder(batch) { (decoder, inputCloseCalls) =>