feat: bound JVM Arrow allocations by Spark's off-heap memory pool

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

@@ -135,9 +135,13 @@ a new subclass needs no new case as long as `getValueVector` returns an Arrow ve
 ## 5. Memory Accounting
 
 An FFI change is usually also a memory change, and the two are easy to review separately and miss
-the interaction. Imported JVM buffers come from `CometArrowAllocator`, which is an unbounded
-`RootAllocator(Long.MaxValue)` that no budget sees. Exported native batches are usually no longer
-pool-reserved by the time the JVM receives them, but they stay resident until the JVM closes them.
+the interaction. The JVM buffers native imports belong to a child of `CometArrowAllocator`, an
+unbounded `RootAllocator(Long.MaxValue)` that no budget sees, and a native operator that retains
+them reserves them itself. Batches the JVM owns within a task come from `CometTaskArrowAllocator`,
+whose listener charges them to Spark and refuses what Spark cannot cover, so a new path that exports
+them without going through `CometArrowStream`, which moves the charge off the task, charges them on
+both sides. Exported native batches are usually no longer pool-reserved by the time the JVM receives
+them, but they stay resident until the JVM closes them.
 
 If the PR changes how long either side holds a batch, use `review-comet-memory-pr` as well.
 

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

@@ -158,15 +158,27 @@ memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction
       `spark.executor.memoryOverhead` is the only real slack in the container, and JVM non-heap
       usage already consumes much of it.
 
-## 6. Unaccounted Allocators
+## 6. JVM-Side Allocators
 
 Two allocators sit outside Comet's pool entirely, and they are easy to confuse:
 
 - **`CometArrowAllocator`** (`spark/src/main/scala/org/apache/comet/package.scala`) is a
-  process-wide `RootAllocator(Long.MaxValue)`. It is unbounded and no budget sees it. Child
-  allocators are cut from it for FFI stream export, broadcast coalescing, and
-  `CometSparkToColumnarExec`. A PR that adds a child allocator, or makes an existing one hold more,
-  is adding uncounted container RSS. Say so even if the volume is small.
+  process-wide `RootAllocator(Long.MaxValue)` with no allocation listener. It is unbounded. Buffers
+  allocated from it, or from a child cut with the three-argument `newChildAllocator`, are seen by no
+  budget: FFI stream export, `NativeUtil`, the JVM UDF result, and the import path's
+  `CometArrowImportAllocator`. JVM-owned allocations inside a task go through
+  `CometTaskArrowAllocator.forCurrentTask()` instead, a per-task child whose
+  `CometArrowAllocationListener` charges what it owns to that task's `TaskMemoryManager` and refuses
+  an allocation Spark cannot cover, so it fails with Arrow's `OutOfMemoryException`. A PR that
+  allocates from the root inside a task, or makes a root-backed allocator hold more, is adding
+  uncounted container RSS. Say so even if the volume is small. A PR that allocates something native
+  will retain from the task allocator is charging it twice, unless the batch reaches native through
+  `CometArrowStream`, whose reader moves the charge off the task allocator. A new call site on the
+  task allocator also has to cope with that refusal: nothing spills these buffers, so it fails the
+  task. Check that the listener stays without a lock in `getUsed` and `spill`, that only
+  `onPreAllocation` throws, and that it sizes the reservation from what the allocator owns rather
+  than from a sum of its callbacks, because an ownership transfer between allocators calls no
+  listener.
 - **JVM shuffle pages** go through `CometShuffleMemoryAllocator.getInstance`, which returns
   `CometUnifiedShuffleMemoryAllocator`, an ordinary Spark `MemoryConsumer` drawing from
   `spark.memory.offHeap.size`. These **are** arbitrated by Spark against its other consumers. Do

.github/workflows/pr_build_linux.yml

@@ -543,6 +543,7 @@ jobs:
               org.apache.spark.CometPluginsSuite
               org.apache.spark.CometRuntimeShutdownSuite
               org.apache.spark.CometTaskMemoryManagerSuite
+              org.apache.spark.comet.CometArrowAllocationListenerSuite
               org.apache.spark.CometExecIteratorLifecycleSuite
               org.apache.spark.CometPluginsDefaultSuite
               org.apache.spark.CometPluginsMemoryOverheadWarningSuite

.github/workflows/pr_build_macos.yml

@@ -191,6 +191,7 @@ jobs:
               org.apache.spark.CometPluginsSuite
               org.apache.spark.CometRuntimeShutdownSuite
               org.apache.spark.CometTaskMemoryManagerSuite
+              org.apache.spark.comet.CometArrowAllocationListenerSuite
               org.apache.spark.CometExecIteratorLifecycleSuite
               org.apache.spark.CometPluginsDefaultSuite
               org.apache.spark.CometPluginsMemoryOverheadWarningSuite