You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/11 14:12:01 UTC
(datafusion-comet) 01/01: perf: read shuffle blocks in 64 KiB pieces instead of through `Channels.newChannel` (#6805)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6805-7321b42f3b4aceb1e86c9bae6f6e3b2276471664
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 0691bd39f75da61e02aa23ce188de141bd623252
Author: Oleks V <co...@users.noreply.github.com>
AuthorDate: Sun Oct 11 14:11:39 2026 +0000
perf: read shuffle blocks in 64 KiB pieces instead of through `Channels.newChannel` (#6805)
* perf: read shuffle blocks in 64 KiB pieces instead of through `Channels.newChannel`
Both shuffle readers, `NativeBatchDecoderIterator` and `CometShuffleBlockIterator`,
read blocks through `Channels.newChannel(in)`. For a stream that is not a plain
`FileInputStream`, the JDK channel copies at most 8 KiB per read and calls
`in.available()` before every read after the first, which costs an fstat and an
lseek on a local shuffle file. Read `in` directly in pieces of up to 64 KiB instead.
* Make the shuffle read buffer size configurable and share the read loop
Add `spark.comet.shuffle.readBufferSize`, 64 KiB by default. Both shuffle readers now
call one `CometShuffleBlockIterator.readFully`, which owns a single buffer per thread,
and one test covers both readers.
* Pass the read buffer size default as a Long
The strict-warnings build rejects the implicit Int to Long widening in
`createWithDefault`, so convert the default explicitly.
* feedback
* Move the local benchmarks out of this PR
`CometHashJoinTaskTimeBenchmark` and `CometShuffleReadKillBenchmark` were local tools for
measuring this change and its review, and came in with the previous commit by accident. They
stay local, and the skew-join benchmark will get its own PR against #6528.
---
.../apache/comet/CometShuffleBlockIterator.java | 63 ++++++++++++++-------
.../main/scala/org/apache/comet/CometConf.scala | 12 ++++
.../shuffle/CometBlockStoreShuffleReader.scala | 21 ++++++-
.../shuffle/CometCelebornShuffleReader.scala | 3 +-
.../execution/shuffle/CometShuffledRowRDD.scala | 5 +-
.../shuffle/NativeBatchDecoderIterator.scala | 20 +++----
.../shuffle/CometCelebornShuffleReaderSuite.scala | 33 +++++++++++
...NativeBatchDecoderIteratorLifecycleChecks.scala | 65 +++++++++++++++++++++-
8 files changed, 182 insertions(+), 40 deletions(-)
diff --git a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java
index 86220f9009..c19ab285e8 100644
--- a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java
+++ b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java
@@ -25,8 +25,6 @@ import java.io.IOException;
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 class CometShuffleBlockIterator implements Closeable {
}
// 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,11 +123,9 @@ public class CometShuffleBlockIterator implements Closeable {
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.
@@ -133,6 +133,29 @@ public class CometShuffleBlockIterator implements Closeable {
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.
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 2beaf8f6a3..c6ac50606f 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -894,6 +894,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")
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 3592095a37..156b50a9c0 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
@@ -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 = {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala
index b7deac55b4..945dbe2c09 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala
+++ b/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) {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala
index edd0332ac5..dddb976776 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala
+++ b/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] = {
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 6227da4bf4..b74732a8c1 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
@@ -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")
}
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 54448ec203..0e618b9608 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
@@ -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()
}
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 52d11d94e3..01e93397f3 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
@@ -27,7 +27,7 @@ import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
-import org.apache.comet.{CometShuffleReadFailureHandler, Native}
+import org.apache.comet.{CometShuffleBlockIterator, CometShuffleReadFailureHandler, Native}
import org.apache.comet.serde.{OperatorOuterClass, QueryPlanSerde}
import org.apache.comet.vector.NativeUtil
@@ -117,6 +117,69 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks {
}
}
+ /** Both shuffle readers read in pieces of their read buffer size, without `available()`. */
+ def readersReadInPiecesOfTheReadBufferSize(): Unit = {
+ val bodyBytes = 1024 * 1024
+ val readBufferSize = 16 * 1024
+ val largeFrame = ByteBuffer
+ .allocate(16 + bodyBytes)
+ .order(ByteOrder.LITTLE_ENDIAN)
+ .putLong(8L + bodyBytes)
+ .putLong(0L)
+ .array()
+ class CountingInput extends ByteArrayInputStream(largeFrame) {
+ var availableCalls = 0
+ var largestRead = 0
+
+ override def available(): Int = {
+ availableCalls += 1
+ super.available()
+ }
+
+ override def read(bytes: Array[Byte], offset: Int, length: Int): Int = {
+ largestRead = math.max(largestRead, length)
+ super.read(bytes, offset, length)
+ }
+
+ def check(): Unit = {
+ assert(availableCalls == 0, "The reader asked the stream for available bytes")
+ assert(largestRead == readBufferSize, s"The largest read was $largestRead bytes")
+ }
+ }
+
+ val decoderInput = new CountingInput
+ val batch = new TrackingBatch()
+ val util = new NativeUtil {
+ override def getNextBatch(
+ numOutputCols: Int,
+ decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = Some(batch)
+ }
+ try {
+ val decoder = NativeBatchDecoderIterator(
+ decoderInput,
+ new SQLMetric("nsTiming", 0L),
+ null,
+ util,
+ tracingEnabled = false,
+ readBufferSize = readBufferSize)
+ assert(decoder.next() eq batch)
+ assert(!decoder.hasNext)
+ } finally {
+ util.close()
+ }
+ decoderInput.check()
+
+ val blockInput = new CountingInput
+ val blocks = new CometShuffleBlockIterator(blockInput, (_: Long) => (), readBufferSize)
+ try {
+ assert(blocks.hasNext() == bodyBytes)
+ assert(blocks.hasNext() == -1)
+ } finally {
+ blocks.close()
+ }
+ blockInput.check()
+ }
+
def closesPrefetchedBatch(): Unit = {
val batch = new TrackingBatch()
withDecoder(batch) { (decoder, inputCloseCalls) =>
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org