perf: read shuffle blocks in 64 KiB pieces instead of through Channels.newChannel

spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java

@@ -25,8 +25,6 @@
 import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
-import java.nio.channels.Channels;
-import java.nio.channels.ReadableByteChannel;
 import java.util.function.LongConsumer;
 
 /**
@@ -44,7 +42,12 @@ public class CometShuffleBlockIterator implements Closeable {
 
   private static final int INITIAL_BUFFER_SIZE = 128 * 1024;
 
-  private final ReadableByteChannel channel;
+  /** Default of {@code spark.comet.shuffle.readBufferSize}. */
+  public static final int DEFAULT_READ_BUFFER_SIZE = 64 * 1024;
+
+  private static final ThreadLocal<byte[]> READ_BUFFER = new ThreadLocal<>();
+
+  private final int readBufferSize;
   private final InputStream inputStream;
   private final LongConsumer recordsReadUpdater;
   private final ByteBuffer headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN);
@@ -53,13 +56,14 @@ public class CometShuffleBlockIterator implements Closeable {
   private int currentBlockLength = 0;
 
   public CometShuffleBlockIterator(InputStream in) {
-    this(in, records -> {});
+    this(in, records -> {}, DEFAULT_READ_BUFFER_SIZE);
   }
 
-  public CometShuffleBlockIterator(InputStream in, LongConsumer recordsReadUpdater) {
+  public CometShuffleBlockIterator(
+      InputStream in, LongConsumer recordsReadUpdater, int readBufferSize) {
     this.inputStream = in;
-    this.channel = Channels.newChannel(in);
     this.recordsReadUpdater = recordsReadUpdater;
+    this.readBufferSize = readBufferSize;
   }
 
   /**
@@ -77,17 +81,15 @@ public int hasNext() throws IOException {
     }
 
     // Read 16-byte header: clear() resets position=0, limit=capacity,
-    // preparing the buffer for channel.read() to fill it
+    // preparing the buffer for readFully() to fill it
     headerBuf.clear();
-    while (headerBuf.hasRemaining()) {
-      int bytesRead = channel.read(headerBuf);
-      if (bytesRead < 0) {
-        if (headerBuf.position() == 0) {
-          close();
-          return -1;
-        }
-        throw new EOFException("Data corrupt: unexpected EOF while reading batch header");
+    readFully(inputStream, headerBuf, readBufferSize);
+    if (headerBuf.hasRemaining()) {
+      if (headerBuf.position() == 0) {
+        close();
+        return -1;
       }
+      throw new EOFException("Data corrupt: unexpected EOF while reading batch header");
     }
     headerBuf.flip();
     long compressedLength = headerBuf.getLong();
@@ -121,18 +123,39 @@ public int hasNext() throws IOException {
 
     dataBuf.clear();
     dataBuf.limit(currentBlockLength);
-    while (dataBuf.hasRemaining()) {
-      int bytesRead = channel.read(dataBuf);
-      if (bytesRead < 0) {
-        throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch");
-      }
+    readFully(inputStream, dataBuf, readBufferSize);
+    if (dataBuf.hasRemaining()) {
+      throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch");
     }
     // Note: native side uses get_direct_buffer_address (base pointer) + currentBlockLength,
     // not the buffer's position/limit. No flip needed.
 
     return currentBlockLength;
   }
 
+  /**
+   * Fills {@code dst} from {@code in}, stopping short only at the end of the stream. Both shuffle
+   * readers use it. The stream is read directly, in pieces of up to {@code readBufferSize}, through
+   * a buffer reused per thread. {@code Channels.newChannel} would read 8 KiB at a time and call
+   * {@code available()} before every piece but the first, which costs an fstat and an lseek on a
+   * local shuffle file.
+   */
+  public static void readFully(InputStream in, ByteBuffer dst, int readBufferSize)
+      throws IOException {
+    byte[] chunk = READ_BUFFER.get();
+    if (chunk == null || chunk.length != readBufferSize) {
+      chunk = new byte[readBufferSize];
+      READ_BUFFER.set(chunk);
+    }
+    while (dst.hasRemaining()) {
+      int read = in.read(chunk, 0, Math.min(dst.remaining(), chunk.length));
+      if (read < 0) {
+        return;
+      }
+      dst.put(chunk, 0, read);
+    }
+  }
+
   /**
    * Reports a native IPC/codec failure to the remote stream that supplied the current block. Called
    * by ShuffleScan via JNI; local streams preserve the original native exception.

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -890,6 +890,18 @@ object CometConf extends ShimCometConf {
         s"Write buffer size must be between 1 and ${Int.MaxValue} bytes")
       .createWithDefault(1024 * 1024)
 
+  val COMET_SHUFFLE_READ_BUFFER_SIZE: ConfigEntry[Long] =
+    conf("spark.comet.shuffle.readBufferSize")
+      .category(CATEGORY_SHUFFLE)
+      .doc(
+        "Size in bytes of each read from a shuffle block stream. Comet reads a shuffle block " +
+          "through a reusable buffer of this size per task thread before decoding it.")
+      .bytesConf(ByteUnit.BYTE)
+      .checkValue(
+        v => v > 0 && v <= Int.MaxValue,
+        s"Read buffer size must be between 1 and ${Int.MaxValue} bytes")
+      .createWithDefault(CometShuffleBlockIterator.DEFAULT_READ_BUFFER_SIZE.toLong)
+
   val COMET_SHUFFLE_JVM_PREFER_DICTIONARY_RATIO: ConfigEntry[Double] = conf(
     "spark.comet.shuffle.jvm.preferDictionary.ratio")
     .withAlternative("spark.comet.shuffle.preferDictionary.ratio")

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala

@@ -19,7 +19,7 @@
 
 package org.apache.spark.sql.comet.execution.shuffle
 
-import java.io.InputStream
+import java.io.{FilterInputStream, InputStream}
 
 import org.apache.spark.{InterruptibleIterator, MapOutputTracker, SparkEnv, TaskContext}
 import org.apache.spark.internal.{config, Logging}
@@ -91,6 +91,7 @@ class CometBlockStoreShuffleReader[K, C](
     val nativeLib = new Native()
     val nativeUtil = new NativeUtil()
     val tracingEnabled = CometConf.COMET_TRACING_ENABLED.get()
+    val readBufferSize = CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt
 
     // Closes last read iterator and shared resources after the task is finished.
     // We need to close read iterator during iterating input streams,
@@ -113,7 +114,8 @@ class CometBlockStoreShuffleReader[K, C](
           dep.decodeTime,
           nativeLib,
           nativeUtil,
-          tracingEnabled)
+          tracingEnabled,
+          readBufferSize = readBufferSize)
         currentReadIterator
       })
       .map(b => (0, b))
@@ -158,10 +160,23 @@ class CometBlockStoreShuffleReader[K, C](
    */
   override def readAsRawStream(): InputStream = {
     val streams = fetchIterator.map(_._2)
-    new java.io.SequenceInputStream(new java.util.Enumeration[InputStream] {
+    val blocks = new java.io.SequenceInputStream(new java.util.Enumeration[InputStream] {
       override def hasMoreElements: Boolean = streams.hasNext
       override def nextElement(): InputStream = streams.next()
     })
+    // Native code reads this stream on the task thread, outside any InterruptibleIterator, so
+    // every read checks whether the task was killed.
+    new FilterInputStream(blocks) {
+      override def read(): Int = {
+        context.killTaskIfInterrupted()
+        in.read()
+      }
+
+      override def read(buffer: Array[Byte], offset: Int, length: Int): Int = {
+        context.killTaskIfInterrupted()
+        in.read(buffer, offset, length)
+      }
+    }
   }
 
   private def fetchContinuousBlocksInBatch: Boolean = {

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala

@@ -122,7 +122,8 @@ private[shuffle] final class CometCelebornShuffleReader[K, C](
           new Native(),
           nativeUtil,
           CometConf.COMET_TRACING_ENABLED.get(),
-          Some(serializedSchema))
+          Some(serializedSchema),
+          CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt)
       } catch {
         case NonFatal(failure) =>
           Utils.tryWithSafeFinally(throw failure) {

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala

@@ -27,7 +27,7 @@ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLShuffleReadMetricsRe
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.CometShuffleBlockIterator
+import org.apache.comet.{CometConf, CometShuffleBlockIterator}
 
 /**
  * Different from [[org.apache.spark.sql.execution.ShuffledRowRDD]], this RDD is specialized for
@@ -176,7 +176,8 @@ class CometShuffledBatchRDD(
     }
     new CometShuffleBlockIterator(
       readerAndMetrics._1.readAsRawStream(),
-      readerAndMetrics._2.incRecordsRead)
+      readerAndMetrics._2.incRecordsRead,
+      CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt)
   }
 
   override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = {

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala

@@ -21,14 +21,13 @@ package org.apache.spark.sql.comet.execution.shuffle
 
 import java.io.{EOFException, InputStream}
 import java.nio.{ByteBuffer, ByteOrder}
-import java.nio.channels.{Channels, ReadableByteChannel}
 
 import scala.util.control.NonFatal
 
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.{CometShuffleReadFailureHandler, Native}
+import org.apache.comet.{CometShuffleBlockIterator, CometShuffleReadFailureHandler, Native}
 import org.apache.comet.vector.NativeUtil
 
 /**
@@ -42,7 +41,8 @@ case class NativeBatchDecoderIterator(
     nativeLib: Native,
     nativeUtil: NativeUtil,
     tracingEnabled: Boolean,
-    expectedSchema: Option[Array[Byte]] = None)
+    expectedSchema: Option[Array[Byte]] = None,
+    readBufferSize: Int = CometShuffleBlockIterator.DEFAULT_READ_BUFFER_SIZE)
     extends Iterator[ColumnarBatch] {
 
   // One consumer reads this iterator, while task completion may close it from another thread.
@@ -60,15 +60,9 @@ case class NativeBatchDecoderIterator(
 
   import NativeBatchDecoderIterator._
 
-  private val channel: ReadableByteChannel = if (in != null) {
-    Channels.newChannel(in)
-  } else {
-    null
-  }
-
   override def hasNext: Boolean = {
     synchronized {
-      if (channel == null || isClosed) {
+      if (in == null || isClosed) {
         return false
       }
       if (batch.isDefined) {
@@ -172,7 +166,7 @@ case class NativeBatchDecoderIterator(
   private def readNextBlock(): Option[(Int, ByteBuffer, Int)] = {
     // read compressed batch size from header
     longBuf.clear()
-    while (longBuf.hasRemaining && channel.read(longBuf) >= 0) {}
+    CometShuffleBlockIterator.readFully(in, longBuf, readBufferSize)
 
     // If we reach the end of the stream, we are done, or if we read partial length
     // then the stream is corrupted.
@@ -189,7 +183,7 @@ case class NativeBatchDecoderIterator(
 
     // read field count from header
     longBuf.clear()
-    while (longBuf.hasRemaining && channel.read(longBuf) >= 0) {}
+    CometShuffleBlockIterator.readFully(in, longBuf, readBufferSize)
     if (longBuf.hasRemaining) {
       throw new EOFException("Data corrupt: unexpected EOF while reading field count")
     }
@@ -214,7 +208,7 @@ case class NativeBatchDecoderIterator(
     }
     dataBuf.clear()
     dataBuf.limit(bytesToRead.toInt)
-    while (dataBuf.hasRemaining && channel.read(dataBuf) >= 0) {}
+    CometShuffleBlockIterator.readFully(in, dataBuf, readBufferSize)
     if (dataBuf.hasRemaining) {
       throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch")
     }

spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala

@@ -931,6 +931,35 @@ class CometCelebornShuffleReaderSuite extends CometTestBase {
     context.markTaskCompleted(None)
   }
 
+  test("a killed task stops at the next read of a local shuffle's direct-read stream") {
+    withSQLConf(
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_SHUFFLE_MODE.key -> "native",
+      CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+      "spark.sql.adaptive.enabled" -> "false") {
+      withParquetTable((0 until 12).map(i => (i, s"row-$i")), "local_direct_read_rows") {
+        val shuffled = sql("SELECT * FROM local_direct_read_rows").repartition(1, $"_1")
+        // Running the query writes the shuffle that the stream below reads again.
+        assert(shuffled.collect().length == 12)
+        val rdd = shuffled.queryExecution.executedPlan
+          .collectFirst { case exchange: CometShuffleExchangeExec =>
+            exchange.executeColumnar().asInstanceOf[CometShuffledBatchRDD]
+          }
+          .getOrElse(fail("Expected a native Comet shuffle exchange"))
+        val context = TaskContext.empty()
+        val blocks = rdd.computeAsShuffleBlockIterator(rdd.partitions.head, context)
+        try {
+          assert(blocks.hasNext() > 0)
+          context.markInterrupted("killed after the first batch")
+          intercept[TaskKilledException](blocks.hasNext())
+        } finally {
+          blocks.close()
+          context.markTaskCompleted(None)
+        }
+      }
+    }
+  }
+
   test("malformed native frames invalidate once before allocating or opening another reducer") {
     def header(length: Long, fields: Long): Array[Byte] = ByteBuffer
       .allocate(16)
@@ -1631,6 +1660,10 @@ class CometCelebornShuffleReaderSuite extends CometTestBase {
     NativeBatchDecoderIteratorLifecycleChecks.closesPrefetchedBatch()
   }
 
+  test("shuffle readers read in pieces of the read buffer size without asking for available") {
+    NativeBatchDecoderIteratorLifecycleChecks.readersReadInPiecesOfTheReadBufferSize()
+  }
+
   test("decoder cleanup releases its delivered batch exactly once") {
     NativeBatchDecoderIteratorLifecycleChecks.closesDeliveredBatch()
   }