comments

Author: sunchao

core/src/execution/memory_pool.rs

@@ -30,8 +30,13 @@ use datafusion::{
     execution::memory_pool::{MemoryPool, MemoryReservation},
 };
 
-use crate::jvm_bridge::{jni_call, JVMClasses};
+use crate::{
+    errors::CometResult,
+    jvm_bridge::{jni_call, JVMClasses},
+};
 
+/// A DataFusion `MemoryPool` implementation for Comet. Internally this is
+/// implemented via delegating calls to [`crate::jvm_bridge::CometTaskMemoryManager`].
 pub struct CometMemoryPool {
     task_memory_manager_handle: Arc<GlobalRef>,
     used: AtomicUsize,
@@ -52,44 +57,56 @@ impl CometMemoryPool {
             used: AtomicUsize::new(0),
         }
     }
+
+    fn acquire(&self, additional: usize) -> CometResult<i64> {
+        let mut env = JVMClasses::get_env();
+        let handle = self.task_memory_manager_handle.as_obj();
+        unsafe {
+            jni_call!(&mut env,
+              comet_task_memory_manager(handle).acquire_memory(additional as i64) -> i64)
+        }
+    }
+
+    fn release(&self, size: usize) -> CometResult<()> {
+        let mut env = JVMClasses::get_env();
+        let handle = self.task_memory_manager_handle.as_obj();
+        unsafe {
+            jni_call!(&mut env, comet_task_memory_manager(handle).release_memory(size as i64) -> ())
+        }
+    }
 }
 
 unsafe impl Send for CometMemoryPool {}
 unsafe impl Sync for CometMemoryPool {}
 
 impl MemoryPool for CometMemoryPool {
     fn grow(&self, _: &MemoryReservation, additional: usize) {
+        self.acquire(additional)
+            .unwrap_or_else(|_| panic!("Failed to acquire {} bytes", additional));
         self.used.fetch_add(additional, Relaxed);
     }
 
     fn shrink(&self, _: &MemoryReservation, size: usize) {
-        let mut env = JVMClasses::get_env();
-        let handle = self.task_memory_manager_handle.as_obj();
-        unsafe {
-            jni_call!(&mut env, comet_task_memory_manager(handle).release_memory(size as i64) -> ())
-                .unwrap();
-        }
+        self.release(size)
+            .unwrap_or_else(|_| panic!("Failed to release {} bytes", size));
         self.used.fetch_sub(size, Relaxed);
     }
 
     fn try_grow(&self, _: &MemoryReservation, additional: usize) -> Result<(), DataFusionError> {
         if additional > 0 {
-            let mut env = JVMClasses::get_env();
-            let handle = self.task_memory_manager_handle.as_obj();
-            unsafe {
-                let acquired = jni_call!(&mut env,
-                  comet_task_memory_manager(handle).acquire_memory(additional as i64) -> i64)?;
+            let acquired = self.acquire(additional)?;
+            // If the number of bytes we acquired is less than the requested, return an error,
+            // and hopefully will trigger spilling from the caller side.
+            if acquired < additional as i64 {
+                // Release the acquired bytes before throwing error
+                self.release(acquired as usize)?;
 
-                // If the number of bytes we acquired is less than the requested, return an error,
-                // and hopefully will trigger spilling from the caller side.
-                if acquired < additional as i64 {
-                    return Err(DataFusionError::Execution(format!(
-                        "Failed to acquire {} bytes, only got {}. Reserved: {}",
-                        additional,
-                        acquired,
-                        self.reserved(),
-                    )));
-                }
+                return Err(DataFusionError::Execution(format!(
+                    "Failed to acquire {} bytes, only got {}. Reserved: {}",
+                    additional,
+                    acquired,
+                    self.reserved(),
+                )));
             }
             self.used.fetch_add(additional, Relaxed);
         }

core/src/jvm_bridge/comet_task_memory_manager.rs

@@ -24,8 +24,8 @@ use jni::{
 
 use crate::jvm_bridge::get_global_jclass;
 
-/// A DataFusion `MemoryPool` implementation for Comet, which delegate to the JVM
-/// side `CometTaskMemoryManager`.
+/// A wrapper which delegate acquire/release memory calls to the
+/// JVM side `CometTaskMemoryManager`.
 #[derive(Debug)]
 pub struct CometTaskMemoryManager<'a> {
     pub class: JClass<'a>,

spark/src/test/scala/org/apache/spark/sql/CometTPCDSQuerySuite.scala

@@ -20,8 +20,7 @@
 package org.apache.spark.sql
 
 import org.apache.spark.SparkConf
-import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE, SHUFFLE_MANAGER}
-import org.apache.spark.sql.internal.StaticSQLConf.SPARK_SESSION_EXTENSIONS
+import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE}
 
 import org.apache.comet.CometConf
 

spark/src/test/scala/org/apache/spark/sql/CometTPCHQuerySuite.scala

@@ -25,11 +25,10 @@ import java.nio.file.{Files, Paths}
 import scala.collection.JavaConverters._
 
 import org.apache.spark.{SparkConf, SparkContext}
-import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE, SHUFFLE_MANAGER}
+import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE}
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.util.{fileToString, resourceToString, stringToFile}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.internal.StaticSQLConf.SPARK_SESSION_EXTENSIONS
 import org.apache.spark.sql.test.{SharedSparkSession, TestSparkSession}
 
 import org.apache.comet.CometConf

spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

@@ -68,6 +68,8 @@ abstract class CometTestBase
     conf.set(SHUFFLE_MANAGER, shuffleManager)
     conf.set(MEMORY_OFFHEAP_ENABLED.key, "true")
     conf.set(MEMORY_OFFHEAP_SIZE.key, "2g")
+    conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "1g")
+    conf.set(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key, "1g")
     conf.set(CometConf.COMET_ENABLED.key, "true")
     conf.set(CometConf.COMET_EXEC_ENABLED.key, "true")
     conf.set(CometConf.COMET_EXEC_ALL_OPERATOR_ENABLED.key, "true")