perf: keep one shuffle block read buffer per reduce task

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

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
+}

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()
   }

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