feat: remove memory accounting from Comet's on-heap mode

.ai/skills/review-comet-memory-pr/SKILL.md

@@ -115,9 +115,11 @@ configuration, from the inside out:
       so it only removes an entry that is still its own, which handles the race where an `acquire`
       observes an expired `Weak` and inserts a replacement first. Do not let that check be
       simplified away.
-- [ ] **Only `fair_unified` and `greedy_unified` are valid in off-heap mode.** Other pool types are
-      on-heap only and belong to `CATEGORY_TESTING`, because on-heap mode exists so the Spark SQL
-      test suite can run against Comet and must not be used in production.
+- [ ] **`fair_unified` and `greedy_unified` are the only pool types.** On-heap mode ignores the
+      pool-type string and always gets `UnboundedMemoryPool`: it exists so the Spark SQL test suite
+      can run against Comet, it accounts for nothing, and it must not be used in production. A PR
+      re-adding a sized on-heap pool is reintroducing a budget that bounds nothing real (see
+      issue #6063).
 
 ## 4. The Spark Bridge
 
@@ -145,8 +147,6 @@ memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction
       native code.
 - [ ] Changing `memoryPool.fraction` semantics affects every deployment that tuned it as a haircut
       for the accounting gap.
-- [ ] `memory_limit_per_task` is read only by the on-heap pool types. A PR wiring it into an
-      off-heap path is probably confused.
 - [ ] On Kubernetes, `spark.memory.offHeap.size` is **part of** the pod limit, not headroom on top
       of it. A PR whose fix is "raise the off-heap size" is asking for fewer executors per node.
       `spark.executor.memoryOverhead` is the only real slack in the container, and JVM non-heap
@@ -182,7 +182,7 @@ what test was added. Ask for at least one of:
   is recoverable at task level.
 
 `spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala` and
-`CometBoundedShuffleMemoryAllocatorSuite.scala` are the existing JVM-side tests. A change to the
+`CometUnboundedShuffleMemoryAllocatorSuite.scala` are the existing JVM-side tests. A change to the
 bridge or an allocator should extend one of them.
 
 ## 8. Does the PR Make `memory_management.md` Stale?

.ai/skills/review-comet-shuffle-pr/SKILL.md

@@ -184,16 +184,16 @@ crates before, and the "Key Classes" tables in the docs are exactly what goes st
 
 ## 7. Tests
 
-| Suite                                                               | Covers                                           |
-| ------------------------------------------------------------------- | ------------------------------------------------ |
-| `org.apache.comet.exec.CometNativeShuffleSuite`                     | Native shuffle end to end                        |
-| `org.apache.comet.exec.CometColumnarShuffleSuite`                   | JVM columnar shuffle end to end                  |
-| `CometShuffle4_0Suite`                                              | Spark 4.x specific behavior                      |
-| `CometDiskBlockWriterSuite`                                         | JVM spill and page handling                      |
-| `NativeBatchDecoderIteratorLifecycleChecks`, `...ConcurrencyChecks` | Reader lifetime and concurrency                  |
-| `CometNativeShuffleInputRDDSuite`                                   | The scheduling-anchor RDD                        |
-| `CometCeleborn*Suite`                                               | The Celeborn path, which is easy to forget       |
-| `CometShuffleBenchmark`                                             | Throughput, needs `-Dspark.comet.memoryOverhead` |
+| Suite                                                               | Covers                                     |
+| ------------------------------------------------------------------- | ------------------------------------------ |
+| `org.apache.comet.exec.CometNativeShuffleSuite`                     | Native shuffle end to end                  |
+| `org.apache.comet.exec.CometColumnarShuffleSuite`                   | JVM columnar shuffle end to end            |
+| `CometShuffle4_0Suite`                                              | Spark 4.x specific behavior                |
+| `CometDiskBlockWriterSuite`                                         | JVM spill and page handling                |
+| `NativeBatchDecoderIteratorLifecycleChecks`, `...ConcurrencyChecks` | Reader lifetime and concurrency            |
+| `CometNativeShuffleInputRDDSuite`                                   | The scheduling-anchor RDD                  |
+| `CometCeleborn*Suite`                                               | The Celeborn path, which is easy to forget |
+| `CometShuffleBenchmark`                                             | Throughput                                 |
 
 Ask specifically:
 

.github/workflows/pr_build_linux.yml

@@ -513,7 +513,7 @@ jobs:
               org.apache.comet.exec.CometAsyncShuffleSuite
               org.apache.comet.exec.DisableAQECometShuffleSuite
               org.apache.comet.exec.DisableAQECometAsyncShuffleSuite
-              org.apache.spark.shuffle.comet.CometBoundedShuffleMemoryAllocatorSuite
+              org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite
               org.apache.spark.shuffle.sort.SpillSorterSuite
           - name: "exec"
             value: |

.github/workflows/pr_build_macos.yml

@@ -161,7 +161,7 @@ jobs:
               org.apache.comet.exec.CometAsyncShuffleSuite
               org.apache.comet.exec.DisableAQECometShuffleSuite
               org.apache.comet.exec.DisableAQECometAsyncShuffleSuite
-              org.apache.spark.shuffle.comet.CometBoundedShuffleMemoryAllocatorSuite
+              org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite
               org.apache.spark.shuffle.sort.SpillSorterSuite
           - name: "exec"
             value: |

dev/diffs/3.4.3.diff

@@ -3053,10 +3053,10 @@ index dd55fcfe42c..d9a3f2df535 100644
 
      spark.internalCreateDataFrame(withoutFilters.execute(), schema)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-index ed2e309fa07..040013dc8ab 100644
+index ed2e309fa07..54d417624ff 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-@@ -74,6 +74,20 @@ trait SharedSparkSessionBase
+@@ -74,6 +74,19 @@ trait SharedSparkSessionBase
        // this rule may potentially block testing of other optimization rules such as
        // ConstantPropagation etc.
        .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
@@ -3072,7 +3072,6 @@ index ed2e309fa07..040013dc8ab 100644
 +        .set("spark.shuffle.manager",
 +          "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
 +        .set("spark.comet.shuffle.enabled", "true")
-+        .set("spark.comet.memoryOverhead", "4g")
 +    }
      conf.set(
        StaticSQLConf.WAREHOUSE_PATH,

dev/diffs/3.5.9.diff

@@ -3064,10 +3064,10 @@ index e937173a590..263934fbe7b 100644
 
      spark.internalCreateDataFrame(withoutFilters.execute(), schema)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-index c23bf4204f7..f91ba524e6e 100644
+index c23bf4204f7..07d215aad2b 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-@@ -97,6 +97,20 @@ trait SharedSparkSessionBase
+@@ -97,6 +97,19 @@ trait SharedSparkSessionBase
        // this rule may potentially block testing of other optimization rules such as
        // ConstantPropagation etc.
        .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
@@ -3083,7 +3083,6 @@ index c23bf4204f7..f91ba524e6e 100644
 +        .set("spark.shuffle.manager",
 +          "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
 +        .set("spark.comet.shuffle.enabled", "true")
-+        .set("spark.comet.memoryOverhead", "2g")
 +    }
      conf.set(
        StaticSQLConf.WAREHOUSE_PATH,

dev/diffs/4.0.4.diff

@@ -3030,7 +3030,7 @@ index 30503af0fab..1491f4bc2d5 100644
 
    import testImplicits._
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
-index 08fd8a9ecb5..06967aec8e1 100644
+index 08fd8a9ecb5..e0b8cada307 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
 @@ -27,6 +27,7 @@ import org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE
@@ -3847,10 +3847,10 @@ index f0f3f94b811..b7d18771314 100644
 
      spark.internalCreateDataFrame(withoutFilters.execute(), schema)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-index 720b13b812e..93221ef4cf5 100644
+index 720b13b812e..e3ac2cebc6e 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-@@ -98,6 +98,21 @@ trait SharedSparkSessionBase
+@@ -98,6 +98,20 @@ trait SharedSparkSessionBase
        // this rule may potentially block testing of other optimization rules such as
        // ConstantPropagation etc.
        .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
@@ -3866,7 +3866,6 @@ index 720b13b812e..93221ef4cf5 100644
 +        .set("spark.shuffle.manager",
 +          "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
 +        .set("spark.comet.shuffle.enabled", "true")
-+        .set("spark.comet.memoryOverhead", "2g")
 +
 +    }
      conf.set(

dev/diffs/4.1.3.diff

@@ -3196,7 +3196,7 @@ index 30503af0fab..1491f4bc2d5 100644
 
    import testImplicits._
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
-index 08fd8a9ecb5..06967aec8e1 100644
+index 08fd8a9ecb5..e0b8cada307 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
 @@ -27,6 +27,7 @@ import org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE
@@ -4141,10 +4141,10 @@ index f0f3f94b811..b7d18771314 100644
 
      spark.internalCreateDataFrame(withoutFilters.execute(), schema)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-index 720b13b812e..93221ef4cf5 100644
+index 720b13b812e..e3ac2cebc6e 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-@@ -98,6 +98,21 @@ trait SharedSparkSessionBase
+@@ -98,6 +98,20 @@ trait SharedSparkSessionBase
        // this rule may potentially block testing of other optimization rules such as
        // ConstantPropagation etc.
        .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
@@ -4160,7 +4160,6 @@ index 720b13b812e..93221ef4cf5 100644
 +        .set("spark.shuffle.manager",
 +          "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
 +        .set("spark.comet.shuffle.enabled", "true")
-+        .set("spark.comet.memoryOverhead", "2g")
 +
 +    }
      conf.set(

docs/source/contributor-guide/jvm_shuffle.md

@@ -190,8 +190,20 @@ writes the same Arrow IPC block format as native shuffle, so direct read applies
 
 ## Memory Management
 
-- `CometShuffleMemoryAllocator`: Custom allocator for off-heap memory pages
-- Memory is allocated in pages; when allocation fails, writers spill to disk
+- `CometShuffleMemoryAllocator.getInstance` returns the allocator for the task. Pages are always
+  `Unsafe`-allocated, because their addresses are handed to native code, but what bounds them
+  depends on Spark's memory mode.
+- Off-heap mode gets `CometUnifiedShuffleMemoryAllocator`, an ordinary Spark `MemoryConsumer`
+  drawing from `spark.memory.offHeap.size`. When it cannot acquire a page the writer spills to
+  disk.
+- On-heap mode gets `CometUnboundedShuffleMemoryAllocator`, which keeps no budget and so never
+  refuses a page. Nothing bounds these allocations, and memory pressure never triggers a spill.
+  That mode exists only so the Spark SQL tests can run against Comet. See
+  [Memory Management](memory_management.md).
+- Row count still triggers spilling in either mode. `CometDiskBlockWriter` spills at
+  `min(spark.comet.shuffle.jvm.spillThreshold, spark.comet.shuffle.jvm.batchSize)`.
+  `CometShuffleExternalSorter` spills at `spark.comet.shuffle.jvm.spillThreshold` alone, which
+  defaults to `Int.MaxValue`, so on the sort path that trigger is effectively off by default.
 - `CometDiskBlockWriter` coordinates spilling across all partition writers (largest first)
 
 ## Configuration

docs/source/contributor-guide/memory_management.md

@@ -26,9 +26,12 @@ anyone debugging an out-of-memory report. For user-facing tuning advice, see the
 
 This page covers off-heap mode (`spark.memory.offHeap.enabled=true`) only. Comet also has an
 on-heap mode, but it exists so that the Spark SQL test suite can run against Comet without changing
-Spark's memory configuration. It must not be used in production, and it is not described here. The
-pool types that only on-heap mode exposes belong to the `CATEGORY_TESTING` config group for the
-same reason.
+Spark's memory configuration. Comet performs no memory accounting in it: the native side gets
+DataFusion's `UnboundedMemoryPool` and the JVM shuffle allocator
+(`CometUnboundedShuffleMemoryAllocator`) hands out `Unsafe` pages against no budget. Native memory
+is not on the JVM heap, so there is no Spark pool it could honestly be charged to, and the
+fixed-size pool that used to stand in for one bounded nothing the container cares about. On-heap
+mode must not be used in production, and it is not described further here.
 
 ## Overview
 
@@ -203,9 +206,6 @@ Comet's under-accounting (see [The accounting gap](#the-accounting-gap)). It hol
 the off-heap pool that Comet is not allowed to reserve, on the assumption that Comet's real usage
 overshoots its reservations by roughly that slice.
 
-A second value, `memory_limit_per_task`, is computed and passed alongside it, but only the on-heap
-pool types read it.
-
 ### Resolving the pool type
 
 `parse_memory_pool_config` (`native/core/src/execution/memory_pools/config.rs`) turns the pool-type
@@ -216,7 +216,8 @@ string and the limit into a `MemoryPoolConfig`. Two pool types are valid in off-
 | `fair_unified` (default) | `memory_limit`      | Delegates to Spark's `TaskMemoryManager`; task-shared |
 | `greedy_unified`         | n/a (pool size `0`) | Spark owns the limit entirely; task-shared            |
 
-Any other pool type is rejected with a configuration error.
+Any other pool type is rejected with a configuration error. In on-heap mode the pool-type string is
+ignored and the pool is always `UnboundedMemoryPool`.
 
 ## The pool stack
 

native/core/src/execution/jni_api.rs

@@ -565,7 +565,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
     off_heap_mode: jboolean,
     memory_pool_type: JString,
     memory_limit: jlong,
-    memory_limit_per_task: jlong,
     task_attempt_id: jlong,
     task_cpus: jlong,
     key_unwrapper_obj: JObject,
@@ -621,7 +620,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
                 off_heap_mode != JNI_FALSE,
                 memory_pool_type,
                 memory_limit,
-                memory_limit_per_task,
             )?;
             let memory_pool =
                 create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id);

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

@@ -21,12 +21,6 @@ use crate::errors::{CometError, CometResult};
 pub(crate) enum MemoryPoolType {
     GreedyUnified,
     FairUnified,
-    Greedy,
-    FairSpill,
-    GreedyTaskShared,
-    FairSpillTaskShared,
-    GreedyGlobal,
-    FairSpillGlobal,
     Unbounded,
 }
 
@@ -48,47 +42,30 @@ pub(crate) fn parse_memory_pool_config(
     off_heap_mode: bool,
     memory_pool_type: String,
     memory_limit: i64,
-    memory_limit_per_task: i64,
 ) -> CometResult<MemoryPoolConfig> {
+    if !off_heap_mode {
+        // On-heap mode exists so that the Spark SQL tests can run against Comet without changing
+        // Spark's memory configuration. Comet's native allocations are not on the JVM heap, so
+        // there is no Spark pool they can honestly be charged to, and the fixed-size pool that
+        // used to stand in for one bounded nothing the container cares about. It is not a
+        // production configuration, so it accounts for nothing.
+        return Ok(MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0));
+    }
+
     let pool_size = memory_limit as usize;
-    let memory_pool_config = if off_heap_mode {
-        match memory_pool_type.as_str() {
-            "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)
-            }
-            _ => {
-                return Err(CometError::Config(format!(
-                    "Unsupported memory pool type for off-heap mode: {memory_pool_type}"
-                )))
-            }
-        }
-    } else {
-        // Use the memory pool from DF
-        let pool_size_per_task = memory_limit_per_task as usize;
-        match memory_pool_type.as_str() {
-            "fair_spill_task_shared" => {
-                MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared, pool_size_per_task)
-            }
-            "greedy_task_shared" => {
-                MemoryPoolConfig::new(MemoryPoolType::GreedyTaskShared, pool_size_per_task)
-            }
-            "fair_spill_global" => {
-                MemoryPoolConfig::new(MemoryPoolType::FairSpillGlobal, pool_size)
-            }
-            "greedy_global" => MemoryPoolConfig::new(MemoryPoolType::GreedyGlobal, pool_size),
-            "fair_spill" => MemoryPoolConfig::new(MemoryPoolType::FairSpill, pool_size_per_task),
-            "greedy" => MemoryPoolConfig::new(MemoryPoolType::Greedy, pool_size_per_task),
-            "unbounded" => MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0),
-            _ => {
-                return Err(CometError::Config(format!(
-                    "Unsupported memory pool type for on-heap mode: {memory_pool_type}"
-                )))
-            }
+    match memory_pool_type.as_str() {
+        "fair_unified" => Ok(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`.
+            Ok(MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0))
         }
-    };
-    Ok(memory_pool_config)
+        _ => Err(CometError::Config(format!(
+            "Unsupported memory pool type: {memory_pool_type}"
+        ))),
+    }
 }