feat: experiment with RealUsagePool

docs/source/user-guide/latest/tuning.md

@@ -60,8 +60,9 @@ The valid pool types are:
 
 - `fair_unified` (default when `spark.memory.offHeap.enabled=true` is set)
 - `greedy_unified`
+- `unbounded`
 
-Both pool types are shared across all native execution contexts within the same Spark task. When
+The two `unified` pool types are shared across all native execution contexts within the same Spark task. When
 Comet executes a shuffle, it runs two native execution contexts concurrently (e.g. one for
 pre-shuffle operators and one for the shuffle writer). The shared pool ensures that the combined
 memory usage stays within the per-task limit.
@@ -74,6 +75,12 @@ when there is sufficient memory in order to leave enough memory for other operat
 The `greedy_unified` pool type implements a greedy first-come first-serve limit. This pool works well for queries that do not
 need to spill or have a single spillable operator.
 
+The `unbounded` pool does no accounting of its own and imposes no limit, so Comet's native memory is capped only by
+the experimental `spark.comet.exec.memoryGuard.enabled`. Enabling that guard in off-heap mode selects `unbounded`
+automatically and ignores this setting: growth is then gated on real allocator usage against the off-heap budget
+rather than on Spark's per-task reservations. Note that Spark's `TaskMemoryManager` no longer sees Comet's off-heap
+usage in that mode, so Spark cannot ask Comet to spill on behalf of its own operators.
+
 [shuffle]: #shuffle
 [Advanced Memory Tuning]: #advanced-memory-tuning
 

native/core/Cargo.toml

@@ -103,9 +103,16 @@ datafusion-functions-nested = { version = "55.0.0" }
 
 [features]
 backtrace = ["datafusion/backtrace"]
-default = ["hdfs-opendal"]
+default = ["hdfs-opendal", "oom-guard"]
 hdfs-opendal = ["opendal", "object_store_opendal", "hdfs-sys"]
 jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
+
+# Allocator-level OOM circuit breaker. Wraps the global allocator to track real
+# allocated bytes and gate/abort over-budget query-worker threads. Enabled by default
+# so `spark.comet.exec.memoryGuard.*` works without a special build; an idle guard is
+# near-free (tracking stays off until a task arms it).
+# Drop it from `default` for a bare allocator.
+oom-guard = []
 # Delta Lake integration. When enabled, links the `comet-contrib-delta` crate
 # into `libcomet` and activates the `OpStruct::DeltaScan` dispatcher arm.
 # Default builds carry zero Delta surface.

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

@@ -30,6 +30,22 @@ pub(crate) enum MemoryPoolType {
     Unbounded,
 }
 
+#[cfg(feature = "oom-guard")]
+impl MemoryPoolType {
+    /// True when this pool's `reserved()` reflects a single task's usage, so a per-task
+    /// fair-share comparison is meaningful (false for process-wide pools). The non-shared
+    /// per-task pools (`Greedy`/`FairSpill`) return true but keep no task registry, so the
+    /// fair-share divisor falls back to `executor_cores` rather than the active-task count.
+    pub(crate) fn has_per_task_budget(&self) -> bool {
+        !matches!(
+            self,
+            MemoryPoolType::GreedyGlobal
+                | MemoryPoolType::FairSpillGlobal
+                | MemoryPoolType::Unbounded
+        )
+    }
+}
+
 pub(crate) struct MemoryPoolConfig {
     pub(crate) pool_type: MemoryPoolType,
     pub(crate) pool_size: usize,
@@ -46,20 +62,27 @@ impl MemoryPoolConfig {
 
 pub(crate) fn parse_memory_pool_config(
     off_heap_mode: bool,
-    memory_pool_type: String,
+    memory_pool_type: &str,
     memory_limit: i64,
     memory_limit_per_task: i64,
 ) -> CometResult<MemoryPoolConfig> {
     let pool_size = memory_limit as usize;
     let memory_pool_config = if off_heap_mode {
-        match memory_pool_type.as_str() {
+        match memory_pool_type {
             "fair_unified" => MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size),
             "greedy_unified" => {
                 // the `unified` memory pool interacts with Spark's memory pool to allocate
                 // memory therefore does not need a size to be explicitly set. The pool size
                 // shared with Spark is set by `spark.memory.offHeap.size`.
                 MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0)
             }
+            "unbounded" => {
+                // No accounting of its own. In off-heap mode this is what
+                // `spark.comet.exec.memoryGuard.enabled` forces, so the real-usage gate
+                // wrapped around it is the only thing rejecting growth, instead of
+                // delegating per-task accounting to Spark's TaskMemoryManager.
+                MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0)
+            }
             _ => {
                 return Err(CometError::Config(format!(
                     "Unsupported memory pool type for off-heap mode: {memory_pool_type}"
@@ -69,7 +92,7 @@ pub(crate) fn parse_memory_pool_config(
     } else {
         // Use the memory pool from DF
         let pool_size_per_task = memory_limit_per_task as usize;
-        match memory_pool_type.as_str() {
+        match memory_pool_type {
             "fair_spill_task_shared" => {
                 MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared, pool_size_per_task)
             }

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

@@ -18,6 +18,10 @@
 mod config;
 mod fair_pool;
 pub mod logging_pool;
+#[cfg(feature = "oom-guard")]
+pub(crate) mod oom_guard;
+#[cfg(feature = "oom-guard")]
+mod real_usage_pool;
 mod task_shared;
 mod unified_pool;
 
@@ -32,6 +36,8 @@ use std::sync::Arc;
 use unified_pool::CometUnifiedMemoryPool;
 
 pub(crate) use config::*;
+#[cfg(feature = "oom-guard")]
+pub(crate) use real_usage_pool::RealUsageMemoryPool;
 pub(crate) use task_shared::*;
 
 /// Creates the memory pool for a native plan.