(arrow-datafusion-comet) branch main updated: fix: Include active spiller when computing peak shuffle memory (#196)

[sunchao]

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/arrow-datafusion-comet.git


The following commit(s) were added to refs/heads/main by this push:
     new e920aa4  fix: Include active spiller when computing peak shuffle 
memory (#196)
e920aa4 is described below

commit e920aa4f956478c7cdcb939a1689779daef04b09
Author: Chao Sun <sunc...@apache.org>
AuthorDate: Tue Mar 12 22:50:54 2024 -0700

    fix: Include active spiller when computing peak shuffle memory (#196)
---
 .../org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java     | 4 ++++
 1 file changed, 4 insertions(+)

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 aa806e2..9fe88ec 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
@@ -257,6 +257,9 @@ public final class CometShuffleExternalSorter implements 
CometShuffleChecksumSup
     for (SpillSorter sorter : spillingSorters) {
       totalPageSize += sorter.getMemoryUsage();
     }
+    if (activeSpillSorter != null) {
+      totalPageSize += activeSpillSorter.getMemoryUsage();
+    }
     return totalPageSize;
   }
 
@@ -274,6 +277,7 @@ public final class CometShuffleExternalSorter implements 
CometShuffleChecksumSup
   }
 
   private long freeMemory() {
+    updatePeakMemoryUsed();
     long memoryFreed = 0;
     if (isAsync) {
       for (SpillSorter sorter : spillingSorters) {