You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/30 09:24:38 UTC
(datafusion-comet) 01/01: fix: count every hash-based JVM shuffle batch as shuffle bytes written (#6369)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6369-0ce2ff2d290b5398e4de665ed63b8c94b98d8118
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit f375c4c345a61c9a41c240f527351b076cbdfffb
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Sep 30 09:24:16 2026 +0000
fix: count every hash-based JVM shuffle batch as shuffle bytes written (#6369)
* fix: count hash-based JVM shuffle batches as shuffle bytes written, not spill
CometDiskBlockWriter appends every batch, whether written at the batch
size, under memory pressure, or on close, to the partition file that
CometBypassMergeSortShuffleWriter concatenates into the map output. It
counted only the final batch as shuffle bytes written, reported the
earlier batches as disk spill, and dropped their write time. Count every
batch with the task's shuffle write metrics, as Spark's
BypassMergeSortShuffleWriter does for its partition files, and report
nothing as spill.
* fix: also report memory-pressure flushes of the JVM hash shuffle as spill
A flush that CometDiskBlockWriter makes to relieve memory pressure now
adds its bytes to the task's diskBytesSpilled and the memory it freed
to memoryBytesSpilled, like a sort-based spill. The flush still counts
as shuffle bytes and records written, since it lands in the partition
file that becomes the map output. Flushes on reaching the batch size
and the final flush on close count only as writes. The log line for a
batch-size flush no longer calls it a spill and moves to debug level.
---
.../execution/shuffle/CometDiskBlockWriter.java | 65 ++++++++++------------
.../spark/sql/comet/CometTaskMetricsSuite.scala | 57 +++++++++++++++++++
.../shuffle/CometDiskBlockWriterSuite.scala | 31 +++++++++--
3 files changed, 111 insertions(+), 42 deletions(-)
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
index 6cda37779a..105452a434 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
@@ -33,7 +33,6 @@ import org.slf4j.LoggerFactory;
import org.apache.spark.SparkConf;
import org.apache.spark.TaskContext;
-import org.apache.spark.executor.ShuffleWriteMetrics;
import org.apache.spark.serializer.SerializationStream;
import org.apache.spark.serializer.SerializerInstance;
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
@@ -173,7 +172,12 @@ public final class CometDiskBlockWriter {
return this.activeWriter.getChecksum();
}
- private void doSpill() throws IOException {
+ /**
+ * Writes the active writer's buffered rows to the partition file as one batch and frees their
+ * memory. The batch always counts as shuffle bytes written. A write that relieves memory pressure
+ * is also reported as a spill, the way a sort-based spill is.
+ */
+ private void doSpill(boolean underMemoryPressure) throws IOException {
// We only allow spilling request from `NativeDiskBlockArrowIPCWriter`.
if (spilling || activeWriter.numRecords() == 0) {
return;
@@ -183,8 +187,13 @@ public final class CometDiskBlockWriter {
spilling = true;
synchronized (CometDiskBlockWriter.this) {
- totalWritten += activeWriter.doSpilling(false);
- activeWriter.freeMemory();
+ final long written = activeWriter.doSpilling();
+ totalWritten += written;
+ final long freed = activeWriter.freeMemory();
+ if (underMemoryPressure) {
+ taskContext.taskMetrics().incDiskBytesSpilled(written);
+ taskContext.taskMetrics().incMemoryBytesSpilled(freed);
+ }
}
spilling = false;
@@ -217,10 +226,9 @@ public final class CometDiskBlockWriter {
if (activeWriter.numRecords() >= numElementsForSpillThreshold
|| activeWriter.numRecords() >= columnarBatchSize) {
int threshold = Math.min(numElementsForSpillThreshold, columnarBatchSize);
- logger.info(
- "Spilling data because number of spilledRecords crossed the threshold " + threshold);
- // Spill the current writer
- doSpill();
+ logger.debug("Writing a batch to the partition file because it reached {} rows", threshold);
+ // Not a spill: the batch is written to the map output because it is full.
+ doSpill(false);
if (activeWriter.numRecords() != 0) {
throw new RuntimeException(
"activeWriter.numRecords()(" + activeWriter.numRecords() + ") != 0");
@@ -241,7 +249,7 @@ public final class CometDiskBlockWriter {
}
FileSegment close() throws IOException {
- totalWritten += activeWriter.doSpilling(true);
+ totalWritten += activeWriter.doSpilling();
if (outputRecords != insertRecords) {
throw new RuntimeException(
@@ -318,20 +326,14 @@ public final class CometDiskBlockWriter {
return rowPartition.getNumRows();
}
- /** Spills the current in-memory records of this `ArrowIPCWriter` to disk. */
- long doSpilling(boolean isLast) throws IOException {
- final ShuffleWriteMetricsReporter writeMetricsToUse;
-
- if (isLast) {
- // We're writing the final non-spill file, so we _do_ want to count this as shuffle bytes.
- writeMetricsToUse = writeMetrics;
- } else {
- // We're spilling, so bytes written should be counted towards spill rather than write.
- // Create a dummy WriteMetrics object to absorb these metrics, since we don't want to count
- // them towards shuffle bytes written.
- writeMetricsToUse = new ShuffleWriteMetrics();
- }
-
+ /**
+ * Writes the current in-memory records of this `ArrowIPCWriter` to the partition file as one
+ * batch. Every batch, whether written on reaching the batch size, under memory pressure, or on
+ * close, is appended to the same file, which is copied as is into the map output. So every
+ * batch counts toward shuffle bytes written, records written, and write time, as in Spark's
+ * `BypassMergeSortShuffleWriter`.
+ */
+ long doSpilling() throws IOException {
final long written;
synchronized (file) {
@@ -341,24 +343,13 @@ public final class CometDiskBlockWriter {
dataTypes,
file,
rowPartition,
- writeMetricsToUse,
+ writeMetrics,
preferDictionaryRatio,
compressionCodec,
compressionLevel,
tracingEnabled);
}
- // Update metrics
- synchronized (writeMetrics) {
- if (!isLast) {
- writeMetrics.incRecordsWritten(
- ((ShuffleWriteMetrics) writeMetricsToUse).recordsWritten());
- taskContext
- .taskMetrics()
- .incDiskBytesSpilled(((ShuffleWriteMetrics) writeMetricsToUse).bytesWritten());
- }
- }
-
return written;
}
@@ -373,7 +364,7 @@ public final class CometDiskBlockWriter {
// initialCurrentPage() requires this writer to have released its current page, even when
// spilling a larger sibling would free enough memory for the allocation on its own.
long totalFreed = getActiveMemoryUsage();
- CometDiskBlockWriter.this.doSpill();
+ CometDiskBlockWriter.this.doSpill(true);
if (totalFreed >= required) {
return;
}
@@ -400,7 +391,7 @@ public final class CometDiskBlockWriter {
}
long used = writer.getActiveMemoryUsage();
- writer.doSpill();
+ writer.doSpill(true);
totalFreed += used;
}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala
index 153a7b422f..29271220f7 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala
@@ -32,6 +32,8 @@ import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocator
import org.apache.spark.shuffle.sort.CometShuffleExternalSorter
import org.apache.spark.sql.CometTestBase
import org.apache.spark.sql.catalyst.expressions.UnsafeRow
+import org.apache.spark.sql.comet.execution.shuffle.CometBypassMergeSortShuffleHandle
+import org.apache.spark.sql.comet.execution.shuffle.CometColumnarShuffle
import org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffle
import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
import org.apache.spark.sql.execution.SparkPlan
@@ -646,6 +648,61 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("JVM hash shuffle counts batch-size writes as shuffle bytes written, not spill") {
+ val expectedRecords = 20000L
+ withParquetTable((0 until expectedRecords.toInt).map(i => (i, s"row-$i")), "tbl") {
+ // A small batch size makes each partition writer flush several batches to its partition
+ // file before the final one.
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
+ CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "100") {
+ val shuffled = sql("SELECT * FROM tbl").repartition(4, $"_1")
+ val store = spark.sparkContext.statusStore
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+ val stagesBefore = store.stageList(null).map(_.stageId).toSet
+
+ assert(shuffled.collect().length == expectedRecords)
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+
+ val exchange = collectFirst(shuffled.queryExecution.executedPlan) {
+ case jvm: CometShuffleExchangeExec if jvm.shuffleType == CometColumnarShuffle => jvm
+ }.getOrElse(fail("Expected a JVM columnar shuffle exchange"))
+ // Four partitions is under spark.shuffle.sort.bypassMergeThreshold, so the hash-based
+ // writer ran.
+ assert(
+ exchange.shuffleDependency.shuffleHandle
+ .isInstanceOf[CometBypassMergeSortShuffleHandle[_, _]])
+
+ // The map output is every batch written to the partition files, concatenated.
+ val shuffleId = exchange.shuffleDependency.shuffleId
+ val dataFiles = SparkEnv.get.blockManager.diskBlockManager
+ .getAllFiles()
+ .filter { file =>
+ file.getName.startsWith(s"shuffle_${shuffleId}_") && file.getName.endsWith(".data")
+ }
+ assert(dataFiles.nonEmpty, s"No map output files found for shuffle $shuffleId")
+ val mapOutputBytes = dataFiles.map(_.length()).sum
+
+ val metrics = exchange.metrics
+ assert(metrics("shuffleRecordsWritten").value == expectedRecords)
+ assert(metrics("shuffleBytesWritten").value == mapOutputBytes)
+ assert(metrics("shuffleWriteTime").value > 0L)
+
+ val shuffleWriteStages = store
+ .stageList(null)
+ .filterNot(stage => stagesBefore.contains(stage.stageId))
+ .filter(_.shuffleWriteRecords > 0L)
+
+ assert(shuffleWriteStages.nonEmpty, "No JVM shuffle write stage was recorded")
+ assert(shuffleWriteStages.map(_.shuffleWriteRecords).sum == expectedRecords)
+ assert(shuffleWriteStages.map(_.shuffleWriteBytes).sum == mapOutputBytes)
+ // Without memory pressure, no batch was spilled.
+ assert(shuffleWriteStages.map(_.memoryBytesSpilled).sum == 0L)
+ assert(shuffleWriteStages.map(_.diskBytesSpilled).sum == 0L)
+ }
+ }
+ }
+
test("native shuffle task metrics include existing child sort spill metrics once") {
val expectedRecords = 20000L
val compressibleValue = "native-child-sort-spill-metrics-" * 8
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala
index 8ce6823999..57da656b28 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala
@@ -85,6 +85,8 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
val writersB = new JLinkedList[CometDiskBlockWriter]()
val serializer = new UnsafeRowSerializer(1).newInstance()
val fileB = new File(tempDir, "taskB-partition0")
+ val writeMetricsA = new ShuffleWriteMetrics
+ val writeMetricsB = new ShuffleWriteMetrics
def newTaskAWriter(partition: Int): CometDiskBlockWriter =
new CometDiskBlockWriter(
new File(tempDir, s"taskA-partition$partition"),
@@ -92,7 +94,7 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
taskContextA,
serializer,
schema,
- new ShuffleWriteMetrics,
+ writeMetricsA,
conf,
false,
writersA)
@@ -105,7 +107,7 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
taskContextB,
serializer,
schema,
- new ShuffleWriteMetrics,
+ writeMetricsB,
conf,
false,
writersB)
@@ -165,14 +167,23 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
assert(writerA1.getOutputRecords > 0)
assert(writerA2.getOutputRecords > 0)
- // Task A resolved its memory pressure by spilling its own data...
- assert(taskContextA.taskMetrics.diskBytesSpilled > 0)
+ // Task A resolved its memory pressure by spilling its own buffered rows to its partition
+ // files. Everything in those files so far was written by a spill...
+ val filesA = Seq(writerA0, writerA1, writerA2).map(_.getFile)
+ val spilledBytesA = filesA.map(_.length()).sum
+ assert(spilledBytesA > 0)
+ assert(taskContextA.taskMetrics.diskBytesSpilled == spilledBytesA)
+ assert(taskContextA.taskMetrics.memoryBytesSpilled > 0)
+ // ... and, as the files become the map output, also counts as shuffle bytes written...
+ assert(writeMetricsA.bytesWritten == spilledBytesA)
assert(writerA0.getOutputRecords > 0)
- // ... and task B's buffered rows were not spilled, not written out, and not charged.
+ // ... while task B's buffered rows were not spilled, not written out, and not charged.
assert(allocatorB.getUsed == 2 * pageSize)
assert(writerB.getActiveMemoryUsage == 2 * pageSize)
assert(writerB.getOutputRecords == 0)
assert(taskContextB.taskMetrics.diskBytesSpilled == 0)
+ assert(taskContextB.taskMetrics.memoryBytesSpilled == 0)
+ assert(writeMetricsB.bytesWritten == 0)
assert(fileB.length() == 0)
val segmentA0 = writerA0.close()
@@ -188,6 +199,16 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
assert(segmentA2.length > 0)
assert(segmentB.length > 0)
assert(writersA.isEmpty && writersB.isEmpty)
+ // Closing writes the remaining rows as final batches, which count as shuffle bytes written
+ // but not as spill.
+ assert(filesA.map(_.length()).sum > spilledBytesA)
+ assert(writeMetricsA.bytesWritten == filesA.map(_.length()).sum)
+ assert(writeMetricsA.recordsWritten == rowsA0 + rowsA1 + rowsA2)
+ assert(taskContextA.taskMetrics.diskBytesSpilled == spilledBytesA)
+ assert(writeMetricsB.bytesWritten == fileB.length())
+ assert(writeMetricsB.recordsWritten == rowsB)
+ assert(taskContextB.taskMetrics.diskBytesSpilled == 0)
+ assert(taskContextB.taskMetrics.memoryBytesSpilled == 0)
} finally {
Utils.deleteRecursively(tempDir)
tmmA.cleanUpAllAllocatedMemory()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org