From 9ad49c4956cb6e9568798f69a8b46e7a3ea77a2a Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 8 Oct 2026 13:06:17 -0700 Subject: [PATCH 1/5] 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. --- .../comet/CometShuffleBlockIterator.java | 50 ++++++++++++------- .../shuffle/NativeBatchDecoderIterator.scala | 37 ++++++++++---- .../comet/exec/CometNativeShuffleSuite.scala | 32 ++++++++++++ .../CometCelebornShuffleReaderSuite.scala | 4 ++ ...eBatchDecoderIteratorLifecycleChecks.scala | 43 ++++++++++++++++ 5 files changed, 137 insertions(+), 29 deletions(-) diff --git a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java index 86220f90093..4b6a4f5ea11 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.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,11 @@ public class CometShuffleBlockIterator implements Closeable { private static final int INITIAL_BUFFER_SIZE = 128 * 1024; - private final ReadableByteChannel channel; + private static final int READ_CHUNK_SIZE = 64 * 1024; + + private static final ThreadLocal READ_CHUNK = + ThreadLocal.withInitial(() -> new byte[READ_CHUNK_SIZE]); + private final InputStream inputStream; private final LongConsumer recordsReadUpdater; private final ByteBuffer headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN); @@ -58,7 +60,6 @@ public CometShuffleBlockIterator(InputStream in) { public CometShuffleBlockIterator(InputStream in, LongConsumer recordsReadUpdater) { this.inputStream = in; - this.channel = Channels.newChannel(in); this.recordsReadUpdater = recordsReadUpdater; } @@ -77,17 +78,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(headerBuf); + 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 +120,9 @@ 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(dataBuf); + 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 +130,23 @@ public int hasNext() throws IOException { return currentBlockLength; } + /** + * Fills {@code dst} from the stream, stopping short only at the end of the stream. The stream is + * read directly, in pieces of up to {@code READ_CHUNK_SIZE}. {@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. + */ + private void readFully(ByteBuffer dst) throws IOException { + byte[] chunk = READ_CHUNK.get(); + while (dst.hasRemaining()) { + int read = inputStream.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/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala index 6227da4bf4f..847b4aa33e4 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,7 +21,6 @@ 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 @@ -60,15 +59,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 +165,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) {} + readFully(longBuf) // If we reach the end of the stream, we are done, or if we read partial length // then the stream is corrupted. @@ -189,7 +182,7 @@ case class NativeBatchDecoderIterator( // read field count from header longBuf.clear() - while (longBuf.hasRemaining && channel.read(longBuf) >= 0) {} + readFully(longBuf) if (longBuf.hasRemaining) { throw new EOFException("Data corrupt: unexpected EOF while reading field count") } @@ -214,7 +207,7 @@ case class NativeBatchDecoderIterator( } dataBuf.clear() dataBuf.limit(bytesToRead.toInt) - while (dataBuf.hasRemaining && channel.read(dataBuf) >= 0) {} + readFully(dataBuf) if (dataBuf.hasRemaining) { throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch") } @@ -222,6 +215,23 @@ case class NativeBatchDecoderIterator( Some((fieldCount, dataBuf, bytesToRead.toInt)) } + /** + * Fills `dst` from `in`, stopping short only at the end of the stream. `in` is read directly, + * in pieces of up to `READ_CHUNK_SIZE`. `Channels.newChannel(in)` would read 8 KiB at a time + * and call `in.available()` before every piece but the first, which costs an fstat and an lseek + * on a local shuffle file. + */ + private def readFully(dst: ByteBuffer): Unit = { + val chunk = threadLocalReadChunk.get() + while (dst.hasRemaining) { + val read = in.read(chunk, 0, math.min(dst.remaining(), chunk.length)) + if (read < 0) { + return + } + dst.put(chunk, 0, read) + } + } + def close(): Unit = { synchronized { if (!isClosed) { @@ -261,6 +271,11 @@ object NativeBatchDecoderIterator { private val INITIAL_BUFFER_SIZE = 128 * 1024 + private val READ_CHUNK_SIZE = 64 * 1024 + + private val threadLocalReadChunk: ThreadLocal[Array[Byte]] = + ThreadLocal.withInitial(() => new Array[Byte](READ_CHUNK_SIZE)) + private val threadLocalDataBuf: ThreadLocal[ByteBuffer] = ThreadLocal.withInitial(() => { ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE) }) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 0c56e11b9ab..ead521d8358 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -119,6 +119,38 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper assert(registration.getReturnType == java.lang.Void.TYPE) } + test("shuffle block iterator reads large blocks in big pieces without asking for available") { + val bodyBytes = 1024 * 1024 + val frame = ByteBuffer + .allocate(16 + bodyBytes) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(8L + bodyBytes) + .putLong(0L) + .array() + var availableCalls = 0 + var largestRead = 0 + val input = new ByteArrayInputStream(frame) { + 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) + } + } + val blocks = new CometShuffleBlockIterator(input) + try { + assert(blocks.hasNext() == bodyBytes) + assert(blocks.hasNext() == -1) + } finally { + blocks.close() + } + assert(availableCalls == 0, "The iterator asked the stream for available bytes") + assert(largestRead > 8 * 1024, s"The largest read was $largestRead bytes") + } + test("native RSS shuffle invokes its task-owned JVM callback with a complete frame") { val planBytes = rssShufflePlanBytes 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..259d8d09a12 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 @@ -1631,6 +1631,10 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { NativeBatchDecoderIteratorLifecycleChecks.closesPrefetchedBatch() } + test("decoder reads large blocks in big pieces without asking for available bytes") { + NativeBatchDecoderIteratorLifecycleChecks.readsLargeBlocksInBigPiecesWithoutAvailable() + } + 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 52d11d94e30..a6a726aabbc 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,49 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } } + def readsLargeBlocksInBigPiecesWithoutAvailable(): Unit = { + val bodyBytes = 1024 * 1024 + val largeFrame = ByteBuffer + .allocate(16 + bodyBytes) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(8L + bodyBytes) + .putLong(0L) + .array() + var availableCalls = 0 + var largestRead = 0 + val input = new ByteArrayInputStream(largeFrame) { + 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) + } + } + 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( + input, + new SQLMetric("nsTiming", 0L), + null, + util, + tracingEnabled = false) + assert(decoder.next() eq batch) + assert(!decoder.hasNext) + assert(availableCalls == 0, "The decoder asked the stream for available bytes") + assert(largestRead > 8 * 1024, s"The largest read was $largestRead bytes") + } finally { + util.close() + } + } + def closesPrefetchedBatch(): Unit = { val batch = new TrackingBatch() withDecoder(batch) { (decoder, inputCloseCalls) => From 4edc57a0dc9c774503be7a4f297d8aa5e9d1ef9b Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 8 Oct 2026 13:53:03 -0700 Subject: [PATCH 2/5] 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. --- .../comet/CometShuffleBlockIterator.java | 37 +++++++++++------- .../scala/org/apache/comet/CometConf.scala | 12 ++++++ .../CometBlockStoreShuffleReader.scala | 4 +- .../shuffle/CometCelebornShuffleReader.scala | 3 +- .../shuffle/CometShuffledRowRDD.scala | 5 ++- .../shuffle/NativeBatchDecoderIterator.scala | 33 +++------------- .../comet/exec/CometNativeShuffleSuite.scala | 32 ---------------- .../CometCelebornShuffleReaderSuite.scala | 4 +- ...eBatchDecoderIteratorLifecycleChecks.scala | 38 ++++++++++++++----- 9 files changed, 80 insertions(+), 88 deletions(-) diff --git a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java index 4b6a4f5ea11..c19ab285e8c 100644 --- a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java +++ b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java @@ -42,11 +42,12 @@ public class CometShuffleBlockIterator implements Closeable { private static final int INITIAL_BUFFER_SIZE = 128 * 1024; - private static final int READ_CHUNK_SIZE = 64 * 1024; + /** Default of {@code spark.comet.shuffle.readBufferSize}. */ + public static final int DEFAULT_READ_BUFFER_SIZE = 64 * 1024; - private static final ThreadLocal READ_CHUNK = - ThreadLocal.withInitial(() -> new byte[READ_CHUNK_SIZE]); + private static final ThreadLocal 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); @@ -55,12 +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.recordsReadUpdater = recordsReadUpdater; + this.readBufferSize = readBufferSize; } /** @@ -80,7 +83,7 @@ public int hasNext() throws IOException { // Read 16-byte header: clear() resets position=0, limit=capacity, // preparing the buffer for readFully() to fill it headerBuf.clear(); - readFully(headerBuf); + readFully(inputStream, headerBuf, readBufferSize); if (headerBuf.hasRemaining()) { if (headerBuf.position() == 0) { close(); @@ -120,7 +123,7 @@ public int hasNext() throws IOException { dataBuf.clear(); dataBuf.limit(currentBlockLength); - readFully(dataBuf); + readFully(inputStream, dataBuf, readBufferSize); if (dataBuf.hasRemaining()) { throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch"); } @@ -131,15 +134,21 @@ public int hasNext() throws IOException { } /** - * Fills {@code dst} from the stream, stopping short only at the end of the stream. The stream is - * read directly, in pieces of up to {@code READ_CHUNK_SIZE}. {@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. + * 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. */ - private void readFully(ByteBuffer dst) throws IOException { - byte[] chunk = READ_CHUNK.get(); + 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 = inputStream.read(chunk, 0, Math.min(dst.remaining(), chunk.length)); + int read = in.read(chunk, 0, Math.min(dst.remaining(), chunk.length)); if (read < 0) { return; } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a187f8fa5f8..1cc93affd4c 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/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) + 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 3592095a37c..e219fe41672 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,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)) 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 b7deac55b49..945dbe2c094 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 edd0332ac56..dddb9767761 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 847b4aa33e4..b74732a8c11 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 @@ -27,7 +27,7 @@ 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 /** @@ -41,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. @@ -165,7 +166,7 @@ case class NativeBatchDecoderIterator( private def readNextBlock(): Option[(Int, ByteBuffer, Int)] = { // read compressed batch size from header longBuf.clear() - readFully(longBuf) + 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. @@ -182,7 +183,7 @@ case class NativeBatchDecoderIterator( // read field count from header longBuf.clear() - readFully(longBuf) + CometShuffleBlockIterator.readFully(in, longBuf, readBufferSize) if (longBuf.hasRemaining) { throw new EOFException("Data corrupt: unexpected EOF while reading field count") } @@ -207,7 +208,7 @@ case class NativeBatchDecoderIterator( } dataBuf.clear() dataBuf.limit(bytesToRead.toInt) - readFully(dataBuf) + CometShuffleBlockIterator.readFully(in, dataBuf, readBufferSize) if (dataBuf.hasRemaining) { throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch") } @@ -215,23 +216,6 @@ case class NativeBatchDecoderIterator( Some((fieldCount, dataBuf, bytesToRead.toInt)) } - /** - * Fills `dst` from `in`, stopping short only at the end of the stream. `in` is read directly, - * in pieces of up to `READ_CHUNK_SIZE`. `Channels.newChannel(in)` would read 8 KiB at a time - * and call `in.available()` before every piece but the first, which costs an fstat and an lseek - * on a local shuffle file. - */ - private def readFully(dst: ByteBuffer): Unit = { - val chunk = threadLocalReadChunk.get() - while (dst.hasRemaining) { - val read = in.read(chunk, 0, math.min(dst.remaining(), chunk.length)) - if (read < 0) { - return - } - dst.put(chunk, 0, read) - } - } - def close(): Unit = { synchronized { if (!isClosed) { @@ -271,11 +255,6 @@ object NativeBatchDecoderIterator { private val INITIAL_BUFFER_SIZE = 128 * 1024 - private val READ_CHUNK_SIZE = 64 * 1024 - - private val threadLocalReadChunk: ThreadLocal[Array[Byte]] = - ThreadLocal.withInitial(() => new Array[Byte](READ_CHUNK_SIZE)) - private val threadLocalDataBuf: ThreadLocal[ByteBuffer] = ThreadLocal.withInitial(() => { ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE) }) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index ead521d8358..0c56e11b9ab 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -119,38 +119,6 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper assert(registration.getReturnType == java.lang.Void.TYPE) } - test("shuffle block iterator reads large blocks in big pieces without asking for available") { - val bodyBytes = 1024 * 1024 - val frame = ByteBuffer - .allocate(16 + bodyBytes) - .order(ByteOrder.LITTLE_ENDIAN) - .putLong(8L + bodyBytes) - .putLong(0L) - .array() - var availableCalls = 0 - var largestRead = 0 - val input = new ByteArrayInputStream(frame) { - 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) - } - } - val blocks = new CometShuffleBlockIterator(input) - try { - assert(blocks.hasNext() == bodyBytes) - assert(blocks.hasNext() == -1) - } finally { - blocks.close() - } - assert(availableCalls == 0, "The iterator asked the stream for available bytes") - assert(largestRead > 8 * 1024, s"The largest read was $largestRead bytes") - } - test("native RSS shuffle invokes its task-owned JVM callback with a complete frame") { val planBytes = rssShufflePlanBytes 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 259d8d09a12..0462bd70421 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 @@ -1631,8 +1631,8 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { NativeBatchDecoderIteratorLifecycleChecks.closesPrefetchedBatch() } - test("decoder reads large blocks in big pieces without asking for available bytes") { - NativeBatchDecoderIteratorLifecycleChecks.readsLargeBlocksInBigPiecesWithoutAvailable() + 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") { 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 a6a726aabbc..01e93397f37 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,17 +117,20 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } } - def readsLargeBlocksInBigPiecesWithoutAvailable(): Unit = { + /** 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() - var availableCalls = 0 - var largestRead = 0 - val input = new ByteArrayInputStream(largeFrame) { + class CountingInput extends ByteArrayInputStream(largeFrame) { + var availableCalls = 0 + var largestRead = 0 + override def available(): Int = { availableCalls += 1 super.available() @@ -137,7 +140,14 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { 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( @@ -146,18 +156,28 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } try { val decoder = NativeBatchDecoderIterator( - input, + decoderInput, new SQLMetric("nsTiming", 0L), null, util, - tracingEnabled = false) + tracingEnabled = false, + readBufferSize = readBufferSize) assert(decoder.next() eq batch) assert(!decoder.hasNext) - assert(availableCalls == 0, "The decoder asked the stream for available bytes") - assert(largestRead > 8 * 1024, s"The largest read was $largestRead bytes") } 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 = { From cf50ca0edccd646b8917946c775348725bdac3c5 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 8 Oct 2026 16:54:42 -0700 Subject: [PATCH 3/5] 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. --- spark/src/main/scala/org/apache/comet/CometConf.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 1cc93affd4c..bf61fecdd58 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -900,7 +900,7 @@ object CometConf extends ShimCometConf { .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) + .createWithDefault(CometShuffleBlockIterator.DEFAULT_READ_BUFFER_SIZE.toLong) val COMET_SHUFFLE_JVM_PREFER_DICTIONARY_RATIO: ConfigEntry[Double] = conf( "spark.comet.shuffle.jvm.preferDictionary.ratio") From 5c025a8acb2278a576895cbff6c781f930267969 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 9 Oct 2026 13:25:32 -0700 Subject: [PATCH 4/5] feedback --- .../CometBlockStoreShuffleReader.scala | 17 +- .../CometHashJoinTaskTimeBenchmark.scala | 1084 +++++++++++++++++ .../CometShuffleReadKillBenchmark.scala | 380 ++++++ .../CometCelebornShuffleReaderSuite.scala | 29 + 4 files changed, 1508 insertions(+), 2 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala 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 e219fe41672..156b50a9c00 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} @@ -160,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/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala new file mode 100644 index 00000000000..c1cc5259f33 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala @@ -0,0 +1,1084 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import java.io.File +import java.nio.charset.StandardCharsets +import java.util.Locale + +import scala.collection.mutable + +import org.apache.spark.{SparkConf, Success} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.rdd.RDD +import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometColumnarToRowExec, CometExec, CometHashJoinExec, CometNativeColumnarToRowExec, CometNativeExec} +import org.apache.spark.sql.execution.{SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Compares Spark and Comet on the one reduce task of a skewed shuffled hash join that receives + * 95% of the probe rows, each row carrying nested and wide columns + * (apache/datafusion-comet#6528). + * + * The join builds `dim` and probes with `fact`, which puts 95% of its rows on one key. AQE is + * off, so the hot partition is not split, and the one reduce task that receives it reads, joins + * and emits about 95% of the rows. Every `fact` row carries a struct holding an array of structs, + * an array of strings, a map, and `WideColumns` flat columns. + * + * Both engines run the same query over the same files. Spark's output rows are counted. Comet's + * plan ends in a columnar-to-row conversion, which has no metric of its own, so the child of that + * conversion runs instead and the rows of its batches are counted. Neither engine writes its + * output anywhere. + * + * Task time is the executor run time of each task, read from task end events with the task's + * shuffle and input metrics and its SQL metric updates. The comparison shows the run with the + * median hot task time of each engine: the hot task, the map stage that scanned `fact` and wrote + * its shuffle, and the whole run. Executor CPU time is shown for reference only: Comet runs part + * of its native work on Tokio threads, which a task's CPU time does not count. After the + * comparison, every metric of the hot task and of the map stage is listed by operator. Metrics of + * one operator can contain each other, so they are listed but not added up. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometHashJoinTaskTimeBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometHashJoinTaskTimeBenchmark-**results.txt". + * + * With the argument `-- profile`, each engine then runs the hot task alone, reusing the shuffle + * of one more run, for two windows of `ProfileSeconds` each. In the first, a thread samples the + * stack of the hot task's thread every millisecond, and the samples are counted by what they were + * doing. A native method on top of a stack stands for all the native code under it. The second + * window is left for an external profiler of native code, such as `sample` on macOS, and starts + * when a line beginning with "PROFILE" names the process id. + * + * With the argument `-- aqe`, AQE is on, with skew join handling and partition coalescing off, so + * that the hot partition stays whole and the join keeps 32 tasks. Comet then reads the shuffle + * directly in native code, which it does only behind AQE query stages. Its conversion to rows + * sits inside the adaptive plan, so each Comet run executes the query once to materialize the + * shuffle stages and then runs the native part of the final plan alone. The run reports the map + * stages of the first execution and the join stage of the second. + * + * A `-Dspark.*` system property, for example in `JAVA_TOOL_OPTIONS`, sets a session config, which + * is how to measure a feature flag of a pull request. + */ +object CometHashJoinTaskTimeBenchmark extends CometBenchmarkBase { + + // `SqlBasedBenchmark` calls `getSparkSession` while it initializes, before the fields of this + // object are assigned. The fields that method reads are therefore `final` with literal values, + // which makes them compile-time constants. + + /** Local cores. The hot task costs wall time only while the other cores sit idle. */ + private final val Cores = 8 + + /** Reduce tasks of the join. The hot task gets the hot key and its share of the other keys. */ + private final val ShufflePartitions = 32 + + /** Rows of `fact`, and so the rows the join emits. */ + private val FactRows = 2 * 1024 * 1024 + + /** The share of `fact` rows on the hot key. */ + private val HotPercent = 95 + + /** Parquet files, and so map tasks, of each table. */ + private val FactFiles = 32 + private val DimFiles = 8 + + /** Rows of `dim`, one per key, so every `fact` row matches exactly one of them. */ + private val DimRows = 256 * 1024 + + /** Flat columns of the payload, next to its three nested ones. */ + private val WideColumns = 100 + + /** Measured runs per engine, set by `runs=N`. Odd, so that the median is one of the runs. */ + private var measuredRuns = 5 + + /** Length of each profiling window. */ + private val ProfileSeconds = 25 + + /** + * Rows of `dim` in the broadcast mode, the largest build side of the earlier broadcast cases. + */ + private val BroadcastDimRows = 1024 * 1024 + + /** Keys of `fact` in the broadcast mode. Each is in `dim`, so every row matches once. */ + private val BroadcastKeys = 64 * 1024 + + private def query: String = + if (broadcast) { + "SELECT /*+ BROADCAST(d) */ f.key, f.f_long, f.f_double, f.f_str, d.d_long, d.d_str " + + "FROM fact f JOIN dim d ON f.key = d.k" + } else { + val payload = Seq("n_order", "n_tags", "n_attrs") ++ (0 until WideColumns).map(i => s"w_$i") + val columns = ("f.key" +: payload.map(c => s"f.$c")) ++ Seq("d.d_long", "d.d_str") + s"SELECT /*+ SHUFFLE_HASH(d) */ ${columns.mkString(", ")} " + + "FROM fact f JOIN dim d ON f.key = d.k" + } + + private case class Arm( + name: String, + comet: Boolean, + extraConfigs: Seq[(String, String)] = Nil) { + def configs: Seq[(String, String)] = Seq( + CometConf.COMET_ENABLED.key -> comet.toString, + CometConf.COMET_EXEC_ENABLED.key -> comet.toString) ++ extraConfigs + } + + private val SparkArm = Arm("Spark", comet = false) + private val CometArm = Arm("Comet", comet = true) + + /** The arms of the run: Spark and Comet, or one Comet arm per read buffer size of a sweep. */ + private var arms = Seq(SparkArm, CometArm) + + /** + * One task of a run, with its task metrics and what it added to each accumulator, keyed by + * accumulator id, which for a SQL metric is `SQLMetric.id`. `partition` is the task's index in + * its stage. + */ + private case class TaskRecord( + stageId: Int, + partition: Int, + runTimeMs: Long, + cpuTimeMs: Long, + gcTimeMs: Long, + inputBytes: Long, + shuffleReadBytes: Long, + shuffleReadRecords: Long, + fetchWaitMs: Long, + shuffleWriteBytes: Long, + shuffleWriteRecords: Long, + shuffleWriteNs: Long, + failed: Boolean, + metricUpdates: Map[Long, Long]) + + /** + * A SQL metric of the executed plan, which a task reports under the metric's accumulator id. + * `column` is the operator's first output column, which tells the two sides of a join apart. + */ + private case class PlanMetric( + id: Long, + nodeName: String, + column: String, + key: String, + name: String, + metricType: String) + + /** + * One run of the query. `rows` is the row count the engine produced, `joinRowsId` the + * accumulator of the join's output row count. A run keeps the plan's metrics rather than the + * plan, because a plan holds its shuffle dependencies, and Spark deletes a shuffle's files only + * once nothing references them. + */ + private case class Run( + tasks: Seq[TaskRecord], + rows: Long, + metrics: Seq[PlanMetric], + joinRowsId: Option[Long], + wholeStageJoin: Boolean) { + private val lastStage = tasks.map(_.stageId).max + + /** The tasks of the last stage of the run, the result stage that runs the join. */ + def joinStage: Seq[TaskRecord] = tasks.filter(_.stageId == lastStage) + + /** The tasks of the map stage that wrote the most shuffle bytes, the one that read `fact`. */ + def mapStage: Seq[TaskRecord] = + tasks + .filter(_.stageId != lastStage) + .groupBy(_.stageId) + .values + .maxBy(_.map(_.shuffleWriteBytes).sum) + + def joinRows(task: TaskRecord): Long = + joinRowsId.flatMap(task.metricUpdates.get).getOrElse(0L) + + /** The join stage task that emitted the most rows, or the longest one without a count. */ + def hotTask: TaskRecord = joinStage.maxBy(task => (joinRows(task), task.runTimeMs)) + + /** + * What `of` reported for the metric `key` of the operators whose name starts with `nodeName`, + * in ms for a timing, MiB for a size and as is otherwise. Empty when the plan has no such + * metric. + */ + def metric(of: Seq[TaskRecord], nodeName: String, key: String): Option[Double] = { + val matching = metrics.filter(m => m.nodeName.startsWith(nodeName) && m.key == key) + matching.headOption.map { first => + val total = of.map(task => matching.flatMap(m => task.metricUpdates.get(m.id)).sum).sum + scaled(first.metricType, total) + } + } + } + + /** A metric value in ms for a timing, MiB for a size and as is otherwise. */ + private def scaled(metricType: String, value: Long): Double = metricType match { + case "nsTiming" => value / 1e6 + case "size" => value / (1024.0 * 1024) + case _ => value.toDouble + } + + /** Collects every task end. The listener bus delivers them on its own thread. */ + private object TaskRecorder extends SparkListener { + private val records = mutable.ArrayBuffer.empty[TaskRecord] + + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + def metric(f: TaskMetrics => Long): Long = Option(taskEnd.taskMetrics).map(f).getOrElse(0L) + val updates = taskEnd.taskInfo.accumulables.flatMap { info => + info.update.collect { case value: Long => info.id -> value } + }.toMap + val record = TaskRecord( + stageId = taskEnd.stageId, + partition = taskEnd.taskInfo.index, + runTimeMs = metric(_.executorRunTime), + cpuTimeMs = metric(_.executorCpuTime) / 1000000, + gcTimeMs = metric(_.jvmGCTime), + inputBytes = metric(_.inputMetrics.bytesRead), + shuffleReadBytes = metric(_.shuffleReadMetrics.totalBytesRead), + shuffleReadRecords = metric(_.shuffleReadMetrics.recordsRead), + fetchWaitMs = metric(_.shuffleReadMetrics.fetchWaitTime), + shuffleWriteBytes = metric(_.shuffleWriteMetrics.bytesWritten), + shuffleWriteRecords = metric(_.shuffleWriteMetrics.recordsWritten), + shuffleWriteNs = metric(_.shuffleWriteMetrics.writeTime), + failed = taskEnd.reason != Success, + metricUpdates = updates) + synchronized { records += record } + } + + def drain(): List[TaskRecord] = synchronized { + val drained = records.toList + records.clear() + drained + } + } + + // Keyed by arm name. + private val runs = mutable.Map.empty[String, mutable.Buffer[Run]] + + /** Whether the run uses AQE, set by the `aqe` argument. */ + private var adaptive = false + + /** Whether the run joins by broadcast, set by the `bhj` argument. */ + private var broadcast = false + + /** Whether profiling follows a median small join task instead of the hot one. */ + private var profileSmallTask = false + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometHashJoinTaskTimeBenchmark") + // Since `spark.master` always exists, overrides this value + .set("spark.master", s"local[$Cores]") + .setIfMissing("spark.driver.memory", "3g") + .setIfMissing("spark.executor.memory", "3g") + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + // Comet requires off-heap memory in production. Both engines then draw from the same pool. + .set("spark.memory.offHeap.enabled", "true") + .setIfMissing("spark.memory.offHeap.size", "8g") + + val sparkSession = SparkSession + .builder() + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + + sparkSession.conf.set(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") + sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.ANSI_ENABLED.key, "false") + // Keeps the hot partition whole. AQE would split it, and Comet then leaves the join to Spark + // (apache/datafusion-comet#6530). + sparkSession.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.SHUFFLE_PARTITIONS.key, ShufflePartitions.toString) + // The join carries a hint, so it cannot become a broadcast join. + sparkSession.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") + // An open cost as large as the largest split closes a scan partition after every file, so + // each Parquet file is exactly one map task. `prepareTables` checks that it worked. + val maxSplitBytes = (128L * 1024 * 1024).toString + sparkSession.conf.set(SQLConf.FILES_MAX_PARTITION_BYTES.key, maxSplitBytes) + sparkSession.conf.set(SQLConf.FILES_OPEN_COST_IN_BYTES.key, maxSplitBytes) + + sparkSession + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + val args = mainArgs.map(_.toLowerCase(Locale.ROOT)) + val profiling = args.exists(_.startsWith("profile")) + profileSmallTask = args.contains("profile=small") + adaptive = args.contains("aqe") + broadcast = args.contains("bhj") + if (adaptive) { + // Comet reads a shuffle directly in native code only behind an AQE query stage. Skew join + // handling stays off, because splitting the hot partition makes Comet leave the join to + // Spark (apache/datafusion-comet#6530), and so does coalescing, which keeps 32 join tasks. + spark.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true") + spark.conf.set(SQLConf.SKEW_JOIN_ENABLED.key, "false") + spark.conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED.key, "false") + } + val readBufferSizes = mainArgs.collectFirst { + case arg if arg.startsWith("readBufferSize=") => + arg.stripPrefix("readBufferSize=").split(",").toSeq + } + readBufferSizes.foreach { sizes => + arms = sizes.map { size => + Arm( + s"Comet, $size", + comet = true, + Seq(CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.key -> size)) + } + } + // An A/B of the kill check in `readAsRawStream` (review of #6805). Only a temporary build + // that reads this local property skips the check. + if (args.contains("killcheck")) { + arms = Seq( + Arm("Comet, kill check", comet = true), + Arm( + "Comet, no kill check", + comet = true, + Seq("spark.comet.benchmark.skipKillCheck" -> "true"))) + } + mainArgs + .collectFirst { case arg if arg.startsWith("runs=") => arg.stripPrefix("runs=").toInt } + .foreach(runs => measuredRuns = runs) + val cometSweep = readBufferSizes.isDefined || args.contains("killcheck") + spark.sparkContext.addSparkListener(TaskRecorder) + try { + withTempPath { dir => + withTempTable("fact", "dim") { + prepareTables(dir) + runBenchmark("Skewed hash join, Spark vs Comet: environment") { + emitEnvironment() + } + runBenchmark("Skewed hash join, Spark vs Comet: runs") { + emit("One untimed run of each engine, checked before the measured runs.") + arms.foreach(inspect) + for (_ <- 1 to measuredRuns; arm <- arms) { + runs.getOrElseUpdate(arm.name, mutable.Buffer.empty) += execute(arm)._1 + } + arms.foreach { arm => + val times = measured(arm).map(_.hotTask.runTimeMs).sorted + emit(s"${arm.name} hot task, every measured run (ms): ${times.mkString(", ")}") + } + } + if (cometSweep) { + runBenchmark("Skewed hash join, Comet: one arm per setting") { + emitReadBufferSweep() + } + } else if (broadcast) { + runBenchmark("Broadcast hash join, Spark vs Comet: probe tasks") { + emitProbeSummary() + } + } else { + runBenchmark("Skewed hash join, Spark vs Comet: comparison") { + emitComparison() + } + runBenchmark("Skewed hash join, Spark vs Comet: metrics by operator") { + emitMetricListings() + } + } + if (profiling && !cometSweep) { + runBenchmark("Skewed hash join, Spark vs Comet: profiling") { + arms.foreach(profile) + } + } + } + } + } catch { + case e: Throwable => + // `BenchmarkBase.main` stops the session only after a run that returned. Under + // `exec:java` the classloader is closed before the JVM's shutdown hooks run, so Spark's + // own cleanup hook fails too, and a failed run would leave its shuffle files on disk. + spark.stop() + throw e + } + } + + /** Writes both tables with Comet disabled and checks that each file is one map task. */ + private def prepareTables(dir: File): Unit = { + // The hot rows are every row whose id falls in the first `HotPercent` of each hundred, so + // every map task contributes to the hot partition. The other rows hash onto keys 1 to + // `DimRows - 1`. + val key = s"CAST(IF(PMOD(id, 100) < $HotPercent, 0, 1 + PMOD(HASH(id), ${DimRows - 1})) " + + "AS BIGINT) AS key" + val order = "NAMED_STRUCT(" + + "'id', id, " + + "'customer', NAMED_STRUCT(" + + "'name', CONCAT('customer-', CAST(PMOD(id, 100000) AS STRING)), " + + "'tier', CAST(PMOD(id, 5) AS INT)), " + + "'items', TRANSFORM(SEQUENCE(0, CAST(PMOD(id, 3) AS INT)), " + + "i -> NAMED_STRUCT('sku', id * 10 + i, 'qty', i + 1)))" + val nested = Seq( + s"$order AS n_order", + "TRANSFORM(SEQUENCE(0, CAST(PMOD(id, 3) AS INT)), " + + "i -> CONCAT('tag-', CAST(PMOD(id + i, 50) AS STRING))) AS n_tags", + "MAP('views', PMOD(id, 1000), 'clicks', PMOD(id, 97)) AS n_attrs") + // BIGINT, DOUBLE, STRING and INT in turn, with values hashed from the row, so that neither + // Parquet nor the shuffle compresses them like a sequence. + val wide = (0 until WideColumns).map { i => + val value = i % 4 match { + case 0 => s"XXHASH64(id, $i)" + case 1 => s"CAST(HASH(id, $i) AS DOUBLE) / 7" + case 2 => s"CONCAT('v-', CAST(PMOD(HASH(id, $i), 1000000) AS STRING))" + case _ => s"HASH(id, $i)" + } + s"$value AS w_$i" + } + val dim = Seq("id AS k", "id * 3 AS d_long", "CONCAT('dim-', CAST(id AS STRING)) AS d_str") + + val broadcastFact = Seq( + s"CAST(PMOD(HASH(id), $BroadcastKeys) AS BIGINT) AS key", + "id AS f_long", + "CAST(id AS DOUBLE) / 7 AS f_double", + "CONCAT('order-', CAST(id AS STRING)) AS f_str") + val (factColumns, dimRows) = + if (broadcast) (broadcastFact, BroadcastDimRows) else (key +: (nested ++ wide), DimRows) + Seq(("fact", FactRows, FactFiles, factColumns), ("dim", dimRows, DimFiles, dim)) + .foreach { case (table, rows, files, columns) => + val path = new File(dir, table).getCanonicalPath + spark + .range(0L, rows.toLong, 1L, files) + .selectExpr(columns: _*) + .write + .option("compression", "snappy") + .parquet(path) + spark.read.parquet(path).createOrReplaceTempView(table) + val tasks = spark.table(table).queryExecution.toRdd.getNumPartitions + require( + tasks == files, + s"$table scans as $tasks map tasks, expected one per file ($files)") + } + } + + /** + * One row count per partition of Spark's output, or per batch of Comet's. Comet's plan ends in + * a columnar-to-row conversion, whose child produces the batches. + */ + private def rowCounts(plan: SparkPlan, arm: Arm): RDD[Long] = + if (arm.comet) { + val nativeChild = collectFirst(plan) { + case conversion: CometColumnarToRowExec => conversion.child + case conversion: CometNativeColumnarToRowExec => conversion.child + }.getOrElse { + throw new IllegalStateException(s"No Comet columnar-to-row conversion in:\n$plan") + } + nativeChild.executeColumnar().map(_.numRows().toLong) + } else { + plan.execute().mapPartitions(rows => Iterator(rows.size.toLong)) + } + + private def joinOf(plan: SparkPlan): Option[SparkPlan] = collectFirst(plan) { + case join: ShuffledHashJoinExec => join + case join: CometHashJoinExec => join + case join: BroadcastHashJoinExec => join + case join: CometBroadcastHashJoinExec => join + } + + /** + * Runs the query once under the arm and returns the tasks of the run with the executed plan. + */ + private def execute(arm: Arm): (Run, SparkPlan) = { + // `ContextCleaner` deletes a shuffle only once its dependency is garbage collected, which a + // large heap can postpone until the shuffle files fill the disk. Collecting before each run + // releases the shuffle of the run before, and no collection lands in a measured task. + System.gc() + spark.sparkContext.listenerBus.waitUntilEmpty() + TaskRecorder.drain() + var plan: Option[SparkPlan] = None + var rows = 0L + var materializing: Seq[TaskRecord] = Nil + withSQLConf(arm.configs: _*) { + withConfsPropagated { + val executed = spark.sql(query).queryExecution.executedPlan + val (counts, firstRun) = countsToRun(executed, arm) + materializing = firstRun + rows = counts.fold(0L)(_ + _) + plan = Some(executed) + } + } + spark.sparkContext.listenerBus.waitUntilEmpty() + val executedPlan = plan.get + val join = joinOf(executedPlan) + val joinRowsId = join + .flatMap(op => op.metrics.get("numOutputRows").orElse(op.metrics.get("output_rows"))) + .map(_.id) + val wholeStageJoin = join.exists { op => + collect(executedPlan) { case stage: WholeStageCodegenExec => stage } + .exists(_.find(_ eq op).isDefined) + } + val tasks = materializing ++ TaskRecorder.drain() + val run = Run(tasks, rows, planMetrics(executedPlan), joinRowsId, wholeStageJoin) + (run, executedPlan) + } + + /** + * The row counts to run for `executed`, after running whatever must come first. Under AQE, + * Comet's conversion to rows sits inside the adaptive plan. The query then runs once, which + * materializes its shuffle stages, and the counts come from the native part of the final plan, + * which reads those stages again. The tasks of that first run are returned without its final + * stage, so that a run has the map stages of the first run and the join stage of the second. + */ + private def countsToRun(executed: SparkPlan, arm: Arm): (RDD[Long], Seq[TaskRecord]) = + executed match { + case adaptivePlan: AdaptiveSparkPlanExec if arm.comet => + adaptivePlan.execute().foreach(_ => ()) + spark.sparkContext.listenerBus.waitUntilEmpty() + val firstRun = TaskRecorder.drain() + val finalStage = firstRun.map(_.stageId).max + (rowCounts(adaptivePlan.executedPlan, arm), firstRun.filter(_.stageId != finalStage)) + case _ => + (rowCounts(executed, arm), Nil) + } + + /** How many shuffle inputs the native plan under Comet's conversion reads in native code. */ + private def directReadInputs(plan: SparkPlan): Int = + collectFirst(plan) { + case conversion: CometColumnarToRowExec => conversion.child + case conversion: CometNativeColumnarToRowExec => conversion.child + }.collect { case native: CometNativeExec => + CometExec.findShuffleScanIndices(native.nativeOp).size + }.getOrElse(0) + + /** + * The first operator that is not Comet native. `findFirstNonCometOperator` does not descend + * into AQE wrappers or query stages, so the final plan and the plan of every stage are checked + * one by one. + */ + private def firstNonCometOperator(plan: SparkPlan): Option[SparkPlan] = { + val fragments = plan +: collect(plan) { + case adaptivePlan: AdaptiveSparkPlanExec => adaptivePlan.executedPlan + case stage: QueryStageExec => stage.plan + } + fragments.iterator + .flatMap(fragment => + findFirstNonCometOperator( + fragment, + classOf[AdaptiveSparkPlanExec], + classOf[QueryStageExec], + classOf[AQEShuffleReadExec])) + .find(_ => true) + } + + private def planMetrics(plan: SparkPlan): Seq[PlanMetric] = + collect(plan) { case op => op }.flatMap { op => + val column = op.output.headOption.fold("")(_.name) + op.metrics.toSeq.sortBy(_._1).map { case (key, metric) => + PlanMetric( + metric.id, + op.nodeName.trim, + column, + key, + metric.name.getOrElse(key), + metric.metricType) + } + } + + /** + * Runs the query once, untimed, prints what it did, and warns when it does not measure the + * case. The plan is dropped on return, so its shuffle is released before the measured runs. + */ + private def inspect(arm: Arm): Unit = { + val (run, plan) = execute(arm) + val join = joinOf(plan) + val hotRows = run.joinRows(run.hotTask) + val share = 100.0 * hotRows / FactRows + val directReads = if (arm.comet) directReadInputs(plan) else 0 + val codegen = + if (arm.comet) s", $directReads shuffle inputs read directly in native code" + else if (run.wholeStageJoin) ", join in whole-stage codegen" + else ", join without whole-stage codegen" + emit( + f"${arm.name}: join ${join.fold("none")(_.nodeName)}, ${run.rows} rows, " + + f"${run.joinStage.size} join stage tasks, the hot task emitted $hotRows rows " + + f"($share%.1f%%)$codegen.") + + if (join.isEmpty) { + warning(s"WARNING: '${arm.name}' ran no shuffled hash join.\n$plan") + } + if (run.rows != FactRows) { + warning(s"WARNING: '${arm.name}' produced ${run.rows} rows, expected $FactRows.") + } + if (!broadcast && share < HotPercent) { + warning( + f"WARNING: the hot task of '${arm.name}' emitted $share%.1f%% of the rows, expected " + + f"$HotPercent%d%%.") + } + if (adaptive && arm.comet && directReads == 0) { + warning("WARNING: under AQE, Comet read no shuffle input directly in native code.") + } + if (arm.comet) { + firstNonCometOperator(plan).foreach { op => + warning( + s"WARNING: the Comet plan is not fully Comet native (first non-Comet operator: " + + s"${op.nodeName}), so it partly measures Spark.\n$plan") + } + } + if (run.joinStage.size != ShufflePartitions) { + warning( + s"WARNING: the join stage of '${arm.name}' ran ${run.joinStage.size} tasks, expected " + + s"$ShufflePartitions.") + } + if (run.tasks.exists(_.failed)) { + warning(s"WARNING: '${arm.name}' had failed tasks, so its task time is not comparable.") + } + } + + private def emitEnvironment(): Unit = { + val conf = spark.sparkContext.getConf + emit( + s"Spark ${spark.version}, Java ${System.getProperty("java.version")}, " + + s"master ${conf.get("spark.master")}") + emit( + if (adaptive) "AQE on, skew join handling and partition coalescing off" + else "AQE off") + emit(s"spark.memory.offHeap.size: ${conf.get("spark.memory.offHeap.size")}") + emit(s"${SQLConf.SHUFFLE_PARTITIONS.key}: $ShufflePartitions") + emit( + s"spark.io.compression.codec (Spark shuffle): " + + conf.get("spark.io.compression.codec", "lz4")) + emit( + s"${CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.key} (Comet shuffle): " + + CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.get()) + emit( + s"${SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key}: " + + spark.conf.get(SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key)) + emit( + s"fact: $FactRows rows in $FactFiles files, $HotPercent% of them on one key, 3 nested " + + s"and $WideColumns flat payload columns") + emit(s"dim: $DimRows rows in $DimFiles files, one per key") + emit(s"measured runs per arm: $measuredRuns, arms: ${arms.map(_.name).mkString("; ")}") + } + + private def measured(arm: Arm): Seq[Run] = + runs.get(arm.name).map(_.toList).getOrElse(Nil) + + /** The measured run with the median hot task time. */ + private def medianRun(arm: Arm): Option[Run] = { + val all = measured(arm).sortBy(_.hotTask.runTimeMs) + all.lift(all.length / 2) + } + + /** A line of the comparison: how to read it from a run, and how many decimals to show. */ + private case class Line(label: String, decimals: Int, value: Run => Option[Double]) + + private def emitComparison(): Unit = { + val medians = arms.map(arm => arm -> medianRun(arm)).toMap + def hot(run: Run): Seq[TaskRecord] = Seq(run.hotTask) + def sum(tasks: Seq[TaskRecord])(f: TaskRecord => Long): Option[Double] = + Some(tasks.map(f).sum.toDouble) + val mib = 1024.0 * 1024 + val sections = Seq( + "Hot task, the join stage task that emitted the most rows" -> Seq( + Line("run time (ms)", 0, r => Some(r.hotTask.runTimeMs.toDouble)), + Line("CPU time (ms)", 0, r => Some(r.hotTask.cpuTimeMs.toDouble)), + Line("GC time (ms)", 0, r => Some(r.hotTask.gcTimeMs.toDouble)), + Line("rows emitted", 0, r => Some(r.joinRows(r.hotTask).toDouble)), + Line("shuffle bytes read (MiB)", 1, r => Some(r.hotTask.shuffleReadBytes / mib)), + Line("shuffle records read", 0, r => Some(r.hotTask.shuffleReadRecords.toDouble)), + Line("shuffle fetch wait (ms)", 0, r => Some(r.hotTask.fetchWaitMs.toDouble)), + Line( + "shuffle read operator (ms)", + 1, + r => r.metric(hot(r), "CometExchange", "elapsed_compute")), + Line( + " decode and decompress (ms)", + 1, + r => r.metric(hot(r), "CometExchange", "decode_time")), + Line( + "join build (ms)", + 1, + r => + r.metric(hot(r), "ShuffledHashJoin", "buildTime") + .orElse(r.metric(hot(r), "CometHashJoin", "build_time"))), + Line( + "join probe and output (ms)", + 1, + r => r.metric(hot(r), "CometHashJoin", "join_time")), + Line( + "whole-stage codegen pipeline (ms)", + 0, + r => r.metric(hot(r), "WholeStageCodegen", "pipelineTime"))), + "Map stage that scanned fact and wrote its shuffle" -> Seq( + Line("tasks", 0, r => Some(r.mapStage.size.toDouble)), + Line("task time, total (ms)", 0, r => sum(r.mapStage)(_.runTimeMs)), + Line("longest task (ms)", 0, r => Some(r.mapStage.map(_.runTimeMs).max.toDouble)), + Line("scan input (MiB)", 1, r => sum(r.mapStage)(_.inputBytes).map(_ / mib)), + Line( + "data size before compression (MiB)", + 1, + r => + r.metric(r.mapStage, "Exchange", "dataSize") + .orElse(r.metric(r.mapStage, "CometExchange", "dataSize"))), + Line( + "shuffle bytes written (MiB)", + 1, + r => sum(r.mapStage)(_.shuffleWriteBytes).map(_ / mib)), + Line("shuffle records written", 0, r => sum(r.mapStage)(_.shuffleWriteRecords)), + Line( + "shuffle bytes written per row", + 0, + r => { + val records = r.mapStage.map(_.shuffleWriteRecords).sum + if (records == 0) None + else Some(r.mapStage.map(_.shuffleWriteBytes).sum.toDouble / records) + }), + Line("shuffle write time (ms)", 0, r => sum(r.mapStage)(_.shuffleWriteNs).map(_ / 1e6)), + Line( + "native shuffle writer (ms)", + 0, + r => r.metric(r.mapStage, "CometExchange", "elapsed_compute")), + Line(" repartition (ms)", 0, r => r.metric(r.mapStage, "CometExchange", "repart_time")), + Line( + " encode and compress (ms)", + 0, + r => r.metric(r.mapStage, "CometExchange", "encode_time"))), + "Whole run" -> Seq( + Line("join stage task time, total (ms)", 0, r => sum(r.joinStage)(_.runTimeMs)), + Line("all tasks, total (ms)", 0, r => sum(r.tasks)(_.runTimeMs)))) + + emit("The run with the median hot task time of each engine. Ratio is Spark divided by Comet.") + emit("A dash means the engine has no such metric.") + val width = sections.flatMap(_._2).map(_.label.length).max + 2 + emit(s"%-${width}s %12s %12s %8s".format("", SparkArm.name, CometArm.name, "ratio")) + sections.foreach { case (title, lines) => + emit("") + emit(title) + lines.foreach { line => + val values = arms.map(arm => medians(arm).flatMap(line.value)) + val cells = values.map(_.fold("-")(v => s"%.${line.decimals}f".format(v))) + val ratio = values match { + case Seq(Some(s), Some(c)) if c > 0 => f"${s / c}%.2f" + case _ => "" + } + emit(s" %-${width - 2}s %12s %12s %8s".format(line.label, cells.head, cells(1), ratio)) + } + } + } + + /** The broadcast mode's probe stage: one task per `fact` file, each probing the whole `dim`. */ + private def emitProbeSummary(): Unit = { + emit( + "The join stage probes, one task per fact file. Times in ms, from the run with the median") + emit("join stage time of each engine. Other stages scan dim for the broadcast.") + emit("") + emit( + f"${"arm"}%-6s ${"tasks"}%6s ${"stage total"}%12s ${"mean task"}%10s ${"median task"}%12s " + + f"${"longest"}%8s ${"other stages"}%13s") + arms.foreach { arm => + val all = measured(arm).sortBy(_.joinStage.map(_.runTimeMs).sum) + all.lift(all.length / 2).foreach { run => + val probe = run.joinStage.map(_.runTimeMs).sorted + val other = run.tasks.map(_.runTimeMs).sum - probe.sum + emit( + f"${arm.name}%-6s ${probe.size}%6d ${probe.sum}%12d ${probe.sum.toDouble / probe.size}%10.1f " + + f"${probe(probe.size / 2)}%12d ${probe.last}%8d $other%13d") + } + } + arms.foreach { arm => + val all = measured(arm).sortBy(_.joinStage.map(_.runTimeMs).sum) + all.lift(all.length / 2).foreach { run => + val tasks = run.joinStage + emit("") + emit(s"${arm.name}, mean per probe task over ${tasks.size} tasks:") + run.metrics.foreach { metric => + val values = tasks.flatMap(_.metricUpdates.get(metric.id)) + if (values.nonEmpty) { + val value = scaled(metric.metricType, values.sum) / tasks.size + val text = metric.metricType match { + case "nsTiming" | "timing" => f"$value%.2f ms" + case "size" => f"$value%.2f MiB" + case _ => f"$value%.0f" + } + val operator = s"${metric.nodeName} [${metric.column}]" + emit(f" $text%14s $operator%-28s ${metric.name}") + } + } + } + } + } + + /** Every metric the hot task and the map stage of each engine's median run reported. */ + private def emitMetricListings(): Unit = { + arms.foreach { arm => + medianRun(arm).foreach { run => + Seq("hot task" -> Seq(run.hotTask), "map stage that read fact" -> run.mapStage).foreach { + case (what, tasks) => + emit("") + emit(s"${arm.name}, $what, ${tasks.size} task(s):") + run.metrics.foreach { metric => + val values = tasks.flatMap(_.metricUpdates.get(metric.id)) + if (values.nonEmpty) { + val value = scaled(metric.metricType, values.sum) + val text = metric.metricType match { + case "nsTiming" | "timing" => f"$value%.1f ms" + case "size" => f"$value%.1f MiB" + case _ => f"$value%.0f" + } + val operator = s"${metric.nodeName} [${metric.column}]" + emit(f" $text%14s $operator%-28s ${metric.name}") + } + } + } + } + } + } + + /** + * Runs the hot task of the arm alone, first while sampling the stack of its thread and then for + * an external profiler. One more run of the query writes the shuffle the hot task reads. + */ + private def profile(arm: Arm): Unit = { + val run = medianRun(arm).getOrElse { + throw new IllegalStateException(s"No measured run of '${arm.name}' to profile") + } + val partition = if (profileSmallTask) { + val small = + run.joinStage.filterNot(_.partition == run.hotTask.partition).sortBy(_.runTimeMs) + small(small.size / 2).partition + } else { + run.hotTask.partition + } + val which = if (profileSmallTask) "a median small join task" else "the hot task" + withSQLConf(arm.configs: _*)(withConfsPropagated { + val (counts, _) = countsToRun(spark.sql(query).queryExecution.executedPlan, arm) + counts.fold(0L)(_ + _) + val sampler = new StackSampler + sampler.start() + val sampled = runHotTask(counts, partition) + sampler.finish() + emit("") + emit( + s"${arm.name}: ran $which (partition $partition) $sampled times, sampling the " + + s"stack of its thread every millisecond, ${sampler.samples} samples.") + sampler.emitSummary() + emit(s"PROFILE ${arm.name}: pid ${ProcessHandle.current().pid()}, $ProfileSeconds s") + val unsampled = runHotTask(counts, partition) + emit(s"${arm.name}: ran $which $unsampled times for the external profiler") + }) + } + + /** + * Runs `body` with the session's configs set as local properties of the jobs it submits, as + * Spark does for a query it executes, so that tasks see them. The benchmark runs plans + * directly, which skips that step, and Comet reads some configs, such as the shuffle read + * buffer size, in tasks. + */ + private def withConfsPropagated[T](body: => T): T = { + val sc = spark.sparkContext + val confs = spark.conf.getAll.filter(_._1.startsWith("spark")).toList + val previous = confs.map { case (key, _) => key -> sc.getLocalProperty(key) } + confs.foreach { case (key, value) => sc.setLocalProperty(key, value) } + try body + finally previous.foreach { case (key, value) => sc.setLocalProperty(key, value) } + } + + /** One line per Comet arm of a sweep over the shuffle read buffer size. */ + private def emitReadBufferSweep(): Unit = { + def median(values: Seq[Double]): Option[Double] = + if (values.isEmpty) None else Some(values.sorted.apply(values.length / 2)) + emit(s"Comet only, one arm per setting, run in turn. Medians over $measuredRuns runs, in ms.") + emit( + "Shuffle read is the hot task's shuffle read operator, which Comet reports with AQE off.") + emit("") + val width = arms.map(_.name.length).max + emit( + s"%-${width}s %9s %6s %6s %13s %11s %10s" + .format("arm", "hot task", "min", "max", "shuffle read", "join stage", "all tasks")) + arms.foreach { arm => + val all = measured(arm) + if (all.nonEmpty) { + val hot = all.map(_.hotTask.runTimeMs.toDouble) + val read = median( + all + .flatMap(run => run.metric(Seq(run.hotTask), "CometExchange", "elapsed_compute")) + .filter(_ > 0)) + emit( + s"%-${width}s %9.0f %6.0f %6.0f %13s %11.0f %10.0f".format( + arm.name, + median(hot).get, + hot.min, + hot.max, + read.fold("-")(ms => f"$ms%.0f"), + median(all.map(_.joinStage.map(_.runTimeMs).sum.toDouble)).get, + median(all.map(_.tasks.map(_.runTimeMs).sum.toDouble)).get)) + } + } + } + + /** + * What a stack sample was doing: the first rule that a frame matches, from the top of the stack + * down, so that generic frames such as a buffer copy take the category of their caller. A + * native method on top is the time spent in native code under it. + */ + private val StackRules: Seq[(String, StackTraceElement => Boolean)] = { + def is(cls: String, method: String)(f: StackTraceElement) = + f.getClassName == cls && f.getMethodName == method + def under(prefix: String)(f: StackTraceElement) = f.getClassName.startsWith(prefix) + Seq( + "Comet native plan (join, batch import, project)" -> is( + "org.apache.comet.Native", + "executePlan"), + "Comet native shuffle decode" -> (f => + f.getClassName == "org.apache.comet.Native" && f.getMethodName.startsWith( + "decodeShuffle")), + "Comet native, other" -> under("org.apache.comet.Native"), + "Arrow C JNI (array release, export)" -> under("org.apache.arrow.c.jni."), + "file read (read syscall)" -> (f => + f.isNativeMethod && Set("readBytes", "read0", "pread0").contains(f.getMethodName)), + "FileInputStream.available (fstat, lseek)" -> (f => f.getMethodName == "available0"), + "decompress (lz4-java)" -> under("net.jpountz."), + "Channels.newChannel copy loop" -> under("java.nio.channels.Channels"), + "Comet shuffle block read (direct read)" -> under( + "org.apache.comet.CometShuffleBlockIterator"), + "Comet shuffle block read" -> is( + "org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIterator", + "readNextBlock"), + "Arrow C import into JVM" -> (f => + under("org.apache.arrow.c.")(f) && (f.getClassName + f.getMethodName).contains("mport")), + "Arrow C export to native" -> (f => + under("org.apache.arrow.c.")(f) && (f.getClassName + f.getMethodName).contains("xport")), + "Comet vectors" -> under("org.apache.comet.vector."), + "Comet iterator glue" -> (f => + under("org.apache.comet.CometBatchIterator")(f) || + under("org.apache.comet.CometExecIterator")(f)), + "row deserialization" -> (f => + f.getClassName.contains("UnsafeRowSerializer") || under("java.io.DataInputStream")(f)), + "shuffle fetch and stream wrappers" -> (f => + under("org.apache.spark.storage.")(f) || under("org.apache.spark.network.")(f) || + under("org.apache.spark.shuffle.")(f)), + "join hash lookup" -> under("org.apache.spark.sql.execution.joins."), + "output projection (UnsafeRow writes)" -> (f => + f.getClassName.contains("SpecificUnsafeProjection") || + under("org.apache.spark.sql.catalyst.expressions.codegen.Unsafe")(f)), + "UnsafeRow, array and map reads" -> under( + "org.apache.spark.sql.catalyst.expressions.Unsafe"), + "row count" -> under("org.apache.spark.sql.benchmark."), + "broadcast decode: Arrow IPC read" -> (f => + under("org.apache.arrow.vector.ipc.")(f) || f.getClassName.endsWith( + "ArrowReaderIterator") || + f.getClassName == "org.apache.comet.vector.StreamReader"), + "broadcast decode, other" -> under("org.apache.spark.sql.comet.CometBatchRDD"), + "whole-stage codegen (scan, probe, project)" -> (f => + f.getClassName.contains("GeneratedIteratorForCodegenStage")), + "Parquet scan (JVM)" -> (f => + under("org.apache.spark.sql.execution.datasources.parquet.")(f) || + under("org.apache.parquet.")(f) || under("org.apache.spark.sql.execution.vectorized.")( + f))) + } + + private def stackCategory(frames: Array[StackTraceElement]): String = + frames.iterator + .flatMap(frame => StackRules.find(_._2(frame)).map(_._1)) + .find(_ => true) + .getOrElse("other Java") + + /** + * Samples the stack of every running executor thread every millisecond until `finish`, and + * counts the samples by category and by top frame. `getStackTrace` waits for the next safepoint + * of a thread running Java code, so Java frames carry safepoint bias. A thread in native code + * is sampled where it is. + */ + private class StackSampler extends Thread("hash-join-benchmark-stack-sampler") { + setDaemon(true) + @volatile private var running = true + private val categories = mutable.Map.empty[String, Long].withDefaultValue(0L) + private val tops = mutable.Map.empty[String, Long].withDefaultValue(0L) + var samples = 0L + + override def run(): Unit = { + var threads = Seq.empty[Thread] + var refreshedAt = 0L + while (running) { + if (System.nanoTime() - refreshedAt > 200000000L) { + threads = executorThreads() + refreshedAt = System.nanoTime() + } + threads.filter(_.getState == Thread.State.RUNNABLE).foreach { thread => + val frames = thread.getStackTrace + if (frames.nonEmpty) { + samples += 1 + categories(stackCategory(frames)) += 1 + val top = frames.head + val native = if (top.isNativeMethod) " (native)" else "" + tops(s"${top.getClassName}.${top.getMethodName}$native") += 1 + } + } + Thread.sleep(1) + } + } + + def finish(): Unit = { + running = false + join() + } + + def emitSummary(): Unit = { + def share(n: Long): String = f"${100.0 * n / math.max(samples, 1L)}%5.1f%%" + categories.toSeq.sortBy(-_._2).foreach { case (category, n) => + emit(s" ${share(n)} $category") + } + emit(" Top frames:") + tops.toSeq.sortBy(-_._2).take(20).foreach { case (frame, n) => + emit(s" ${share(n)} $frame") + } + } + } + + private def executorThreads(): Seq[Thread] = { + val threads = Thread.getAllStackTraces.keySet.iterator() + val found = mutable.Buffer.empty[Thread] + while (threads.hasNext) { + val thread = threads.next() + if (thread.getName.startsWith("Executor task launch worker")) { + found += thread + } + } + found.toList + } + + /** Runs the partition of `counts` alone, again and again for `ProfileSeconds`. */ + private def runHotTask(counts: RDD[Long], partition: Int): Int = { + val end = System.nanoTime() + ProfileSeconds * 1000000000L + var times = 0 + while (System.nanoTime() < end) { + spark.sparkContext.runJob(counts, (rows: Iterator[Long]) => rows.sum, Seq(partition)) + times += 1 + } + times + } + + /** Sets a warning off from the lines around it, in the console and the results file. */ + private def warning(message: String): Unit = { + val border = "=" * 80 + emit(s"\n$border\n$message\n$border") + } + + /** Writes a line to the console and, when results files are generated, to the results file. */ + private def emit(line: String): Unit = { + // scalastyle:off println + println(line) + // scalastyle:on println + output.foreach(_.write(s"$line\n".getBytes(StandardCharsets.UTF_8))) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala new file mode 100644 index 00000000000..c833e21721d --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala @@ -0,0 +1,380 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import java.io.File +import java.nio.charset.StandardCharsets +import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.locks.LockSupport + +import scala.collection.mutable + +import org.apache.spark.{SparkConf, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometNativeExec} +import org.apache.spark.sql.execution.{InputAdapter, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Measures how long a Comet reduce task keeps running after it is killed while it reads its + * shuffle (review of apache/datafusion-comet#6805). + * + * One map task writes the whole shuffle as one block, and one reduce task reads it with a native + * project on top. On the `aqe` path the project reads the shuffle directly in native code + * (`ShuffleScan` over `CometShuffleBlockIterator`), which Comet does only behind an AQE query + * stage. On the `noaqe` path it reads it through `NativeBatchDecoderIterator`. + * + * Each trial runs the reduce task alone over the same shuffle. A kill trial cancels the job + * `KillAfterMs` after the task starts, with or without interrupting the task's thread, as + * `spark.job.interruptOnCancel` selects, and measures how long the task runs after the cancel, + * until its completion listeners run. A trial without a kill measures the whole task. + * + * With `variants=fix,head`, the default, the variants run in turn in one JVM. `head` sets the + * local property `spark.comet.benchmark.skipKillCheck`, which only a temporary benchmark build + * reads: it makes `readAsRawStream` skip its kill check, as #6805 did before the fix. Any other + * variant name only labels the classes on the classpath. `paths=aqe,noaqe` selects the read + * paths, `rows=N` the rows of the shuffle and `trials=N` the trials. + */ +object CometShuffleReadKillBenchmark extends CometBenchmarkBase { + + // `SqlBasedBenchmark` calls `getSparkSession` while it initializes, before the fields of this + // object are assigned, so the field that method reads is a compile-time constant. + private final val Cores = 4 + + /** Rows of the shuffle. Its strings repeat, so it is small on disk but slow to decode. */ + private var rows = 16L * 1024 * 1024 + + /** Trials per variant and mode, set by `trials=N`. */ + private var trials = 3 + + private val KillAfterMs = 1000L + + private val SkipKillCheckKey = "spark.comet.benchmark.skipKillCheck" + + private case class Mode(name: String, kill: Boolean, interrupt: Boolean) + + private val Modes = Seq( + Mode("not killed: task time", kill = false, interrupt = false), + Mode("killed, thread interrupted: time after the kill", kill = true, interrupt = true), + Mode("killed, thread not interrupted: time after the kill", kill = true, interrupt = false)) + + /** Counts task ends and keeps the shuffle sizes and the end reason of the last one. */ + private object TaskRecorder extends SparkListener { + val ends = new AtomicLong(0L) + @volatile var lastReason = "" + @volatile var shuffleWritten = 0L + @volatile var shuffleRead = 0L + + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + Option(taskEnd.taskMetrics).foreach { metrics => + if (metrics.shuffleWriteMetrics.bytesWritten > 0) { + shuffleWritten = metrics.shuffleWriteMetrics.bytesWritten + } + if (metrics.shuffleReadMetrics.totalBytesRead > 0) { + shuffleRead = metrics.shuffleReadMetrics.totalBytesRead + } + } + lastReason = taskEnd.reason.toString.takeWhile(_ != '(') + ends.incrementAndGet() + } + } + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometShuffleReadKillBenchmark") + // Since `spark.master` always exists, overrides this value + .set("spark.master", s"local[$Cores]") + .setIfMissing("spark.driver.memory", "3g") + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.memory.offHeap.enabled", "true") + .setIfMissing("spark.memory.offHeap.size", "8g") + + val sparkSession = SparkSession + .builder() + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + + sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.SHUFFLE_PARTITIONS.key, "1") + sparkSession.conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED.key, "false") + // An open cost as large as the largest split keeps the one file in one map task. + val maxSplitBytes = (1024L * 1024 * 1024).toString + sparkSession.conf.set(SQLConf.FILES_MAX_PARTITION_BYTES.key, maxSplitBytes) + sparkSession.conf.set(SQLConf.FILES_OPEN_COST_IN_BYTES.key, maxSplitBytes) + sparkSession + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + def listArg(name: String, default: Seq[String]): Seq[String] = + mainArgs + .collectFirst { + case arg if arg.startsWith(s"$name=") => arg.stripPrefix(s"$name=").split(",").toSeq + } + .getOrElse(default) + val variants = listArg("variants", Seq("fix", "head")) + val paths = listArg("paths", Seq("aqe", "noaqe")) + listArg("rows", Nil).headOption.foreach(value => rows = value.toLong) + listArg("trials", Nil).headOption.foreach(value => trials = value.toInt) + spark.sparkContext.addSparkListener(TaskRecorder) + try { + withTempPath { dir => + withTempTable("src") { + writeSource(dir) + runBenchmark("Killing a Comet reduce task while it reads its shuffle") { + val conf = spark.sparkContext.getConf + emit( + s"Spark ${spark.version}, Java ${System.getProperty("java.version")}, " + + s"master ${conf.get("spark.master")}") + emit( + s"src: $rows rows in 1 Parquet file, so 1 map task writes the shuffle as 1 block " + + "for 1 reduce task") + emit( + s"variants: ${variants.mkString(", ")}. $trials trials per variant and mode, " + + s"kills $KillAfterMs ms after the task starts. Times in ms, median (trials).") + paths.foreach(path => measure(path, variants)) + } + } + } + } catch { + case e: Throwable => + // `BenchmarkBase.main` stops the session only after a run that returned. + spark.stop() + throw e + } + } + + /** Writes `src` with Comet disabled and checks that it scans as one map task. */ + private def writeSource(dir: File): Unit = { + val path = new File(dir, "src").getCanonicalPath + spark + .range(0L, rows, 1L, 1) + .selectExpr( + "id", + "XXHASH64(id) AS h", + "REPEAT(CONCAT('k', CAST(PMOD(id, 97) AS STRING), '|'), 24) AS s1", + "REPEAT(CONCAT('v', CAST(PMOD(id, 89) AS STRING), ';'), 24) AS s2") + .write + .option("compression", "snappy") + .parquet(path) + spark.read.parquet(path).createOrReplaceTempView("src") + val tasks = spark.table("src").queryExecution.toRdd.getNumPartitions + require(tasks == 1, s"src scans as $tasks map tasks, expected 1") + } + + /** Writes the shuffle of one read path, then runs and reports its trials. */ + private def measure(path: String, variants: Seq[String]): Unit = { + val adaptive = path == "aqe" + val results = mutable.LinkedHashMap.empty[(String, String), mutable.Buffer[Option[Double]]] + val reasons = mutable.LinkedHashMap.empty[(String, String), mutable.Set[String]] + var description = "" + System.gc() + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true") { + withConfsPropagated { + val executed = spark + .table("src") + .repartition(1, col("id")) + .selectExpr("id + 1 AS a", "h", "LENGTH(s1) + LENGTH(s2) AS l") + .queryExecution + .executedPlan + // Under AQE, Comet's conversion to rows sits inside the adaptive plan. The query runs + // once to materialize its shuffle stage, and the native part of the final plan then runs + // alone. Without AQE, the native part runs alone from the start. + val finalPlan = executed match { + case adaptivePlan: AdaptiveSparkPlanExec => + adaptivePlan.execute().foreach(_ => ()) + adaptivePlan.executedPlan + case plan => plan + } + val native = nativeChild(finalPlan) + val direct = CometExec.findShuffleScanIndices(native.nativeOp).size + require(direct == (if (adaptive) 1 else 0), s"$direct direct reads in:\n$finalPlan") + val counts = native.executeColumnar().map(_.numRows().toLong) + // Without AQE, this run also writes the shuffle. + val counted = counts.fold(0L)(_ + _) + require(counted == rows, s"counted $counted rows, expected $rows") + spark.sparkContext.listenerBus.waitUntilEmpty() + val mib = 1024.0 * 1024 + description = + (if (adaptive) "AQE on, the project reads the shuffle directly (ShuffleScan)" + else "AQE off, the project reads the shuffle through NativeBatchDecoderIterator") + + f", shuffle ${TaskRecorder.shuffleWritten / mib}%.0f MiB written, " + + f"${TaskRecorder.shuffleRead / mib}%.0f MiB read by the reduce task" + emit("") + emit(description) + emit(native.treeString.linesIterator.map(" " + _).mkString("\n")) + for (_ <- 1 to trials; mode <- Modes; variant <- variants) { + val (result, reason) = trial(counts, variant, mode) + results.getOrElseUpdate((variant, mode.name), mutable.Buffer.empty) += result + reasons.getOrElseUpdate((variant, mode.name), mutable.Set.empty) += reason + } + } + } + val width = Modes.map(_.name.length).max + emit("") + emit(s"%-${width}s".format("") + variants.map(v => "%24s".format(v)).mkString) + Modes.foreach { mode => + val cells = variants.map { variant => + val values = results((variant, mode.name)).toList + val measured = values.flatten.sorted + val median = measured.lift(measured.size / 2).fold("-")(ms => f"$ms%.0f") + val each = values.map(_.fold("ended first")(ms => f"$ms%.0f")).mkString(", ") + "%24s".format(s"$median ($each)") + } + emit(s"%-${width}s".format(mode.name) + cells.mkString) + } + emit("Task end reasons:") + reasons.foreach { case ((variant, mode), seen) => + emit(s" $variant, $mode: ${seen.mkString(", ")}") + } + } + + /** + * The native plan under the Comet columnar-to-row conversion of `plan`. `collectFirst` of + * `AdaptiveSparkPlanHelper` looks into query stages, such as AQE's final `ResultQueryStage`. + * Under whole-stage codegen an `InputAdapter`, which tree strings do not show, sits between. + */ + private def nativeChild(plan: SparkPlan): CometNativeExec = + collectFirst(plan) { + case conversion: CometColumnarToRowExec => conversion.child + case conversion: CometNativeColumnarToRowExec => conversion.child + } + .map { + case adapter: InputAdapter => adapter.child + case child => child + } + .collect { case native: CometNativeExec => native } + .getOrElse { + throw new IllegalStateException(s"No native plan under a conversion to rows in:\n$plan") + } + + /** + * Runs the reduce task once. Returns how long it ran after the kill, or in total without one, + * in ms, with the reason its task ended. The time is empty when the task ended before the kill. + */ + private def trial(counts: RDD[Long], variant: String, mode: Mode): (Option[Double], String) = { + val sc = spark.sparkContext + // A collection now keeps one out of the time measured after a kill. + System.gc() + CometShuffleReadKillTaskClock.started.set(0L) + CometShuffleReadKillTaskClock.ended.set(0L) + val endsBefore = TaskRecorder.ends.get() + val group = s"kill-benchmark-${System.nanoTime()}" + val job = new Thread( + () => { + sc.setJobGroup(group, s"$variant, ${mode.name}", interruptOnCancel = mode.interrupt) + sc.setLocalProperty(SkipKillCheckKey, if (variant == "head") "true" else null) + try { + sc.runJob( + counts, + (context: TaskContext, batches: Iterator[Long]) => { + context.addTaskCompletionListener[Unit] { _ => + CometShuffleReadKillTaskClock.ended.set(System.nanoTime()) + } + CometShuffleReadKillTaskClock.started.set(System.nanoTime()) + batches.sum + }, + Seq(0)) + } catch { + // A cancelled job fails here, as expected. + case _: Throwable => + } + }, + "kill-benchmark-job") + job.setDaemon(true) + job.start() + val start = await(CometShuffleReadKillTaskClock.started) + val result = if (mode.kill) { + val killAt = start + KillAfterMs * 1000000L + while (System.nanoTime() < killAt) { + LockSupport.parkNanos(math.min(killAt - System.nanoTime(), 1000000L)) + } + val killedAt = System.nanoTime() + if (CometShuffleReadKillTaskClock.ended.get() != 0L) { + None + } else { + sc.cancelJobGroup(group) + Some((await(CometShuffleReadKillTaskClock.ended) - killedAt) / 1e6) + } + } else { + Some((await(CometShuffleReadKillTaskClock.ended) - start) / 1e6) + } + job.join() + // The task end event follows the executor's final status update, so the next trial does not + // overlap the cleanup of this one. + while (TaskRecorder.ends.get() == endsBefore) { + LockSupport.parkNanos(1000000L) + } + (result, TaskRecorder.lastReason) + } + + /** Waits until `clock` is set and returns it. */ + private def await(clock: AtomicLong): Long = { + val deadline = System.nanoTime() + 300L * 1000000000L + while (clock.get() == 0L) { + require(System.nanoTime() < deadline, "the task did not start or end within 300 s") + LockSupport.parkNanos(100000L) + } + clock.get() + } + + /** + * Runs `body` with the session's configs set as local properties of the jobs it submits, as + * Spark does for a query it executes, so that tasks see them. The benchmark runs plans + * directly, which skips that step. + */ + private def withConfsPropagated[T](body: => T): T = { + val sc = spark.sparkContext + val confs = spark.conf.getAll.filter(_._1.startsWith("spark")).toList + val previous = confs.map { case (key, _) => key -> sc.getLocalProperty(key) } + confs.foreach { case (key, value) => sc.setLocalProperty(key, value) } + try body + finally previous.foreach { case (key, value) => sc.setLocalProperty(key, value) } + } + + /** Writes a line to the console and, when results files are generated, to the results file. */ + private def emit(line: String): Unit = { + // scalastyle:off println + println(line) + // scalastyle:on println + output.foreach(_.write(s"$line\n".getBytes(StandardCharsets.UTF_8))) + } +} + +/** + * When the task of a trial started, and when its completion listeners ran, in `nanoTime`. A + * top-level object, so that the task's closure reads it without capturing the benchmark. + */ +private object CometShuffleReadKillTaskClock { + val started = new AtomicLong(0L) + val ended = new AtomicLong(0L) +} 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 0462bd70421..0e618b96088 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) From 1346697f8c83680c738f8270181a9e5996b178d3 Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 10 Oct 2026 11:00:38 -0700 Subject: [PATCH 5/5] 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. --- .../CometHashJoinTaskTimeBenchmark.scala | 1084 ----------------- .../CometShuffleReadKillBenchmark.scala | 380 ------ 2 files changed, 1464 deletions(-) delete mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala delete mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala deleted file mode 100644 index c1cc5259f33..00000000000 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinTaskTimeBenchmark.scala +++ /dev/null @@ -1,1084 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.spark.sql.benchmark - -import java.io.File -import java.nio.charset.StandardCharsets -import java.util.Locale - -import scala.collection.mutable - -import org.apache.spark.{SparkConf, Success} -import org.apache.spark.executor.TaskMetrics -import org.apache.spark.rdd.RDD -import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometColumnarToRowExec, CometExec, CometHashJoinExec, CometNativeColumnarToRowExec, CometNativeExec} -import org.apache.spark.sql.execution.{SparkPlan, WholeStageCodegenExec} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.comet.{CometConf, CometSparkSessionExtensions} - -/** - * Compares Spark and Comet on the one reduce task of a skewed shuffled hash join that receives - * 95% of the probe rows, each row carrying nested and wide columns - * (apache/datafusion-comet#6528). - * - * The join builds `dim` and probes with `fact`, which puts 95% of its rows on one key. AQE is - * off, so the hot partition is not split, and the one reduce task that receives it reads, joins - * and emits about 95% of the rows. Every `fact` row carries a struct holding an array of structs, - * an array of strings, a map, and `WideColumns` flat columns. - * - * Both engines run the same query over the same files. Spark's output rows are counted. Comet's - * plan ends in a columnar-to-row conversion, which has no metric of its own, so the child of that - * conversion runs instead and the rows of its batches are counted. Neither engine writes its - * output anywhere. - * - * Task time is the executor run time of each task, read from task end events with the task's - * shuffle and input metrics and its SQL metric updates. The comparison shows the run with the - * median hot task time of each engine: the hot task, the map stage that scanned `fact` and wrote - * its shuffle, and the whole run. Executor CPU time is shown for reference only: Comet runs part - * of its native work on Tokio threads, which a task's CPU time does not count. After the - * comparison, every metric of the hot task and of the map stage is listed by operator. Metrics of - * one operator can contain each other, so they are listed but not added up. - * - * To run this benchmark: - * {{{ - * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometHashJoinTaskTimeBenchmark - * }}} - * Results will be written to "spark/benchmarks/CometHashJoinTaskTimeBenchmark-**results.txt". - * - * With the argument `-- profile`, each engine then runs the hot task alone, reusing the shuffle - * of one more run, for two windows of `ProfileSeconds` each. In the first, a thread samples the - * stack of the hot task's thread every millisecond, and the samples are counted by what they were - * doing. A native method on top of a stack stands for all the native code under it. The second - * window is left for an external profiler of native code, such as `sample` on macOS, and starts - * when a line beginning with "PROFILE" names the process id. - * - * With the argument `-- aqe`, AQE is on, with skew join handling and partition coalescing off, so - * that the hot partition stays whole and the join keeps 32 tasks. Comet then reads the shuffle - * directly in native code, which it does only behind AQE query stages. Its conversion to rows - * sits inside the adaptive plan, so each Comet run executes the query once to materialize the - * shuffle stages and then runs the native part of the final plan alone. The run reports the map - * stages of the first execution and the join stage of the second. - * - * A `-Dspark.*` system property, for example in `JAVA_TOOL_OPTIONS`, sets a session config, which - * is how to measure a feature flag of a pull request. - */ -object CometHashJoinTaskTimeBenchmark extends CometBenchmarkBase { - - // `SqlBasedBenchmark` calls `getSparkSession` while it initializes, before the fields of this - // object are assigned. The fields that method reads are therefore `final` with literal values, - // which makes them compile-time constants. - - /** Local cores. The hot task costs wall time only while the other cores sit idle. */ - private final val Cores = 8 - - /** Reduce tasks of the join. The hot task gets the hot key and its share of the other keys. */ - private final val ShufflePartitions = 32 - - /** Rows of `fact`, and so the rows the join emits. */ - private val FactRows = 2 * 1024 * 1024 - - /** The share of `fact` rows on the hot key. */ - private val HotPercent = 95 - - /** Parquet files, and so map tasks, of each table. */ - private val FactFiles = 32 - private val DimFiles = 8 - - /** Rows of `dim`, one per key, so every `fact` row matches exactly one of them. */ - private val DimRows = 256 * 1024 - - /** Flat columns of the payload, next to its three nested ones. */ - private val WideColumns = 100 - - /** Measured runs per engine, set by `runs=N`. Odd, so that the median is one of the runs. */ - private var measuredRuns = 5 - - /** Length of each profiling window. */ - private val ProfileSeconds = 25 - - /** - * Rows of `dim` in the broadcast mode, the largest build side of the earlier broadcast cases. - */ - private val BroadcastDimRows = 1024 * 1024 - - /** Keys of `fact` in the broadcast mode. Each is in `dim`, so every row matches once. */ - private val BroadcastKeys = 64 * 1024 - - private def query: String = - if (broadcast) { - "SELECT /*+ BROADCAST(d) */ f.key, f.f_long, f.f_double, f.f_str, d.d_long, d.d_str " + - "FROM fact f JOIN dim d ON f.key = d.k" - } else { - val payload = Seq("n_order", "n_tags", "n_attrs") ++ (0 until WideColumns).map(i => s"w_$i") - val columns = ("f.key" +: payload.map(c => s"f.$c")) ++ Seq("d.d_long", "d.d_str") - s"SELECT /*+ SHUFFLE_HASH(d) */ ${columns.mkString(", ")} " + - "FROM fact f JOIN dim d ON f.key = d.k" - } - - private case class Arm( - name: String, - comet: Boolean, - extraConfigs: Seq[(String, String)] = Nil) { - def configs: Seq[(String, String)] = Seq( - CometConf.COMET_ENABLED.key -> comet.toString, - CometConf.COMET_EXEC_ENABLED.key -> comet.toString) ++ extraConfigs - } - - private val SparkArm = Arm("Spark", comet = false) - private val CometArm = Arm("Comet", comet = true) - - /** The arms of the run: Spark and Comet, or one Comet arm per read buffer size of a sweep. */ - private var arms = Seq(SparkArm, CometArm) - - /** - * One task of a run, with its task metrics and what it added to each accumulator, keyed by - * accumulator id, which for a SQL metric is `SQLMetric.id`. `partition` is the task's index in - * its stage. - */ - private case class TaskRecord( - stageId: Int, - partition: Int, - runTimeMs: Long, - cpuTimeMs: Long, - gcTimeMs: Long, - inputBytes: Long, - shuffleReadBytes: Long, - shuffleReadRecords: Long, - fetchWaitMs: Long, - shuffleWriteBytes: Long, - shuffleWriteRecords: Long, - shuffleWriteNs: Long, - failed: Boolean, - metricUpdates: Map[Long, Long]) - - /** - * A SQL metric of the executed plan, which a task reports under the metric's accumulator id. - * `column` is the operator's first output column, which tells the two sides of a join apart. - */ - private case class PlanMetric( - id: Long, - nodeName: String, - column: String, - key: String, - name: String, - metricType: String) - - /** - * One run of the query. `rows` is the row count the engine produced, `joinRowsId` the - * accumulator of the join's output row count. A run keeps the plan's metrics rather than the - * plan, because a plan holds its shuffle dependencies, and Spark deletes a shuffle's files only - * once nothing references them. - */ - private case class Run( - tasks: Seq[TaskRecord], - rows: Long, - metrics: Seq[PlanMetric], - joinRowsId: Option[Long], - wholeStageJoin: Boolean) { - private val lastStage = tasks.map(_.stageId).max - - /** The tasks of the last stage of the run, the result stage that runs the join. */ - def joinStage: Seq[TaskRecord] = tasks.filter(_.stageId == lastStage) - - /** The tasks of the map stage that wrote the most shuffle bytes, the one that read `fact`. */ - def mapStage: Seq[TaskRecord] = - tasks - .filter(_.stageId != lastStage) - .groupBy(_.stageId) - .values - .maxBy(_.map(_.shuffleWriteBytes).sum) - - def joinRows(task: TaskRecord): Long = - joinRowsId.flatMap(task.metricUpdates.get).getOrElse(0L) - - /** The join stage task that emitted the most rows, or the longest one without a count. */ - def hotTask: TaskRecord = joinStage.maxBy(task => (joinRows(task), task.runTimeMs)) - - /** - * What `of` reported for the metric `key` of the operators whose name starts with `nodeName`, - * in ms for a timing, MiB for a size and as is otherwise. Empty when the plan has no such - * metric. - */ - def metric(of: Seq[TaskRecord], nodeName: String, key: String): Option[Double] = { - val matching = metrics.filter(m => m.nodeName.startsWith(nodeName) && m.key == key) - matching.headOption.map { first => - val total = of.map(task => matching.flatMap(m => task.metricUpdates.get(m.id)).sum).sum - scaled(first.metricType, total) - } - } - } - - /** A metric value in ms for a timing, MiB for a size and as is otherwise. */ - private def scaled(metricType: String, value: Long): Double = metricType match { - case "nsTiming" => value / 1e6 - case "size" => value / (1024.0 * 1024) - case _ => value.toDouble - } - - /** Collects every task end. The listener bus delivers them on its own thread. */ - private object TaskRecorder extends SparkListener { - private val records = mutable.ArrayBuffer.empty[TaskRecord] - - override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { - def metric(f: TaskMetrics => Long): Long = Option(taskEnd.taskMetrics).map(f).getOrElse(0L) - val updates = taskEnd.taskInfo.accumulables.flatMap { info => - info.update.collect { case value: Long => info.id -> value } - }.toMap - val record = TaskRecord( - stageId = taskEnd.stageId, - partition = taskEnd.taskInfo.index, - runTimeMs = metric(_.executorRunTime), - cpuTimeMs = metric(_.executorCpuTime) / 1000000, - gcTimeMs = metric(_.jvmGCTime), - inputBytes = metric(_.inputMetrics.bytesRead), - shuffleReadBytes = metric(_.shuffleReadMetrics.totalBytesRead), - shuffleReadRecords = metric(_.shuffleReadMetrics.recordsRead), - fetchWaitMs = metric(_.shuffleReadMetrics.fetchWaitTime), - shuffleWriteBytes = metric(_.shuffleWriteMetrics.bytesWritten), - shuffleWriteRecords = metric(_.shuffleWriteMetrics.recordsWritten), - shuffleWriteNs = metric(_.shuffleWriteMetrics.writeTime), - failed = taskEnd.reason != Success, - metricUpdates = updates) - synchronized { records += record } - } - - def drain(): List[TaskRecord] = synchronized { - val drained = records.toList - records.clear() - drained - } - } - - // Keyed by arm name. - private val runs = mutable.Map.empty[String, mutable.Buffer[Run]] - - /** Whether the run uses AQE, set by the `aqe` argument. */ - private var adaptive = false - - /** Whether the run joins by broadcast, set by the `bhj` argument. */ - private var broadcast = false - - /** Whether profiling follows a median small join task instead of the hot one. */ - private var profileSmallTask = false - - override def getSparkSession: SparkSession = { - val conf = new SparkConf() - .setAppName("CometHashJoinTaskTimeBenchmark") - // Since `spark.master` always exists, overrides this value - .set("spark.master", s"local[$Cores]") - .setIfMissing("spark.driver.memory", "3g") - .setIfMissing("spark.executor.memory", "3g") - .set( - "spark.shuffle.manager", - "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") - // Comet requires off-heap memory in production. Both engines then draw from the same pool. - .set("spark.memory.offHeap.enabled", "true") - .setIfMissing("spark.memory.offHeap.size", "8g") - - val sparkSession = SparkSession - .builder() - .config(conf) - .withExtensions(new CometSparkSessionExtensions) - .getOrCreate() - - sparkSession.conf.set(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key, "true") - sparkSession.conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") - sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") - sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") - sparkSession.conf.set(SQLConf.ANSI_ENABLED.key, "false") - // Keeps the hot partition whole. AQE would split it, and Comet then leaves the join to Spark - // (apache/datafusion-comet#6530). - sparkSession.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") - sparkSession.conf.set(SQLConf.SHUFFLE_PARTITIONS.key, ShufflePartitions.toString) - // The join carries a hint, so it cannot become a broadcast join. - sparkSession.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") - // An open cost as large as the largest split closes a scan partition after every file, so - // each Parquet file is exactly one map task. `prepareTables` checks that it worked. - val maxSplitBytes = (128L * 1024 * 1024).toString - sparkSession.conf.set(SQLConf.FILES_MAX_PARTITION_BYTES.key, maxSplitBytes) - sparkSession.conf.set(SQLConf.FILES_OPEN_COST_IN_BYTES.key, maxSplitBytes) - - sparkSession - } - - override def runCometBenchmark(mainArgs: Array[String]): Unit = { - val args = mainArgs.map(_.toLowerCase(Locale.ROOT)) - val profiling = args.exists(_.startsWith("profile")) - profileSmallTask = args.contains("profile=small") - adaptive = args.contains("aqe") - broadcast = args.contains("bhj") - if (adaptive) { - // Comet reads a shuffle directly in native code only behind an AQE query stage. Skew join - // handling stays off, because splitting the hot partition makes Comet leave the join to - // Spark (apache/datafusion-comet#6530), and so does coalescing, which keeps 32 join tasks. - spark.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true") - spark.conf.set(SQLConf.SKEW_JOIN_ENABLED.key, "false") - spark.conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED.key, "false") - } - val readBufferSizes = mainArgs.collectFirst { - case arg if arg.startsWith("readBufferSize=") => - arg.stripPrefix("readBufferSize=").split(",").toSeq - } - readBufferSizes.foreach { sizes => - arms = sizes.map { size => - Arm( - s"Comet, $size", - comet = true, - Seq(CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.key -> size)) - } - } - // An A/B of the kill check in `readAsRawStream` (review of #6805). Only a temporary build - // that reads this local property skips the check. - if (args.contains("killcheck")) { - arms = Seq( - Arm("Comet, kill check", comet = true), - Arm( - "Comet, no kill check", - comet = true, - Seq("spark.comet.benchmark.skipKillCheck" -> "true"))) - } - mainArgs - .collectFirst { case arg if arg.startsWith("runs=") => arg.stripPrefix("runs=").toInt } - .foreach(runs => measuredRuns = runs) - val cometSweep = readBufferSizes.isDefined || args.contains("killcheck") - spark.sparkContext.addSparkListener(TaskRecorder) - try { - withTempPath { dir => - withTempTable("fact", "dim") { - prepareTables(dir) - runBenchmark("Skewed hash join, Spark vs Comet: environment") { - emitEnvironment() - } - runBenchmark("Skewed hash join, Spark vs Comet: runs") { - emit("One untimed run of each engine, checked before the measured runs.") - arms.foreach(inspect) - for (_ <- 1 to measuredRuns; arm <- arms) { - runs.getOrElseUpdate(arm.name, mutable.Buffer.empty) += execute(arm)._1 - } - arms.foreach { arm => - val times = measured(arm).map(_.hotTask.runTimeMs).sorted - emit(s"${arm.name} hot task, every measured run (ms): ${times.mkString(", ")}") - } - } - if (cometSweep) { - runBenchmark("Skewed hash join, Comet: one arm per setting") { - emitReadBufferSweep() - } - } else if (broadcast) { - runBenchmark("Broadcast hash join, Spark vs Comet: probe tasks") { - emitProbeSummary() - } - } else { - runBenchmark("Skewed hash join, Spark vs Comet: comparison") { - emitComparison() - } - runBenchmark("Skewed hash join, Spark vs Comet: metrics by operator") { - emitMetricListings() - } - } - if (profiling && !cometSweep) { - runBenchmark("Skewed hash join, Spark vs Comet: profiling") { - arms.foreach(profile) - } - } - } - } - } catch { - case e: Throwable => - // `BenchmarkBase.main` stops the session only after a run that returned. Under - // `exec:java` the classloader is closed before the JVM's shutdown hooks run, so Spark's - // own cleanup hook fails too, and a failed run would leave its shuffle files on disk. - spark.stop() - throw e - } - } - - /** Writes both tables with Comet disabled and checks that each file is one map task. */ - private def prepareTables(dir: File): Unit = { - // The hot rows are every row whose id falls in the first `HotPercent` of each hundred, so - // every map task contributes to the hot partition. The other rows hash onto keys 1 to - // `DimRows - 1`. - val key = s"CAST(IF(PMOD(id, 100) < $HotPercent, 0, 1 + PMOD(HASH(id), ${DimRows - 1})) " + - "AS BIGINT) AS key" - val order = "NAMED_STRUCT(" + - "'id', id, " + - "'customer', NAMED_STRUCT(" + - "'name', CONCAT('customer-', CAST(PMOD(id, 100000) AS STRING)), " + - "'tier', CAST(PMOD(id, 5) AS INT)), " + - "'items', TRANSFORM(SEQUENCE(0, CAST(PMOD(id, 3) AS INT)), " + - "i -> NAMED_STRUCT('sku', id * 10 + i, 'qty', i + 1)))" - val nested = Seq( - s"$order AS n_order", - "TRANSFORM(SEQUENCE(0, CAST(PMOD(id, 3) AS INT)), " + - "i -> CONCAT('tag-', CAST(PMOD(id + i, 50) AS STRING))) AS n_tags", - "MAP('views', PMOD(id, 1000), 'clicks', PMOD(id, 97)) AS n_attrs") - // BIGINT, DOUBLE, STRING and INT in turn, with values hashed from the row, so that neither - // Parquet nor the shuffle compresses them like a sequence. - val wide = (0 until WideColumns).map { i => - val value = i % 4 match { - case 0 => s"XXHASH64(id, $i)" - case 1 => s"CAST(HASH(id, $i) AS DOUBLE) / 7" - case 2 => s"CONCAT('v-', CAST(PMOD(HASH(id, $i), 1000000) AS STRING))" - case _ => s"HASH(id, $i)" - } - s"$value AS w_$i" - } - val dim = Seq("id AS k", "id * 3 AS d_long", "CONCAT('dim-', CAST(id AS STRING)) AS d_str") - - val broadcastFact = Seq( - s"CAST(PMOD(HASH(id), $BroadcastKeys) AS BIGINT) AS key", - "id AS f_long", - "CAST(id AS DOUBLE) / 7 AS f_double", - "CONCAT('order-', CAST(id AS STRING)) AS f_str") - val (factColumns, dimRows) = - if (broadcast) (broadcastFact, BroadcastDimRows) else (key +: (nested ++ wide), DimRows) - Seq(("fact", FactRows, FactFiles, factColumns), ("dim", dimRows, DimFiles, dim)) - .foreach { case (table, rows, files, columns) => - val path = new File(dir, table).getCanonicalPath - spark - .range(0L, rows.toLong, 1L, files) - .selectExpr(columns: _*) - .write - .option("compression", "snappy") - .parquet(path) - spark.read.parquet(path).createOrReplaceTempView(table) - val tasks = spark.table(table).queryExecution.toRdd.getNumPartitions - require( - tasks == files, - s"$table scans as $tasks map tasks, expected one per file ($files)") - } - } - - /** - * One row count per partition of Spark's output, or per batch of Comet's. Comet's plan ends in - * a columnar-to-row conversion, whose child produces the batches. - */ - private def rowCounts(plan: SparkPlan, arm: Arm): RDD[Long] = - if (arm.comet) { - val nativeChild = collectFirst(plan) { - case conversion: CometColumnarToRowExec => conversion.child - case conversion: CometNativeColumnarToRowExec => conversion.child - }.getOrElse { - throw new IllegalStateException(s"No Comet columnar-to-row conversion in:\n$plan") - } - nativeChild.executeColumnar().map(_.numRows().toLong) - } else { - plan.execute().mapPartitions(rows => Iterator(rows.size.toLong)) - } - - private def joinOf(plan: SparkPlan): Option[SparkPlan] = collectFirst(plan) { - case join: ShuffledHashJoinExec => join - case join: CometHashJoinExec => join - case join: BroadcastHashJoinExec => join - case join: CometBroadcastHashJoinExec => join - } - - /** - * Runs the query once under the arm and returns the tasks of the run with the executed plan. - */ - private def execute(arm: Arm): (Run, SparkPlan) = { - // `ContextCleaner` deletes a shuffle only once its dependency is garbage collected, which a - // large heap can postpone until the shuffle files fill the disk. Collecting before each run - // releases the shuffle of the run before, and no collection lands in a measured task. - System.gc() - spark.sparkContext.listenerBus.waitUntilEmpty() - TaskRecorder.drain() - var plan: Option[SparkPlan] = None - var rows = 0L - var materializing: Seq[TaskRecord] = Nil - withSQLConf(arm.configs: _*) { - withConfsPropagated { - val executed = spark.sql(query).queryExecution.executedPlan - val (counts, firstRun) = countsToRun(executed, arm) - materializing = firstRun - rows = counts.fold(0L)(_ + _) - plan = Some(executed) - } - } - spark.sparkContext.listenerBus.waitUntilEmpty() - val executedPlan = plan.get - val join = joinOf(executedPlan) - val joinRowsId = join - .flatMap(op => op.metrics.get("numOutputRows").orElse(op.metrics.get("output_rows"))) - .map(_.id) - val wholeStageJoin = join.exists { op => - collect(executedPlan) { case stage: WholeStageCodegenExec => stage } - .exists(_.find(_ eq op).isDefined) - } - val tasks = materializing ++ TaskRecorder.drain() - val run = Run(tasks, rows, planMetrics(executedPlan), joinRowsId, wholeStageJoin) - (run, executedPlan) - } - - /** - * The row counts to run for `executed`, after running whatever must come first. Under AQE, - * Comet's conversion to rows sits inside the adaptive plan. The query then runs once, which - * materializes its shuffle stages, and the counts come from the native part of the final plan, - * which reads those stages again. The tasks of that first run are returned without its final - * stage, so that a run has the map stages of the first run and the join stage of the second. - */ - private def countsToRun(executed: SparkPlan, arm: Arm): (RDD[Long], Seq[TaskRecord]) = - executed match { - case adaptivePlan: AdaptiveSparkPlanExec if arm.comet => - adaptivePlan.execute().foreach(_ => ()) - spark.sparkContext.listenerBus.waitUntilEmpty() - val firstRun = TaskRecorder.drain() - val finalStage = firstRun.map(_.stageId).max - (rowCounts(adaptivePlan.executedPlan, arm), firstRun.filter(_.stageId != finalStage)) - case _ => - (rowCounts(executed, arm), Nil) - } - - /** How many shuffle inputs the native plan under Comet's conversion reads in native code. */ - private def directReadInputs(plan: SparkPlan): Int = - collectFirst(plan) { - case conversion: CometColumnarToRowExec => conversion.child - case conversion: CometNativeColumnarToRowExec => conversion.child - }.collect { case native: CometNativeExec => - CometExec.findShuffleScanIndices(native.nativeOp).size - }.getOrElse(0) - - /** - * The first operator that is not Comet native. `findFirstNonCometOperator` does not descend - * into AQE wrappers or query stages, so the final plan and the plan of every stage are checked - * one by one. - */ - private def firstNonCometOperator(plan: SparkPlan): Option[SparkPlan] = { - val fragments = plan +: collect(plan) { - case adaptivePlan: AdaptiveSparkPlanExec => adaptivePlan.executedPlan - case stage: QueryStageExec => stage.plan - } - fragments.iterator - .flatMap(fragment => - findFirstNonCometOperator( - fragment, - classOf[AdaptiveSparkPlanExec], - classOf[QueryStageExec], - classOf[AQEShuffleReadExec])) - .find(_ => true) - } - - private def planMetrics(plan: SparkPlan): Seq[PlanMetric] = - collect(plan) { case op => op }.flatMap { op => - val column = op.output.headOption.fold("")(_.name) - op.metrics.toSeq.sortBy(_._1).map { case (key, metric) => - PlanMetric( - metric.id, - op.nodeName.trim, - column, - key, - metric.name.getOrElse(key), - metric.metricType) - } - } - - /** - * Runs the query once, untimed, prints what it did, and warns when it does not measure the - * case. The plan is dropped on return, so its shuffle is released before the measured runs. - */ - private def inspect(arm: Arm): Unit = { - val (run, plan) = execute(arm) - val join = joinOf(plan) - val hotRows = run.joinRows(run.hotTask) - val share = 100.0 * hotRows / FactRows - val directReads = if (arm.comet) directReadInputs(plan) else 0 - val codegen = - if (arm.comet) s", $directReads shuffle inputs read directly in native code" - else if (run.wholeStageJoin) ", join in whole-stage codegen" - else ", join without whole-stage codegen" - emit( - f"${arm.name}: join ${join.fold("none")(_.nodeName)}, ${run.rows} rows, " + - f"${run.joinStage.size} join stage tasks, the hot task emitted $hotRows rows " + - f"($share%.1f%%)$codegen.") - - if (join.isEmpty) { - warning(s"WARNING: '${arm.name}' ran no shuffled hash join.\n$plan") - } - if (run.rows != FactRows) { - warning(s"WARNING: '${arm.name}' produced ${run.rows} rows, expected $FactRows.") - } - if (!broadcast && share < HotPercent) { - warning( - f"WARNING: the hot task of '${arm.name}' emitted $share%.1f%% of the rows, expected " + - f"$HotPercent%d%%.") - } - if (adaptive && arm.comet && directReads == 0) { - warning("WARNING: under AQE, Comet read no shuffle input directly in native code.") - } - if (arm.comet) { - firstNonCometOperator(plan).foreach { op => - warning( - s"WARNING: the Comet plan is not fully Comet native (first non-Comet operator: " + - s"${op.nodeName}), so it partly measures Spark.\n$plan") - } - } - if (run.joinStage.size != ShufflePartitions) { - warning( - s"WARNING: the join stage of '${arm.name}' ran ${run.joinStage.size} tasks, expected " + - s"$ShufflePartitions.") - } - if (run.tasks.exists(_.failed)) { - warning(s"WARNING: '${arm.name}' had failed tasks, so its task time is not comparable.") - } - } - - private def emitEnvironment(): Unit = { - val conf = spark.sparkContext.getConf - emit( - s"Spark ${spark.version}, Java ${System.getProperty("java.version")}, " + - s"master ${conf.get("spark.master")}") - emit( - if (adaptive) "AQE on, skew join handling and partition coalescing off" - else "AQE off") - emit(s"spark.memory.offHeap.size: ${conf.get("spark.memory.offHeap.size")}") - emit(s"${SQLConf.SHUFFLE_PARTITIONS.key}: $ShufflePartitions") - emit( - s"spark.io.compression.codec (Spark shuffle): " + - conf.get("spark.io.compression.codec", "lz4")) - emit( - s"${CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.key} (Comet shuffle): " + - CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.get()) - emit( - s"${SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key}: " + - spark.conf.get(SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key)) - emit( - s"fact: $FactRows rows in $FactFiles files, $HotPercent% of them on one key, 3 nested " + - s"and $WideColumns flat payload columns") - emit(s"dim: $DimRows rows in $DimFiles files, one per key") - emit(s"measured runs per arm: $measuredRuns, arms: ${arms.map(_.name).mkString("; ")}") - } - - private def measured(arm: Arm): Seq[Run] = - runs.get(arm.name).map(_.toList).getOrElse(Nil) - - /** The measured run with the median hot task time. */ - private def medianRun(arm: Arm): Option[Run] = { - val all = measured(arm).sortBy(_.hotTask.runTimeMs) - all.lift(all.length / 2) - } - - /** A line of the comparison: how to read it from a run, and how many decimals to show. */ - private case class Line(label: String, decimals: Int, value: Run => Option[Double]) - - private def emitComparison(): Unit = { - val medians = arms.map(arm => arm -> medianRun(arm)).toMap - def hot(run: Run): Seq[TaskRecord] = Seq(run.hotTask) - def sum(tasks: Seq[TaskRecord])(f: TaskRecord => Long): Option[Double] = - Some(tasks.map(f).sum.toDouble) - val mib = 1024.0 * 1024 - val sections = Seq( - "Hot task, the join stage task that emitted the most rows" -> Seq( - Line("run time (ms)", 0, r => Some(r.hotTask.runTimeMs.toDouble)), - Line("CPU time (ms)", 0, r => Some(r.hotTask.cpuTimeMs.toDouble)), - Line("GC time (ms)", 0, r => Some(r.hotTask.gcTimeMs.toDouble)), - Line("rows emitted", 0, r => Some(r.joinRows(r.hotTask).toDouble)), - Line("shuffle bytes read (MiB)", 1, r => Some(r.hotTask.shuffleReadBytes / mib)), - Line("shuffle records read", 0, r => Some(r.hotTask.shuffleReadRecords.toDouble)), - Line("shuffle fetch wait (ms)", 0, r => Some(r.hotTask.fetchWaitMs.toDouble)), - Line( - "shuffle read operator (ms)", - 1, - r => r.metric(hot(r), "CometExchange", "elapsed_compute")), - Line( - " decode and decompress (ms)", - 1, - r => r.metric(hot(r), "CometExchange", "decode_time")), - Line( - "join build (ms)", - 1, - r => - r.metric(hot(r), "ShuffledHashJoin", "buildTime") - .orElse(r.metric(hot(r), "CometHashJoin", "build_time"))), - Line( - "join probe and output (ms)", - 1, - r => r.metric(hot(r), "CometHashJoin", "join_time")), - Line( - "whole-stage codegen pipeline (ms)", - 0, - r => r.metric(hot(r), "WholeStageCodegen", "pipelineTime"))), - "Map stage that scanned fact and wrote its shuffle" -> Seq( - Line("tasks", 0, r => Some(r.mapStage.size.toDouble)), - Line("task time, total (ms)", 0, r => sum(r.mapStage)(_.runTimeMs)), - Line("longest task (ms)", 0, r => Some(r.mapStage.map(_.runTimeMs).max.toDouble)), - Line("scan input (MiB)", 1, r => sum(r.mapStage)(_.inputBytes).map(_ / mib)), - Line( - "data size before compression (MiB)", - 1, - r => - r.metric(r.mapStage, "Exchange", "dataSize") - .orElse(r.metric(r.mapStage, "CometExchange", "dataSize"))), - Line( - "shuffle bytes written (MiB)", - 1, - r => sum(r.mapStage)(_.shuffleWriteBytes).map(_ / mib)), - Line("shuffle records written", 0, r => sum(r.mapStage)(_.shuffleWriteRecords)), - Line( - "shuffle bytes written per row", - 0, - r => { - val records = r.mapStage.map(_.shuffleWriteRecords).sum - if (records == 0) None - else Some(r.mapStage.map(_.shuffleWriteBytes).sum.toDouble / records) - }), - Line("shuffle write time (ms)", 0, r => sum(r.mapStage)(_.shuffleWriteNs).map(_ / 1e6)), - Line( - "native shuffle writer (ms)", - 0, - r => r.metric(r.mapStage, "CometExchange", "elapsed_compute")), - Line(" repartition (ms)", 0, r => r.metric(r.mapStage, "CometExchange", "repart_time")), - Line( - " encode and compress (ms)", - 0, - r => r.metric(r.mapStage, "CometExchange", "encode_time"))), - "Whole run" -> Seq( - Line("join stage task time, total (ms)", 0, r => sum(r.joinStage)(_.runTimeMs)), - Line("all tasks, total (ms)", 0, r => sum(r.tasks)(_.runTimeMs)))) - - emit("The run with the median hot task time of each engine. Ratio is Spark divided by Comet.") - emit("A dash means the engine has no such metric.") - val width = sections.flatMap(_._2).map(_.label.length).max + 2 - emit(s"%-${width}s %12s %12s %8s".format("", SparkArm.name, CometArm.name, "ratio")) - sections.foreach { case (title, lines) => - emit("") - emit(title) - lines.foreach { line => - val values = arms.map(arm => medians(arm).flatMap(line.value)) - val cells = values.map(_.fold("-")(v => s"%.${line.decimals}f".format(v))) - val ratio = values match { - case Seq(Some(s), Some(c)) if c > 0 => f"${s / c}%.2f" - case _ => "" - } - emit(s" %-${width - 2}s %12s %12s %8s".format(line.label, cells.head, cells(1), ratio)) - } - } - } - - /** The broadcast mode's probe stage: one task per `fact` file, each probing the whole `dim`. */ - private def emitProbeSummary(): Unit = { - emit( - "The join stage probes, one task per fact file. Times in ms, from the run with the median") - emit("join stage time of each engine. Other stages scan dim for the broadcast.") - emit("") - emit( - f"${"arm"}%-6s ${"tasks"}%6s ${"stage total"}%12s ${"mean task"}%10s ${"median task"}%12s " + - f"${"longest"}%8s ${"other stages"}%13s") - arms.foreach { arm => - val all = measured(arm).sortBy(_.joinStage.map(_.runTimeMs).sum) - all.lift(all.length / 2).foreach { run => - val probe = run.joinStage.map(_.runTimeMs).sorted - val other = run.tasks.map(_.runTimeMs).sum - probe.sum - emit( - f"${arm.name}%-6s ${probe.size}%6d ${probe.sum}%12d ${probe.sum.toDouble / probe.size}%10.1f " + - f"${probe(probe.size / 2)}%12d ${probe.last}%8d $other%13d") - } - } - arms.foreach { arm => - val all = measured(arm).sortBy(_.joinStage.map(_.runTimeMs).sum) - all.lift(all.length / 2).foreach { run => - val tasks = run.joinStage - emit("") - emit(s"${arm.name}, mean per probe task over ${tasks.size} tasks:") - run.metrics.foreach { metric => - val values = tasks.flatMap(_.metricUpdates.get(metric.id)) - if (values.nonEmpty) { - val value = scaled(metric.metricType, values.sum) / tasks.size - val text = metric.metricType match { - case "nsTiming" | "timing" => f"$value%.2f ms" - case "size" => f"$value%.2f MiB" - case _ => f"$value%.0f" - } - val operator = s"${metric.nodeName} [${metric.column}]" - emit(f" $text%14s $operator%-28s ${metric.name}") - } - } - } - } - } - - /** Every metric the hot task and the map stage of each engine's median run reported. */ - private def emitMetricListings(): Unit = { - arms.foreach { arm => - medianRun(arm).foreach { run => - Seq("hot task" -> Seq(run.hotTask), "map stage that read fact" -> run.mapStage).foreach { - case (what, tasks) => - emit("") - emit(s"${arm.name}, $what, ${tasks.size} task(s):") - run.metrics.foreach { metric => - val values = tasks.flatMap(_.metricUpdates.get(metric.id)) - if (values.nonEmpty) { - val value = scaled(metric.metricType, values.sum) - val text = metric.metricType match { - case "nsTiming" | "timing" => f"$value%.1f ms" - case "size" => f"$value%.1f MiB" - case _ => f"$value%.0f" - } - val operator = s"${metric.nodeName} [${metric.column}]" - emit(f" $text%14s $operator%-28s ${metric.name}") - } - } - } - } - } - } - - /** - * Runs the hot task of the arm alone, first while sampling the stack of its thread and then for - * an external profiler. One more run of the query writes the shuffle the hot task reads. - */ - private def profile(arm: Arm): Unit = { - val run = medianRun(arm).getOrElse { - throw new IllegalStateException(s"No measured run of '${arm.name}' to profile") - } - val partition = if (profileSmallTask) { - val small = - run.joinStage.filterNot(_.partition == run.hotTask.partition).sortBy(_.runTimeMs) - small(small.size / 2).partition - } else { - run.hotTask.partition - } - val which = if (profileSmallTask) "a median small join task" else "the hot task" - withSQLConf(arm.configs: _*)(withConfsPropagated { - val (counts, _) = countsToRun(spark.sql(query).queryExecution.executedPlan, arm) - counts.fold(0L)(_ + _) - val sampler = new StackSampler - sampler.start() - val sampled = runHotTask(counts, partition) - sampler.finish() - emit("") - emit( - s"${arm.name}: ran $which (partition $partition) $sampled times, sampling the " + - s"stack of its thread every millisecond, ${sampler.samples} samples.") - sampler.emitSummary() - emit(s"PROFILE ${arm.name}: pid ${ProcessHandle.current().pid()}, $ProfileSeconds s") - val unsampled = runHotTask(counts, partition) - emit(s"${arm.name}: ran $which $unsampled times for the external profiler") - }) - } - - /** - * Runs `body` with the session's configs set as local properties of the jobs it submits, as - * Spark does for a query it executes, so that tasks see them. The benchmark runs plans - * directly, which skips that step, and Comet reads some configs, such as the shuffle read - * buffer size, in tasks. - */ - private def withConfsPropagated[T](body: => T): T = { - val sc = spark.sparkContext - val confs = spark.conf.getAll.filter(_._1.startsWith("spark")).toList - val previous = confs.map { case (key, _) => key -> sc.getLocalProperty(key) } - confs.foreach { case (key, value) => sc.setLocalProperty(key, value) } - try body - finally previous.foreach { case (key, value) => sc.setLocalProperty(key, value) } - } - - /** One line per Comet arm of a sweep over the shuffle read buffer size. */ - private def emitReadBufferSweep(): Unit = { - def median(values: Seq[Double]): Option[Double] = - if (values.isEmpty) None else Some(values.sorted.apply(values.length / 2)) - emit(s"Comet only, one arm per setting, run in turn. Medians over $measuredRuns runs, in ms.") - emit( - "Shuffle read is the hot task's shuffle read operator, which Comet reports with AQE off.") - emit("") - val width = arms.map(_.name.length).max - emit( - s"%-${width}s %9s %6s %6s %13s %11s %10s" - .format("arm", "hot task", "min", "max", "shuffle read", "join stage", "all tasks")) - arms.foreach { arm => - val all = measured(arm) - if (all.nonEmpty) { - val hot = all.map(_.hotTask.runTimeMs.toDouble) - val read = median( - all - .flatMap(run => run.metric(Seq(run.hotTask), "CometExchange", "elapsed_compute")) - .filter(_ > 0)) - emit( - s"%-${width}s %9.0f %6.0f %6.0f %13s %11.0f %10.0f".format( - arm.name, - median(hot).get, - hot.min, - hot.max, - read.fold("-")(ms => f"$ms%.0f"), - median(all.map(_.joinStage.map(_.runTimeMs).sum.toDouble)).get, - median(all.map(_.tasks.map(_.runTimeMs).sum.toDouble)).get)) - } - } - } - - /** - * What a stack sample was doing: the first rule that a frame matches, from the top of the stack - * down, so that generic frames such as a buffer copy take the category of their caller. A - * native method on top is the time spent in native code under it. - */ - private val StackRules: Seq[(String, StackTraceElement => Boolean)] = { - def is(cls: String, method: String)(f: StackTraceElement) = - f.getClassName == cls && f.getMethodName == method - def under(prefix: String)(f: StackTraceElement) = f.getClassName.startsWith(prefix) - Seq( - "Comet native plan (join, batch import, project)" -> is( - "org.apache.comet.Native", - "executePlan"), - "Comet native shuffle decode" -> (f => - f.getClassName == "org.apache.comet.Native" && f.getMethodName.startsWith( - "decodeShuffle")), - "Comet native, other" -> under("org.apache.comet.Native"), - "Arrow C JNI (array release, export)" -> under("org.apache.arrow.c.jni."), - "file read (read syscall)" -> (f => - f.isNativeMethod && Set("readBytes", "read0", "pread0").contains(f.getMethodName)), - "FileInputStream.available (fstat, lseek)" -> (f => f.getMethodName == "available0"), - "decompress (lz4-java)" -> under("net.jpountz."), - "Channels.newChannel copy loop" -> under("java.nio.channels.Channels"), - "Comet shuffle block read (direct read)" -> under( - "org.apache.comet.CometShuffleBlockIterator"), - "Comet shuffle block read" -> is( - "org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIterator", - "readNextBlock"), - "Arrow C import into JVM" -> (f => - under("org.apache.arrow.c.")(f) && (f.getClassName + f.getMethodName).contains("mport")), - "Arrow C export to native" -> (f => - under("org.apache.arrow.c.")(f) && (f.getClassName + f.getMethodName).contains("xport")), - "Comet vectors" -> under("org.apache.comet.vector."), - "Comet iterator glue" -> (f => - under("org.apache.comet.CometBatchIterator")(f) || - under("org.apache.comet.CometExecIterator")(f)), - "row deserialization" -> (f => - f.getClassName.contains("UnsafeRowSerializer") || under("java.io.DataInputStream")(f)), - "shuffle fetch and stream wrappers" -> (f => - under("org.apache.spark.storage.")(f) || under("org.apache.spark.network.")(f) || - under("org.apache.spark.shuffle.")(f)), - "join hash lookup" -> under("org.apache.spark.sql.execution.joins."), - "output projection (UnsafeRow writes)" -> (f => - f.getClassName.contains("SpecificUnsafeProjection") || - under("org.apache.spark.sql.catalyst.expressions.codegen.Unsafe")(f)), - "UnsafeRow, array and map reads" -> under( - "org.apache.spark.sql.catalyst.expressions.Unsafe"), - "row count" -> under("org.apache.spark.sql.benchmark."), - "broadcast decode: Arrow IPC read" -> (f => - under("org.apache.arrow.vector.ipc.")(f) || f.getClassName.endsWith( - "ArrowReaderIterator") || - f.getClassName == "org.apache.comet.vector.StreamReader"), - "broadcast decode, other" -> under("org.apache.spark.sql.comet.CometBatchRDD"), - "whole-stage codegen (scan, probe, project)" -> (f => - f.getClassName.contains("GeneratedIteratorForCodegenStage")), - "Parquet scan (JVM)" -> (f => - under("org.apache.spark.sql.execution.datasources.parquet.")(f) || - under("org.apache.parquet.")(f) || under("org.apache.spark.sql.execution.vectorized.")( - f))) - } - - private def stackCategory(frames: Array[StackTraceElement]): String = - frames.iterator - .flatMap(frame => StackRules.find(_._2(frame)).map(_._1)) - .find(_ => true) - .getOrElse("other Java") - - /** - * Samples the stack of every running executor thread every millisecond until `finish`, and - * counts the samples by category and by top frame. `getStackTrace` waits for the next safepoint - * of a thread running Java code, so Java frames carry safepoint bias. A thread in native code - * is sampled where it is. - */ - private class StackSampler extends Thread("hash-join-benchmark-stack-sampler") { - setDaemon(true) - @volatile private var running = true - private val categories = mutable.Map.empty[String, Long].withDefaultValue(0L) - private val tops = mutable.Map.empty[String, Long].withDefaultValue(0L) - var samples = 0L - - override def run(): Unit = { - var threads = Seq.empty[Thread] - var refreshedAt = 0L - while (running) { - if (System.nanoTime() - refreshedAt > 200000000L) { - threads = executorThreads() - refreshedAt = System.nanoTime() - } - threads.filter(_.getState == Thread.State.RUNNABLE).foreach { thread => - val frames = thread.getStackTrace - if (frames.nonEmpty) { - samples += 1 - categories(stackCategory(frames)) += 1 - val top = frames.head - val native = if (top.isNativeMethod) " (native)" else "" - tops(s"${top.getClassName}.${top.getMethodName}$native") += 1 - } - } - Thread.sleep(1) - } - } - - def finish(): Unit = { - running = false - join() - } - - def emitSummary(): Unit = { - def share(n: Long): String = f"${100.0 * n / math.max(samples, 1L)}%5.1f%%" - categories.toSeq.sortBy(-_._2).foreach { case (category, n) => - emit(s" ${share(n)} $category") - } - emit(" Top frames:") - tops.toSeq.sortBy(-_._2).take(20).foreach { case (frame, n) => - emit(s" ${share(n)} $frame") - } - } - } - - private def executorThreads(): Seq[Thread] = { - val threads = Thread.getAllStackTraces.keySet.iterator() - val found = mutable.Buffer.empty[Thread] - while (threads.hasNext) { - val thread = threads.next() - if (thread.getName.startsWith("Executor task launch worker")) { - found += thread - } - } - found.toList - } - - /** Runs the partition of `counts` alone, again and again for `ProfileSeconds`. */ - private def runHotTask(counts: RDD[Long], partition: Int): Int = { - val end = System.nanoTime() + ProfileSeconds * 1000000000L - var times = 0 - while (System.nanoTime() < end) { - spark.sparkContext.runJob(counts, (rows: Iterator[Long]) => rows.sum, Seq(partition)) - times += 1 - } - times - } - - /** Sets a warning off from the lines around it, in the console and the results file. */ - private def warning(message: String): Unit = { - val border = "=" * 80 - emit(s"\n$border\n$message\n$border") - } - - /** Writes a line to the console and, when results files are generated, to the results file. */ - private def emit(line: String): Unit = { - // scalastyle:off println - println(line) - // scalastyle:on println - output.foreach(_.write(s"$line\n".getBytes(StandardCharsets.UTF_8))) - } -} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala deleted file mode 100644 index c833e21721d..00000000000 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleReadKillBenchmark.scala +++ /dev/null @@ -1,380 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.spark.sql.benchmark - -import java.io.File -import java.nio.charset.StandardCharsets -import java.util.concurrent.atomic.AtomicLong -import java.util.concurrent.locks.LockSupport - -import scala.collection.mutable - -import org.apache.spark.{SparkConf, TaskContext} -import org.apache.spark.rdd.RDD -import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometNativeExec} -import org.apache.spark.sql.execution.{InputAdapter, SparkPlan} -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec -import org.apache.spark.sql.functions.col -import org.apache.spark.sql.internal.SQLConf - -import org.apache.comet.{CometConf, CometSparkSessionExtensions} - -/** - * Measures how long a Comet reduce task keeps running after it is killed while it reads its - * shuffle (review of apache/datafusion-comet#6805). - * - * One map task writes the whole shuffle as one block, and one reduce task reads it with a native - * project on top. On the `aqe` path the project reads the shuffle directly in native code - * (`ShuffleScan` over `CometShuffleBlockIterator`), which Comet does only behind an AQE query - * stage. On the `noaqe` path it reads it through `NativeBatchDecoderIterator`. - * - * Each trial runs the reduce task alone over the same shuffle. A kill trial cancels the job - * `KillAfterMs` after the task starts, with or without interrupting the task's thread, as - * `spark.job.interruptOnCancel` selects, and measures how long the task runs after the cancel, - * until its completion listeners run. A trial without a kill measures the whole task. - * - * With `variants=fix,head`, the default, the variants run in turn in one JVM. `head` sets the - * local property `spark.comet.benchmark.skipKillCheck`, which only a temporary benchmark build - * reads: it makes `readAsRawStream` skip its kill check, as #6805 did before the fix. Any other - * variant name only labels the classes on the classpath. `paths=aqe,noaqe` selects the read - * paths, `rows=N` the rows of the shuffle and `trials=N` the trials. - */ -object CometShuffleReadKillBenchmark extends CometBenchmarkBase { - - // `SqlBasedBenchmark` calls `getSparkSession` while it initializes, before the fields of this - // object are assigned, so the field that method reads is a compile-time constant. - private final val Cores = 4 - - /** Rows of the shuffle. Its strings repeat, so it is small on disk but slow to decode. */ - private var rows = 16L * 1024 * 1024 - - /** Trials per variant and mode, set by `trials=N`. */ - private var trials = 3 - - private val KillAfterMs = 1000L - - private val SkipKillCheckKey = "spark.comet.benchmark.skipKillCheck" - - private case class Mode(name: String, kill: Boolean, interrupt: Boolean) - - private val Modes = Seq( - Mode("not killed: task time", kill = false, interrupt = false), - Mode("killed, thread interrupted: time after the kill", kill = true, interrupt = true), - Mode("killed, thread not interrupted: time after the kill", kill = true, interrupt = false)) - - /** Counts task ends and keeps the shuffle sizes and the end reason of the last one. */ - private object TaskRecorder extends SparkListener { - val ends = new AtomicLong(0L) - @volatile var lastReason = "" - @volatile var shuffleWritten = 0L - @volatile var shuffleRead = 0L - - override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { - Option(taskEnd.taskMetrics).foreach { metrics => - if (metrics.shuffleWriteMetrics.bytesWritten > 0) { - shuffleWritten = metrics.shuffleWriteMetrics.bytesWritten - } - if (metrics.shuffleReadMetrics.totalBytesRead > 0) { - shuffleRead = metrics.shuffleReadMetrics.totalBytesRead - } - } - lastReason = taskEnd.reason.toString.takeWhile(_ != '(') - ends.incrementAndGet() - } - } - - override def getSparkSession: SparkSession = { - val conf = new SparkConf() - .setAppName("CometShuffleReadKillBenchmark") - // Since `spark.master` always exists, overrides this value - .set("spark.master", s"local[$Cores]") - .setIfMissing("spark.driver.memory", "3g") - .set( - "spark.shuffle.manager", - "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") - .set("spark.memory.offHeap.enabled", "true") - .setIfMissing("spark.memory.offHeap.size", "8g") - - val sparkSession = SparkSession - .builder() - .config(conf) - .withExtensions(new CometSparkSessionExtensions) - .getOrCreate() - - sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") - sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") - sparkSession.conf.set(SQLConf.SHUFFLE_PARTITIONS.key, "1") - sparkSession.conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED.key, "false") - // An open cost as large as the largest split keeps the one file in one map task. - val maxSplitBytes = (1024L * 1024 * 1024).toString - sparkSession.conf.set(SQLConf.FILES_MAX_PARTITION_BYTES.key, maxSplitBytes) - sparkSession.conf.set(SQLConf.FILES_OPEN_COST_IN_BYTES.key, maxSplitBytes) - sparkSession - } - - override def runCometBenchmark(mainArgs: Array[String]): Unit = { - def listArg(name: String, default: Seq[String]): Seq[String] = - mainArgs - .collectFirst { - case arg if arg.startsWith(s"$name=") => arg.stripPrefix(s"$name=").split(",").toSeq - } - .getOrElse(default) - val variants = listArg("variants", Seq("fix", "head")) - val paths = listArg("paths", Seq("aqe", "noaqe")) - listArg("rows", Nil).headOption.foreach(value => rows = value.toLong) - listArg("trials", Nil).headOption.foreach(value => trials = value.toInt) - spark.sparkContext.addSparkListener(TaskRecorder) - try { - withTempPath { dir => - withTempTable("src") { - writeSource(dir) - runBenchmark("Killing a Comet reduce task while it reads its shuffle") { - val conf = spark.sparkContext.getConf - emit( - s"Spark ${spark.version}, Java ${System.getProperty("java.version")}, " + - s"master ${conf.get("spark.master")}") - emit( - s"src: $rows rows in 1 Parquet file, so 1 map task writes the shuffle as 1 block " + - "for 1 reduce task") - emit( - s"variants: ${variants.mkString(", ")}. $trials trials per variant and mode, " + - s"kills $KillAfterMs ms after the task starts. Times in ms, median (trials).") - paths.foreach(path => measure(path, variants)) - } - } - } - } catch { - case e: Throwable => - // `BenchmarkBase.main` stops the session only after a run that returned. - spark.stop() - throw e - } - } - - /** Writes `src` with Comet disabled and checks that it scans as one map task. */ - private def writeSource(dir: File): Unit = { - val path = new File(dir, "src").getCanonicalPath - spark - .range(0L, rows, 1L, 1) - .selectExpr( - "id", - "XXHASH64(id) AS h", - "REPEAT(CONCAT('k', CAST(PMOD(id, 97) AS STRING), '|'), 24) AS s1", - "REPEAT(CONCAT('v', CAST(PMOD(id, 89) AS STRING), ';'), 24) AS s2") - .write - .option("compression", "snappy") - .parquet(path) - spark.read.parquet(path).createOrReplaceTempView("src") - val tasks = spark.table("src").queryExecution.toRdd.getNumPartitions - require(tasks == 1, s"src scans as $tasks map tasks, expected 1") - } - - /** Writes the shuffle of one read path, then runs and reports its trials. */ - private def measure(path: String, variants: Seq[String]): Unit = { - val adaptive = path == "aqe" - val results = mutable.LinkedHashMap.empty[(String, String), mutable.Buffer[Option[Double]]] - val reasons = mutable.LinkedHashMap.empty[(String, String), mutable.Set[String]] - var description = "" - System.gc() - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - CometConf.COMET_ENABLED.key -> "true", - CometConf.COMET_EXEC_ENABLED.key -> "true") { - withConfsPropagated { - val executed = spark - .table("src") - .repartition(1, col("id")) - .selectExpr("id + 1 AS a", "h", "LENGTH(s1) + LENGTH(s2) AS l") - .queryExecution - .executedPlan - // Under AQE, Comet's conversion to rows sits inside the adaptive plan. The query runs - // once to materialize its shuffle stage, and the native part of the final plan then runs - // alone. Without AQE, the native part runs alone from the start. - val finalPlan = executed match { - case adaptivePlan: AdaptiveSparkPlanExec => - adaptivePlan.execute().foreach(_ => ()) - adaptivePlan.executedPlan - case plan => plan - } - val native = nativeChild(finalPlan) - val direct = CometExec.findShuffleScanIndices(native.nativeOp).size - require(direct == (if (adaptive) 1 else 0), s"$direct direct reads in:\n$finalPlan") - val counts = native.executeColumnar().map(_.numRows().toLong) - // Without AQE, this run also writes the shuffle. - val counted = counts.fold(0L)(_ + _) - require(counted == rows, s"counted $counted rows, expected $rows") - spark.sparkContext.listenerBus.waitUntilEmpty() - val mib = 1024.0 * 1024 - description = - (if (adaptive) "AQE on, the project reads the shuffle directly (ShuffleScan)" - else "AQE off, the project reads the shuffle through NativeBatchDecoderIterator") + - f", shuffle ${TaskRecorder.shuffleWritten / mib}%.0f MiB written, " + - f"${TaskRecorder.shuffleRead / mib}%.0f MiB read by the reduce task" - emit("") - emit(description) - emit(native.treeString.linesIterator.map(" " + _).mkString("\n")) - for (_ <- 1 to trials; mode <- Modes; variant <- variants) { - val (result, reason) = trial(counts, variant, mode) - results.getOrElseUpdate((variant, mode.name), mutable.Buffer.empty) += result - reasons.getOrElseUpdate((variant, mode.name), mutable.Set.empty) += reason - } - } - } - val width = Modes.map(_.name.length).max - emit("") - emit(s"%-${width}s".format("") + variants.map(v => "%24s".format(v)).mkString) - Modes.foreach { mode => - val cells = variants.map { variant => - val values = results((variant, mode.name)).toList - val measured = values.flatten.sorted - val median = measured.lift(measured.size / 2).fold("-")(ms => f"$ms%.0f") - val each = values.map(_.fold("ended first")(ms => f"$ms%.0f")).mkString(", ") - "%24s".format(s"$median ($each)") - } - emit(s"%-${width}s".format(mode.name) + cells.mkString) - } - emit("Task end reasons:") - reasons.foreach { case ((variant, mode), seen) => - emit(s" $variant, $mode: ${seen.mkString(", ")}") - } - } - - /** - * The native plan under the Comet columnar-to-row conversion of `plan`. `collectFirst` of - * `AdaptiveSparkPlanHelper` looks into query stages, such as AQE's final `ResultQueryStage`. - * Under whole-stage codegen an `InputAdapter`, which tree strings do not show, sits between. - */ - private def nativeChild(plan: SparkPlan): CometNativeExec = - collectFirst(plan) { - case conversion: CometColumnarToRowExec => conversion.child - case conversion: CometNativeColumnarToRowExec => conversion.child - } - .map { - case adapter: InputAdapter => adapter.child - case child => child - } - .collect { case native: CometNativeExec => native } - .getOrElse { - throw new IllegalStateException(s"No native plan under a conversion to rows in:\n$plan") - } - - /** - * Runs the reduce task once. Returns how long it ran after the kill, or in total without one, - * in ms, with the reason its task ended. The time is empty when the task ended before the kill. - */ - private def trial(counts: RDD[Long], variant: String, mode: Mode): (Option[Double], String) = { - val sc = spark.sparkContext - // A collection now keeps one out of the time measured after a kill. - System.gc() - CometShuffleReadKillTaskClock.started.set(0L) - CometShuffleReadKillTaskClock.ended.set(0L) - val endsBefore = TaskRecorder.ends.get() - val group = s"kill-benchmark-${System.nanoTime()}" - val job = new Thread( - () => { - sc.setJobGroup(group, s"$variant, ${mode.name}", interruptOnCancel = mode.interrupt) - sc.setLocalProperty(SkipKillCheckKey, if (variant == "head") "true" else null) - try { - sc.runJob( - counts, - (context: TaskContext, batches: Iterator[Long]) => { - context.addTaskCompletionListener[Unit] { _ => - CometShuffleReadKillTaskClock.ended.set(System.nanoTime()) - } - CometShuffleReadKillTaskClock.started.set(System.nanoTime()) - batches.sum - }, - Seq(0)) - } catch { - // A cancelled job fails here, as expected. - case _: Throwable => - } - }, - "kill-benchmark-job") - job.setDaemon(true) - job.start() - val start = await(CometShuffleReadKillTaskClock.started) - val result = if (mode.kill) { - val killAt = start + KillAfterMs * 1000000L - while (System.nanoTime() < killAt) { - LockSupport.parkNanos(math.min(killAt - System.nanoTime(), 1000000L)) - } - val killedAt = System.nanoTime() - if (CometShuffleReadKillTaskClock.ended.get() != 0L) { - None - } else { - sc.cancelJobGroup(group) - Some((await(CometShuffleReadKillTaskClock.ended) - killedAt) / 1e6) - } - } else { - Some((await(CometShuffleReadKillTaskClock.ended) - start) / 1e6) - } - job.join() - // The task end event follows the executor's final status update, so the next trial does not - // overlap the cleanup of this one. - while (TaskRecorder.ends.get() == endsBefore) { - LockSupport.parkNanos(1000000L) - } - (result, TaskRecorder.lastReason) - } - - /** Waits until `clock` is set and returns it. */ - private def await(clock: AtomicLong): Long = { - val deadline = System.nanoTime() + 300L * 1000000000L - while (clock.get() == 0L) { - require(System.nanoTime() < deadline, "the task did not start or end within 300 s") - LockSupport.parkNanos(100000L) - } - clock.get() - } - - /** - * Runs `body` with the session's configs set as local properties of the jobs it submits, as - * Spark does for a query it executes, so that tasks see them. The benchmark runs plans - * directly, which skips that step. - */ - private def withConfsPropagated[T](body: => T): T = { - val sc = spark.sparkContext - val confs = spark.conf.getAll.filter(_._1.startsWith("spark")).toList - val previous = confs.map { case (key, _) => key -> sc.getLocalProperty(key) } - confs.foreach { case (key, value) => sc.setLocalProperty(key, value) } - try body - finally previous.foreach { case (key, value) => sc.setLocalProperty(key, value) } - } - - /** Writes a line to the console and, when results files are generated, to the results file. */ - private def emit(line: String): Unit = { - // scalastyle:off println - println(line) - // scalastyle:on println - output.foreach(_.write(s"$line\n".getBytes(StandardCharsets.UTF_8))) - } -} - -/** - * When the task of a trial started, and when its completion listeners ran, in `nanoTime`. A - * top-level object, so that the task's closure reads it without capturing the benchmark. - */ -private object CometShuffleReadKillTaskClock { - val started = new AtomicLong(0L) - val ended = new AtomicLong(0L) -}