fix: Peak memory usage should include active spiller

apache · Mar 12, 2024 · 7d97c5b · 7d97c5b
1 parent 77ed549
commit 7d97c5b
Showing 1 changed file with 2 additions and 0 deletions.
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
@@ -257,6 +257,7 @@ private long getMemoryUsage() {
     for (SpillSorter sorter : spillingSorters) {
       totalPageSize += sorter.getMemoryUsage();
     }
+    totalPageSize += activeSpillSorter.getMemoryUsage();
     return totalPageSize;
   }
 
@@ -274,6 +275,7 @@ public long getPeakMemoryUsedBytes() {
   }
 
   private long freeMemory() {
+    updatePeakMemoryUsed();
     long memoryFreed = 0;
     if (isAsync) {
       for (SpillSorter sorter : spillingSorters) {