add memory_limit to expose the info in MemoryPool.

Author: Rachelint

datafusion/execution/src/memory_pool/mod.rs

@@ -141,6 +141,21 @@ pub trait MemoryPool: Send + Sync + std::fmt::Debug {
 
     /// Return the total amount of memory reserved
     fn reserved(&self) -> usize;
+
+    /// Return the memory limit of the pool
+    ///
+    /// It is useful to know the limit of the pool sometime (e.g. we can know
+    /// if spilling will be triggered, it only we be triggered when the memory
+    /// limit isn't infinite)
+    fn memory_limit(&self) -> MemoryLimit {
+        MemoryLimit::Unknown
+    }
+}
+
+pub enum MemoryLimit {
+    Infinite,
+    Finite(usize),
+    Unknown,
 }
 
 /// A memory consumer is a named allocation traced by a particular

datafusion/execution/src/memory_pool/pool.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation};
+use crate::memory_pool::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation};
 use datafusion_common::HashMap;
 use datafusion_common::{resources_datafusion_err, DataFusionError, Result};
 use log::debug;
@@ -48,6 +48,10 @@ impl MemoryPool for UnboundedMemoryPool {
     fn reserved(&self) -> usize {
         self.used.load(Ordering::Relaxed)
     }
+
+    fn memory_limit(&self) -> MemoryLimit {
+        MemoryLimit::Infinite
+    }
 }
 
 /// A [`MemoryPool`] that implements a greedy first-come first-serve limit.
@@ -100,6 +104,10 @@ impl MemoryPool for GreedyMemoryPool {
     fn reserved(&self) -> usize {
         self.used.load(Ordering::Relaxed)
     }
+
+    fn memory_limit(&self) -> MemoryLimit {
+        MemoryLimit::Finite(self.pool_size)
+    }
 }
 
 /// A [`MemoryPool`] that prevents spillable reservations from using more than
@@ -233,6 +241,10 @@ impl MemoryPool for FairSpillPool {
         let state = self.state.lock();
         state.spillable + state.unspillable
     }
+
+    fn memory_limit(&self) -> MemoryLimit {
+        MemoryLimit::Finite(self.pool_size)
+    }
 }
 
 /// Constructs a resources error based upon the individual [`MemoryReservation`].
@@ -408,6 +420,10 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
     fn reserved(&self) -> usize {
         self.inner.reserved()
     }
+
+    fn memory_limit(&self) -> MemoryLimit {
+        self.inner.memory_limit()
+    }
 }
 
 fn provide_top_memory_consumers_to_error_msg(

datafusion/execution/src/runtime_env.rs

@@ -253,24 +253,12 @@ impl RuntimeEnvBuilder {
             cache_manager,
             object_store_registry,
         } = self;
-
-        // If `memory_pool` is not set, it represents spilling is disabled.
-        // Because `UnboundedMemoryPool` will be used in this case, and
-        // it will never limit memory usage.
-        // And we disable the `disk_manager`(use `DiskManagerConfig::Disabled`)
-        // to tell we have disabled spilling.
-        let (memory_pool, disk_manager_config) = if let Some(pool) = memory_pool {
-            (pool, disk_manager)
-        } else {
-            (
-                Arc::new(UnboundedMemoryPool::default()) as _,
-                DiskManagerConfig::Disabled,
-            )
-        };
+        let memory_pool =
+            memory_pool.unwrap_or_else(|| Arc::new(UnboundedMemoryPool::default()));
 
         Ok(RuntimeEnv {
             memory_pool,
-            disk_manager: DiskManager::try_new(disk_manager_config)?,
+            disk_manager: DiskManager::try_new(disk_manager)?,
             cache_manager: CacheManager::try_new(&cache_manager)?,
             object_store_registry,
         })