You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/08/25 20:26:01 UTC
(datafusion-comet) branch main updated: fix: track peak memory before JVM shuffle spill (#5463)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new c881acf4b3 fix: track peak memory before JVM shuffle spill (#5463)
c881acf4b3 is described below
commit c881acf4b3b9e1cc5038159b61cb54974b2f6b32
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Wed Aug 26 04:25:56 2026 +0800
fix: track peak memory before JVM shuffle spill (#5463)
---
.../shuffle/sort/CometShuffleExternalSorter.java | 1 +
.../spark/sql/comet/CometTaskMetricsSuite.scala | 71 +++++++++++++++++++++-
2 files changed, 71 insertions(+), 1 deletion(-)
diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
index 13ea96dfcb..1a2ba0f96c 100644
--- a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
+++ b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
@@ -196,6 +196,7 @@ public final class CometShuffleExternalSorter implements CometShuffleChecksumSup
activeSpillSorter.setSpillInfo(spillInfo);
activeSpillSorter.writeSortedFileNative(false, tracingEnabled);
+ updatePeakMemoryUsed();
final long spillSize = activeSpillSorter.freeMemory();
activeSpillSorter.reset();
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 86fab4f349..bd3123447e 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
@@ -23,13 +23,17 @@ import java.io.File
import scala.collection.mutable
-import org.apache.spark.{SparkConf, SparkContext, Success}
+import org.apache.spark.{SparkConf, SparkContext, SparkEnv, Success, TaskContext}
import org.apache.spark.executor.ShuffleReadMetrics
import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
import org.apache.spark.scheduler.SparkListener
import org.apache.spark.scheduler.SparkListenerJobStart
import org.apache.spark.scheduler.SparkListenerTaskEnd
+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.CometNativeShuffle
import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
import org.apache.spark.sql.execution.SparkPlan
@@ -37,6 +41,8 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.command.DataWritingCommandExec
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructType}
+import org.apache.spark.unsafe.Platform
import org.apache.comet.CometConf
import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus
@@ -78,6 +84,69 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper {
assert(metricTree.sumMetricValues("memory_spilled_bytes") == 20L)
}
+ test("JVM shuffle peak memory includes a larger earlier spill") {
+ withSQLConf(CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> Int.MaxValue.toString) {
+ val pageSize = 256L
+ val initialSize = 128
+ val conf = new SparkConf(false)
+ .set("spark.memory.offHeap.enabled", "true")
+ .set("spark.memory.offHeap.size", "10m")
+ val memoryManager = new TestMemoryManager(conf)
+ memoryManager.limit(10 * 1024 * 1024)
+ val taskMemoryManager = new TaskMemoryManager(memoryManager, 0)
+ val allocator =
+ CometShuffleMemoryAllocator.getInstance(conf, taskMemoryManager, pageSize)
+ val taskContext = TaskContext.empty()
+ val writeMetrics = taskContext.taskMetrics.shuffleWriteMetrics
+ val sorter = new CometShuffleExternalSorter(
+ allocator,
+ SparkEnv.get.blockManager,
+ taskContext,
+ initialSize,
+ 2,
+ conf,
+ writeMetrics,
+ new StructType().add("id", IntegerType))
+
+ def insert(value: Int): Unit = {
+ val bytes = new Array[Byte](4 + 16)
+ Platform.putInt(bytes, Platform.BYTE_ARRAY_OFFSET, value)
+ val row = new UnsafeRow(1)
+ row.pointTo(bytes, Platform.BYTE_ARRAY_OFFSET + 4, 16)
+ row.setInt(0, value)
+ sorter.insertRecord(bytes, Platform.BYTE_ARRAY_OFFSET, bytes.length, value % 2)
+ }
+
+ try {
+ val initialMemory = allocator.getUsed
+ assert(initialMemory == initialSize * 8L)
+
+ (0 until 21).foreach(insert)
+ val firstHighWater = allocator.getUsed
+ sorter.spill()
+ val firstSpillBytes = taskContext.taskMetrics.memoryBytesSpilled
+
+ (21 until 26).foreach(insert)
+ val secondHighWater = allocator.getUsed
+ sorter.spill()
+ val secondSpillBytes = taskContext.taskMetrics.memoryBytesSpilled - firstSpillBytes
+ val spills = sorter.closeAndGetSpills()
+
+ assert(firstHighWater == initialMemory + firstSpillBytes)
+ assert(secondHighWater == initialMemory + secondSpillBytes)
+ assert(firstHighWater > secondHighWater)
+ assert(spills.length == 2)
+ assert(writeMetrics.recordsWritten == 26L)
+ assert(spills.map(_.file.length()).sum == taskContext.taskMetrics.diskBytesSpilled)
+ assert(spills.map(_.partitionLengths.sum).sum == taskContext.taskMetrics.diskBytesSpilled)
+ assert(sorter.getPeakMemoryUsedBytes() == firstHighWater)
+ } finally {
+ sorter.cleanupResources()
+ assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L)
+ }
+ }
+ }
+
test("per-task native shuffle metrics") {
withParquetTable((0 until 10000).map(i => (i, (i + 1).toLong)), "tbl") {
val df = sql("SELECT * FROM tbl").sortWithinPartitions($"_1".desc)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org