Memory pool fix for Arrow 7.0

zhztheplayer · Apr 8, 2022 · 1543aff · 1543aff
1 parent c5eddea
commit 1543aff
Showing 1 changed file with 20 additions and 18 deletions.
diff --git a/...src/main/scala/org/apache/spark/sql/execution/datasources/v2/arrow/SparkMemoryUtils.scala b/...src/main/scala/org/apache/spark/sql/execution/datasources/v2/arrow/SparkMemoryUtils.scala
@@ -109,23 +109,24 @@ object SparkMemoryUtils extends Logging {
         .newChildAllocator("CHILD-ALLOC-BUFFER-IMPORT", allocListenerForBufferImport, 0L,
           Long.MaxValue)
 
-    val defaultMemoryPool: NativeMemoryPoolWrapper = {
-      val rl = new SparkManagedReservationListener(
-        new NativeSQLMemoryConsumer(getTaskMemoryManager(), Spiller.NO_OP),
-        sharedMetrics)
-      val pool = NativeMemoryPoolWrapper(NativeMemoryPool.createListenable(rl), rl,
-        collectStackForDebug)
-      memoryPools.add(pool)
-      pool
-    }
+//    val defaultMemoryPool: NativeMemoryPoolWrapper = {
+//      val rl = new SparkManagedReservationListener(
+//        new NativeSQLMemoryConsumer(getTaskMemoryManager(), Spiller.NO_OP),
+//        sharedMetrics)
+//      val pool = NativeMemoryPoolWrapper(NativeMemoryPool.createListenable(rl), rl,
+//        collectStackForDebug)
+//      memoryPools.add(pool)
+//      pool
+//    }
 
     def createSpillableMemoryPool(spiller: Spiller): NativeMemoryPool = {
-      val rl = new SparkManagedReservationListener(
-        new NativeSQLMemoryConsumer(getTaskMemoryManager(), spiller),
-        sharedMetrics)
-      val pool = NativeMemoryPool.createListenable(rl)
-      memoryPools.add(NativeMemoryPoolWrapper(pool, rl, collectStackForDebug))
-      pool
+      return NativeMemoryPool.getDefault
+//      val rl = new SparkManagedReservationListener(
+//        new NativeSQLMemoryConsumer(getTaskMemoryManager(), spiller),
+//        sharedMetrics)
+//      val pool = NativeMemoryPool.createListenable(rl)
+//      memoryPools.add(NativeMemoryPoolWrapper(pool, rl, collectStackForDebug))
+//      pool
     }
 
     def createSpillableAllocator(spiller: Spiller): BufferAllocator = {
@@ -301,10 +302,11 @@ object SparkMemoryUtils extends Logging {
   }
 
   def contextMemoryPool(): NativeMemoryPool = {
-    if (!inSparkTask()) {
       return globalMemoryPool()
-    }
-    getTaskMemoryResources().defaultMemoryPool.pool
+//    if (!inSparkTask()) {
+//      return globalMemoryPool()
+//    }
+//    getTaskMemoryResources().defaultMemoryPool.pool
   }
 
   def getLeakedAllocators(): List[BufferAllocator] = {