save

Author: andygrove

native/core/src/execution/memory_pools/config.rs

@@ -21,7 +21,6 @@ use crate::errors::{CometError, CometResult};
 pub(crate) enum MemoryPoolType {
     GreedyUnified,
     FairUnified,
-    FairUnifiedGlobal,
     Greedy,
     FairSpill,
     GreedyTaskShared,
@@ -65,7 +64,7 @@ pub(crate) fn parse_memory_pool_config(
     let memory_pool_config = if off_heap_mode {
         match memory_pool_type.as_str() {
             "fair_unified_global" => {
-                MemoryPoolConfig::new(MemoryPoolType::FairUnifiedGlobal, pool_size_global)
+                MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size_global)
             }
             "fair_unified" => {
                 MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size_per_task)

native/core/src/execution/memory_pools/fair_pool.rs

@@ -36,7 +36,7 @@ use parking_lot::Mutex;
 
 /// A DataFusion fair `MemoryPool` implementation for Comet. Internally this is
 /// implemented via delegating calls to [`crate::jvm_bridge::CometTaskMemoryManager`].
-pub struct CometFairMemoryPool {
+pub struct CometFairUnifiedMemoryPool {
     task_memory_manager_handle: Arc<GlobalRef>,
     pool_size: usize,
     state: Mutex<CometFairPoolState>,
@@ -47,7 +47,7 @@ struct CometFairPoolState {
     num: usize,
 }
 
-impl Debug for CometFairMemoryPool {
+impl Debug for CometFairUnifiedMemoryPool {
     fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
         let state = self.state.lock();
         f.debug_struct("CometFairMemoryPool")
@@ -58,19 +58,19 @@ impl Debug for CometFairMemoryPool {
     }
 }
 
-impl CometFairMemoryPool {
+impl CometFairUnifiedMemoryPool {
     pub fn new(
         task_memory_manager_handle: Arc<GlobalRef>,
         pool_size: usize,
-    ) -> CometFairMemoryPool {
+    ) -> CometFairUnifiedMemoryPool {
         Self {
             task_memory_manager_handle,
             pool_size,
             state: Mutex::new(CometFairPoolState { used: 0, num: 0 }),
         }
     }
 
-    fn acquire(&self, additional: usize) -> CometResult<i64> {
+    fn acquire_from_spark(&self, additional: usize) -> CometResult<i64> {
         let mut env = JVMClasses::get_env()?;
         let handle = self.task_memory_manager_handle.as_obj();
         unsafe {
@@ -79,7 +79,7 @@ impl CometFairMemoryPool {
         }
     }
 
-    fn release(&self, size: usize) -> CometResult<()> {
+    fn release_to_spark(&self, size: usize) -> CometResult<()> {
         let mut env = JVMClasses::get_env()?;
         let handle = self.task_memory_manager_handle.as_obj();
         unsafe {
@@ -88,10 +88,10 @@ impl CometFairMemoryPool {
     }
 }
 
-unsafe impl Send for CometFairMemoryPool {}
-unsafe impl Sync for CometFairMemoryPool {}
+unsafe impl Send for CometFairUnifiedMemoryPool {}
+unsafe impl Sync for CometFairUnifiedMemoryPool {}
 
-impl MemoryPool for CometFairMemoryPool {
+impl MemoryPool for CometFairUnifiedMemoryPool {
     fn register(&self, _: &MemoryConsumer) {
         let mut state = self.state.lock();
         state.num = state
@@ -119,7 +119,7 @@ impl MemoryPool for CometFairMemoryPool {
             if size < subtractive {
                 panic!("Failed to release {subtractive} bytes where only {size} bytes reserved")
             }
-            self.release(subtractive)
+            self.release_to_spark(subtractive)
                 .unwrap_or_else(|_| panic!("Failed to release {subtractive} bytes"));
             state.used = state.used.checked_sub(subtractive).unwrap();
         }
@@ -133,20 +133,23 @@ impl MemoryPool for CometFairMemoryPool {
         if additional > 0 {
             let mut state = self.state.lock();
             let num = state.num;
-            let limit = self.pool_size.checked_div(num).unwrap();
+            let limit = self
+                .pool_size
+                .checked_div(num)
+                .expect("overflow in checked_div");
             let size = reservation.size();
             if limit < size + additional {
                 return resources_err!(
                     "Failed to acquire {additional} bytes where {size} bytes already reserved and the fair limit is {limit} bytes, {num} registered"
                 );
             }
 
-            let acquired = self.acquire(additional)?;
+            let acquired = self.acquire_from_spark(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)?;
+                self.release_to_spark(acquired as usize)?;
 
                 return resources_err!(
                     "Failed to acquire {} bytes, only got {} bytes. Reserved: {} bytes",
@@ -155,7 +158,10 @@ impl MemoryPool for CometFairMemoryPool {
                     state.used
                 );
             }
-            state.used = state.used.checked_add(additional).unwrap();
+            state.used = state
+                .used
+                .checked_add(additional)
+                .expect("overflow in checked_add");
         }
         Ok(())
     }

native/core/src/execution/memory_pools/mod.rs

@@ -23,7 +23,7 @@ mod unified_pool;
 use datafusion::execution::memory_pool::{
     FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool,
 };
-use fair_pool::CometFairMemoryPool;
+use fair_pool::CometFairUnifiedMemoryPool;
 use jni::objects::GlobalRef;
 use once_cell::sync::OnceCell;
 use std::num::NonZeroUsize;
@@ -51,8 +51,10 @@ pub(crate) fn create_memory_pool(
         }
         MemoryPoolType::FairUnified => {
             // Set Comet fair memory pool for native
-            let memory_pool =
-                CometFairMemoryPool::new(comet_task_memory_manager, memory_pool_config.pool_size);
+            let memory_pool = CometFairUnifiedMemoryPool::new(
+                comet_task_memory_manager,
+                memory_pool_config.pool_size,
+            );
             Arc::new(TrackConsumersPool::new(
                 memory_pool,
                 NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),