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) =>