fix: Deallocate row addresses and size arrays after exporting

apache · Apr 7, 2024 · 4317796 · 4317796
1 parent a86626a
commit 4317796
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/spark/src/main/scala/org/apache/spark/shuffle/sort/RowPartition.scala b/spark/src/main/scala/org/apache/spark/shuffle/sort/RowPartition.scala
@@ -32,8 +32,17 @@ class RowPartition(initialSize: Int) {
 
   def getNumRows: Int = rowAddresses.size
 
-  def getRowAddresses: Array[Long] = rowAddresses.toArray
-  def getRowSizes: Array[Int] = rowSizes.toArray
+  def getRowAddresses: Array[Long] = {
+    val array = rowAddresses.toArray
+    rowAddresses = null
+    array
+  }
+
+  def getRowSizes: Array[Int] = {
+    val array = rowSizes.toArray
+    rowSizes = null
+    array
+  }
 
   def reset(): Unit = {
     rowAddresses = new ArrayBuffer[Long](initialSize)