From fda7bd63407d6970345491baaef7f130e734b946 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 08:52:40 -0600
Subject: [PATCH 01/18] feat: report JVM Arrow allocations to Spark's memory
 manager

CometArrowAllocator is a single process-wide RootAllocator with no limit, so the off-heap bytes it hands out were counted by neither Spark's TaskMemoryManager nor Comet's native memory pool, despite being resident in the container and contributing to container OOM kills. Attach an AllocationListener that charges each allocation to a Spark MemoryConsumer belonging to the task that made it.

Reservations are grown and shrunk in whole blocks, 1 MiB by default, because Arrow allocates per buffer and acquireExecutionMemory takes locks. Arrow passes a parent's listener down to child allocators, so a single attachment point covers every call site including those that cut child allocators.

Report without enforcing. A short grant from Spark is logged and the allocation proceeds, because Arrow allocation on these paths cannot fail today and making it fail is a behavioural change that belongs in its own commit.

Do nothing in three cases: when there is no active task, since broadcast coalescing and the cached batch serializer can allocate off a task thread; in on-heap mode, which exists so the Spark SQL suite can run without off-heap memory configured; and for buffers released after their allocating task has finished, since the allocator is process-wide precisely because buffers can outlive their task.

Part of #5997.
---
 .github/workflows/pr_build_linux.yml          |   1 +
 .github/workflows/pr_build_macos.yml          |   1 +
 .../main/scala/org/apache/comet/package.scala |   8 +-
 .../comet/CometArrowAllocationListener.scala  | 221 ++++++++++++++++++
 .../CometArrowAllocationListenerSuite.scala   | 132 +++++++++++
 5 files changed, 362 insertions(+), 1 deletion(-)
 create mode 100644 spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
 create mode 100644 spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala

diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index fd1b0a6c0a7..89c45083839 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -544,6 +544,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.CometPluginsNonOverrideSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 1d1aa4ff262..c50dcc5827f 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -183,6 +183,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.CometPluginsNonOverrideSuite
diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala
index 0eb65c9ba6b..55982b3645e 100644
--- a/spark/src/main/scala/org/apache/comet/package.scala
+++ b/spark/src/main/scala/org/apache/comet/package.scala
@@ -22,6 +22,7 @@ package org.apache
 import java.util.Properties
 
 import org.apache.arrow.memory.RootAllocator
+import org.apache.spark.comet.CometArrowAllocationListener
 import org.apache.spark.internal.Logging
 
 package object comet {
@@ -32,8 +33,13 @@ package object comet {
    * Until the reference count is zero, the memory will not be released. If the consumer side is
    * finished later than the close of the allocator, the allocator will think the memory is
    * leaked. To avoid this, we use a single allocator for the whole execution process.
+   *
+   * The allocator itself is unlimited, but [[CometArrowAllocationListener]] reports every
+   * allocation to Spark's memory manager so that these off-heap bytes are no longer invisible to
+   * Spark's accounting. It reports without enforcing, so allocation here still cannot fail.
    */
-  val CometArrowAllocator = new RootAllocator(Long.MaxValue)
+  val CometArrowAllocator =
+    new RootAllocator(new CometArrowAllocationListener, Long.MaxValue)
 
   /**
    * Provides access to build information about the Comet libraries. This will be used by the
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
new file mode 100644
index 00000000000..544b0503ad4
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.comet
+
+import java.util.concurrent.ConcurrentHashMap
+
+import org.apache.arrow.memory.AllocationListener
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager}
+
+/**
+ * Reports JVM-side Arrow allocations to Spark's memory manager.
+ *
+ * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, so until now the
+ * off-heap bytes it hands out were counted by nobody: not Spark's `TaskMemoryManager`, and not
+ * Comet's native memory pool. They are still resident in the container, which makes them a blind
+ * spot when an executor is killed for exceeding its memory limit.
+ *
+ * This listener closes the reporting half of that gap. Every allocation is charged to a
+ * [[MemoryConsumer]] belonging to the task that made it, so the bytes appear in
+ * `TaskMemoryManager.showMemoryUsage` and are arbitrated against Spark's other off-heap
+ * consumers.
+ *
+ * It deliberately does not enforce. A short grant from Spark is logged and the allocation
+ * proceeds, because Arrow allocation on these paths cannot fail today and making it fail is a
+ * behavioural change that belongs in its own commit. See
+ * [[https://github.com/apache/datafusion-comet/issues/5997]].
+ *
+ * Three cases are handled by doing nothing, each for a different reason:
+ *   - No active task. Broadcast coalescing and the cached batch serializer can allocate from the
+ *     driver or a non-task thread, where there is no task to charge.
+ *   - On-heap mode. Comet's on-heap mode exists so the Spark SQL suite can run without off-heap
+ *     memory configured; charging an off-heap consumer there would be wrong.
+ *   - A buffer released after its allocating task has finished. The allocator is process-wide
+ *     precisely because buffers can outlive the task that created them, so the task's reservation
+ *     is dropped at task end and later releases are ignored rather than double-counted.
+ */
+class CometArrowAllocationListener extends AllocationListener with Logging {
+
+  import CometArrowAllocationListener._
+
+  private val reservations = new ConcurrentHashMap[Long, TaskReservation]()
+
+  @volatile private var configResolved = false
+  @volatile private var accountingEnabled = true
+  @volatile private var blockSize = DEFAULT_BLOCK_SIZE
+
+  override def onAllocation(size: Long): Unit = {
+    val reservation = reservationForCurrentTask()
+    if (reservation != null) {
+      reservation.allocated(size)
+    }
+  }
+
+  override def onRelease(size: Long): Unit = {
+    val reservation = reservationForCurrentTask()
+    if (reservation != null) {
+      reservation.released(size)
+    }
+  }
+
+  /** Bytes currently reserved with Spark on behalf of the given task. Visible for testing. */
+  private[comet] def reservedBytesForTask(taskAttemptId: Long): Long = {
+    val reservation = reservations.get(taskAttemptId)
+    if (reservation == null) 0L else reservation.reservedBytes
+  }
+
+  private[comet] def trackedTaskCount: Int = reservations.size()
+
+  /**
+   * Resolves configuration from the `SparkConf` rather than a `SQLConf` entry. This listener is
+   * attached to a `val` in a package object, so it is constructed on first touch of
+   * `CometArrowAllocator`, which can happen before any `SparkSession` exists and on executors
+   * where `SQLConf` does not carry Comet's settings. `SparkEnv` is absent until the executor is
+   * up, so the read is retried until it succeeds rather than cached from a null environment.
+   */
+  private def resolveConfig(): Unit = {
+    if (!configResolved) {
+      val env = SparkEnv.get
+      if (env != null) {
+        accountingEnabled = env.conf.getBoolean(ACCOUNTING_ENABLED_KEY, defaultValue = true)
+        blockSize = env.conf.getSizeAsBytes(BLOCK_SIZE_KEY, DEFAULT_BLOCK_SIZE_STRING)
+        configResolved = true
+      }
+    }
+  }
+
+  private def reservationForCurrentTask(): TaskReservation = {
+    resolveConfig()
+    if (!accountingEnabled) return null
+
+    val taskContext = TaskContext.get()
+    if (taskContext == null) return null
+
+    val taskMemoryManager = taskContext.taskMemoryManager()
+    if (taskMemoryManager == null ||
+      taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) {
+      return null
+    }
+
+    val taskAttemptId = taskContext.taskAttemptId()
+    val existing = reservations.get(taskAttemptId)
+    if (existing != null) return existing
+
+    val created = new TaskReservation(taskMemoryManager, blockSize, this)
+    val previous = reservations.putIfAbsent(taskAttemptId, created)
+    if (previous != null) return previous
+
+    taskContext.addTaskCompletionListener[Unit] { _ =>
+      val finished = reservations.remove(taskAttemptId)
+      if (finished != null) {
+        finished.close()
+      }
+    }
+    created
+  }
+
+  private[comet] def warnOnShortGrant(requested: Long, granted: Long): Unit = {
+    if (!shortGrantLogged) {
+      shortGrantLogged = true
+      logWarning(
+        s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " +
+          "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " +
+          s"Set $ACCOUNTING_ENABLED_KEY=false to stop reporting these allocations to Spark.")
+    }
+  }
+
+  @volatile private var shortGrantLogged = false
+}
+
+object CometArrowAllocationListener {
+
+  val ACCOUNTING_ENABLED_KEY = "spark.comet.arrowAllocator.accounting.enabled"
+  val BLOCK_SIZE_KEY = "spark.comet.arrowAllocator.accounting.blockSize"
+
+  private val DEFAULT_BLOCK_SIZE_STRING = "1m"
+  private val DEFAULT_BLOCK_SIZE = 1024L * 1024L
+
+  /**
+   * One task's reservation against Spark's off-heap pool.
+   *
+   * Arrow allocates per buffer, and `acquireExecutionMemory` takes locks, so reserving for every
+   * buffer would be needlessly chatty. Instead the reservation is grown and shrunk in whole
+   * blocks and only block-crossing changes reach Spark.
+   */
+  private class TaskReservation(
+      taskMemoryManager: TaskMemoryManager,
+      blockSize: Long,
+      listener: CometArrowAllocationListener)
+      extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) {
+
+    // Named `usedBytes` rather than `used` on purpose: `MemoryConsumer` already declares a
+    // `protected long used`, and a private field of that name narrows the inherited member, which
+    // the compiler rejects as weaker access privileges in overriding.
+    private var usedBytes: Long = 0L
+    private var reserved: Long = 0L
+
+    /** Comet's native operators cannot be made to spill from here. See issue #5997. */
+    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+
+    /**
+     * Reports our own tally. The inherited `used` counter stays at zero because this consumer
+     * never calls `acquireMemory` or `allocatePage`; Arrow has already obtained the memory and we
+     * are only accounting for it.
+     */
+    override def getUsed: Long = synchronized(usedBytes)
+
+    def reservedBytes: Long = synchronized(reserved)
+
+    def allocated(size: Long): Unit = synchronized {
+      usedBytes += size
+      while (reserved < usedBytes) {
+        val request = math.max(blockSize, usedBytes - reserved)
+        val granted = taskMemoryManager.acquireExecutionMemory(request, this)
+        if (granted <= 0L) {
+          listener.warnOnShortGrant(request, granted)
+          return
+        }
+        reserved += granted
+        if (granted < request) {
+          listener.warnOnShortGrant(request, granted)
+          return
+        }
+      }
+    }
+
+    def released(size: Long): Unit = synchronized {
+      usedBytes = math.max(0L, usedBytes - size)
+      while (reserved - usedBytes >= blockSize) {
+        taskMemoryManager.releaseExecutionMemory(blockSize, this)
+        reserved -= blockSize
+      }
+    }
+
+    def close(): Unit = synchronized {
+      if (reserved > 0L) {
+        taskMemoryManager.releaseExecutionMemory(reserved, this)
+        reserved = 0L
+      }
+      usedBytes = 0L
+    }
+  }
+}
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
new file mode 100644
index 00000000000..cde69e0b848
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.comet
+
+import java.util.Properties
+
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+
+/**
+ * Tests that JVM Arrow allocations are reported to Spark, and, just as importantly, that the
+ * paths where they cannot be reported fail quietly rather than throwing. Arrow allocation on
+ * these paths cannot fail today and this listener must not change that.
+ */
+class CometArrowAllocationListenerSuite extends AnyFunSuite {
+
+  private val blockSize = 1024L * 1024L
+
+  test("allocations are charged to the current task in whole blocks") {
+    withOffHeapTask(taskAttemptId = 1L) { listener =>
+      // Far smaller than a block, so the reservation should round up to exactly one block.
+      listener.onAllocation(128L)
+      assert(listener.reservedBytesForTask(1L) == blockSize)
+
+      // Still inside the first block, so Spark is not asked again.
+      listener.onAllocation(1024L)
+      assert(listener.reservedBytesForTask(1L) == blockSize)
+    }
+  }
+
+  test("a request larger than a block reserves enough to cover it") {
+    withOffHeapTask(taskAttemptId = 2L) { listener =>
+      listener.onAllocation(blockSize * 3 + 7L)
+      assert(listener.reservedBytesForTask(2L) >= blockSize * 3 + 7L)
+    }
+  }
+
+  test("releasing returns whole blocks to Spark") {
+    withOffHeapTask(taskAttemptId = 3L) { listener =>
+      listener.onAllocation(blockSize * 2)
+      val afterAllocation = listener.reservedBytesForTask(3L)
+      assert(afterAllocation >= blockSize * 2)
+
+      listener.onRelease(blockSize * 2)
+      assert(listener.reservedBytesForTask(3L) == 0L)
+    }
+  }
+
+  test("no active task is a no-op rather than an error") {
+    TaskContext.unset()
+    val listener = new CometArrowAllocationListener
+    // Broadcast coalescing and the cached batch serializer can allocate off a task thread.
+    listener.onAllocation(4096L)
+    listener.onRelease(4096L)
+    assert(listener.trackedTaskCount == 0)
+  }
+
+  test("on-heap mode is not accounted") {
+    val memoryManager = new TestMemoryManager(new SparkConf(false))
+    memoryManager.limit(64L * 1024 * 1024)
+    val taskMemoryManager = new TaskMemoryManager(memoryManager, 4L)
+    withTaskContext(taskMemoryManager, taskAttemptId = 4L) {
+      val listener = new CometArrowAllocationListener
+      listener.onAllocation(blockSize)
+      // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory.
+      // Charging an off-heap consumer there would be wrong, so nothing is tracked.
+      assert(listener.trackedTaskCount == 0)
+      assert(listener.reservedBytesForTask(4L) == 0L)
+    }
+  }
+
+  private def withOffHeapTask(taskAttemptId: Long)(
+      f: CometArrowAllocationListener => Unit): Unit = {
+    val conf = new SparkConf(false)
+      .set("spark.memory.offHeap.enabled", "true")
+      .set("spark.memory.offHeap.size", "64m")
+    val memoryManager = new TestMemoryManager(conf)
+    memoryManager.limit(64L * 1024 * 1024)
+    val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId)
+    withTaskContext(taskMemoryManager, taskAttemptId) {
+      f(new CometArrowAllocationListener)
+    }
+  }
+
+  private def withTaskContext(taskMemoryManager: TaskMemoryManager, taskAttemptId: Long)(
+      body: => Unit): Unit = {
+    val taskContext = new TaskContextImpl(
+      stageId = 0,
+      stageAttemptNumber = 0,
+      partitionId = 0,
+      numPartitions = 1,
+      taskAttemptId = taskAttemptId,
+      attemptNumber = 0,
+      taskMemoryManager = taskMemoryManager,
+      localProperties = new Properties,
+      metricsSystem = null,
+      taskMetrics = TaskMetrics.empty,
+      cpus = 1,
+      resources = Map.empty)
+
+    TaskContext.setTaskContext(taskContext)
+    try {
+      body
+    } finally {
+      try {
+        taskMemoryManager.cleanUpAllAllocatedMemory()
+      } finally {
+        TaskContext.unset()
+      }
+    }
+  }
+}

From 0e3dfa7392957d3a6f946dfc387c8da5ca6ee2a3 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 09:10:57 -0600
Subject: [PATCH 02/18] refactor: reduce lock traffic in Arrow allocation
 accounting

Release reserved blocks in a single call instead of one per block. releaseExecutionMemory synchronizes on the executor-wide pool, so freeing a 64 MiB buffer at the 1 MiB block size took 64 acquisitions of that lock rather than one, which defeated the point of batching reservations into blocks.

Round the grow request up to a block multiple. Requesting the bare deficit left reserved exactly equal to usedBytes for any buffer at or above the block size, so the next allocation of any size went straight back into Spark's lock. Reserved is now always a block multiple and growth always leaves headroom.

Check for an active task before resolving config. Resolving first meant re-reading SparkEnv on every allocation in any process that never has one, and the task check is both cheaper and eliminates more callers.

Declare the accounting toggle in CometConf so it reaches the generated configuration docs and follows the conventions every other spark.comet.* key follows, and demote the block size to a private constant rather than an undocumented key that would become a permanent compatibility surface.

Collapse the loop in allocated() that could never iterate twice, move warn-once to the companion so TaskReservation no longer holds a back-reference to its parent, and route the on-heap test through the shared fixture.
---
 .../scala/org/apache/comet/CometConf.scala    |  12 ++
 .../comet/CometArrowAllocationListener.scala  | 136 +++++++++---------
 .../CometArrowAllocationListenerSuite.scala   |  54 ++++---
 3 files changed, 104 insertions(+), 98 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index fc858a75a18..56d192ebc4c 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -850,6 +850,18 @@ object CometConf extends ShimCometConf {
       .doubleConf
       .createWithDefault(1.0)
 
+  val COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.arrowAllocator.accounting.enabled")
+      .category(CATEGORY_TUNING)
+      .doc("When enabled, JVM-side Arrow allocations are reported to Spark's memory manager so that " +
+        "they are visible to Spark's off-heap accounting rather than being invisible to every " +
+        "budget. Reporting only: an allocation is never failed by this setting, but the bytes " +
+        "do consume the off-heap pool, so other consumers see correspondingly less headroom. " +
+        "Disable to restore the previous behaviour of not accounting for these allocations. " +
+        s"$TUNING_GUIDE.")
+      .booleanConf
+      .createWithDefault(true)
+
   val COMET_NATIVE_LOAD_REQUIRED: ConfigEntry[Boolean] = conf("spark.comet.nativeLoadRequired")
     .category(CATEGORY_EXEC)
     .doc(
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index 544b0503ad4..6132e954757 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -20,12 +20,15 @@
 package org.apache.spark.comet
 
 import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicBoolean
 
 import org.apache.arrow.memory.AllocationListener
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager}
 
+import org.apache.comet.CometConf
+
 /**
  * Reports JVM-side Arrow allocations to Spark's memory manager.
  *
@@ -41,8 +44,9 @@ import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager}
  *
  * It deliberately does not enforce. A short grant from Spark is logged and the allocation
  * proceeds, because Arrow allocation on these paths cannot fail today and making it fail is a
- * behavioural change that belongs in its own commit. See
- * [[https://github.com/apache/datafusion-comet/issues/5997]].
+ * behavioural change that belongs in its own commit. Note that enforcement belongs in
+ * `onPreAllocation`, the only callback permitted to throw, and `onFailedAllocation`, not here.
+ * See [[https://github.com/apache/datafusion-comet/issues/5997]].
  *
  * Three cases are handled by doing nothing, each for a different reason:
  *   - No active task. Broadcast coalescing and the cached batch serializer can allocate from the
@@ -53,16 +57,12 @@ import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager}
  *     precisely because buffers can outlive the task that created them, so the task's reservation
  *     is dropped at task end and later releases are ignored rather than double-counted.
  */
-class CometArrowAllocationListener extends AllocationListener with Logging {
+class CometArrowAllocationListener extends AllocationListener {
 
   import CometArrowAllocationListener._
 
   private val reservations = new ConcurrentHashMap[Long, TaskReservation]()
 
-  @volatile private var configResolved = false
-  @volatile private var accountingEnabled = true
-  @volatile private var blockSize = DEFAULT_BLOCK_SIZE
-
   override def onAllocation(size: Long): Unit = {
     val reservation = reservationForCurrentTask()
     if (reservation != null) {
@@ -85,30 +85,14 @@ class CometArrowAllocationListener extends AllocationListener with Logging {
 
   private[comet] def trackedTaskCount: Int = reservations.size()
 
-  /**
-   * Resolves configuration from the `SparkConf` rather than a `SQLConf` entry. This listener is
-   * attached to a `val` in a package object, so it is constructed on first touch of
-   * `CometArrowAllocator`, which can happen before any `SparkSession` exists and on executors
-   * where `SQLConf` does not carry Comet's settings. `SparkEnv` is absent until the executor is
-   * up, so the read is retried until it succeeds rather than cached from a null environment.
-   */
-  private def resolveConfig(): Unit = {
-    if (!configResolved) {
-      val env = SparkEnv.get
-      if (env != null) {
-        accountingEnabled = env.conf.getBoolean(ACCOUNTING_ENABLED_KEY, defaultValue = true)
-        blockSize = env.conf.getSizeAsBytes(BLOCK_SIZE_KEY, DEFAULT_BLOCK_SIZE_STRING)
-        configResolved = true
-      }
-    }
-  }
-
   private def reservationForCurrentTask(): TaskReservation = {
-    resolveConfig()
-    if (!accountingEnabled) return null
-
+    // Cheapest check first, and the one that eliminates the most callers: the driver, broadcast
+    // coalescing and the cached batch serializer all allocate with no task in scope. Reading the
+    // config before this would also mean re-reading `SparkEnv` on every allocation in a process
+    // that never has one.
     val taskContext = TaskContext.get()
     if (taskContext == null) return null
+    if (!accountingEnabled) return null
 
     val taskMemoryManager = taskContext.taskMemoryManager()
     if (taskMemoryManager == null ||
@@ -120,7 +104,10 @@ class CometArrowAllocationListener extends AllocationListener with Logging {
     val existing = reservations.get(taskAttemptId)
     if (existing != null) return existing
 
-    val created = new TaskReservation(taskMemoryManager, blockSize, this)
+    // Deliberately not `computeIfAbsent`: `addTaskCompletionListener` runs the callback inline if
+    // the task has already completed, and that callback removes from this same map, which is a
+    // recursive update inside a mapping function. Registering outside the map operation avoids it.
+    val created = new TaskReservation(taskMemoryManager)
     val previous = reservations.putIfAbsent(taskAttemptId, created)
     if (previous != null) return previous
 
@@ -132,39 +119,48 @@ class CometArrowAllocationListener extends AllocationListener with Logging {
     }
     created
   }
-
-  private[comet] def warnOnShortGrant(requested: Long, granted: Long): Unit = {
-    if (!shortGrantLogged) {
-      shortGrantLogged = true
-      logWarning(
-        s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " +
-          "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " +
-          s"Set $ACCOUNTING_ENABLED_KEY=false to stop reporting these allocations to Spark.")
-    }
-  }
-
-  @volatile private var shortGrantLogged = false
 }
 
-object CometArrowAllocationListener {
+object CometArrowAllocationListener extends Logging {
 
-  val ACCOUNTING_ENABLED_KEY = "spark.comet.arrowAllocator.accounting.enabled"
-  val BLOCK_SIZE_KEY = "spark.comet.arrowAllocator.accounting.blockSize"
+  /**
+   * Batching granularity for reservations. Arrow allocates per buffer and
+   * `acquireExecutionMemory` takes an executor-wide lock, so the reservation is grown and shrunk
+   * in whole blocks and only block-crossing changes reach Spark. Deliberately not configurable:
+   * it trades lock chatter against reservation slack and has no plausible per-workload tuning.
+   */
+  private val BLOCK_SIZE = 1024L * 1024L
 
-  private val DEFAULT_BLOCK_SIZE_STRING = "1m"
-  private val DEFAULT_BLOCK_SIZE = 1024L * 1024L
+  private val shortGrantLogged = new AtomicBoolean(false)
 
   /**
-   * One task's reservation against Spark's off-heap pool.
-   *
-   * Arrow allocates per buffer, and `acquireExecutionMemory` takes locks, so reserving for every
-   * buffer would be needlessly chatty. Instead the reservation is grown and shrunk in whole
-   * blocks and only block-crossing changes reach Spark.
+   * Resolved once per JVM. The listener is attached to a `val` in a package object, so it is
+   * constructed on first touch of `CometArrowAllocator`, which can happen before any
+   * `SparkSession` exists and on executors where `SQLConf` does not carry Comet's settings. This
+   * is only read once a `TaskContext` exists, by which point an executor has a `SparkEnv`; the
+   * `Option` guard covers tests that install a task context without one.
    */
-  private class TaskReservation(
-      taskMemoryManager: TaskMemoryManager,
-      blockSize: Long,
-      listener: CometArrowAllocationListener)
+  private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env =>
+    env.conf.getBoolean(
+      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key,
+      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.defaultValue.get)
+  }
+
+  private def roundUpToBlock(bytes: Long): Long =
+    ((bytes + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE
+
+  private def warnOnShortGrant(requested: Long, granted: Long): Unit = {
+    if (shortGrantLogged.compareAndSet(false, true)) {
+      logWarning(
+        s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " +
+          "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " +
+          s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " +
+          "reporting these allocations to Spark.")
+    }
+  }
+
+  /** One task's reservation against Spark's off-heap pool. */
+  private class TaskReservation(taskMemoryManager: TaskMemoryManager)
       extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) {
 
     // Named `usedBytes` rather than `used` on purpose: `MemoryConsumer` already declares a
@@ -177,9 +173,10 @@ object CometArrowAllocationListener {
     override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
 
     /**
-     * Reports our own tally. The inherited `used` counter stays at zero because this consumer
-     * never calls `acquireMemory` or `allocatePage`; Arrow has already obtained the memory and we
-     * are only accounting for it.
+     * Reports our own tally. Spark reads this for spill-victim ordering, `showMemoryUsage` and
+     * end-of-task leak reporting. The inherited `used` counter stays at zero because this
+     * consumer never calls `acquireMemory` or `allocatePage`; Arrow has already obtained the
+     * memory and we are only accounting for it.
      */
     override def getUsed: Long = synchronized(usedBytes)
 
@@ -187,26 +184,27 @@ object CometArrowAllocationListener {
 
     def allocated(size: Long): Unit = synchronized {
       usedBytes += size
-      while (reserved < usedBytes) {
-        val request = math.max(blockSize, usedBytes - reserved)
+      if (reserved < usedBytes) {
+        // Round up so `reserved` stays a block multiple and growth always leaves headroom.
+        // Requesting the bare deficit would land exactly on `usedBytes` for any buffer at or above
+        // the block size, sending the very next allocation straight back into Spark's lock.
+        val request = roundUpToBlock(usedBytes - reserved)
         val granted = taskMemoryManager.acquireExecutionMemory(request, this)
-        if (granted <= 0L) {
-          listener.warnOnShortGrant(request, granted)
-          return
-        }
         reserved += granted
         if (granted < request) {
-          listener.warnOnShortGrant(request, granted)
-          return
+          warnOnShortGrant(request, granted)
         }
       }
     }
 
     def released(size: Long): Unit = synchronized {
       usedBytes = math.max(0L, usedBytes - size)
-      while (reserved - usedBytes >= blockSize) {
-        taskMemoryManager.releaseExecutionMemory(blockSize, this)
-        reserved -= blockSize
+      // Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes on
+      // the executor-wide pool, so a per-block loop would take that lock once per megabyte freed.
+      val excess = ((reserved - usedBytes) / BLOCK_SIZE) * BLOCK_SIZE
+      if (excess > 0L) {
+        taskMemoryManager.releaseExecutionMemory(excess, this)
+        reserved -= excess
       }
     }
 
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index cde69e0b848..79d532ee002 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -35,34 +35,37 @@ import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
 class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
   private val blockSize = 1024L * 1024L
+  private val poolBytes = 64L * 1024 * 1024
+  private val taskAttemptId = 1L
 
   test("allocations are charged to the current task in whole blocks") {
-    withOffHeapTask(taskAttemptId = 1L) { listener =>
+    withTask() { listener =>
       // Far smaller than a block, so the reservation should round up to exactly one block.
       listener.onAllocation(128L)
-      assert(listener.reservedBytesForTask(1L) == blockSize)
+      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize)
 
       // Still inside the first block, so Spark is not asked again.
       listener.onAllocation(1024L)
-      assert(listener.reservedBytesForTask(1L) == blockSize)
+      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize)
     }
   }
 
-  test("a request larger than a block reserves enough to cover it") {
-    withOffHeapTask(taskAttemptId = 2L) { listener =>
+  test("a request larger than a block rounds up to a block multiple") {
+    withTask() { listener =>
       listener.onAllocation(blockSize * 3 + 7L)
-      assert(listener.reservedBytesForTask(2L) >= blockSize * 3 + 7L)
+      // Rounded up rather than sized to the exact deficit, so growth leaves headroom and the next
+      // small allocation does not go straight back into Spark.
+      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize * 4)
     }
   }
 
   test("releasing returns whole blocks to Spark") {
-    withOffHeapTask(taskAttemptId = 3L) { listener =>
+    withTask() { listener =>
       listener.onAllocation(blockSize * 2)
-      val afterAllocation = listener.reservedBytesForTask(3L)
-      assert(afterAllocation >= blockSize * 2)
+      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize * 2)
 
       listener.onRelease(blockSize * 2)
-      assert(listener.reservedBytesForTask(3L) == 0L)
+      assert(listener.reservedBytesForTask(taskAttemptId) == 0L)
     }
   }
 
@@ -76,34 +79,27 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   }
 
   test("on-heap mode is not accounted") {
-    val memoryManager = new TestMemoryManager(new SparkConf(false))
-    memoryManager.limit(64L * 1024 * 1024)
-    val taskMemoryManager = new TaskMemoryManager(memoryManager, 4L)
-    withTaskContext(taskMemoryManager, taskAttemptId = 4L) {
-      val listener = new CometArrowAllocationListener
+    withTask(offHeap = false) { listener =>
       listener.onAllocation(blockSize)
       // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory.
-      // Charging an off-heap consumer there would be wrong, so nothing is tracked.
+      // Charging an off-heap consumer there would be wrong, so nothing is tracked. Distinguishing
+      // "not tracked" from "tracked with zero reserved" is why trackedTaskCount is asserted here.
       assert(listener.trackedTaskCount == 0)
-      assert(listener.reservedBytesForTask(4L) == 0L)
+      assert(listener.reservedBytesForTask(taskAttemptId) == 0L)
     }
   }
 
-  private def withOffHeapTask(taskAttemptId: Long)(
-      f: CometArrowAllocationListener => Unit): Unit = {
+  private def withTask(offHeap: Boolean = true)(f: CometArrowAllocationListener => Unit): Unit = {
     val conf = new SparkConf(false)
-      .set("spark.memory.offHeap.enabled", "true")
-      .set("spark.memory.offHeap.size", "64m")
+    if (offHeap) {
+      conf
+        .set("spark.memory.offHeap.enabled", "true")
+        .set("spark.memory.offHeap.size", poolBytes.toString)
+    }
     val memoryManager = new TestMemoryManager(conf)
-    memoryManager.limit(64L * 1024 * 1024)
+    memoryManager.limit(poolBytes)
     val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId)
-    withTaskContext(taskMemoryManager, taskAttemptId) {
-      f(new CometArrowAllocationListener)
-    }
-  }
 
-  private def withTaskContext(taskMemoryManager: TaskMemoryManager, taskAttemptId: Long)(
-      body: => Unit): Unit = {
     val taskContext = new TaskContextImpl(
       stageId = 0,
       stageAttemptNumber = 0,
@@ -120,7 +116,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
     TaskContext.setTaskContext(taskContext)
     try {
-      body
+      f(new CometArrowAllocationListener)
     } finally {
       try {
         taskMemoryManager.cleanUpAllAllocatedMemory()

From 55903fd9361ef63a952dc884fd525d0db4d04048 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 09:21:54 -0600
Subject: [PATCH 03/18] fix: exclude FFI-imported buffers from Arrow allocation
 accounting

Arrow's BaseAllocator.wrapForeignAllocation reports an imported buffer to the allocator's listener at full capacity even though no JVM-side allocation happened. Comet imports every batch from native execution through that path, so charging the root allocator's listener counted memory the native side owns and frees as though it were JVM Arrow memory.

That double counted whatever an operator had already reserved in Comet's native pool, and the error grew in proportion to batch throughput, which is the opposite of what this accounting is for.

Import through CometImportedArrowAllocator, a child allocator with no listener. Arrow notifies only the allocating allocator's own listener and never its ancestors, so the imported bytes stay out of Spark's accounting while reference counting and lifetime are unchanged. The two import sites are NativeUtil and CometUdfBridge; export, IPC and materialisation paths keep the charging root.

Add a test pinning both halves of the mechanism: a child with no listener is not charged, and the same root still charges for allocations made directly against it.
---
 .../org/apache/comet/udf/CometUdfBridge.java  |  7 ++++-
 .../main/scala/org/apache/comet/package.scala | 22 ++++++++++++-
 .../org/apache/comet/vector/NativeUtil.scala  |  7 +++--
 .../comet/CometArrowAllocationListener.scala  |  6 ++++
 .../CometArrowAllocationListenerSuite.scala   | 31 +++++++++++++++++++
 5 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
index d8dea731353..77f3936b7f1 100644
--- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
+++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
@@ -210,6 +210,8 @@ private static void evaluateInternal(
     assert udf != null : "reflective instantiation returned null for " + udfClassName;
 
     BufferAllocator allocator = org.apache.comet.package$.MODULE$.CometArrowAllocator();
+    BufferAllocator importAllocator =
+        org.apache.comet.package$.MODULE$.CometImportedArrowAllocator();
 
     ValueVector[] inputs = new ValueVector[inputArrayPtrs.length];
     ValueVector result = null;
@@ -217,7 +219,10 @@ private static void evaluateInternal(
       for (int i = 0; i < inputArrayPtrs.length; i++) {
         ArrowArray inArr = ArrowArray.wrap(inputArrayPtrs[i]);
         ArrowSchema inSch = ArrowSchema.wrap(inputSchemaPtrs[i]);
-        inputs[i] = Data.importVector(allocator, inArr, inSch, null);
+        // Imported from native memory that the native side owns and frees, so it is deliberately
+        // not reported to Spark's memory manager. The export below still uses the root allocator,
+        // whose allocations are JVM-owned.
+        inputs[i] = Data.importVector(importAllocator, inArr, inSch, null);
       }
 
       result = udf.evaluate(inputs, numRows);
diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala
index 55982b3645e..7c378544b93 100644
--- a/spark/src/main/scala/org/apache/comet/package.scala
+++ b/spark/src/main/scala/org/apache/comet/package.scala
@@ -21,7 +21,7 @@ package org.apache
 
 import java.util.Properties
 
-import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.memory.{AllocationListener, RootAllocator}
 import org.apache.spark.comet.CometArrowAllocationListener
 import org.apache.spark.internal.Logging
 
@@ -41,6 +41,26 @@ package object comet {
   val CometArrowAllocator =
     new RootAllocator(new CometArrowAllocationListener, Long.MaxValue)
 
+  /**
+   * The allocator for buffers imported over the Arrow C Data Interface.
+   *
+   * An imported buffer wraps memory that the native side owns and frees, but Arrow's
+   * `wrapForeignAllocation` still reports it to the allocator's listener at full buffer capacity,
+   * as though a JVM-side allocation had happened. Importing through [[CometArrowAllocator]] would
+   * therefore charge Spark for native bytes, double counting whatever an operator has already
+   * reserved in Comet's native pool, and the error would grow with batch throughput.
+   *
+   * Arrow notifies only the allocating allocator's own listener, never its ancestors, so a child
+   * with no listener keeps these buffers out of Spark's accounting. It stays a child of the root
+   * so that reference counting and lifetime are unchanged.
+   */
+  val CometImportedArrowAllocator =
+    CometArrowAllocator.newChildAllocator(
+      "comet-imported-ffi",
+      AllocationListener.NOOP,
+      0,
+      Long.MaxValue)
+
   /**
    * Provides access to build information about the Comet libraries. This will be used by the
    * benchmarking software to provide the source revision and repository. In addition, the build
diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
index 173086d2fd7..ecd36cd720d 100644
--- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
+++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution.vectorized.ConstantColumnVector
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.CometArrowAllocator
+import org.apache.comet.{CometArrowAllocator, CometImportedArrowAllocator}
 
 /**
  * Provides functionality for importing Arrow vectors from native code and wrapping them as
@@ -51,7 +51,10 @@ class NativeUtil extends AutoCloseable {
   private val allocator = CometArrowAllocator
 
   /** ArrowImporter does not hold any state and does not need to be closed */
-  private val importer = new ArrowImporter(allocator)
+  // Imported buffers wrap memory the native side owns and frees, so they are charged to Comet's
+  // native pool, not to Spark. Importing through the root allocator would report them to Spark as
+  // if they were JVM Arrow bytes and double count them. See CometImportedArrowAllocator.
+  private val importer = new ArrowImporter(CometImportedArrowAllocator)
 
   /**
    * Dictionary provider to use for the lifetime of this instance of NativeUtil. The dictionary
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index 6132e954757..09060742199 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -48,6 +48,12 @@ import org.apache.comet.CometConf
  * `onPreAllocation`, the only callback permitted to throw, and `onFailedAllocation`, not here.
  * See [[https://github.com/apache/datafusion-comet/issues/5997]].
  *
+ * Buffers imported over the C Data Interface never reach this listener at all. They wrap memory
+ * the native side owns, so Comet imports them through `CometImportedArrowAllocator`, a child with
+ * no listener. Charging them here would double count bytes already reserved in Comet's native
+ * pool. Arrow notifies only the allocating allocator's own listener, which is what makes that
+ * separation work.
+ *
  * Three cases are handled by doing nothing, each for a different reason:
  *   - No active task. Broadcast coalescing and the cached batch serializer can allocate from the
  *     driver or a non-task thread, where there is no task to charge.
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index 79d532ee002..11eb4fc807e 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -23,6 +23,7 @@ import java.util.Properties
 
 import org.scalatest.funsuite.AnyFunSuite
 
+import org.apache.arrow.memory.{AllocationListener, RootAllocator}
 import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
@@ -69,6 +70,36 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
+  test("a child allocator with no listener is not charged, but the root still is") {
+    withTask() { listener =>
+      val root = new RootAllocator(listener, Long.MaxValue)
+      val imported =
+        root.newChildAllocator("imported", AllocationListener.NOOP, 0, Long.MaxValue)
+      try {
+        // This is how FFI-imported buffers, which wrap memory the native side owns, stay out of
+        // Spark's accounting: Arrow notifies only the allocating allocator's own listener, never
+        // its ancestors, so the root's listener does not see the child's allocation.
+        val importedBuf = imported.buffer(blockSize)
+        try {
+          assert(listener.trackedTaskCount == 0)
+
+          // The same root does still charge for JVM-owned allocations made directly against it.
+          val jvmBuf = root.buffer(blockSize)
+          try {
+            assert(listener.reservedBytesForTask(taskAttemptId) >= blockSize)
+          } finally {
+            jvmBuf.close()
+          }
+        } finally {
+          importedBuf.close()
+        }
+      } finally {
+        imported.close()
+        root.close()
+      }
+    }
+  }
+
   test("no active task is a no-op rather than an error") {
     TaskContext.unset()
     val listener = new CometArrowAllocationListener

From 2d531ecad093645f7bf8137b9931927590f24f04 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 09:45:25 -0600
Subject: [PATCH 04/18] docs: record that JVM Arrow allocations are now
 reported to Spark

The inventory table and the observation below it both stated that CometArrowAllocator was accounted by nobody, which the listener has changed. Update the row, rewrite the observation to say that allocations are reported to Spark but the allocator is still unbounded, and soften the Open problems entry to match: no longer invisible, still uncapped.

Add the imported-buffer distinction as its own observation and its own inventory row. A reader who learns the allocator is accounted would otherwise reasonably assume FFI imports are too, and they are deliberately excluded because charging them would double count memory Comet's native pool already reserved.
---
 .../contributor-guide/memory_management.md    | 41 +++++++++++++------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index c252afedc45..8cf0cbca263 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -64,23 +64,36 @@ Enabling Comet does not add one new memory consumer, it adds several, and they a
 accounted by the same party. This inventory is worth internalizing before reading the rest of the
 page:
 
-| Allocator                               | Lives in    | Bounded by                                                    | Visible to Spark? |
-| --------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
-| Spark execution + storage (on-heap)     | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
-| Spark Tungsten (off-heap)               | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
-| Comet native (Rust global allocator)    | Native heap | `memory_limit` (see below), enforced only via the memory pool | No                |
-| Comet JVM Arrow (`CometArrowAllocator`) | Off-heap    | **Nothing**: a `RootAllocator(Long.MaxValue)`                 | No                |
-| Comet JVM shuffle pages                 | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+| Allocator                                                  | Lives in    | Bounded by                                                    | Visible to Spark? |
+| ---------------------------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
+| Spark execution + storage (on-heap)                        | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
+| Spark Tungsten (off-heap)                                  | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+| Comet native (Rust global allocator)                       | Native heap | `memory_limit` (see below), enforced only via the memory pool | No                |
+| Comet JVM Arrow (`CometArrowAllocator`)                    | Off-heap    | **Nothing**, but reported to `TaskMemoryManager`              | Yes               |
+| Comet imported FFI buffers (`CometImportedArrowAllocator`) | Native heap | Comet's native memory pool, when an operator reserved them    | No                |
+| Comet JVM shuffle pages                                    | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
 
 Two observations follow.
 
-**Comet's JVM-side Arrow allocator is unbounded and accounted by nobody.** `CometArrowAllocator`
+**Comet's JVM-side Arrow allocator is reported to Spark, but still unbounded.** `CometArrowAllocator`
 (`spark/src/main/scala/org/apache/comet/package.scala`) is a single process-wide
 `new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream export
-(`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. These are real
-off-heap bytes in container RSS that neither Spark's `TaskMemoryManager` nor Comet's native memory
-pool sees. In practice the volume is modest, a batch at a time per stream, but there is no
-ceiling and no backpressure.
+(`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. A
+`CometArrowAllocationListener` on the root charges each allocation to a `MemoryConsumer` for the
+task that made it, in whole blocks, so these bytes now appear in `showMemoryUsage` and are
+arbitrated against Spark's other off-heap consumers. The limit is still `Long.MaxValue`: the
+listener reports without enforcing, so an allocation here cannot fail, though the bytes do consume
+the off-heap pool and other consumers see correspondingly less headroom. Set
+`spark.comet.arrowAllocator.accounting.enabled=false` to stop reporting them.
+
+**Buffers imported over the C Data Interface are deliberately excluded from that accounting.**
+Arrow's `wrapForeignAllocation` reports an imported buffer to the allocator's listener at full
+capacity even though no JVM-side allocation happened, so charging them would count native memory
+that Comet's own pool has already reserved, with the error growing in proportion to batch
+throughput. Imports therefore go through `CometImportedArrowAllocator`, a child allocator with no
+listener; Arrow notifies only the allocating allocator's own listener and never its ancestors. The
+import sites are `NativeUtil` and `CometUdfBridge`. Export, IPC and materialisation keep the
+charging root.
 
 **The JVM shuffle allocator is an ordinary Spark consumer.** `CometShuffleMemoryAllocator.getInstance`
 returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from
@@ -321,7 +334,9 @@ much they matter:
   the fact. There is no runtime value that an operator, a metric, or a policy could read.
 - **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to
   guess a per-workload haircut rather than measuring anything.
-- **`CometArrowAllocator` is unbounded** and participates in no budget.
+- **`CometArrowAllocator` is unbounded.** Its allocations are now reported to Spark's memory
+  manager, so they are no longer invisible, but nothing caps them: the listener reports without
+  enforcing, and enforcing would mean failing allocations on paths that cannot fail today.
 - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be
   resident on either side with no reservation covering it, because reservations are made and
   withdrawn by individual operators while the bytes outlive them.

From 4d7657d4ceb4d9e9aa6cbd23d34211053511944c Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 09:45:25 -0600
Subject: [PATCH 05/18] style: keep the new config doc within the 100 character
 line limit

The doc string for spark.comet.arrowAllocator.accounting.enabled left a 103 character line, which failed scalastyle and took every downstream CI job with it. scalafmt cannot split a string literal, so the segments themselves are shortened rather than reindented.
---
 .../main/scala/org/apache/comet/CometConf.scala    | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 56d192ebc4c..7044904159d 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -853,12 +853,14 @@ object CometConf extends ShimCometConf {
   val COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.arrowAllocator.accounting.enabled")
       .category(CATEGORY_TUNING)
-      .doc("When enabled, JVM-side Arrow allocations are reported to Spark's memory manager so that " +
-        "they are visible to Spark's off-heap accounting rather than being invisible to every " +
-        "budget. Reporting only: an allocation is never failed by this setting, but the bytes " +
-        "do consume the off-heap pool, so other consumers see correspondingly less headroom. " +
-        "Disable to restore the previous behaviour of not accounting for these allocations. " +
-        s"$TUNING_GUIDE.")
+      .doc(
+        "When enabled, JVM-side Arrow allocations are reported to Spark's memory " +
+          "manager, so they are visible to Spark's off-heap accounting rather than to no " +
+          "budget at all. Reporting only: this setting never fails an allocation, but the " +
+          "bytes do consume the off-heap pool, so other consumers see correspondingly less " +
+          "headroom. Disable to restore the previous behaviour of not accounting for these " +
+          "allocations. " +
+          s"$TUNING_GUIDE.")
       .booleanConf
       .createWithDefault(true)
 

From a2518d97da374dc8c20c2444559dd92a74ca223f Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 10:32:05 -0600
Subject: [PATCH 06/18] style: give the imported Arrow allocator an explicit
 result type

scalafix's ExplicitResultTypes rule requires a declared type on new public members, and CometImportedArrowAllocator was inferred. The rule's own suggested patch was to annotate it as BufferAllocator and import that type, which is what this does; the import is merged into the existing Arrow import rather than added separately, which scalafix accepts since the rule only requires the annotation to exist.

This escaped local checks because scalafix runs only in the Lint Java jobs, under -Psemanticdb for Spark 3.4, 3.5 and 4.0 on JDK 17, and is not part of the default build or of spotless. Verified by running the CI invocation locally: ./mvnw package -DskipTests scalafix:scalafix -Dscalafix.mode=CHECK -Psemanticdb -Pspark-3.5 -Pscala-2.12
---
 spark/src/main/scala/org/apache/comet/package.scala | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala
index 7c378544b93..ffaad7ac293 100644
--- a/spark/src/main/scala/org/apache/comet/package.scala
+++ b/spark/src/main/scala/org/apache/comet/package.scala
@@ -21,7 +21,7 @@ package org.apache
 
 import java.util.Properties
 
-import org.apache.arrow.memory.{AllocationListener, RootAllocator}
+import org.apache.arrow.memory.{AllocationListener, BufferAllocator, RootAllocator}
 import org.apache.spark.comet.CometArrowAllocationListener
 import org.apache.spark.internal.Logging
 
@@ -54,7 +54,7 @@ package object comet {
    * with no listener keeps these buffers out of Spark's accounting. It stays a child of the root
    * so that reference counting and lifetime are unchanged.
    */
-  val CometImportedArrowAllocator =
+  val CometImportedArrowAllocator: BufferAllocator =
     CometArrowAllocator.newChildAllocator(
       "comet-imported-ffi",
       AllocationListener.NOOP,

From e58b177a4ce2c24bfcf0d08e17e538ce42d9b6ce Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 14:01:57 -0600
Subject: [PATCH 07/18] fix: bind JVM Arrow accounting to a per-task allocator

Arrow hands `AllocationListener` nothing but a size, and it reports a
release on whichever thread drops the last reference. Reading
`TaskContext` inside the callbacks therefore lost every release that
happens off the allocating task's thread, the JVM UDF export path being
the clear case: it exports a JVM-owned vector and Rust drops it later
from a Tokio worker with no task context installed. Batch after batch,
that left the task charged for memory it had already freed, and a
release arriving under a different task subtracted from that task
instead.

Bind one listener to one task's allocator instead, the way Gluten does.
`CometTaskArrowAllocator.forCurrentTask()` cuts a child of the root per
task; Arrow passes the listener down to children, so the paths that make
their own children are covered too. The root keeps no listener, which is
what FFI imports want, so the separate imported-buffer allocator is no
longer needed. Off a task, and in Comet's on-heap mode, callers get the
root and nothing is reported, exactly as before.

Two further fixes to the listener itself:

- `getUsed` and `spill` are now lock-free. Spark calls both while holding
  the `TaskMemoryManager` monitor, and the listener holds its own monitor
  across `acquireExecutionMemory`, so a native reservation arriving
  through `CometTaskMemoryManager` could hold Spark's monitor and wait
  for ours while an Arrow allocation on the same task did the reverse.

- Neither callback propagates an exception. Arrow documents that they may
  not, and `BaseAllocator.buffer` marks the allocation successful before
  calling `onAllocation`, so a throw loses the buffer Arrow has already
  created. Spark's acquisition is fallible: it runs other consumers'
  `spill`, which turns a task interrupt into a `RuntimeException` and an
  I/O failure into a `SparkOutOfMemoryError`.

A task allocator cannot simply be closed at task end, because the
process-wide allocator exists precisely so buffers can outlive their
task and Arrow treats closing a non-empty allocator as a leak. At
completion the reservation is dropped and the allocator is closed if it
is drained; otherwise it is parked and reaped by a later task.
---
 .../contributor-guide/memory_management.md    |  77 +++--
 .../org/apache/comet/udf/CometUdfBridge.java  |  17 +-
 .../scala/org/apache/comet/CometConf.scala    |  13 +-
 .../CometBatchKernelCodegenOutput.scala       |  12 +-
 .../main/scala/org/apache/comet/package.scala |  36 +--
 .../org/apache/comet/vector/NativeUtil.scala  |  15 +-
 .../apache/comet/vector/StreamReader.scala    |  10 +-
 .../comet/CometArrowAllocationListener.scala  | 276 +++++++++---------
 .../spark/comet/CometTaskArrowAllocator.scala | 194 ++++++++++++
 .../arrow/ArrowCachedBatchSerializer.scala    |  10 +-
 .../arrow/CometNativeArrowSource.scala        |   8 +-
 .../apache/spark/sql/comet/util/Utils.scala   |   5 +-
 .../python/CometArrowPythonRunnerBase.scala   |  10 +-
 13 files changed, 444 insertions(+), 239 deletions(-)
 create mode 100644 spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala

diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 8cf0cbca263..81deaef4953 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -64,36 +64,51 @@ Enabling Comet does not add one new memory consumer, it adds several, and they a
 accounted by the same party. This inventory is worth internalizing before reading the rest of the
 page:
 
-| Allocator                                                  | Lives in    | Bounded by                                                    | Visible to Spark? |
-| ---------------------------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
-| Spark execution + storage (on-heap)                        | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
-| Spark Tungsten (off-heap)                                  | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
-| Comet native (Rust global allocator)                       | Native heap | `memory_limit` (see below), enforced only via the memory pool | No                |
-| Comet JVM Arrow (`CometArrowAllocator`)                    | Off-heap    | **Nothing**, but reported to `TaskMemoryManager`              | Yes               |
-| Comet imported FFI buffers (`CometImportedArrowAllocator`) | Native heap | Comet's native memory pool, when an operator reserved them    | No                |
-| Comet JVM shuffle pages                                    | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
-
-Two observations follow.
-
-**Comet's JVM-side Arrow allocator is reported to Spark, but still unbounded.** `CometArrowAllocator`
-(`spark/src/main/scala/org/apache/comet/package.scala`) is a single process-wide
-`new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream export
-(`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. A
-`CometArrowAllocationListener` on the root charges each allocation to a `MemoryConsumer` for the
-task that made it, in whole blocks, so these bytes now appear in `showMemoryUsage` and are
+| Allocator                                    | Lives in    | Bounded by                                                    | Visible to Spark? |
+| -------------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
+| Spark execution + storage (on-heap)          | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
+| Spark Tungsten (off-heap)                    | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+| Comet native (Rust global allocator)         | Native heap | `memory_limit` (see below), enforced only via the memory pool | No                |
+| Comet JVM Arrow, in a task (task allocator)  | Off-heap    | **Nothing**, but reported to `TaskMemoryManager`              | Yes               |
+| Comet JVM Arrow, off a task, and FFI imports | Off-heap    | **Nothing**: a `RootAllocator(Long.MaxValue)`                 | No                |
+| Comet JVM shuffle pages                      | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+
+Several observations follow.
+
+**Comet's JVM-side Arrow allocations inside a task are reported to Spark, but still unbounded.**
+`CometArrowAllocator` (`spark/src/main/scala/org/apache/comet/package.scala`) is a single
+process-wide `new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream
+export (`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. JVM-owned
+allocations do not use the root directly: `CometTaskArrowAllocator.forCurrentTask()` hands each
+task a child carrying its own `CometArrowAllocationListener`, which charges the bytes to a
+`MemoryConsumer` for that task in whole blocks, so they appear in `showMemoryUsage` and are
 arbitrated against Spark's other off-heap consumers. The limit is still `Long.MaxValue`: the
 listener reports without enforcing, so an allocation here cannot fail, though the bytes do consume
 the off-heap pool and other consumers see correspondingly less headroom. Set
 `spark.comet.arrowAllocator.accounting.enabled=false` to stop reporting them.
 
-**Buffers imported over the C Data Interface are deliberately excluded from that accounting.**
+**The owner is the allocator, not the calling thread.** Arrow hands `AllocationListener` nothing
+but a size, and it reports a release on whichever thread drops the last reference, which for
+anything exported over the C Data Interface is a Tokio worker with no task context installed.
+Attributing by `TaskContext` would therefore drop those releases and leave a task charged for
+memory it had already freed. Binding one listener to one task's allocator is what makes the release
+land on the task that allocated. Children cut from a task allocator inherit its listener, so the
+paths that make their own children are covered without knowing about any of this.
+
+**Anything allocated from the root itself is not accounted, which is what FFI imports want.**
 Arrow's `wrapForeignAllocation` reports an imported buffer to the allocator's listener at full
 capacity even though no JVM-side allocation happened, so charging them would count native memory
 that Comet's own pool has already reserved, with the error growing in proportion to batch
-throughput. Imports therefore go through `CometImportedArrowAllocator`, a child allocator with no
-listener; Arrow notifies only the allocating allocator's own listener and never its ancestors. The
-import sites are `NativeUtil` and `CometUdfBridge`. Export, IPC and materialisation keep the
-charging root.
+throughput. `NativeUtil` and `CometUdfBridge` therefore import through the listener-less root,
+while export, IPC and materialisation use the task allocator. The driver, broadcast coalescing and
+the cached batch serializer also fall back to the root when there is no task to charge, as does
+Comet's on-heap mode.
+
+**A task allocator is closed only once it has been drained.** The process-wide allocator exists
+because Arrow buffers can outlive the task that created them, and Arrow treats closing an allocator
+that still owns bytes as a leak. At task completion the Spark reservation is dropped and the
+allocator is closed if it is empty; otherwise it is parked and closed by a later task once the
+stragglers are released, since `BaseAllocator` keeps every child in a map until it closes.
 
 **The JVM shuffle allocator is an ordinary Spark consumer.** `CometShuffleMemoryAllocator.getInstance`
 returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from
@@ -218,13 +233,13 @@ accounting (if any) saw the bytes appear. Whether the buffer is _reserved_ in Co
 separate decision, made by whichever operator holds it, and that operator neither knows nor cares
 which side of the boundary the bytes came from.
 
-**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of
-`CometArrowAllocator` and exports the whole per-partition iterator once as an `ArrowArrayStream`.
+**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of the task
+allocator and exports the whole per-partition iterator once as an `ArrowArrayStream`.
 `ScanExec` imports each batch through `AlignedArrowStreamReader` with `CopyMode::UnpackOrClone`:
 dictionary columns are unpacked into new native arrays, everything else is an `Arc` clone of the
 imported buffers. Those bytes stay where Java Arrow put them and are pinned for as long as any native
-reference survives. They are invisible to Spark's `TaskMemoryManager`, and `CometArrowAllocator` is
-unbounded, so nobody charged for them at allocation time. Whether they are charged _later_ depends
+reference survives. They are reported to Spark's `TaskMemoryManager` at allocation time, but nothing
+caps them, because the task allocator reports without enforcing. Whether they are charged _later_ depends
 on who holds them. DataFusion's `ExternalSorter` reserves `get_reserved_bytes_for_record_batch` for
 every batch it retains, and the hash join build side reserves `get_record_batch_memory_size` for
 each incoming batch; both read buffer sizes off the `ArrayData` and apply equally to imported
@@ -299,7 +314,7 @@ hard ceiling on the sum of everything in the container. That cgroup counts, amon
 - JVM non-heap: metaspace, code cache, thread stacks, GC structures, Netty direct buffers,
 - Spark's own off-heap allocations,
 - **all of Comet's native allocations**,
-- Comet's JVM-side Arrow buffers (`CometArrowAllocator`),
+- Comet's JVM-side Arrow buffers (`CometArrowAllocator` and the per-task allocators cut from it),
 - page cache charged to the cgroup by the container's file I/O, including spill files.
 
 Only the first and a portion of the third are visible to Spark's accounting. When the total crosses
@@ -334,9 +349,11 @@ much they matter:
   the fact. There is no runtime value that an operator, a metric, or a policy could read.
 - **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to
   guess a per-workload haircut rather than measuring anything.
-- **`CometArrowAllocator` is unbounded.** Its allocations are now reported to Spark's memory
-  manager, so they are no longer invisible, but nothing caps them: the listener reports without
-  enforcing, and enforcing would mean failing allocations on paths that cannot fail today.
+- **`CometArrowAllocator` is unbounded.** Allocations made inside a task are now reported to
+  Spark's memory manager, so they are no longer invisible, but nothing caps them: the listener
+  reports without enforcing, and enforcing would mean failing allocations on paths that cannot fail
+  today. Allocations made off a task, and buffers imported over the C Data Interface, are not
+  reported at all.
 - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be
   resident on either side with no reservation covering it, because reservations are made and
   withdrawn by individual operators while the bytes outlive them.
diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
index 77f3936b7f1..3a7ae0e53a4 100644
--- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
+++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
@@ -209,9 +209,11 @@ private static void evaluateInternal(
             });
     assert udf != null : "reflective instantiation returned null for " + udfClassName;
 
-    BufferAllocator allocator = org.apache.comet.package$.MODULE$.CometArrowAllocator();
-    BufferAllocator importAllocator =
-        org.apache.comet.package$.MODULE$.CometImportedArrowAllocator();
+    // The result vector below is JVM-owned, so it is allocated from this task's accounted
+    // allocator. Rust retains it through `from_ffi` and drops it later from a Tokio worker with no
+    // task context installed, which is exactly why the accounting is bound to the allocator rather
+    // than to the releasing thread.
+    BufferAllocator allocator = org.apache.spark.comet.CometTaskArrowAllocator.forCurrentTask();
 
     ValueVector[] inputs = new ValueVector[inputArrayPtrs.length];
     ValueVector result = null;
@@ -219,10 +221,11 @@ private static void evaluateInternal(
       for (int i = 0; i < inputArrayPtrs.length; i++) {
         ArrowArray inArr = ArrowArray.wrap(inputArrayPtrs[i]);
         ArrowSchema inSch = ArrowSchema.wrap(inputSchemaPtrs[i]);
-        // Imported from native memory that the native side owns and frees, so it is deliberately
-        // not reported to Spark's memory manager. The export below still uses the root allocator,
-        // whose allocations are JVM-owned.
-        inputs[i] = Data.importVector(importAllocator, inArr, inSch, null);
+        // Imported from native memory that the native side owns and frees, already charged to
+        // Comet's native pool, so it goes through the unaccounted root allocator.
+        inputs[i] =
+            Data.importVector(
+                org.apache.comet.package$.MODULE$.CometArrowAllocator(), inArr, inSch, null);
       }
 
       result = udf.evaluate(inputs, numRows);
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 7044904159d..31004add087 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -854,12 +854,13 @@ object CometConf extends ShimCometConf {
     conf("spark.comet.arrowAllocator.accounting.enabled")
       .category(CATEGORY_TUNING)
       .doc(
-        "When enabled, JVM-side Arrow allocations are reported to Spark's memory " +
-          "manager, so they are visible to Spark's off-heap accounting rather than to no " +
-          "budget at all. Reporting only: this setting never fails an allocation, but the " +
-          "bytes do consume the off-heap pool, so other consumers see correspondingly less " +
-          "headroom. Disable to restore the previous behaviour of not accounting for these " +
-          "allocations. " +
+        "When enabled, JVM-side Arrow allocations made inside a Spark task are reported to " +
+          "Spark's memory manager, so they are visible to Spark's off-heap accounting rather " +
+          "than to no budget at all. Reporting only: this setting never fails an allocation, " +
+          "but the bytes do consume the off-heap pool, so other consumers see correspondingly " +
+          "less headroom. Buffers imported over the Arrow C Data Interface are never reported, " +
+          "because the memory belongs to the native side. Disable to restore the previous " +
+          "behaviour of not accounting for these allocations. " +
           s"$TUNING_GUIDE.")
       .booleanConf
       .createWithDefault(true)
diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
index 33e6c0c0355..159ef8fd1d5 100644
--- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
+++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
@@ -26,11 +26,11 @@ import org.apache.arrow.memory.BufferAllocator
 import org.apache.arrow.vector._
 import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field}
+import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.types._
 
-import org.apache.comet.CometArrowAllocator
 import org.apache.comet.shims.CometTypeShim
 
 /**
@@ -87,21 +87,23 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
    * Closes the vector on any failure so a partially-initialized tree doesn't leak buffers.
    */
   def allocateOutput(field: Field, numRows: Int, estimatedBytes: Int): FieldVector = {
+    // JVM-owned codegen output, accounted to the task that is running the kernel.
+    val allocator = CometTaskArrowAllocator.forCurrentTask()
     val vec: FieldVector = field.getType match {
       case _: ArrowType.List | _: ArrowType.LargeList | _: ArrowType.FixedSizeList =>
-        val v = new RenamedListVector(field, CometArrowAllocator)
+        val v = new RenamedListVector(field, allocator)
         v.initializeChildrenFromFields(field.getChildren)
         v
       case _: ArrowType.Map =>
-        val v = new RenamedMapVector(field, CometArrowAllocator)
+        val v = new RenamedMapVector(field, allocator)
         v.initializeChildrenFromFields(field.getChildren)
         v
       case _: ArrowType.Struct =>
-        val v = new RenamedStructVector(field, CometArrowAllocator)
+        val v = new RenamedStructVector(field, allocator)
         v.initializeChildrenFromFields(field.getChildren)
         v
       case _ =>
-        field.createVector(CometArrowAllocator).asInstanceOf[FieldVector]
+        field.createVector(allocator).asInstanceOf[FieldVector]
     }
     try {
       vec.setInitialCapacity(numRows)
diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala
index ffaad7ac293..b26cca4f4cc 100644
--- a/spark/src/main/scala/org/apache/comet/package.scala
+++ b/spark/src/main/scala/org/apache/comet/package.scala
@@ -21,8 +21,7 @@ package org.apache
 
 import java.util.Properties
 
-import org.apache.arrow.memory.{AllocationListener, BufferAllocator, RootAllocator}
-import org.apache.spark.comet.CometArrowAllocationListener
+import org.apache.arrow.memory.RootAllocator
 import org.apache.spark.internal.Logging
 
 package object comet {
@@ -34,32 +33,17 @@ package object comet {
    * finished later than the close of the allocator, the allocator will think the memory is
    * leaked. To avoid this, we use a single allocator for the whole execution process.
    *
-   * The allocator itself is unlimited, but [[CometArrowAllocationListener]] reports every
-   * allocation to Spark's memory manager so that these off-heap bytes are no longer invisible to
-   * Spark's accounting. It reports without enforcing, so allocation here still cannot fail.
-   */
-  val CometArrowAllocator =
-    new RootAllocator(new CometArrowAllocationListener, Long.MaxValue)
-
-  /**
-   * The allocator for buffers imported over the Arrow C Data Interface.
-   *
-   * An imported buffer wraps memory that the native side owns and frees, but Arrow's
-   * `wrapForeignAllocation` still reports it to the allocator's listener at full buffer capacity,
-   * as though a JVM-side allocation had happened. Importing through [[CometArrowAllocator]] would
-   * therefore charge Spark for native bytes, double counting whatever an operator has already
-   * reserved in Comet's native pool, and the error would grow with batch throughput.
+   * It carries no allocation listener, so allocating from it directly is not reported to Spark's
+   * memory manager. That is what buffers imported over the Arrow C Data Interface want: they wrap
+   * memory the native side owns and frees, already charged to Comet's native pool, and Arrow's
+   * `wrapForeignAllocation` would otherwise report the full buffer capacity as though a JVM-side
+   * allocation had happened.
    *
-   * Arrow notifies only the allocating allocator's own listener, never its ancestors, so a child
-   * with no listener keeps these buffers out of Spark's accounting. It stays a child of the root
-   * so that reference counting and lifetime are unchanged.
+   * JVM-owned allocations should go through `CometTaskArrowAllocator.forCurrentTask()` instead,
+   * which cuts a per-task child whose listener reports the bytes to Spark. Off a task it hands
+   * back this allocator, so the driver-side paths are unchanged.
    */
-  val CometImportedArrowAllocator: BufferAllocator =
-    CometArrowAllocator.newChildAllocator(
-      "comet-imported-ffi",
-      AllocationListener.NOOP,
-      0,
-      Long.MaxValue)
+  val CometArrowAllocator = new RootAllocator(Long.MaxValue)
 
   /**
    * Provides access to build information about the Comet libraries. This will be used by the
diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
index ecd36cd720d..88df5f5b1ca 100644
--- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
+++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
@@ -26,12 +26,13 @@ import org.apache.arrow.util.AutoCloseables
 import org.apache.arrow.vector.VectorSchemaRoot
 import org.apache.arrow.vector.dictionary.DictionaryProvider
 import org.apache.spark.SparkException
+import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.sql.comet.execution.arrow.ConstantColumnVectors
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution.vectorized.ConstantColumnVector
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.{CometArrowAllocator, CometImportedArrowAllocator}
+import org.apache.comet.CometArrowAllocator
 
 /**
  * Provides functionality for importing Arrow vectors from native code and wrapping them as
@@ -47,14 +48,14 @@ import org.apache.comet.{CometArrowAllocator, CometImportedArrowAllocator}
 class NativeUtil extends AutoCloseable {
   import Utils._
 
-  /** Use the global allocator */
-  private val allocator = CometArrowAllocator
+  /** Accounted to this task, because the structs allocated below are JVM-owned. */
+  private val allocator = CometTaskArrowAllocator.forCurrentTask()
 
   /** ArrowImporter does not hold any state and does not need to be closed */
-  // Imported buffers wrap memory the native side owns and frees, so they are charged to Comet's
-  // native pool, not to Spark. Importing through the root allocator would report them to Spark as
-  // if they were JVM Arrow bytes and double count them. See CometImportedArrowAllocator.
-  private val importer = new ArrowImporter(CometImportedArrowAllocator)
+  // Imported buffers wrap memory the native side owns and frees, already charged to Comet's native
+  // pool, so they go through the unaccounted root rather than this task's allocator. Importing
+  // through the latter would report them to Spark as if they were JVM Arrow bytes.
+  private val importer = new ArrowImporter(CometArrowAllocator)
 
   /**
    * Dictionary provider to use for the lifetime of this instance of NativeUtil. The dictionary
diff --git a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala
index 805eae988e8..8a47fe44557 100644
--- a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala
+++ b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala
@@ -26,17 +26,19 @@ import scala.util.control.NonFatal
 import org.apache.arrow.vector.VectorSchemaRoot
 import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel}
 import org.apache.arrow.vector.ipc.message.MessageChannelReader
+import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.CometArrowAllocator
-
 /**
  * A reader that consumes Arrow data from an input channel, and produces Comet batches.
  */
 case class StreamReader(channel: ReadableByteChannel, source: String) extends AutoCloseable {
+  // Decoded into JVM-owned buffers, so it is accounted to the task doing the read. Held in a val
+  // because the reader and the channel reader must share one allocator.
+  private val allocator = CometTaskArrowAllocator.forCurrentTask()
   private val channelReader =
-    new MessageChannelReader(new ReadChannel(channel), CometArrowAllocator)
-  private var arrowReader = new ArrowStreamReader(channelReader, CometArrowAllocator)
+    new MessageChannelReader(new ReadChannel(channel), allocator)
+  private var arrowReader = new ArrowStreamReader(channelReader, allocator)
 
   // Reading the schema allocates the root's vectors, so it can fail with buffers already taken.
   // No caller holds this reader until its constructor returns, so close it here or nothing will.
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index 09060742199..9db366c02e6 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -19,111 +19,161 @@
 
 package org.apache.spark.comet
 
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
+
+import scala.util.control.NonFatal
 
 import org.apache.arrow.memory.AllocationListener
-import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.internal.Logging
-import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager}
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager}
 
 import org.apache.comet.CometConf
 
 /**
- * Reports JVM-side Arrow allocations to Spark's memory manager.
+ * Accounts one task's JVM-side Arrow allocations against Spark's off-heap execution pool.
  *
  * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, so until now the
  * off-heap bytes it hands out were counted by nobody: not Spark's `TaskMemoryManager`, and not
  * Comet's native memory pool. They are still resident in the container, which makes them a blind
- * spot when an executor is killed for exceeding its memory limit.
+ * spot when an executor is killed for exceeding its memory limit. This closes the reporting half
+ * of that gap: the bytes appear in `TaskMemoryManager.showMemoryUsage` and are arbitrated against
+ * Spark's other off-heap consumers.
  *
- * This listener closes the reporting half of that gap. Every allocation is charged to a
- * [[MemoryConsumer]] belonging to the task that made it, so the bytes appear in
- * `TaskMemoryManager.showMemoryUsage` and are arbitrated against Spark's other off-heap
- * consumers.
+ * '''Ownership.''' One instance is created per task and attached to that task's Arrow allocator
+ * by [[CometTaskArrowAllocator]]. Arrow reports an allocation and its matching release to the
+ * listener of the allocator that '''owns''' the buffer, on whichever thread happens to drop the
+ * last reference, and `AllocationListener` is handed nothing but a size. Binding the listener to
+ * an allocator is therefore the only way to attribute a release, and reading `TaskContext` inside
+ * the callbacks would get it wrong: the JVM UDF path exports a JVM-owned vector to native, which
+ * drops it later from a Tokio worker with no task context installed. That release would be lost,
+ * leaving the task charged for memory it had already freed, batch after batch.
  *
- * It deliberately does not enforce. A short grant from Spark is logged and the allocation
- * proceeds, because Arrow allocation on these paths cannot fail today and making it fail is a
- * behavioural change that belongs in its own commit. Note that enforcement belongs in
- * `onPreAllocation`, the only callback permitted to throw, and `onFailedAllocation`, not here.
- * See [[https://github.com/apache/datafusion-comet/issues/5997]].
+ * '''Reporting only.''' A short grant is logged and the allocation proceeds, because Arrow
+ * allocation on these paths cannot fail today and making it fail is a behavioural change that
+ * belongs in its own commit. Enforcement belongs in `onPreAllocation`, the only callback
+ * permitted to throw, and in `onFailedAllocation`, not here. See
+ * [[https://github.com/apache/datafusion-comet/issues/5997]].
  *
- * Buffers imported over the C Data Interface never reach this listener at all. They wrap memory
- * the native side owns, so Comet imports them through `CometImportedArrowAllocator`, a child with
- * no listener. Charging them here would double count bytes already reserved in Comet's native
- * pool. Arrow notifies only the allocating allocator's own listener, which is what makes that
- * separation work.
+ * '''Neither callback may throw.''' Arrow's `AllocationListener` documents that, and
+ * `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so
+ * throwing from here loses the buffer Arrow has already created and never hands back. Spark's
+ * acquisition is fallible -- it runs other consumers' `spill`, which turns a task interrupt into
+ * a `RuntimeException` and an I/O failure into a `SparkOutOfMemoryError` -- so every call into
+ * the memory manager is wrapped and reported rather than propagated.
  *
- * Three cases are handled by doing nothing, each for a different reason:
- *   - No active task. Broadcast coalescing and the cached batch serializer can allocate from the
- *     driver or a non-task thread, where there is no task to charge.
- *   - On-heap mode. Comet's on-heap mode exists so the Spark SQL suite can run without off-heap
- *     memory configured; charging an off-heap consumer there would be wrong.
- *   - A buffer released after its allocating task has finished. The allocator is process-wide
- *     precisely because buffers can outlive the task that created them, so the task's reservation
- *     is dropped at task end and later releases are ignored rather than double-counted.
+ * '''Lock order.''' [[getUsed]] and [[spill]] must stay lock-free, because Spark calls both while
+ * holding the `TaskMemoryManager` monitor, and [[adjust]] holds this listener's monitor across
+ * `acquireExecutionMemory`, which takes that monitor. Were the snapshot to take this monitor
+ * instead, a native reservation arriving through `CometTaskMemoryManager` on a Comet Tokio thread
+ * could hold Spark's monitor and wait for ours while an Arrow allocation on the same task held
+ * ours and waited for Spark's.
  */
-class CometArrowAllocationListener extends AllocationListener {
+private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryManager)
+    extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP)
+    with AllocationListener {
 
   import CometArrowAllocationListener._
 
-  private val reservations = new ConcurrentHashMap[Long, TaskReservation]()
+  /**
+   * Bytes Arrow currently holds on this task's behalf. An atomic rather than a guarded field so
+   * that [[getUsed]] can read it without taking this listener's monitor; see the lock order note
+   * above.
+   */
+  private val live = new AtomicLong(0L)
+
+  /** Bytes currently reserved with Spark. Guarded by this listener's monitor. */
+  private var reserved = 0L
+
+  /** Set once the owning task has finished. Volatile so [[getUsed]] can read it lock-free. */
+  @volatile private var completed = false
 
   override def onAllocation(size: Long): Unit = {
-    val reservation = reservationForCurrentTask()
-    if (reservation != null) {
-      reservation.allocated(size)
-    }
+    live.addAndGet(size)
+    adjustQuietly()
   }
 
   override def onRelease(size: Long): Unit = {
-    val reservation = reservationForCurrentTask()
-    if (reservation != null) {
-      reservation.released(size)
-    }
+    live.addAndGet(-size)
+    adjustQuietly()
   }
 
-  /** Bytes currently reserved with Spark on behalf of the given task. Visible for testing. */
-  private[comet] def reservedBytesForTask(taskAttemptId: Long): Long = {
-    val reservation = reservations.get(taskAttemptId)
-    if (reservation == null) 0L else reservation.reservedBytes
+  /**
+   * Reports our own tally. Spark reads this for spill-victim ordering, `showMemoryUsage` and
+   * end-of-task leak reporting. The inherited `used` counter stays at zero because this consumer
+   * never calls `acquireMemory` or `allocatePage`; Arrow has already obtained the memory and we
+   * are only accounting for it.
+   *
+   * Reports zero once the task has finished, so that buffers deliberately allowed to outlive
+   * their task are not reported by `cleanUpAllAllocatedMemory` as a Spark memory leak.
+   */
+  override def getUsed: Long = if (completed) 0L else math.max(0L, live.get())
+
+  /** Comet's native operators cannot be made to spill from here. See issue #5997. */
+  override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+
+  /**
+   * Drops the whole reservation and stops accounting.
+   *
+   * Called from the owning task's completion listener. Anything still alive afterwards is a
+   * buffer that outlives its task, which the process-wide allocator exists to allow; those
+   * releases are ignored rather than charged to whichever task happens to be running by then.
+   */
+  private[comet] def taskCompleted(): Unit = {
+    try {
+      synchronized {
+        completed = true
+        if (reserved > 0L) {
+          taskMemoryManager.releaseExecutionMemory(reserved, this)
+          reserved = 0L
+        }
+      }
+    } catch {
+      case NonFatal(e) => warnOnMemoryManagerFailure(e)
+      case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
+    }
   }
 
-  private[comet] def trackedTaskCount: Int = reservations.size()
-
-  private def reservationForCurrentTask(): TaskReservation = {
-    // Cheapest check first, and the one that eliminates the most callers: the driver, broadcast
-    // coalescing and the cached batch serializer all allocate with no task in scope. Reading the
-    // config before this would also mean re-reading `SparkEnv` on every allocation in a process
-    // that never has one.
-    val taskContext = TaskContext.get()
-    if (taskContext == null) return null
-    if (!accountingEnabled) return null
-
-    val taskMemoryManager = taskContext.taskMemoryManager()
-    if (taskMemoryManager == null ||
-      taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) {
-      return null
+  /** Bytes Arrow currently holds on this task's behalf. Visible for testing. */
+  private[comet] def liveBytes: Long = live.get()
+
+  /** Bytes currently reserved with Spark on this task's behalf. Visible for testing. */
+  private[comet] def reservedBytes: Long = synchronized(reserved)
+
+  private def adjustQuietly(): Unit = {
+    try {
+      adjust()
+    } catch {
+      // Both of these are reachable: `acquireExecutionMemory` runs other consumers' `spill`, and
+      // `TaskMemoryManager` rethrows an interrupt as a RuntimeException and an IOException as a
+      // SparkOutOfMemoryError, which is an Error and so slips past NonFatal.
+      case NonFatal(e) => warnOnMemoryManagerFailure(e)
+      case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
     }
+  }
 
-    val taskAttemptId = taskContext.taskAttemptId()
-    val existing = reservations.get(taskAttemptId)
-    if (existing != null) return existing
-
-    // Deliberately not `computeIfAbsent`: `addTaskCompletionListener` runs the callback inline if
-    // the task has already completed, and that callback removes from this same map, which is a
-    // recursive update inside a mapping function. Registering outside the map operation avoids it.
-    val created = new TaskReservation(taskMemoryManager)
-    val previous = reservations.putIfAbsent(taskAttemptId, created)
-    if (previous != null) return previous
-
-    taskContext.addTaskCompletionListener[Unit] { _ =>
-      val finished = reservations.remove(taskAttemptId)
-      if (finished != null) {
-        finished.close()
+  private def adjust(): Unit = synchronized {
+    if (!completed) {
+      val liveBytes = math.max(0L, live.get())
+      if (reserved < liveBytes) {
+        // Round up so `reserved` stays a block multiple and growth always leaves headroom.
+        // Requesting the bare deficit would land exactly on `liveBytes` for any buffer at or above
+        // the block size, sending the very next allocation straight back into Spark's lock.
+        val request = roundUpToBlock(liveBytes - reserved)
+        val granted = taskMemoryManager.acquireExecutionMemory(request, this)
+        reserved += granted
+        if (granted < request) {
+          warnOnShortGrant(request, granted)
+        }
+      } else {
+        // Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes on
+        // the executor-wide pool, so a per-block loop would take that lock once per megabyte freed.
+        val excess = ((reserved - liveBytes) / BLOCK_SIZE) * BLOCK_SIZE
+        if (excess > 0L) {
+          taskMemoryManager.releaseExecutionMemory(excess, this)
+          reserved -= excess
+        }
       }
     }
-    created
   }
 }
 
@@ -135,22 +185,10 @@ object CometArrowAllocationListener extends Logging {
    * in whole blocks and only block-crossing changes reach Spark. Deliberately not configurable:
    * it trades lock chatter against reservation slack and has no plausible per-workload tuning.
    */
-  private val BLOCK_SIZE = 1024L * 1024L
+  private[comet] val BLOCK_SIZE = 1024L * 1024L
 
   private val shortGrantLogged = new AtomicBoolean(false)
-
-  /**
-   * Resolved once per JVM. The listener is attached to a `val` in a package object, so it is
-   * constructed on first touch of `CometArrowAllocator`, which can happen before any
-   * `SparkSession` exists and on executors where `SQLConf` does not carry Comet's settings. This
-   * is only read once a `TaskContext` exists, by which point an executor has a `SparkEnv`; the
-   * `Option` guard covers tests that install a task context without one.
-   */
-  private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env =>
-    env.conf.getBoolean(
-      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key,
-      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.defaultValue.get)
-  }
+  private val memoryManagerFailureLogged = new AtomicBoolean(false)
 
   private def roundUpToBlock(bytes: Long): Long =
     ((bytes + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE
@@ -165,61 +203,15 @@ object CometArrowAllocationListener extends Logging {
     }
   }
 
-  /** One task's reservation against Spark's off-heap pool. */
-  private class TaskReservation(taskMemoryManager: TaskMemoryManager)
-      extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) {
-
-    // Named `usedBytes` rather than `used` on purpose: `MemoryConsumer` already declares a
-    // `protected long used`, and a private field of that name narrows the inherited member, which
-    // the compiler rejects as weaker access privileges in overriding.
-    private var usedBytes: Long = 0L
-    private var reserved: Long = 0L
-
-    /** Comet's native operators cannot be made to spill from here. See issue #5997. */
-    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
-
-    /**
-     * Reports our own tally. Spark reads this for spill-victim ordering, `showMemoryUsage` and
-     * end-of-task leak reporting. The inherited `used` counter stays at zero because this
-     * consumer never calls `acquireMemory` or `allocatePage`; Arrow has already obtained the
-     * memory and we are only accounting for it.
-     */
-    override def getUsed: Long = synchronized(usedBytes)
-
-    def reservedBytes: Long = synchronized(reserved)
-
-    def allocated(size: Long): Unit = synchronized {
-      usedBytes += size
-      if (reserved < usedBytes) {
-        // Round up so `reserved` stays a block multiple and growth always leaves headroom.
-        // Requesting the bare deficit would land exactly on `usedBytes` for any buffer at or above
-        // the block size, sending the very next allocation straight back into Spark's lock.
-        val request = roundUpToBlock(usedBytes - reserved)
-        val granted = taskMemoryManager.acquireExecutionMemory(request, this)
-        reserved += granted
-        if (granted < request) {
-          warnOnShortGrant(request, granted)
-        }
-      }
-    }
-
-    def released(size: Long): Unit = synchronized {
-      usedBytes = math.max(0L, usedBytes - size)
-      // Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes on
-      // the executor-wide pool, so a per-block loop would take that lock once per megabyte freed.
-      val excess = ((reserved - usedBytes) / BLOCK_SIZE) * BLOCK_SIZE
-      if (excess > 0L) {
-        taskMemoryManager.releaseExecutionMemory(excess, this)
-        reserved -= excess
-      }
-    }
-
-    def close(): Unit = synchronized {
-      if (reserved > 0L) {
-        taskMemoryManager.releaseExecutionMemory(reserved, this)
-        reserved = 0L
-      }
-      usedBytes = 0L
+  private def warnOnMemoryManagerFailure(e: Throwable): Unit = {
+    if (memoryManagerFailureLogged.compareAndSet(false, true)) {
+      logWarning(
+        "Failed to report a JVM Arrow allocation to Spark's memory manager. The allocation " +
+          "itself is unaffected, so this is a reporting gap rather than a failure, but Spark's " +
+          "view of these bytes will be short until the task ends. " +
+          s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " +
+          "reporting these allocations to Spark.",
+        e)
     }
   }
 }
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
new file mode 100644
index 00000000000..bcedff62e20
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
@@ -0,0 +1,194 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.comet
+
+import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue}
+
+import scala.util.control.NonFatal
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.memory.MemoryMode
+
+import org.apache.comet.{CometArrowAllocator, CometConf}
+
+/**
+ * Hands out the Arrow allocator that JVM-owned allocations should use, one per Spark task.
+ *
+ * Each task gets a child of `CometArrowAllocator` carrying its own
+ * [[CometArrowAllocationListener]], so the bytes it hands out are reported to that task's
+ * `TaskMemoryManager`. The allocator, not the calling thread, is what identifies the owner:
+ * Arrow's `AllocationListener` is given only a size, and a buffer is released on whichever thread
+ * drops the last reference, which for anything exported over the C Data Interface is a Comet
+ * Tokio worker with no task context installed. Child allocators cut from the returned allocator
+ * inherit its listener, so the paths that make their own children are covered too.
+ *
+ * Callers that get no task -- the driver, broadcast coalescing, the cached batch serializer --
+ * and callers running with Comet's on-heap mode get the unaccounted process-wide root instead,
+ * which is what they used before this existed. Buffers imported over the C Data Interface must
+ * also use the root: they wrap memory the native side owns and frees, and charging them to Spark
+ * would double count bytes an operator has already reserved in Comet's native pool.
+ *
+ * '''Lifetime.''' The task allocator cannot simply be closed when the task ends. The process-wide
+ * allocator exists precisely because Arrow buffers can outlive the task that created them, and
+ * Arrow treats closing an allocator that still owns bytes as a leak. So at task completion the
+ * Spark reservation is dropped and the allocator is closed only if it has been drained; otherwise
+ * it is parked and closed by a later task once the stragglers are released. Leaving it open
+ * indefinitely is not an option: `BaseAllocator` keeps every child in a map until it closes.
+ */
+object CometTaskArrowAllocator extends Logging {
+
+  private class TaskAllocator(
+      val allocator: BufferAllocator,
+      val listener: CometArrowAllocationListener)
+
+  private val perTask = new ConcurrentHashMap[Long, TaskAllocator]()
+
+  /** Allocators whose task has ended but which still own bytes. Guarded by [[closeLock]]. */
+  private val awaitingClose = new ConcurrentLinkedQueue[BufferAllocator]()
+
+  private val closeLock = new Object
+
+  /**
+   * Resolved once per JVM. Read from the `SparkConf` rather than `SQLConf`, because this is
+   * reached from executor threads where `SQLConf` does not carry Comet's settings. The `Option`
+   * guard covers tests that install a task context without a `SparkEnv`.
+   */
+  private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env =>
+    env.conf.getBoolean(
+      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key,
+      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.defaultValue.get)
+  }
+
+  /**
+   * The allocator to use for JVM-owned Arrow buffers on the calling thread. Never null, and never
+   * an allocator belonging to a task other than the current one.
+   */
+  def forCurrentTask(): BufferAllocator = {
+    // Cheapest check first, and the one that eliminates the most callers: the driver, broadcast
+    // coalescing and the cached batch serializer all allocate with no task in scope.
+    val taskContext = TaskContext.get()
+    if (taskContext == null) {
+      CometArrowAllocator
+    } else {
+      val existing = perTask.get(taskContext.taskAttemptId())
+      if (existing != null) existing.allocator else create(taskContext)
+    }
+  }
+
+  private def create(taskContext: TaskContext): BufferAllocator = {
+    val taskMemoryManager = taskContext.taskMemoryManager()
+    // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory
+    // configured, and charging an off-heap consumer there would be wrong. Note this differs from
+    // CometUnifiedShuffleMemoryAllocator, which throws in that situation; throwing here would
+    // break those tests.
+    if (!accountingEnabled || taskMemoryManager == null ||
+      taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) {
+      return CometArrowAllocator
+    }
+
+    val taskAttemptId = taskContext.taskAttemptId()
+    closeDrained()
+    val listener = new CometArrowAllocationListener(taskMemoryManager)
+    val allocator = CometArrowAllocator
+      .newChildAllocator(s"comet-task-$taskAttemptId", listener, 0L, Long.MaxValue)
+    val created = new TaskAllocator(allocator, listener)
+    val previous = perTask.putIfAbsent(taskAttemptId, created)
+    if (previous != null) {
+      // Lost the race with another thread in the same task; the loser's allocator is untouched.
+      closeQuietly(allocator)
+      previous.allocator
+    } else {
+      // Deliberately outside the map operation: `addTaskCompletionListener` runs the callback
+      // inline if the task has already completed, and that callback removes from this same map,
+      // which would be a recursive update inside a mapping function.
+      taskContext.addTaskCompletionListener[Unit](_ => taskCompleted(taskAttemptId))
+      // ...and if it did run inline, the allocator just created is already closed, so hand back
+      // the root rather than something the caller cannot allocate from.
+      if (perTask.containsKey(taskAttemptId)) allocator else CometArrowAllocator
+    }
+  }
+
+  private def taskCompleted(taskAttemptId: Long): Unit = {
+    val finished = perTask.remove(taskAttemptId)
+    if (finished != null) {
+      finished.listener.taskCompleted()
+      closeLock.synchronized {
+        if (!tryClose(finished.allocator)) {
+          awaitingClose.add(finished.allocator)
+        }
+      }
+    }
+    closeDrained()
+  }
+
+  /** Closes any parked allocator whose stragglers have since been released. */
+  private def closeDrained(): Unit = {
+    if (!awaitingClose.isEmpty) {
+      closeLock.synchronized {
+        val parked = awaitingClose.iterator()
+        while (parked.hasNext) {
+          if (tryClose(parked.next())) {
+            parked.remove()
+          }
+        }
+      }
+    }
+  }
+
+  /** Closes the allocator if it has been drained. Returns false if it must stay open. */
+  private def tryClose(allocator: BufferAllocator): Boolean = {
+    if (allocator.getAllocatedMemory != 0L) {
+      false
+    } else {
+      closeQuietly(allocator)
+      true
+    }
+  }
+
+  private def closeQuietly(allocator: BufferAllocator): Unit = {
+    try {
+      allocator.close()
+    } catch {
+      case NonFatal(e) =>
+        // Closing is bookkeeping: the bytes are already gone and the Spark reservation is already
+        // released, so a failure here must not propagate into a task completion listener.
+        logWarning(s"Failed to close Arrow allocator ${allocator.getName}", e)
+    }
+  }
+
+  /** Number of tasks currently holding an accounted allocator. Visible for testing. */
+  private[comet] def trackedTaskCount: Int = perTask.size()
+
+  /** Number of finished tasks whose allocator is still draining. Visible for testing. */
+  private[comet] def awaitingCloseCount: Int = awaitingClose.size()
+
+  /** The listener accounting for the given task, if it has one. Visible for testing. */
+  private[comet] def listenerForTask(
+      taskAttemptId: Long): Option[CometArrowAllocationListener] = {
+    Option(perTask.get(taskAttemptId)).map(_.listener)
+  }
+
+  /** Bytes currently reserved with Spark on behalf of the given task. Visible for testing. */
+  private[comet] def reservedBytesForTask(taskAttemptId: Long): Long = {
+    listenerForTask(taskAttemptId).map(_.reservedBytes).getOrElse(0L)
+  }
+}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala
index 821f84e0c2b..9051c16e901 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala
@@ -25,6 +25,7 @@ import scala.collection.JavaConverters._
 import scala.util.control.NonFatal
 
 import org.apache.spark.TaskContext
+import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, UnsafeProjection}
@@ -39,8 +40,6 @@ import org.apache.spark.storage.StorageLevel
 import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.io.ChunkedByteBuffer
 
-import org.apache.comet.CometArrowAllocator
-
 /**
  * Cached batch format used when Comet writes Spark in-memory cache data.
  *
@@ -356,7 +355,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
         Utils.serializeBatchColumns(batch)
       } else {
         val arrowBatch =
-          CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator)
+          CometArrowConverters.columnarBatchToArrowBatch(
+            batch,
+            arrowSchema,
+            CometTaskArrowAllocator.forCurrentTask())
         try Utils.serializeBatchColumns(arrowBatch)
         finally arrowBatch.close()
       }
@@ -624,7 +626,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
           // the Unix epoch regardless of session timezone, so no values are converted. It also
           // matches Comet's native schema, avoiding a cast at the native boundary.
           CometArrowStream.NATIVE_TIMEZONE,
-          CometArrowAllocator)
+          CometTaskArrowAllocator.forCurrentTask())
 
         encodeBatches(iter, schema)
       }
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
index a60e2811ecb..8f4ed2bef5e 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
@@ -27,6 +27,7 @@ import org.apache.arrow.memory.BufferAllocator
 import org.apache.arrow.vector.ipc.ArrowReader
 import org.apache.arrow.vector.types.pojo.{Field, FieldType, Schema}
 import org.apache.spark.TaskContext
+import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.comet.util.Utils
@@ -34,7 +35,6 @@ import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.CometArrowAllocator
 import org.apache.comet.vector.{CometDictionaryVector, CometVector, NativeUtil}
 
 /**
@@ -230,7 +230,8 @@ object CometArrowStream extends Logging {
       name: String,
       readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = {
     val context = TaskContext.get()
-    val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue)
+    val allocator =
+      CometTaskArrowAllocator.forCurrentTask().newChildAllocator(name, 0, Long.MaxValue)
     var reader: ArrowReader = null
     var arrowStream: ArrowArrayStream = null
     try {
@@ -274,7 +275,8 @@ object CometArrowStream extends Logging {
       name: String,
       readerFactory: BufferAllocator => ArrowReader): Iterator[ColumnarBatch] = {
     val context = TaskContext.get()
-    val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue)
+    val allocator =
+      CometTaskArrowAllocator.forCurrentTask().newChildAllocator(name, 0, Long.MaxValue)
     val reader =
       try readerFactory(allocator)
       catch {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
index d70fdab35e6..ff95daef47c 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
@@ -373,7 +373,8 @@ object Utils extends CometTypeShim with Logging {
       return (Array.empty, 0L, 0L)
     }
 
-    val allocator = org.apache.comet.CometArrowAllocator
+    val allocator = org.apache.spark.comet.CometTaskArrowAllocator
+      .forCurrentTask()
       .newChildAllocator("broadcast-coalesce", 0, Long.MaxValue)
     try {
       var targetRoot: VectorSchemaRoot = null
@@ -567,7 +568,7 @@ object Utils extends CometTypeShim with Logging {
             cv.dataType(),
             rows,
             s"_const_$index",
-            org.apache.comet.CometArrowAllocator,
+            org.apache.spark.comet.CometTaskArrowAllocator.forCurrentTask(),
             "UTC")
           (materialized, None)
 
diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
index fecf781baac..3229f11d7c3 100644
--- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
+++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
@@ -34,13 +34,13 @@ import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch, Me
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
 
-import org.apache.comet.CometArrowAllocator
 import org.apache.comet.vector.{CometDecodedVector, CometDictionaryVector, CometVector, CometVectorUtils}
 
 /**
@@ -124,7 +124,9 @@ private[python] trait CometArrowPythonRunnerBase
     new Writer(env, worker, inputIterator, partitionIndex, context) {
 
       private val allocator =
-        CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+        CometTaskArrowAllocator
+          .forCurrentTask()
+          .newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
       private var batches = inputIterator.flatten
       // Upstream owns this batch. Even hasNext may close it and reuse its native buffers, so
       // leave the upstream iterators untouched until all ranges have been serialized.
@@ -264,7 +266,9 @@ private[python] trait CometArrowPythonRunnerBase
     new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
 
       private val allocator =
-        CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+        CometTaskArrowAllocator
+          .forCurrentTask()
+          .newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
       private var reader: ArrowStreamReader = _
       private var root: VectorSchemaRoot = _
       private var batchLoaded = true

From 094863ff9968edc9ebed5eb82cd59d7d5a5fc156 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 14:02:15 -0600
Subject: [PATCH 08/18] test: cover release ownership, task completion and
 failed acquisitions

The old fixture called `cleanUpAllAllocatedMemory` directly and never ran
the registered completion listener, so nothing exercised what happens at
the end of a task. It now goes through `markTaskCompleted`, and the suite
adds the cases the per-task binding exists for:

- a release on a thread with no task context, and a release arriving
  under a different task, are both charged to the task that allocated
- completion drops the whole reservation and closes the allocator
- a buffer outliving its task parks the allocator until it is released,
  then a later task reaps it
- a spill that fails with an I/O error or an interrupt neither fails the
  Arrow allocation nor leaks the buffer Arrow already created
- `getUsed` and `spill` do not wait on the reservation monitor, probed by
  holding that monitor and calling both from another thread
- concurrent Arrow allocations and native-style reservations both finish

Also covers what the root allocator is for now: allocating from it inside
a task is not reported, which is what keeps FFI-imported buffers out of
Spark's accounting, and a child of a task allocator is reported to that
task, which is what covers the Python runner and the native Arrow source.
---
 .../CometArrowAllocationListenerSuite.scala   | 424 +++++++++++++++---
 1 file changed, 367 insertions(+), 57 deletions(-)

diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index 11eb4fc807e..990f3b2f017 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -19,119 +19,374 @@
 
 package org.apache.spark.comet
 
+import java.io.{InterruptedIOException, IOException}
 import java.util.Properties
+import java.util.concurrent.atomic.{AtomicLong, AtomicReference}
 
 import org.scalatest.funsuite.AnyFunSuite
 
-import org.apache.arrow.memory.{AllocationListener, RootAllocator}
+import org.apache.arrow.memory.BufferAllocator
 import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
 import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager}
+
+import org.apache.comet.CometArrowAllocator
 
 /**
- * Tests that JVM Arrow allocations are reported to Spark, and, just as importantly, that the
- * paths where they cannot be reported fail quietly rather than throwing. Arrow allocation on
- * these paths cannot fail today and this listener must not change that.
+ * Tests that JVM Arrow allocations are reported to Spark, that they are reported against the task
+ * that made them rather than whichever task happens to be on the releasing thread, and, just as
+ * importantly, that the paths where they cannot be reported fail quietly rather than throwing.
+ * Arrow allocation on these paths cannot fail today and this listener must not change that.
  */
 class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
-  private val blockSize = 1024L * 1024L
+  private val blockSize = CometArrowAllocationListener.BLOCK_SIZE
   private val poolBytes = 64L * 1024 * 1024
-  private val taskAttemptId = 1L
+
+  /** Task attempt ids are keys in a process-wide map, so no two tests may share one. */
+  private val nextTaskAttemptId = new AtomicLong(1000L)
+
+  // ---------------------------------------------------------------------------------------------
+  // Reservation arithmetic. Driven through the listener directly, since Arrow's rounding policy
+  // would otherwise decide the sizes under test.
+  // ---------------------------------------------------------------------------------------------
 
   test("allocations are charged to the current task in whole blocks") {
-    withTask() { listener =>
+    withTask() { task =>
+      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
+
       // Far smaller than a block, so the reservation should round up to exactly one block.
       listener.onAllocation(128L)
-      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize)
+      assert(listener.reservedBytes == blockSize)
 
       // Still inside the first block, so Spark is not asked again.
       listener.onAllocation(1024L)
-      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize)
+      assert(listener.reservedBytes == blockSize)
+
+      listener.taskCompleted()
     }
   }
 
   test("a request larger than a block rounds up to a block multiple") {
-    withTask() { listener =>
+    withTask() { task =>
+      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
       listener.onAllocation(blockSize * 3 + 7L)
       // Rounded up rather than sized to the exact deficit, so growth leaves headroom and the next
       // small allocation does not go straight back into Spark.
-      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize * 4)
+      assert(listener.reservedBytes == blockSize * 4)
+      listener.taskCompleted()
     }
   }
 
   test("releasing returns whole blocks to Spark") {
-    withTask() { listener =>
+    withTask() { task =>
+      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
       listener.onAllocation(blockSize * 2)
-      assert(listener.reservedBytesForTask(taskAttemptId) == blockSize * 2)
+      assert(listener.reservedBytes == blockSize * 2)
 
       listener.onRelease(blockSize * 2)
-      assert(listener.reservedBytesForTask(taskAttemptId) == 0L)
+      assert(listener.reservedBytes == 0L)
+      listener.taskCompleted()
+    }
+  }
+
+  // ---------------------------------------------------------------------------------------------
+  // Which allocator is handed out, and what it charges.
+  // ---------------------------------------------------------------------------------------------
+
+  test("a real Arrow allocation from the task allocator is charged to that task") {
+    withTask() { task =>
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      assert(allocator ne CometArrowAllocator)
+      val buf = allocator.buffer(blockSize)
+      try {
+        assert(reservedFor(task) == blockSize)
+        assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize)
+      } finally {
+        buf.close()
+      }
+      assert(reservedFor(task) == 0L)
     }
   }
 
-  test("a child allocator with no listener is not charged, but the root still is") {
-    withTask() { listener =>
-      val root = new RootAllocator(listener, Long.MaxValue)
-      val imported =
-        root.newChildAllocator("imported", AllocationListener.NOOP, 0, Long.MaxValue)
+  test("a child of the task allocator is charged to the same task") {
+    withTask() { task =>
+      // The Python runner and the native Arrow source cut their own children. Arrow passes the
+      // parent's listener down, so they are accounted without knowing anything about it.
+      val child =
+        CometTaskArrowAllocator.forCurrentTask().newChildAllocator("probe", 0L, Long.MaxValue)
       try {
-        // This is how FFI-imported buffers, which wrap memory the native side owns, stay out of
-        // Spark's accounting: Arrow notifies only the allocating allocator's own listener, never
-        // its ancestors, so the root's listener does not see the child's allocation.
-        val importedBuf = imported.buffer(blockSize)
+        val buf = child.buffer(blockSize)
         try {
-          assert(listener.trackedTaskCount == 0)
-
-          // The same root does still charge for JVM-owned allocations made directly against it.
-          val jvmBuf = root.buffer(blockSize)
-          try {
-            assert(listener.reservedBytesForTask(taskAttemptId) >= blockSize)
-          } finally {
-            jvmBuf.close()
-          }
+          assert(reservedFor(task) == blockSize)
         } finally {
-          importedBuf.close()
+          buf.close()
         }
       } finally {
-        imported.close()
-        root.close()
+        child.close()
       }
+      assert(reservedFor(task) == 0L)
     }
   }
 
-  test("no active task is a no-op rather than an error") {
+  test("the process-wide root is not accounted, which is what FFI imports rely on") {
+    withTask() { task =>
+      // Establish the task allocator first, so this asserts "not charged" rather than "no task".
+      CometTaskArrowAllocator.forCurrentTask()
+      // Buffers imported over the C Data Interface wrap memory the native side owns and frees,
+      // already charged to Comet's native pool, so they are allocated from the listener-less root.
+      val buf = CometArrowAllocator.buffer(blockSize)
+      try {
+        assert(reservedFor(task) == 0L)
+      } finally {
+        buf.close()
+      }
+    }
+  }
+
+  test("no active task uses the unaccounted root allocator") {
     TaskContext.unset()
-    val listener = new CometArrowAllocationListener
     // Broadcast coalescing and the cached batch serializer can allocate off a task thread.
-    listener.onAllocation(4096L)
-    listener.onRelease(4096L)
-    assert(listener.trackedTaskCount == 0)
+    assert(CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator)
   }
 
-  test("on-heap mode is not accounted") {
-    withTask(offHeap = false) { listener =>
-      listener.onAllocation(blockSize)
+  test("on-heap mode uses the unaccounted root allocator") {
+    withTask(offHeap = false) { task =>
       // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory.
-      // Charging an off-heap consumer there would be wrong, so nothing is tracked. Distinguishing
-      // "not tracked" from "tracked with zero reserved" is why trackedTaskCount is asserted here.
-      assert(listener.trackedTaskCount == 0)
-      assert(listener.reservedBytesForTask(taskAttemptId) == 0L)
+      // Charging an off-heap consumer there would be wrong.
+      assert(CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator)
+      assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty)
+    }
+  }
+
+  // ---------------------------------------------------------------------------------------------
+  // Ownership: a release is attributed to the task that allocated, not to the releasing thread.
+  // ---------------------------------------------------------------------------------------------
+
+  test("a release on a thread with no task context is charged to the allocating task") {
+    withTask() { task =>
+      val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize)
+      assert(reservedFor(task) == blockSize)
+
+      // This is what the JVM UDF path does: it exports a JVM-owned vector, and Rust drops it later
+      // from a Tokio worker that has no task context installed. Reading TaskContext in onRelease
+      // would ignore this release and leave the task charged for memory it had already freed.
+      onDetachedThread(buf.close())
+
+      assert(reservedFor(task) == 0L)
+      assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
     }
   }
 
-  private def withTask(offHeap: Boolean = true)(f: CometArrowAllocationListener => Unit): Unit = {
+  test("a release under a different task does not touch that task's accounting") {
+    withTask() { taskA =>
+      val bufA = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize)
+      withTask() { taskB =>
+        val bufB = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize)
+        assert(reservedFor(taskB) == blockSize)
+
+        // A's buffer, released while B's context is on the thread. B must not pay for it.
+        bufA.close()
+
+        assert(reservedFor(taskB) == blockSize)
+        assert(reservedFor(taskA) == 0L)
+        bufB.close()
+      }
+    }
+  }
+
+  // ---------------------------------------------------------------------------------------------
+  // Task completion, and buffers that outlive their task.
+  // ---------------------------------------------------------------------------------------------
+
+  test("task completion releases the whole reservation and closes the allocator") {
+    val task = newTask()
+    val allocatorName = withInstalledTask(task) {
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val buf = allocator.buffer(blockSize)
+      assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize)
+      buf.close()
+
+      task.context.markTaskCompleted(None)
+
+      assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty)
+      assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
+      allocator.getName
+    }
+    assert(!rootChildNames().contains(allocatorName))
+  }
+
+  test("a buffer outliving its task parks the allocator until it is released") {
+    val task = newTask()
+    withInstalledTask(task) {
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val buf = allocator.buffer(blockSize)
+
+      task.context.markTaskCompleted(None)
+
+      // The reservation goes back to Spark even though the buffer is still alive: the task is
+      // over, and leaving it charged would be reported as a Spark memory leak.
+      assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
+      assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty)
+      // Arrow treats closing an allocator that still owns bytes as a leak, so it has to stay open.
+      assert(rootChildNames().contains(allocator.getName))
+
+      // A late release is ignored rather than charged to whoever is running by then...
+      buf.close()
+      assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
+
+      // ...and the drained allocator is reaped by the next task, so the root does not accumulate
+      // one child per task attempt.
+      withTask() { _ => CometTaskArrowAllocator.forCurrentTask() }
+      assert(!rootChildNames().contains(allocator.getName))
+    }
+  }
+
+  // ---------------------------------------------------------------------------------------------
+  // Failure containment: Spark's acquisition is fallible, Arrow's callbacks are not allowed to be.
+  // ---------------------------------------------------------------------------------------------
+
+  for ((label, failure) <- Seq(
+      "an I/O failure" -> new IOException("spill failed"),
+      "an interrupt" -> new InterruptedIOException("task killed"))) {
+    test(s"$label while spilling does not fail or leak the Arrow allocation") {
+      // Exactly one block of budget, already taken by a consumer that refuses to spill, so the
+      // listener's acquisition has to go through Spark's spill path and comes back throwing.
+      withTask(pool = blockSize) { task =>
+        val hostile = new FailingSpillConsumer(task.taskMemoryManager, failure)
+        assert(hostile.take(blockSize) == blockSize)
+
+        val allocator = CometTaskArrowAllocator.forCurrentTask()
+        // `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so
+        // throwing from the listener would lose this buffer: Arrow neither returns nor frees it.
+        val buf = allocator.buffer(blockSize)
+        try {
+          assert(allocator.getAllocatedMemory == blockSize)
+          // Nothing was reserved, which is what says the acquisition really did go down the spill
+          // path and throw rather than quietly succeeding and making this test vacuous.
+          assert(reservedFor(task) == 0L)
+        } finally {
+          buf.close()
+        }
+        // Zero here is the leak check: a buffer Arrow created but never handed back would still
+        // be counted.
+        assert(allocator.getAllocatedMemory == 0L)
+      }
+    }
+  }
+
+  // ---------------------------------------------------------------------------------------------
+  // Lock order. Spark calls getUsed and spill while holding the TaskMemoryManager monitor, and the
+  // listener holds its own monitor while waiting for that one.
+  // ---------------------------------------------------------------------------------------------
+
+  test("the usage snapshot does not take the reservation monitor") {
+    withTask() { task =>
+      val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize)
+      try {
+        val listener = CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).get
+        val used = new AtomicLong(-1L)
+        val spilled = new AtomicLong(-1L)
+        listener.synchronized {
+          // A native reservation arriving through CometTaskMemoryManager on a Tokio thread holds
+          // Spark's monitor here. If either call waited on this one, it would deadlock against an
+          // Arrow allocation on the same task that already holds this monitor and wants Spark's.
+          val probe = new Thread(() => {
+            used.set(listener.getUsed)
+            spilled.set(listener.spill(blockSize, listener))
+          })
+          probe.setDaemon(true)
+          probe.setName("lock-order-probe")
+          probe.start()
+          probe.join(30000L)
+          assert(!probe.isAlive, "getUsed or spill blocked on the reservation monitor")
+        }
+        assert(used.get == blockSize)
+        assert(spilled.get == 0L)
+      } finally {
+        buf.close()
+      }
+    }
+  }
+
+  test("concurrent Arrow and native reservations make progress") {
+    // Two blocks of budget shared by both consumers, so most requests are short and Spark walks
+    // its consumer list, calling getUsed on the Arrow listener while holding its own monitor.
+    withTask(pool = blockSize * 2) { task =>
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val native = new NativeLikeConsumer(task.taskMemoryManager)
+      val failure = new AtomicReference[Throwable]()
+
+      val arrowThread = loopingThread("arrow-allocations", failure) {
+        val buf = allocator.buffer(blockSize)
+        buf.close()
+      }
+      val nativeThread = loopingThread("native-reservations", failure) {
+        native.release(native.reserve(blockSize))
+      }
+
+      Seq(arrowThread, nativeThread).foreach(_.start())
+      Seq(arrowThread, nativeThread).foreach { t =>
+        t.join(60000L)
+        assert(!t.isAlive, s"${t.getName} did not finish; concurrent reservations deadlocked")
+      }
+
+      Option(failure.get).foreach(e => fail("a worker failed", e))
+      assert(allocator.getAllocatedMemory == 0L)
+      native.release(native.used())
+    }
+  }
+
+  // ---------------------------------------------------------------------------------------------
+  // Fixtures.
+  // ---------------------------------------------------------------------------------------------
+
+  private class FailingSpillConsumer(tmm: TaskMemoryManager, failure: IOException)
+      extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
+    def take(bytes: Long): Long = acquireMemory(bytes)
+    override def spill(size: Long, trigger: MemoryConsumer): Long = throw failure
+  }
+
+  /** Stands in for `CometTaskMemoryManager`: reserves from Spark directly and never spills. */
+  private class NativeLikeConsumer(tmm: TaskMemoryManager)
+      extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
+    private val reserved = new AtomicLong(0L)
+    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+    override def getUsed: Long = reserved.get()
+    def used(): Long = reserved.get()
+    def reserve(bytes: Long): Long = {
+      val granted = tmm.acquireExecutionMemory(bytes, this)
+      reserved.addAndGet(granted)
+      granted
+    }
+    def release(bytes: Long): Unit = {
+      if (bytes > 0L) {
+        reserved.addAndGet(-bytes)
+        tmm.releaseExecutionMemory(bytes, this)
+      }
+    }
+  }
+
+  private case class TaskFixture(
+      taskAttemptId: Long,
+      context: TaskContextImpl,
+      taskMemoryManager: TaskMemoryManager)
+
+  private def reservedFor(task: TaskFixture): Long =
+    CometTaskArrowAllocator.reservedBytesForTask(task.taskAttemptId)
+
+  private def newTask(offHeap: Boolean = true, pool: Long = poolBytes): TaskFixture = {
     val conf = new SparkConf(false)
     if (offHeap) {
       conf
         .set("spark.memory.offHeap.enabled", "true")
-        .set("spark.memory.offHeap.size", poolBytes.toString)
+        .set("spark.memory.offHeap.size", pool.toString)
     }
     val memoryManager = new TestMemoryManager(conf)
-    memoryManager.limit(poolBytes)
+    memoryManager.limit(pool)
+    val taskAttemptId = nextTaskAttemptId.getAndIncrement()
     val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId)
-
-    val taskContext = new TaskContextImpl(
+    val context = new TaskContextImpl(
       stageId = 0,
       stageAttemptNumber = 0,
       partitionId = 0,
@@ -144,16 +399,71 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       taskMetrics = TaskMetrics.empty,
       cpus = 1,
       resources = Map.empty)
+    TaskFixture(taskAttemptId, context, taskMemoryManager)
+  }
 
-    TaskContext.setTaskContext(taskContext)
+  /** Installs the task on this thread, restoring whatever was there before. */
+  private def withInstalledTask[T](task: TaskFixture)(f: => T): T = {
+    val previous = TaskContext.get()
+    TaskContext.setTaskContext(task.context)
     try {
-      f(new CometArrowAllocationListener)
+      f
     } finally {
       try {
-        taskMemoryManager.cleanUpAllAllocatedMemory()
+        // Fires the completion listener that drops the reservation; harmless if already run.
+        task.context.markTaskCompleted(None)
+        task.taskMemoryManager.cleanUpAllAllocatedMemory()
       } finally {
-        TaskContext.unset()
+        if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous)
       }
     }
   }
+
+  private def withTask(offHeap: Boolean = true, pool: Long = poolBytes)(
+      f: TaskFixture => Unit): Unit = {
+    val task = newTask(offHeap, pool)
+    withInstalledTask(task)(f(task))
+  }
+
+  /** Runs the body on a fresh thread, which by construction carries no task context. */
+  private def onDetachedThread(body: => Unit): Unit = {
+    val failure = new AtomicReference[Throwable]()
+    val thread = new Thread(() => {
+      try body
+      catch { case t: Throwable => failure.set(t) }
+    })
+    thread.setDaemon(true)
+    thread.setName("detached-release")
+    thread.start()
+    thread.join(30000L)
+    assert(!thread.isAlive, "the detached release did not finish")
+    Option(failure.get).foreach(t => throw t)
+  }
+
+  private def loopingThread(name: String, failure: AtomicReference[Throwable])(
+      body: => Unit): Thread = {
+    val thread = new Thread(() => {
+      try {
+        var i = 0
+        while (i < 500) {
+          body
+          i += 1
+        }
+      } catch {
+        case t: Throwable => failure.compareAndSet(null, t)
+      }
+    })
+    thread.setDaemon(true)
+    thread.setName(name)
+    thread
+  }
+
+  private def rootChildNames(): Set[String] = {
+    val names = Set.newBuilder[String]
+    val children = CometArrowAllocator.getChildAllocators.iterator()
+    while (children.hasNext) {
+      names += children.next().asInstanceOf[BufferAllocator].getName
+    }
+    names.result()
+  }
 }

From d591b8f2ba0206e3172f64a1f8ef57573ad9cca6 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 14:02:15 -0600
Subject: [PATCH 09/18] test: benchmark the cost of reporting Arrow allocations
 to Spark

Accounting is on by default, so the overhead needs a number rather than
an assurance. Each case runs the same allocate/release loop against a
plain `RootAllocator` and against a task allocator, over the three buffer
shapes Comet actually produces, and once more with a second thread
reserving from the same constrained pool the way the native side does.
Reservation call counts are printed per iteration, because how often a
buffer crosses a block boundary matters more than the per-buffer
bookkeeping.
---
 ...ometArrowAllocationListenerBenchmark.scala | 311 ++++++++++++++++++
 1 file changed, 311 insertions(+)
 create mode 100644 spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala

diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
new file mode 100644
index 00000000000..b719b4c462e
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
@@ -0,0 +1,311 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.comet
+
+import java.util.Properties
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
+
+import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, RootAllocator}
+import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
+import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, TaskMemoryManager, TestMemoryManager}
+
+/**
+ * Measures what reporting JVM Arrow allocations to Spark costs, since it is on by default.
+ *
+ * Each case runs the same allocate/release loop twice: once against a plain `RootAllocator`,
+ * which is what Comet did before [[CometArrowAllocationListener]] existed and what
+ * `spark.comet.arrowAllocator.accounting.enabled=false` restores, and once against a task
+ * allocator carrying the listener. Reservation call counts are printed under each table, because
+ * the interesting variable is not the per-buffer bookkeeping but how often a buffer size crosses
+ * a block boundary and has to go into `TaskMemoryManager` at all.
+ *
+ * To run this benchmark:
+ * {{{
+ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.comet.CometArrowAllocationListenerBenchmark
+ * }}}
+ */
+object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
+
+  private val blockSize = 1024L * 1024L
+  private val poolBytes = 1024L * 1024L * 1024L
+
+  override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+    runBenchmark("JVM Arrow allocations reported to Spark") {
+      // The FFI struct shape: NativeUtil allocates two of these per column per batch, and they
+      // never come close to a block, so the listener should never reach Spark after the first one.
+      allocateAndRelease("small FFI structs", bufferSize = 128L, buffersPerIteration = 512)
+      // A wide batch: many medium buffers alive at once, one block boundary crossed per few
+      // buffers on the way up and the same on the way down.
+      allocateAndRelease(
+        "wide batch buffers",
+        bufferSize = 64L * 1024L,
+        buffersPerIteration = 128)
+      // Worst case for the block batching: every allocation crosses a boundary, so every one of
+      // them takes the executor-wide lock in `acquireExecutionMemory`.
+      allocateAndRelease("block-sized buffers", bufferSize = blockSize, buffersPerIteration = 8)
+      // Same, with a second thread reserving from the same constrained pool the way Comet's native
+      // side does through CometTaskMemoryManager.
+      allocateUnderNativePressure()
+      // Per call site rather than per buffer, but it is the cost the root allocator `val` did not
+      // have: a TaskContext lookup and a concurrent map read.
+      allocatorLookup()
+    }
+  }
+
+  private def allocateAndRelease(
+      name: String,
+      bufferSize: Long,
+      buffersPerIteration: Int): Unit = {
+    val benchmark =
+      new Benchmark(
+        s"$name (${buffersPerIteration}x$bufferSize)",
+        buffersPerIteration,
+        output = output)
+
+    // Both allocators are built once, outside the timed body, so what is measured is the
+    // steady-state cost of allocating and releasing rather than the cost of standing a task up.
+    val root = new RootAllocator(Long.MaxValue)
+    try {
+      withTaskAllocator() { (accounted, memory) =>
+        benchmark.addCase("not accounted") { _ =>
+          churn(root, bufferSize, buffersPerIteration)
+        }
+        benchmark.addCase("accounted") { _ =>
+          churn(accounted, bufferSize, buffersPerIteration)
+        }
+        benchmark.run()
+
+        // One more round with the counters zeroed, to report how often a single iteration reaches
+        // the memory manager. That, rather than the per-buffer bookkeeping, is the cost that
+        // scales with buffer size.
+        memory.reset()
+        churn(accounted, bufferSize, buffersPerIteration)
+        writeLine(s"  accounted: ${memory.summary(buffersPerIteration)} per iteration")
+      }
+    } finally {
+      root.close()
+    }
+  }
+
+  private def allocateUnderNativePressure(): Unit = {
+    val buffersPerIteration = 8
+    val benchmark = new Benchmark(
+      s"block-sized buffers under native pressure (${buffersPerIteration}x$blockSize)",
+      buffersPerIteration,
+      output = output)
+
+    val root = new RootAllocator(Long.MaxValue)
+    try {
+      // Two blocks of pool for both arms, with the pressure thread running throughout. In the
+      // unaccounted arm only that thread touches the pool, which is the point: the delta is what
+      // the Arrow side adds once it competes for the same budget.
+      withTaskAllocator(poolBytes = blockSize * 2) { (accounted, memory) =>
+        withNativePressure(memory) {
+          benchmark.addCase("not accounted") { _ =>
+            churn(root, blockSize, buffersPerIteration)
+          }
+          benchmark.addCase("accounted") { _ =>
+            churn(accounted, blockSize, buffersPerIteration)
+          }
+          benchmark.run()
+
+          memory.reset()
+          churn(accounted, blockSize, buffersPerIteration)
+          writeLine(s"  accounted: ${memory.summary(buffersPerIteration)} per iteration")
+        }
+      }
+    } finally {
+      root.close()
+    }
+  }
+
+  private def allocatorLookup(): Unit = {
+    val lookupsPerIteration = 100000
+    val benchmark =
+      new Benchmark("allocator lookup", lookupsPerIteration.toLong, output = output)
+
+    // The "before" shape: call sites read a package-object `val`, which the JIT folds away
+    // entirely. The interesting number is therefore the absolute cost of the second case.
+    benchmark.addCase("process-wide val") { _ =>
+      var i = 0
+      var sink = 0
+      while (i < lookupsPerIteration) {
+        sink += System.identityHashCode(org.apache.comet.CometArrowAllocator)
+        i += 1
+      }
+      assert(sink != Int.MinValue)
+    }
+    benchmark.addCase("forCurrentTask()") { _ =>
+      withTaskAllocator() { (_, _) =>
+        var i = 0
+        var sink = 0
+        while (i < lookupsPerIteration) {
+          sink += System.identityHashCode(CometTaskArrowAllocator.forCurrentTask())
+          i += 1
+        }
+        assert(sink != Int.MinValue)
+      }
+    }
+
+    benchmark.run()
+  }
+
+  /**
+   * Allocates the whole set, then releases it, so the peak is what the reservation has to cover.
+   */
+  private def churn(allocator: BufferAllocator, bufferSize: Long, count: Int): Unit = {
+    val buffers = new Array[ArrowBuf](count)
+    var i = 0
+    while (i < count) {
+      buffers(i) = allocator.buffer(bufferSize)
+      i += 1
+    }
+    i = 0
+    while (i < count) {
+      buffers(i).close()
+      i += 1
+    }
+  }
+
+  /**
+   * Runs the body against a task allocator, torn down afterwards, so that repeated iterations do
+   * not accumulate reservations or allocators.
+   */
+  private def withTaskAllocator[T](poolBytes: Long = poolBytes)(
+      f: (BufferAllocator, CountingTaskMemoryManager) => T): T = {
+    val memory = newTaskMemoryManager(poolBytes)
+    val context = newTaskContext(memory)
+    val previous = TaskContext.get()
+    TaskContext.setTaskContext(context)
+    try {
+      f(CometTaskArrowAllocator.forCurrentTask(), memory)
+    } finally {
+      try {
+        context.markTaskCompleted(None)
+        memory.cleanUpAllAllocatedMemory()
+      } finally {
+        if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous)
+      }
+    }
+  }
+
+  /** Hammers the same pool from another thread, the way native reservations do. */
+  private def withNativePressure[T](memory: TaskMemoryManager)(f: => T): T = {
+    val stop = new AtomicBoolean(false)
+    val consumer = new NativeLikeConsumer(memory)
+    val thread = new Thread(() => {
+      while (!stop.get()) {
+        val granted = consumer.reserve(blockSize)
+        consumer.release(granted)
+      }
+    })
+    thread.setDaemon(true)
+    thread.setName("native-reservations")
+    thread.start()
+    try f
+    finally {
+      stop.set(true)
+      thread.join()
+    }
+  }
+
+  private val nextTaskAttemptId = new AtomicLong(1L)
+
+  private def newTaskMemoryManager(poolBytes: Long): CountingTaskMemoryManager = {
+    val conf = new SparkConf(false)
+      .set("spark.memory.offHeap.enabled", "true")
+      .set("spark.memory.offHeap.size", poolBytes.toString)
+    val memoryManager = new TestMemoryManager(conf)
+    memoryManager.limit(poolBytes)
+    new CountingTaskMemoryManager(memoryManager, nextTaskAttemptId.getAndIncrement())
+  }
+
+  private def newTaskContext(memory: CountingTaskMemoryManager): TaskContextImpl = {
+    new TaskContextImpl(
+      stageId = 0,
+      stageAttemptNumber = 0,
+      partitionId = 0,
+      numPartitions = 1,
+      taskAttemptId = memory.getTaskAttemptId,
+      attemptNumber = 0,
+      taskMemoryManager = memory,
+      localProperties = new Properties,
+      metricsSystem = null,
+      taskMetrics = TaskMetrics.empty,
+      cpus = 1,
+      resources = Map.empty)
+  }
+
+  private def writeLine(line: String): Unit = {
+    // scalastyle:off println
+    println(line)
+    // scalastyle:on println
+    output.foreach(_.write(s"$line\n".getBytes("UTF-8")))
+  }
+
+  private class CountingTaskMemoryManager(memoryManager: MemoryManager, taskAttemptId: Long)
+      extends TaskMemoryManager(memoryManager, taskAttemptId) {
+    private val acquires = new AtomicLong(0L)
+    private val releases = new AtomicLong(0L)
+
+    def getTaskAttemptId: Long = taskAttemptId
+
+    // Only the Arrow listener's own calls are counted, so the pressure thread's traffic does not
+    // land in the reported figure.
+    override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = {
+      if (consumer.isInstanceOf[CometArrowAllocationListener]) acquires.incrementAndGet()
+      super.acquireExecutionMemory(required, consumer)
+    }
+
+    override def releaseExecutionMemory(size: Long, consumer: MemoryConsumer): Unit = {
+      if (consumer.isInstanceOf[CometArrowAllocationListener]) releases.incrementAndGet()
+      super.releaseExecutionMemory(size, consumer)
+    }
+
+    def reset(): Unit = {
+      acquires.set(0L)
+      releases.set(0L)
+    }
+
+    def summary(buffers: Int): String =
+      s"${acquires.get()} acquire and ${releases.get()} release calls for $buffers buffers"
+  }
+
+  /** Stands in for `CometTaskMemoryManager`: reserves from Spark directly and never spills. */
+  private class NativeLikeConsumer(memory: TaskMemoryManager)
+      extends MemoryConsumer(memory, 0L, MemoryMode.OFF_HEAP) {
+    private val reserved = new AtomicLong(0L)
+    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+    override def getUsed: Long = reserved.get()
+    def reserve(bytes: Long): Long = {
+      val granted = memory.acquireExecutionMemory(bytes, this)
+      reserved.addAndGet(granted)
+      granted
+    }
+    def release(bytes: Long): Unit = {
+      if (bytes > 0L) {
+        reserved.addAndGet(-bytes)
+        memory.releaseExecutionMemory(bytes, this)
+      }
+    }
+  }
+}

From 054fe9a278ca88cd952a287c98a5568b799ebbdd Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 15:54:58 -0600
Subject: [PATCH 10/18] fix: contain acquisition failures and stop
 double-charging FFI buffers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two more findings from review.

**Spark's acquisition can fail in ways `NonFatal` does not catch, and can
lose memory when it does.** `ExecutionMemoryPool.acquireMemory` parks in
`lock.wait()` when a task is below its fair share, so killing a task
raises a plain `InterruptedException` out of `acquireExecutionMemory` —
and `NonFatal` deliberately excludes it, so it escaped `onAllocation` and
Arrow lost the buffer it had already created. It is now caught, with the
thread's interrupt flag re-armed rather than swallowed, so the
cancellation is still there for the task to observe.

Separately, `acquireExecutionMemory` takes its first grant from the pool
and only then asks other consumers to spill, so when a spill throws it
has already charged the task for bytes it never reports back. Nothing
would have released them before the final task cleanup. `acquire` now
adopts them, measured as the change in what the pool says the task holds
— an estimate, but one that can only come out too small, because
`acquireExecutionMemory` holds the `TaskMemoryManager` monitor for its
whole duration so no concurrent acquire can inflate it.

**Buffers crossing the FFI boundary are no longer charged on the JVM
side.** Native's pool is the authority for bytes native holds: whichever
DataFusion operator retains a batch reserves those same buffers through
Comet's unified pool, which charges the same Spark task. Charging them
here as well reserved the same memory twice and could reject an
allocation that fits — a hard failure this PR has no business
introducing. So `NativeUtil`, the JVM UDF result and
`CometNativeArrowSource.stream` join the imports on the unaccounted root,
while IPC reads, codegen output, the cached batch serializer and
`CometNativeArrowSource.readerBatchIter` keep the task allocator.

The split is by allocation site, so it is not exact: a buffer used in the
JVM and only later handed to native, a shuffle-read batch feeding a
native operator being the common shape, is still charged on both sides.
Closing that needs reservation ownership handed over at the boundary,
which is a change on both sides of it. Recorded in the guide's open
problems.
---
 .../contributor-guide/memory_management.md    | 66 +++++++++------
 .../org/apache/comet/udf/CometUdfBridge.java  | 16 ++--
 .../main/scala/org/apache/comet/package.scala | 17 ++--
 .../org/apache/comet/vector/NativeUtil.scala  | 15 ++--
 .../comet/CometArrowAllocationListener.scala  | 81 ++++++++++++++++---
 .../spark/comet/CometTaskArrowAllocator.scala | 26 ++++--
 .../arrow/CometNativeArrowSource.scala        | 10 ++-
 ...ometArrowAllocationListenerBenchmark.scala |  7 +-
 .../CometArrowAllocationListenerSuite.scala   | 65 +++++++++++++--
 9 files changed, 224 insertions(+), 79 deletions(-)

diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 81deaef4953..1ac9e40b8a3 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -64,14 +64,14 @@ Enabling Comet does not add one new memory consumer, it adds several, and they a
 accounted by the same party. This inventory is worth internalizing before reading the rest of the
 page:
 
-| Allocator                                    | Lives in    | Bounded by                                                    | Visible to Spark? |
-| -------------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
-| Spark execution + storage (on-heap)          | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
-| Spark Tungsten (off-heap)                    | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
-| Comet native (Rust global allocator)         | Native heap | `memory_limit` (see below), enforced only via the memory pool | No                |
-| Comet JVM Arrow, in a task (task allocator)  | Off-heap    | **Nothing**, but reported to `TaskMemoryManager`              | Yes               |
-| Comet JVM Arrow, off a task, and FFI imports | Off-heap    | **Nothing**: a `RootAllocator(Long.MaxValue)`                 | No                |
-| Comet JVM shuffle pages                      | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+| Allocator                                 | Lives in    | Bounded by                                                    | Visible to Spark? |
+| ----------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
+| Spark execution + storage (on-heap)       | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
+| Spark Tungsten (off-heap)                 | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+| Comet native (Rust global allocator)      | Native heap | `memory_limit` (see below), enforced only via the memory pool | No                |
+| Comet JVM Arrow that stays in the JVM     | Off-heap    | **Nothing**, but reported to `TaskMemoryManager`              | Yes               |
+| Comet JVM Arrow crossing the FFI boundary | Off-heap    | **Nothing**: a `RootAllocator(Long.MaxValue)`                 | No                |
+| Comet JVM shuffle pages                   | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
 
 Several observations follow.
 
@@ -95,14 +95,25 @@ memory it had already freed. Binding one listener to one task's allocator is wha
 land on the task that allocated. Children cut from a task allocator inherit its listener, so the
 paths that make their own children are covered without knowing about any of this.
 
-**Anything allocated from the root itself is not accounted, which is what FFI imports want.**
-Arrow's `wrapForeignAllocation` reports an imported buffer to the allocator's listener at full
-capacity even though no JVM-side allocation happened, so charging them would count native memory
-that Comet's own pool has already reserved, with the error growing in proportion to batch
-throughput. `NativeUtil` and `CometUdfBridge` therefore import through the listener-less root,
-while export, IPC and materialisation use the task allocator. The driver, broadcast coalescing and
-the cached batch serializer also fall back to the root when there is no task to charge, as does
-Comet's on-heap mode.
+**Anything crossing the FFI boundary uses the root, and is not accounted.** Both directions, for
+the same reason: native's pool is the authority for bytes native holds. Coming in, Arrow's
+`wrapForeignAllocation` reports an imported buffer to the allocator's listener at full capacity
+even though no JVM-side allocation happened. Going out, whichever DataFusion operator retains the
+batch reserves those buffers itself -- `ExternalSorter` through `get_reserved_bytes_for_record_batch`,
+the hash join build side through `get_record_batch_memory_size`, both of which read sizes off the
+`ArrayData` and so count imported buffers -- and through a unified pool that charges the same Spark
+task. Reporting either direction on the JVM side would reserve the same memory twice and could
+reject an allocation that fits. So `NativeUtil`, the JVM UDF result in `CometUdfBridge`, and
+`CometNativeArrowSource.stream` all use the listener-less root; IPC reads, codegen output, the
+cached batch serializer and `CometNativeArrowSource.readerBatchIter` use the task allocator. The
+driver, broadcast coalescing and the cached batch serializer also fall back to the root when there
+is no task to charge, as does Comet's on-heap mode.
+
+**The split is by allocation site, so it is not exact.** A buffer that is used in the JVM and only
+later handed to native -- a shuffle-read batch feeding a native operator is the common shape -- is
+allocated somewhere that cannot know its future, so it stays charged on the JVM side while native
+may also reserve it. Closing that needs reservation ownership to be handed over at the boundary,
+which is a change on both sides of it.
 
 **A task allocator is closed only once it has been drained.** The process-wide allocator exists
 because Arrow buffers can outlive the task that created them, and Arrow treats closing an allocator
@@ -233,13 +244,13 @@ accounting (if any) saw the bytes appear. Whether the buffer is _reserved_ in Co
 separate decision, made by whichever operator holds it, and that operator neither knows nor cares
 which side of the boundary the bytes came from.
 
-**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of the task
-allocator and exports the whole per-partition iterator once as an `ArrowArrayStream`.
+**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of the
+unaccounted root and exports the whole per-partition iterator once as an `ArrowArrayStream`.
 `ScanExec` imports each batch through `AlignedArrowStreamReader` with `CopyMode::UnpackOrClone`:
 dictionary columns are unpacked into new native arrays, everything else is an `Arc` clone of the
 imported buffers. Those bytes stay where Java Arrow put them and are pinned for as long as any native
-reference survives. They are reported to Spark's `TaskMemoryManager` at allocation time, but nothing
-caps them, because the task allocator reports without enforcing. Whether they are charged _later_ depends
+reference survives. They are invisible to Spark's `TaskMemoryManager` at allocation time, deliberately:
+this is the export path, so nobody charged for them yet. Whether they are charged _later_ depends
 on who holds them. DataFusion's `ExternalSorter` reserves `get_reserved_bytes_for_record_batch` for
 every batch it retains, and the hash join build side reserves `get_record_batch_memory_size` for
 each incoming batch; both read buffer sizes off the `ArrayData` and apply equally to imported
@@ -349,11 +360,16 @@ much they matter:
   the fact. There is no runtime value that an operator, a metric, or a policy could read.
 - **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to
   guess a per-workload haircut rather than measuring anything.
-- **`CometArrowAllocator` is unbounded.** Allocations made inside a task are now reported to
-  Spark's memory manager, so they are no longer invisible, but nothing caps them: the listener
-  reports without enforcing, and enforcing would mean failing allocations on paths that cannot fail
-  today. Allocations made off a task, and buffers imported over the C Data Interface, are not
-  reported at all.
+- **`CometArrowAllocator` is unbounded.** Allocations that live and die inside a task are now
+  reported to Spark's memory manager, so they are no longer invisible, but nothing caps them: the
+  listener reports without enforcing, and enforcing would mean failing allocations on paths that
+  cannot fail today. Allocations made off a task, and anything crossing the FFI boundary in either
+  direction, are not reported at all.
+- **Reservation ownership is not handed over at the FFI boundary.** Which side accounts for a
+  buffer is decided by where it was allocated, not by who holds it, so a buffer allocated for JVM
+  use and later handed to native is charged on the JVM side while a native operator that retains it
+  charges the same task again. Buffers allocated for export dodge this only because their
+  allocation site knows where they are going.
 - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be
   resident on either side with no reservation covering it, because reservations are made and
   withdrawn by individual operators while the bytes outlive them.
diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
index 3a7ae0e53a4..5c70bf370ca 100644
--- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
+++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
@@ -209,11 +209,11 @@ private static void evaluateInternal(
             });
     assert udf != null : "reflective instantiation returned null for " + udfClassName;
 
-    // The result vector below is JVM-owned, so it is allocated from this task's accounted
-    // allocator. Rust retains it through `from_ffi` and drops it later from a Tokio worker with no
-    // task context installed, which is exactly why the accounting is bound to the allocator rather
-    // than to the releasing thread.
-    BufferAllocator allocator = org.apache.spark.comet.CometTaskArrowAllocator.forCurrentTask();
+    // The unaccounted root on both sides. The inputs wrap native memory that the native side owns
+    // and frees, and the result below is exported straight back to native, where whichever
+    // operator retains the batch reserves those same buffers in Comet's pool and so charges the
+    // same Spark task. Reporting either direction to Spark would count native memory a second time.
+    BufferAllocator allocator = org.apache.comet.package$.MODULE$.CometArrowAllocator();
 
     ValueVector[] inputs = new ValueVector[inputArrayPtrs.length];
     ValueVector result = null;
@@ -221,11 +221,7 @@ private static void evaluateInternal(
       for (int i = 0; i < inputArrayPtrs.length; i++) {
         ArrowArray inArr = ArrowArray.wrap(inputArrayPtrs[i]);
         ArrowSchema inSch = ArrowSchema.wrap(inputSchemaPtrs[i]);
-        // Imported from native memory that the native side owns and frees, already charged to
-        // Comet's native pool, so it goes through the unaccounted root allocator.
-        inputs[i] =
-            Data.importVector(
-                org.apache.comet.package$.MODULE$.CometArrowAllocator(), inArr, inSch, null);
+        inputs[i] = Data.importVector(allocator, inArr, inSch, null);
       }
 
       result = udf.evaluate(inputs, numRows);
diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala
index b26cca4f4cc..b67aa301364 100644
--- a/spark/src/main/scala/org/apache/comet/package.scala
+++ b/spark/src/main/scala/org/apache/comet/package.scala
@@ -34,14 +34,17 @@ package object comet {
    * leaked. To avoid this, we use a single allocator for the whole execution process.
    *
    * It carries no allocation listener, so allocating from it directly is not reported to Spark's
-   * memory manager. That is what buffers imported over the Arrow C Data Interface want: they wrap
-   * memory the native side owns and frees, already charged to Comet's native pool, and Arrow's
-   * `wrapForeignAllocation` would otherwise report the full buffer capacity as though a JVM-side
-   * allocation had happened.
+   * memory manager. That is what memory on either side of the C Data Interface wants. Imported
+   * buffers wrap memory the native side owns and frees, already charged to Comet's native pool,
+   * and Arrow's `wrapForeignAllocation` would otherwise report the full buffer capacity as though
+   * a JVM-side allocation had happened. Buffers allocated to be exported are the mirror image:
+   * whichever native operator retains the batch reserves them through Comet's unified pool, which
+   * charges the same Spark task, so reporting them here too would reserve the same memory twice.
    *
-   * JVM-owned allocations should go through `CometTaskArrowAllocator.forCurrentTask()` instead,
-   * which cuts a per-task child whose listener reports the bytes to Spark. Off a task it hands
-   * back this allocator, so the driver-side paths are unchanged.
+   * Allocations that live and die in the JVM should go through
+   * `CometTaskArrowAllocator.forCurrentTask()` instead, which cuts a per-task child whose
+   * listener reports the bytes to Spark. Off a task it hands back this allocator, so the
+   * driver-side paths are unchanged.
    */
   val CometArrowAllocator = new RootAllocator(Long.MaxValue)
 
diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
index 88df5f5b1ca..9ce5a6ed18e 100644
--- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
+++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
@@ -26,7 +26,6 @@ import org.apache.arrow.util.AutoCloseables
 import org.apache.arrow.vector.VectorSchemaRoot
 import org.apache.arrow.vector.dictionary.DictionaryProvider
 import org.apache.spark.SparkException
-import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.sql.comet.execution.arrow.ConstantColumnVectors
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution.vectorized.ConstantColumnVector
@@ -48,14 +47,16 @@ import org.apache.comet.CometArrowAllocator
 class NativeUtil extends AutoCloseable {
   import Utils._
 
-  /** Accounted to this task, because the structs allocated below are JVM-owned. */
-  private val allocator = CometTaskArrowAllocator.forCurrentTask()
+  // The unaccounted root, on both sides of the boundary this class straddles. Imported buffers
+  // wrap memory the native side owns and frees, already charged to Comet's native pool. Everything
+  // allocated here goes the other way -- the FFI structs, the materialised constant vectors, the
+  // exported vectors -- and whichever native operator retains a batch reserves those same buffers
+  // in Comet's pool, which charges the same Spark task. Reporting either direction to Spark would
+  // count native memory a second time.
+  private val allocator = CometArrowAllocator
 
   /** ArrowImporter does not hold any state and does not need to be closed */
-  // Imported buffers wrap memory the native side owns and frees, already charged to Comet's native
-  // pool, so they go through the unaccounted root rather than this task's allocator. Importing
-  // through the latter would report them to Spark as if they were JVM Arrow bytes.
-  private val importer = new ArrowImporter(CometArrowAllocator)
+  private val importer = new ArrowImporter(allocator)
 
   /**
    * Dictionary provider to use for the lifetime of this instance of NativeUtil. The dictionary
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index 9db366c02e6..bc0bb92ea48 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -44,9 +44,9 @@ import org.apache.comet.CometConf
  * listener of the allocator that '''owns''' the buffer, on whichever thread happens to drop the
  * last reference, and `AllocationListener` is handed nothing but a size. Binding the listener to
  * an allocator is therefore the only way to attribute a release, and reading `TaskContext` inside
- * the callbacks would get it wrong: the JVM UDF path exports a JVM-owned vector to native, which
- * drops it later from a Tokio worker with no task context installed. That release would be lost,
- * leaving the task charged for memory it had already freed, batch after batch.
+ * the callbacks would get it wrong: a shuffle-read batch handed on to a native operator is pinned
+ * by native and dropped later from a Tokio worker with no task context installed. That release
+ * would be lost, leaving the task charged for memory it had already freed, batch after batch.
  *
  * '''Reporting only.''' A short grant is logged and the allocation proceeds, because Arrow
  * allocation on these paths cannot fail today and making it fail is a behavioural change that
@@ -57,9 +57,13 @@ import org.apache.comet.CometConf
  * '''Neither callback may throw.''' Arrow's `AllocationListener` documents that, and
  * `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so
  * throwing from here loses the buffer Arrow has already created and never hands back. Spark's
- * acquisition is fallible -- it runs other consumers' `spill`, which turns a task interrupt into
- * a `RuntimeException` and an I/O failure into a `SparkOutOfMemoryError` -- so every call into
- * the memory manager is wrapped and reported rather than propagated.
+ * acquisition is fallible in three ways, and only the first is caught by `NonFatal`: it runs
+ * other consumers' `spill`, which turns a task interrupt into a `RuntimeException` and an I/O
+ * failure into a `SparkOutOfMemoryError`, and the execution pool itself parks in `lock.wait()`,
+ * so killing a task can raise a plain `InterruptedException` here. Every call into the memory
+ * manager is wrapped and reported rather than propagated, and an interrupt additionally re-arms
+ * the thread's flag so the cancellation is not swallowed. A failed acquisition can also leave the
+ * task charged for bytes Spark never reported back; see [[acquire]].
  *
  * '''Lock order.''' [[getUsed]] and [[spill]] must stay lock-free, because Spark calls both while
  * holding the `TaskMemoryManager` monitor, and [[adjust]] holds this listener's monitor across
@@ -128,6 +132,7 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
         }
       }
     } catch {
+      case e: InterruptedException => reportAndReinterrupt(e)
       case NonFatal(e) => warnOnMemoryManagerFailure(e)
       case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
     }
@@ -143,9 +148,13 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
     try {
       adjust()
     } catch {
-      // Both of these are reachable: `acquireExecutionMemory` runs other consumers' `spill`, and
-      // `TaskMemoryManager` rethrows an interrupt as a RuntimeException and an IOException as a
-      // SparkOutOfMemoryError, which is an Error and so slips past NonFatal.
+      // Growth handles its own failures in `acquire`, so this is the net for the release path and
+      // for anything unforeseen. All three are reachable from the memory manager:
+      // `acquireExecutionMemory` runs other consumers' `spill`, `TaskMemoryManager` turns an
+      // interrupted spill into a RuntimeException and an IOException into a SparkOutOfMemoryError,
+      // and the execution pool itself parks in `lock.wait()`. The last two slip past NonFatal,
+      // which excludes Errors and InterruptedException.
+      case e: InterruptedException => reportAndReinterrupt(e)
       case NonFatal(e) => warnOnMemoryManagerFailure(e)
       case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
     }
@@ -159,7 +168,7 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
         // Requesting the bare deficit would land exactly on `liveBytes` for any buffer at or above
         // the block size, sending the very next allocation straight back into Spark's lock.
         val request = roundUpToBlock(liveBytes - reserved)
-        val granted = taskMemoryManager.acquireExecutionMemory(request, this)
+        val granted = acquire(request)
         reserved += granted
         if (granted < request) {
           warnOnShortGrant(request, granted)
@@ -175,6 +184,58 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
       }
     }
   }
+
+  /**
+   * Asks Spark for `request` bytes and returns what this consumer ends up holding, which is not
+   * always what Spark returns.
+   *
+   * `acquireExecutionMemory` takes its first grant from the pool and only then asks other
+   * consumers to spill, so when a spill throws it has already charged the task for bytes it never
+   * reports back. Nothing would release them: [[taskCompleted]] only knows about `reserved`, and
+   * Spark itself only reclaims them in `cleanUpAllAllocatedMemory` at the very end of the task,
+   * so until then they are headroom nobody can use. They are adopted here instead, measured as
+   * the change in what the pool says this task holds.
+   *
+   * That measurement is an estimate, but a safe one. Another consumer in the same task cannot
+   * acquire concurrently, because `acquireExecutionMemory` holds the `TaskMemoryManager` monitor
+   * throughout, so the only interference is a concurrent release, which makes the figure too
+   * small rather than too large; and `request` bounds it from above either way. Too small
+   * degrades to what would have happened anyway.
+   */
+  private def acquire(request: Long): Long = {
+    val heldBefore = taskMemoryManager.getMemoryConsumptionForThisTask
+    try {
+      taskMemoryManager.acquireExecutionMemory(request, this)
+    } catch {
+      case e: InterruptedException =>
+        reportAndReinterrupt(e)
+        adoptOrphanedGrant(heldBefore, request)
+      case NonFatal(e) =>
+        warnOnMemoryManagerFailure(e)
+        adoptOrphanedGrant(heldBefore, request)
+      case e: SparkOutOfMemoryError =>
+        warnOnMemoryManagerFailure(e)
+        adoptOrphanedGrant(heldBefore, request)
+    }
+  }
+
+  private def adoptOrphanedGrant(heldBefore: Long, request: Long): Long = {
+    val orphaned = taskMemoryManager.getMemoryConsumptionForThisTask - heldBefore
+    math.max(0L, math.min(orphaned, request))
+  }
+
+  /**
+   * An interrupt cannot be allowed out of an Arrow callback any more than anything else can, but
+   * swallowing the cancellation would be wrong too. Spark's execution pool parks in `lock.wait()`
+   * when a task is below its fair share, so killing a task lands here, and `NonFatal`
+   * deliberately excludes `InterruptedException`. Re-arming the flag leaves the cancellation for
+   * the task to observe at its next interruptible point, which is the only place it can act on it
+   * anyway.
+   */
+  private def reportAndReinterrupt(e: InterruptedException): Unit = {
+    Thread.currentThread().interrupt()
+    warnOnMemoryManagerFailure(e)
+  }
 }
 
 object CometArrowAllocationListener extends Logging {
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
index bcedff62e20..fc5074e5a11 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
@@ -37,15 +37,25 @@ import org.apache.comet.{CometArrowAllocator, CometConf}
  * [[CometArrowAllocationListener]], so the bytes it hands out are reported to that task's
  * `TaskMemoryManager`. The allocator, not the calling thread, is what identifies the owner:
  * Arrow's `AllocationListener` is given only a size, and a buffer is released on whichever thread
- * drops the last reference, which for anything exported over the C Data Interface is a Comet
- * Tokio worker with no task context installed. Child allocators cut from the returned allocator
- * inherit its listener, so the paths that make their own children are covered too.
+ * drops the last reference, which for anything that reaches native over the C Data Interface is a
+ * Comet Tokio worker with no task context installed. Child allocators cut from the returned
+ * allocator inherit its listener, so the paths that make their own children are covered too.
  *
- * Callers that get no task -- the driver, broadcast coalescing, the cached batch serializer --
- * and callers running with Comet's on-heap mode get the unaccounted process-wide root instead,
- * which is what they used before this existed. Buffers imported over the C Data Interface must
- * also use the root: they wrap memory the native side owns and frees, and charging them to Spark
- * would double count bytes an operator has already reserved in Comet's native pool.
+ * '''This is for buffers the JVM owns.''' Anything allocated to be handed straight to native --
+ * `NativeUtil`, the JVM UDF result, `CometNativeArrowSource.stream` -- uses the unaccounted root
+ * instead, and so does anything imported from native. Native's pool is the authority for bytes
+ * native holds: whichever DataFusion operator retains the batch reserves those buffers through
+ * Comet's unified pool, which charges the same Spark task, so reporting them here as well would
+ * reserve the same memory twice and could reject an allocation that fits. The same fallback
+ * covers callers with no task to charge -- the driver, broadcast coalescing, the cached batch
+ * serializer
+ * -- and Comet's on-heap mode, where charging an off-heap consumer would be wrong.
+ *
+ * What this does not fix is a buffer that is used in the JVM and only later handed to native, a
+ * shuffle-read batch feeding a native operator being the common shape. Its allocation site cannot
+ * know, so it stays charged here while native may also reserve it. Coordinating reservation
+ * ownership across the FFI boundary needs changes on both sides; see
+ * [[https://github.com/apache/datafusion-comet/issues/5997]].
  *
  * '''Lifetime.''' The task allocator cannot simply be closed when the task ends. The process-wide
  * allocator exists precisely because Arrow buffers can outlive the task that created them, and
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
index 8f4ed2bef5e..d5bc9c41a33 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
@@ -35,6 +35,7 @@ import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
+import org.apache.comet.CometArrowAllocator
 import org.apache.comet.vector.{CometDictionaryVector, CometVector, NativeUtil}
 
 /**
@@ -230,8 +231,11 @@ object CometArrowStream extends Logging {
       name: String,
       readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = {
     val context = TaskContext.get()
-    val allocator =
-      CometTaskArrowAllocator.forCurrentTask().newChildAllocator(name, 0, Long.MaxValue)
+    // A child of the unaccounted root, not of the task allocator: every byte allocated here is
+    // exported to native and exists only for native to consume. Whichever native operator retains
+    // the batch reserves those same buffers in Comet's pool, which charges the same Spark task, so
+    // reporting them on this side too would reserve the same memory twice.
+    val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue)
     var reader: ArrowReader = null
     var arrowStream: ArrowArrayStream = null
     try {
@@ -275,6 +279,8 @@ object CometArrowStream extends Logging {
       name: String,
       readerFactory: BufferAllocator => ArrowReader): Iterator[ColumnarBatch] = {
     val context = TaskContext.get()
+    // Accounted, unlike `stream`: these batches are consumed in the JVM, so nothing on the native
+    // side reserves them.
     val allocator =
       CometTaskArrowAllocator.forCurrentTask().newChildAllocator(name, 0, Long.MaxValue)
     val reader =
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
index b719b4c462e..07e2896f2f5 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
@@ -50,9 +50,10 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
 
   override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
     runBenchmark("JVM Arrow allocations reported to Spark") {
-      // The FFI struct shape: NativeUtil allocates two of these per column per batch, and they
-      // never come close to a block, so the listener should never reach Spark after the first one.
-      allocateAndRelease("small FFI structs", bufferSize = 128L, buffersPerIteration = 512)
+      // Many sub-block buffers, the shape a codegen output vector's validity and offset buffers
+      // take. None of them comes close to a block, so the listener should never reach Spark after
+      // the first one.
+      allocateAndRelease("small buffers", bufferSize = 128L, buffersPerIteration = 512)
       // A wide batch: many medium buffers alive at once, one block boundary crossed per few
       // buffers on the way up and the same on the way down.
       allocateAndRelease(
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index 990f3b2f017..43a885e6c94 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -129,12 +129,13 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
-  test("the process-wide root is not accounted, which is what FFI imports rely on") {
+  test("the process-wide root is not accounted, which is what the FFI paths rely on") {
     withTask() { task =>
       // Establish the task allocator first, so this asserts "not charged" rather than "no task".
       CometTaskArrowAllocator.forCurrentTask()
-      // Buffers imported over the C Data Interface wrap memory the native side owns and frees,
-      // already charged to Comet's native pool, so they are allocated from the listener-less root.
+      // Both directions across the C Data Interface use the listener-less root: imported buffers
+      // wrap memory the native side owns, and buffers allocated for export are reserved again by
+      // whichever native operator retains the batch, through a pool that charges the same task.
       val buf = CometArrowAllocator.buffer(blockSize)
       try {
         assert(reservedFor(task) == 0L)
@@ -168,9 +169,10 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize)
       assert(reservedFor(task) == blockSize)
 
-      // This is what the JVM UDF path does: it exports a JVM-owned vector, and Rust drops it later
-      // from a Tokio worker that has no task context installed. Reading TaskContext in onRelease
-      // would ignore this release and leave the task charged for memory it had already freed.
+      // This is what happens when a shuffle-read batch is handed on to a native operator: native
+      // pins it and drops it later from a Tokio worker with no task context installed. Reading
+      // TaskContext in onRelease would ignore this release and leave the task charged for memory
+      // it had already freed.
       onDetachedThread(buf.close())
 
       assert(reservedFor(task) == 0L)
@@ -275,6 +277,54 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
+  test("an interrupt while spilling is re-armed rather than thrown or swallowed") {
+    // Spark's execution pool parks in `lock.wait()` when a task is below its fair share, so a task
+    // kill raises a plain InterruptedException out of `acquireExecutionMemory`. `NonFatal` excludes
+    // it, so before this it escaped `onAllocation` and Arrow lost the buffer it had just created.
+    // TestMemoryManager never parks, so the interrupt is injected through a failing spill instead.
+    withTask(pool = blockSize) { task =>
+      val hostile =
+        new FailingSpillConsumer(task.taskMemoryManager, new InterruptedException("task killed"))
+      assert(hostile.take(blockSize) == blockSize)
+
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val buf = allocator.buffer(blockSize)
+      try {
+        assert(allocator.getAllocatedMemory == blockSize)
+        assert(reservedFor(task) == 0L)
+      } finally {
+        buf.close()
+      }
+      assert(allocator.getAllocatedMemory == 0L)
+      // Cleared here as well as asserted, so the flag does not leak into the next test.
+      assert(Thread.interrupted(), "the interrupt was swallowed instead of being re-armed")
+    }
+  }
+
+  test("a partial grant lost to a failing spill is adopted rather than stranded") {
+    // One block already taken, one still in the pool, and a two-block request: Spark hands over the
+    // block it has and only then asks the other consumer to spill, which throws. It never reports
+    // the block it already took, so nothing would release it before the task ended.
+    withTask(pool = blockSize * 2) { task =>
+      val hostile =
+        new FailingSpillConsumer(task.taskMemoryManager, new IOException("spill failed"))
+      assert(hostile.take(blockSize) == blockSize)
+
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val buf = allocator.buffer(blockSize * 2)
+      try {
+        assert(reservedFor(task) == blockSize)
+        assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize * 2)
+      } finally {
+        buf.close()
+      }
+      assert(reservedFor(task) == 0L)
+      // Only the other consumer's block is left. Without adopting the orphan this would still be
+      // two blocks, with one of them charged to the task and owned by nobody.
+      assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize)
+    }
+  }
+
   // ---------------------------------------------------------------------------------------------
   // Lock order. Spark calls getUsed and spill while holding the TaskMemoryManager monitor, and the
   // listener holds its own monitor while waiting for that one.
@@ -341,7 +391,8 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   // Fixtures.
   // ---------------------------------------------------------------------------------------------
 
-  private class FailingSpillConsumer(tmm: TaskMemoryManager, failure: IOException)
+  /** Holds memory and refuses to give it back, so `trySpillAndAcquire` throws on its behalf. */
+  private class FailingSpillConsumer(tmm: TaskMemoryManager, failure: Exception)
       extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
     def take(bytes: Long): Long = acquireMemory(bytes)
     override def spill(size: Long, trigger: MemoryConsumer): Long = throw failure

From c2a542d693dca674294f8d9d7dffcc06b5fe0c7d Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 17 Sep 2026 21:45:20 -0600
Subject: [PATCH 11/18] fix: measure a lost grant under the task memory manager
 monitor

The orphaned grant adopted after a failed acquisition was measured with a snapshot taken before acquireExecutionMemory took the TaskMemoryManager monitor, and a second one read after that monitor had been released on the exception path. Another consumer in the same task could acquire in either window, so the listener could adopt bytes it never received and later hand them back, leaving two consumers holding the same bytes between them.

Both snapshots and the acquisition now run as one transaction under that monitor. It is the same monitor acquireExecutionMemory itself takes and holds for its whole duration, and it is reentrant, so taking it here only widens that window to cover the two reads. Every acquisition in the task funnels through that method, so no other consumer can take memory in between and have it adopted here. A release does not take the monitor, so a concurrent release, or a spill that frees some bytes before throwing, can still deflate the figure, which is the safe direction.

The regression test drives that interleaving deterministically. A hooked TaskMemoryManager runs another consumer immediately after the real snapshot, and the test asserts that what the listener adopted plus what the other two consumers hold equals what the task actually holds. Without the transaction it fails, adopting 2099200 bytes against a task total of 2098176.
---
 .../comet/CometArrowAllocationListener.scala  |  34 ++++--
 .../CometArrowAllocationListenerSuite.scala   | 113 +++++++++++++++++-
 2 files changed, 131 insertions(+), 16 deletions(-)

diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index bc0bb92ea48..de8e88a7404 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -65,12 +65,15 @@ import org.apache.comet.CometConf
  * the thread's flag so the cancellation is not swallowed. A failed acquisition can also leave the
  * task charged for bytes Spark never reported back; see [[acquire]].
  *
- * '''Lock order.''' [[getUsed]] and [[spill]] must stay lock-free, because Spark calls both while
- * holding the `TaskMemoryManager` monitor, and [[adjust]] holds this listener's monitor across
- * `acquireExecutionMemory`, which takes that monitor. Were the snapshot to take this monitor
- * instead, a native reservation arriving through `CometTaskMemoryManager` on a Comet Tokio thread
- * could hold Spark's monitor and wait for ours while an Arrow allocation on the same task held
- * ours and waited for Spark's.
+ * '''Lock order.''' This listener's monitor is taken before Spark's and never the other way
+ * round: [[adjust]] holds ours across `acquireExecutionMemory`, and [[acquire]] additionally
+ * takes the `TaskMemoryManager` monitor itself, so that the two usage snapshots either side of
+ * that call cannot be split by another consumer. [[getUsed]] and [[spill]] must therefore stay
+ * lock-free, because Spark calls both while holding its own monitor: were either to take ours, a
+ * native reservation arriving through `CometTaskMemoryManager` on a Comet Tokio thread could hold
+ * Spark's monitor and wait for ours while an Arrow allocation on the same task held ours and
+ * waited for Spark's. For the same reason, nothing reachable from a `spill` callback may allocate
+ * JVM Arrow memory.
  */
 private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryManager)
     extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP)
@@ -196,13 +199,20 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
    * so until then they are headroom nobody can use. They are adopted here instead, measured as
    * the change in what the pool says this task holds.
    *
-   * That measurement is an estimate, but a safe one. Another consumer in the same task cannot
-   * acquire concurrently, because `acquireExecutionMemory` holds the `TaskMemoryManager` monitor
-   * throughout, so the only interference is a concurrent release, which makes the figure too
-   * small rather than too large; and `request` bounds it from above either way. Too small
-   * degrades to what would have happened anyway.
+   * Spark reports that figure per task rather than per consumer, so it only measures '''our'''
+   * grant if nothing else in the task can move it while we are looking. Both snapshots and the
+   * acquisition therefore run as one transaction under the `TaskMemoryManager` monitor. That is
+   * the same monitor `acquireExecutionMemory` takes and holds for its whole duration, spills
+   * included, and it is reentrant, so taking it here only widens that window to cover the two
+   * reads. Every acquisition in the task funnels through that method, so with it held no other
+   * consumer can take memory between a snapshot and the call and have it adopted here.
+   *
+   * What the monitor does not cover is a release, which reaches the pool without it. Another
+   * consumer returning memory, or a spill that frees some bytes before throwing, makes the figure
+   * too small, which is the safe direction: too small degrades to what would have happened
+   * anyway. `request` bounds it from above.
    */
-  private def acquire(request: Long): Long = {
+  private def acquire(request: Long): Long = taskMemoryManager.synchronized {
     val heldBefore = taskMemoryManager.getMemoryConsumptionForThisTask
     try {
       taskMemoryManager.acquireExecutionMemory(request, this)
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index 43a885e6c94..c5bbe055b3d 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -28,7 +28,7 @@ import org.scalatest.funsuite.AnyFunSuite
 import org.apache.arrow.memory.BufferAllocator
 import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
 import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager}
 
 import org.apache.comet.CometArrowAllocator
 
@@ -325,6 +325,57 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
+  test("a concurrent acquisition is not adopted as this listener's lost grant") {
+    // The orphan is measured as a change in the task's total consumption, which Spark reports per
+    // task rather than per consumer. If another consumer could acquire between the snapshot taken
+    // before the acquisition and the acquisition itself, its bytes would be adopted here and handed
+    // back when this listener next shrank, leaving two consumers holding the same bytes between
+    // them. Both snapshots and the call therefore run as one transaction under the
+    // TaskMemoryManager monitor. This forces the interleaving that transaction exists to exclude.
+    val spare = 1024L
+    val task = newTask(pool = blockSize * 2 + spare)
+    withInstalledTask(task) {
+      val hostile =
+        new FailingSpillConsumer(task.taskMemoryManager, new IOException("spill failed"))
+      assert(hostile.take(blockSize) == blockSize)
+
+      val interloper = new PlainConsumer(task.taskMemoryManager)
+      // Once the acquisition has taken what was left, the pool is empty and Spark answers the
+      // interloper by asking the hostile consumer to spill, which throws. That is a legitimate
+      // outcome for the interloper and not what is under test here; what it ends up holding is.
+      val interloperThread = daemonThread("interloper") {
+        try interloper.take(spare)
+        catch { case _: SparkOutOfMemoryError => }
+      }
+
+      // Fires once, in place of the snapshot taken before the acquisition: exactly the window the
+      // transaction has to close. The interloper either runs to completion here, which is the bug,
+      // or blocks on the monitor the acquisition is holding, which is the fix.
+      task.snapshotHook.set(() => {
+        interloperThread.start()
+        awaitBlockedOrFinished(interloperThread)
+      })
+
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val buf = allocator.buffer(blockSize * 2)
+      try {
+        interloperThread.join(30000L)
+        assert(
+          !interloperThread.isAlive,
+          "the interloper never finished; the transaction deadlocked")
+        // Nothing is claimed twice: what the listener adopted has to fit alongside what the other
+        // two consumers hold. Without the transaction the listener adopts the interloper's bytes on
+        // top of its own grant, and this sum comes out over what the task actually holds.
+        assert(
+          listenerFor(task).reservedBytes + hostile.getUsed + interloper.getUsed ==
+            task.taskMemoryManager.getMemoryConsumptionForThisTask,
+          "the listener adopted bytes belonging to another consumer")
+      } finally {
+        buf.close()
+      }
+    }
+  }
+
   // ---------------------------------------------------------------------------------------------
   // Lock order. Spark calls getUsed and spill while holding the TaskMemoryManager monitor, and the
   // listener holds its own monitor while waiting for that one.
@@ -398,6 +449,32 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     override def spill(size: Long, trigger: MemoryConsumer): Long = throw failure
   }
 
+  /** Any other consumer in the task: takes memory, and cannot give it back. */
+  private class PlainConsumer(tmm: TaskMemoryManager)
+      extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
+    def take(bytes: Long): Long = acquireMemory(bytes)
+    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+  }
+
+  /**
+   * Runs one action immediately after the usage snapshot the listener takes before asking Spark
+   * for memory, so a test can drive what happens in the window between that snapshot and the
+   * acquisition. The real value is read first, which is what makes the window the one under test:
+   * running the action before the read would fold whatever it does into the snapshot itself.
+   */
+  private class HookedTaskMemoryManager(
+      memoryManager: MemoryManager,
+      taskAttemptId: Long,
+      hook: AtomicReference[Runnable])
+      extends TaskMemoryManager(memoryManager, taskAttemptId) {
+    override def getMemoryConsumptionForThisTask(): Long = {
+      val held = super.getMemoryConsumptionForThisTask()
+      val pending = hook.getAndSet(null)
+      if (pending != null) pending.run()
+      held
+    }
+  }
+
   /** Stands in for `CometTaskMemoryManager`: reserves from Spark directly and never spills. */
   private class NativeLikeConsumer(tmm: TaskMemoryManager)
       extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
@@ -421,11 +498,15 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   private case class TaskFixture(
       taskAttemptId: Long,
       context: TaskContextImpl,
-      taskMemoryManager: TaskMemoryManager)
+      taskMemoryManager: TaskMemoryManager,
+      snapshotHook: AtomicReference[Runnable])
 
   private def reservedFor(task: TaskFixture): Long =
     CometTaskArrowAllocator.reservedBytesForTask(task.taskAttemptId)
 
+  private def listenerFor(task: TaskFixture): CometArrowAllocationListener =
+    CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).get
+
   private def newTask(offHeap: Boolean = true, pool: Long = poolBytes): TaskFixture = {
     val conf = new SparkConf(false)
     if (offHeap) {
@@ -436,7 +517,9 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     val memoryManager = new TestMemoryManager(conf)
     memoryManager.limit(pool)
     val taskAttemptId = nextTaskAttemptId.getAndIncrement()
-    val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId)
+    val snapshotHook = new AtomicReference[Runnable]()
+    val taskMemoryManager =
+      new HookedTaskMemoryManager(memoryManager, taskAttemptId, snapshotHook)
     val context = new TaskContextImpl(
       stageId = 0,
       stageAttemptNumber = 0,
@@ -450,7 +533,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       taskMetrics = TaskMetrics.empty,
       cpus = 1,
       resources = Map.empty)
-    TaskFixture(taskAttemptId, context, taskMemoryManager)
+    TaskFixture(taskAttemptId, context, taskMemoryManager, snapshotHook)
   }
 
   /** Installs the task on this thread, restoring whatever was there before. */
@@ -491,6 +574,28 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     Option(failure.get).foreach(t => throw t)
   }
 
+  /** An unstarted daemon thread, so a test can choose the moment it runs. */
+  private def daemonThread(name: String)(body: => Unit): Thread = {
+    val thread = new Thread(() => body)
+    thread.setDaemon(true)
+    thread.setName(name)
+    thread
+  }
+
+  /**
+   * Waits until the thread is either blocked on a monitor or finished, whichever happens first,
+   * so that a test can tell the two orderings apart without depending on timing.
+   */
+  private def awaitBlockedOrFinished(thread: Thread): Unit = {
+    val deadline = System.currentTimeMillis() + 30000L
+    var state = thread.getState
+    while (state != Thread.State.BLOCKED && state != Thread.State.TERMINATED &&
+      System.currentTimeMillis() < deadline) {
+      Thread.sleep(1L)
+      state = thread.getState
+    }
+  }
+
   private def loopingThread(name: String, failure: AtomicReference[Throwable])(
       body: => Unit): Thread = {
     val thread = new Thread(() => {

From 9c9e1476037d81f6d96e8c382fea5ed5162aa426 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Fri, 18 Sep 2026 08:24:00 -0600
Subject: [PATCH 12/18] fix: allocate codegen UDF output from the unaccounted
 root

CometBatchKernelCodegenOutput.allocateOutput charged its vector to the task
allocator, but every buffer it hands out is exported to native. Its only
production caller is CometScalaUDFCodegen.evaluate, whose result CometUdfBridge
exports over the C Data Interface before closing its own reference, so the
buffers stay alive on the native side. Whichever DataFusion operator retains
that batch reserves the same buffers through Comet's unified pool, which charges
the same Spark task, so the JVM-side charge was reserving the same memory twice
and could reject an allocation that fits.

This is the rule the rest of the PR already follows: native's pool is the
authority for bytes native holds, so anything allocated to be handed straight to
native uses the listener-less root. Bring the contributor guide in line.
---
 docs/source/contributor-guide/memory_management.md     |  7 ++++---
 .../comet/codegen/CometBatchKernelCodegenOutput.scala  | 10 +++++++---
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 1ac9e40b8a3..ed963070d4f 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -103,9 +103,10 @@ batch reserves those buffers itself -- `ExternalSorter` through `get_reserved_by
 the hash join build side through `get_record_batch_memory_size`, both of which read sizes off the
 `ArrayData` and so count imported buffers -- and through a unified pool that charges the same Spark
 task. Reporting either direction on the JVM side would reserve the same memory twice and could
-reject an allocation that fits. So `NativeUtil`, the JVM UDF result in `CometUdfBridge`, and
-`CometNativeArrowSource.stream` all use the listener-less root; IPC reads, codegen output, the
-cached batch serializer and `CometNativeArrowSource.readerBatchIter` use the task allocator. The
+reject an allocation that fits. So `NativeUtil`, the JVM UDF result in `CometUdfBridge`
+together with the codegen output vector it exports, and `CometNativeArrowSource.stream` all use
+the listener-less root; IPC reads, the cached batch serializer and
+`CometNativeArrowSource.readerBatchIter` use the task allocator. The
 driver, broadcast coalescing and the cached batch serializer also fall back to the root when there
 is no task to charge, as does Comet's on-heap mode.
 
diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
index 159ef8fd1d5..e277f00e4e7 100644
--- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
+++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
@@ -26,11 +26,11 @@ import org.apache.arrow.memory.BufferAllocator
 import org.apache.arrow.vector._
 import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field}
-import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.types._
 
+import org.apache.comet.CometArrowAllocator
 import org.apache.comet.shims.CometTypeShim
 
 /**
@@ -87,8 +87,12 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
    * Closes the vector on any failure so a partially-initialized tree doesn't leak buffers.
    */
   def allocateOutput(field: Field, numRows: Int, estimatedBytes: Int): FieldVector = {
-    // JVM-owned codegen output, accounted to the task that is running the kernel.
-    val allocator = CometTaskArrowAllocator.forCurrentTask()
+    // The unaccounted root, because this vector is allocated to be handed to native. Its only
+    // caller is `CometScalaUDFCodegen.evaluate`, whose result `CometUdfBridge` exports over the C
+    // Data Interface before closing its own reference. Whichever DataFusion operator retains the
+    // batch reserves these same buffers through Comet's unified pool, which charges the same Spark
+    // task, so reporting them here as well would reserve the same memory twice.
+    val allocator: BufferAllocator = CometArrowAllocator
     val vec: FieldVector = field.getType match {
       case _: ArrowType.List | _: ArrowType.LargeList | _: ArrowType.FixedSizeList =>
         val v = new RenamedListVector(field, allocator)

From fc4e6e5ddf20f78a154e20ba0ea8b7a0a475f911 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sat, 26 Sep 2026 01:49:44 -0600
Subject: [PATCH 13/18] fix: size the Arrow reservation from what the task
 allocator owns

The listener kept its own running total, adding on onAllocation and
subtracting on onRelease. That total drifts from what the allocator owns
whenever a buffer's ownership moves between allocators, because Arrow
moves the charge with transferBalance and calls no listener: when two
allocators hold references to one buffer and the owner releases first,
the other becomes the owner, and the final onRelease goes to it.

ColumnarBatchArrowReader does that to every batch it streams to native.
It retains the source batch's buffers into the stream's allocator, a
child of the listener-less root, and closes the source. So any batch
read into a task allocator and then handed to a native operator through
wrapColumnarBatchRDD, which covers broadcast reads and cached-table
scans feeding native, was charged to the task on allocation and never
released, and the task's reservation grew by every such batch until the
task ended. In a broadcast hash join into native over a 1M-row build
side, each task finished holding 16 MiB for a broadcast it had already
handed to native, which the native join reserves as well. The reverse
case, a buffer adopted into a task allocator from elsewhere, drove the
total below zero and hid later allocations.

The listener now treats a callback as a prompt only, and sizes the
reservation from the task allocator's own accountant, which Arrow keeps
correct across transfers and which is an atomic, so getUsed stays
lock-free. CometTaskArrowAllocator binds each listener to its allocator
once the allocator exists.

A transfer is still invisible until something prompts the listener, and
the last batch a task streams has nothing after it, which for a
broadcast coalesced into one batch is the whole broadcast. So
ColumnarBatchArrowReader, having closed a source, calls the new
CometTaskArrowAllocator.reconcile for each allocator the source's
columns came from, which re-syncs a task allocator's listener and does
nothing for any other allocator. The same join now finishes each task
with nothing reserved. Any other transfer is picked up at the next
callback on the task's allocator, and task completion returns whatever
is left.

This also removes most of what the branch recorded as double counting:
a batch the JVM reads and then streams to native stops being charged on
the JVM side as soon as the stream has it.

Three tests cover it: an ownership move out of the task allocator, the
real CometArrowStream.fromColumnarBatchIter export over 16 batches
drained by a JVM importer standing in for native, asserting the task
holds nothing once each batch is handed over, and a buffer adopted into
the task allocator. Against the previous listener all three fail.
Without the reader's reconcile the second still fails, with a block held
after the first batch is handed over. The arithmetic tests drive the
listener through a stand-in for the accountant, and a further test pins
that CometArrowAllocator and CometArrowImportAllocator both stay
unaccounted.

The scaladocs of both classes are corrected along the way. The example
of a release on a native thread was a shuffle-read batch, but shuffle
reads decode natively and import through CometArrowImportAllocator.
---
 .../comet/CometArrowAllocationListener.scala  |  83 ++++----
 .../spark/comet/CometTaskArrowAllocator.scala |  56 ++---
 .../arrow/ColumnarBatchArrowReader.scala      |  11 +
 .../arrow/CometNativeArrowSource.scala        |   8 +-
 .../CometArrowAllocationListenerSuite.scala   | 192 ++++++++++++++++--
 5 files changed, 268 insertions(+), 82 deletions(-)

diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index de8e88a7404..c5cfbe5cdb0 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -19,7 +19,8 @@
 
 package org.apache.spark.comet
 
-import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.function.LongSupplier
 
 import scala.util.control.NonFatal
 
@@ -40,13 +41,27 @@ import org.apache.comet.CometConf
  * Spark's other off-heap consumers.
  *
  * '''Ownership.''' One instance is created per task and attached to that task's Arrow allocator
- * by [[CometTaskArrowAllocator]]. Arrow reports an allocation and its matching release to the
- * listener of the allocator that '''owns''' the buffer, on whichever thread happens to drop the
- * last reference, and `AllocationListener` is handed nothing but a size. Binding the listener to
- * an allocator is therefore the only way to attribute a release, and reading `TaskContext` inside
- * the callbacks would get it wrong: a shuffle-read batch handed on to a native operator is pinned
- * by native and dropped later from a Tokio worker with no task context installed. That release
- * would be lost, leaving the task charged for memory it had already freed, batch after batch.
+ * by [[CometTaskArrowAllocator]]. Arrow calls `onAllocation` on the listener of the allocator
+ * that made a buffer, and `onRelease` on the listener of the allocator that '''owns''' it once
+ * the last reference drops, on whichever thread that happens to be. `AllocationListener` is
+ * handed nothing but a size, so binding the listener to an allocator is the only way to attribute
+ * a release. Reading `TaskContext` inside the callbacks would get it wrong: a buffer exported
+ * over the C Data Interface is released by whichever native thread drops it last, and that thread
+ * has no task context installed.
+ *
+ * '''What is charged is what the allocator owns, not a running total of the callbacks.''' The two
+ * come apart because ownership moves between allocators without either listener hearing about it.
+ * When two allocators hold references to one buffer and the owner lets go first, Arrow hands the
+ * charge to the other through `transferBalance`, which calls no listener, and the final
+ * `onRelease` goes to the new owner. That is the everyday path rather than a corner:
+ * `ColumnarBatchArrowReader` retains every batch it streams to native into its own allocator, a
+ * child of the unaccounted root, and then closes the source. Summing the callbacks would charge
+ * the task for each such batch until the task ended. So the callbacks only say when to look, and
+ * the size comes from the task allocator's accountant, which Arrow keeps right across transfers.
+ * `ColumnarBatchArrowReader` calls [[reconcile]] as soon as it closes a source, so a task stops
+ * paying for a batch once native has it. Any other transfer is reflected at the next allocation
+ * or release on the task's allocator, and whatever is still reserved when the task ends goes back
+ * to Spark then.
  *
  * '''Reporting only.''' A short grant is logged and the allocation proceeds, because Arrow
  * allocation on these paths cannot fail today and making it fail is a behavioural change that
@@ -82,11 +97,12 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
   import CometArrowAllocationListener._
 
   /**
-   * Bytes Arrow currently holds on this task's behalf. An atomic rather than a guarded field so
-   * that [[getUsed]] can read it without taking this listener's monitor; see the lock order note
-   * above.
+   * Bytes the task's allocator currently owns, read from Arrow's own accountant, which is an
+   * atomic, so [[getUsed]] can read it without taking this listener's monitor; see the lock order
+   * note above. Set by [[bind]], because the allocator is built with this listener and so cannot
+   * exist before it.
    */
-  private val live = new AtomicLong(0L)
+  @volatile private var owned: LongSupplier = NothingOwned
 
   /** Bytes currently reserved with Spark. Guarded by this listener's monitor. */
   private var reserved = 0L
@@ -94,26 +110,26 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
   /** Set once the owning task has finished. Volatile so [[getUsed]] can read it lock-free. */
   @volatile private var completed = false
 
-  override def onAllocation(size: Long): Unit = {
-    live.addAndGet(size)
-    adjustQuietly()
-  }
+  /** Attaches the listener to what it accounts for. Called once, before the allocator is used. */
+  private[comet] def bind(owned: LongSupplier): Unit = this.owned = owned
 
-  override def onRelease(size: Long): Unit = {
-    live.addAndGet(-size)
-    adjustQuietly()
-  }
+  override def onAllocation(size: Long): Unit = adjustQuietly()
+
+  override def onRelease(size: Long): Unit = adjustQuietly()
+
+  /** Brings the reservation up to date after buffers moved out without a callback. */
+  private[comet] def reconcile(): Unit = adjustQuietly()
 
   /**
-   * Reports our own tally. Spark reads this for spill-victim ordering, `showMemoryUsage` and
-   * end-of-task leak reporting. The inherited `used` counter stays at zero because this consumer
-   * never calls `acquireMemory` or `allocatePage`; Arrow has already obtained the memory and we
-   * are only accounting for it.
+   * Reports what the task's allocator owns. Spark reads this for spill-victim ordering,
+   * `showMemoryUsage` and end-of-task leak reporting. The inherited `used` counter stays at zero
+   * because this consumer never calls `acquireMemory` or `allocatePage`; Arrow has already
+   * obtained the memory and we are only accounting for it.
    *
    * Reports zero once the task has finished, so that buffers deliberately allowed to outlive
    * their task are not reported by `cleanUpAllAllocatedMemory` as a Spark memory leak.
    */
-  override def getUsed: Long = if (completed) 0L else math.max(0L, live.get())
+  override def getUsed: Long = if (completed) 0L else owned.getAsLong
 
   /** Comet's native operators cannot be made to spill from here. See issue #5997. */
   override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
@@ -141,9 +157,6 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
     }
   }
 
-  /** Bytes Arrow currently holds on this task's behalf. Visible for testing. */
-  private[comet] def liveBytes: Long = live.get()
-
   /** Bytes currently reserved with Spark on this task's behalf. Visible for testing. */
   private[comet] def reservedBytes: Long = synchronized(reserved)
 
@@ -165,12 +178,12 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
 
   private def adjust(): Unit = synchronized {
     if (!completed) {
-      val liveBytes = math.max(0L, live.get())
-      if (reserved < liveBytes) {
+      val ownedBytes = owned.getAsLong
+      if (reserved < ownedBytes) {
         // Round up so `reserved` stays a block multiple and growth always leaves headroom.
-        // Requesting the bare deficit would land exactly on `liveBytes` for any buffer at or above
-        // the block size, sending the very next allocation straight back into Spark's lock.
-        val request = roundUpToBlock(liveBytes - reserved)
+        // Requesting the bare deficit would land exactly on `ownedBytes` for any buffer at or
+        // above the block size, sending the very next allocation straight back into Spark's lock.
+        val request = roundUpToBlock(ownedBytes - reserved)
         val granted = acquire(request)
         reserved += granted
         if (granted < request) {
@@ -179,7 +192,7 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
       } else {
         // Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes on
         // the executor-wide pool, so a per-block loop would take that lock once per megabyte freed.
-        val excess = ((reserved - liveBytes) / BLOCK_SIZE) * BLOCK_SIZE
+        val excess = ((reserved - ownedBytes) / BLOCK_SIZE) * BLOCK_SIZE
         if (excess > 0L) {
           taskMemoryManager.releaseExecutionMemory(excess, this)
           reserved -= excess
@@ -258,6 +271,8 @@ object CometArrowAllocationListener extends Logging {
    */
   private[comet] val BLOCK_SIZE = 1024L * 1024L
 
+  private val NothingOwned: LongSupplier = () => 0L
+
   private val shortGrantLogged = new AtomicBoolean(false)
   private val memoryManagerFailureLogged = new AtomicBoolean(false)
 
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
index fc5074e5a11..4d120bc6119 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
@@ -34,28 +34,26 @@ import org.apache.comet.{CometArrowAllocator, CometConf}
  * Hands out the Arrow allocator that JVM-owned allocations should use, one per Spark task.
  *
  * Each task gets a child of `CometArrowAllocator` carrying its own
- * [[CometArrowAllocationListener]], so the bytes it hands out are reported to that task's
+ * [[CometArrowAllocationListener]], so what that allocator owns is reported to the task's
  * `TaskMemoryManager`. The allocator, not the calling thread, is what identifies the owner:
  * Arrow's `AllocationListener` is given only a size, and a buffer is released on whichever thread
  * drops the last reference, which for anything that reaches native over the C Data Interface is a
  * Comet Tokio worker with no task context installed. Child allocators cut from the returned
- * allocator inherit its listener, so the paths that make their own children are covered too.
+ * allocator inherit its listener and roll their bytes up into it, so the paths that make their
+ * own children are covered too.
  *
- * '''This is for buffers the JVM owns.''' Anything allocated to be handed straight to native --
- * `NativeUtil`, the JVM UDF result, `CometNativeArrowSource.stream` -- uses the unaccounted root
- * instead, and so does anything imported from native. Native's pool is the authority for bytes
- * native holds: whichever DataFusion operator retains the batch reserves those buffers through
- * Comet's unified pool, which charges the same Spark task, so reporting them here as well would
- * reserve the same memory twice and could reject an allocation that fits. The same fallback
- * covers callers with no task to charge -- the driver, broadcast coalescing, the cached batch
- * serializer
- * -- and Comet's on-heap mode, where charging an off-heap consumer would be wrong.
- *
- * What this does not fix is a buffer that is used in the JVM and only later handed to native, a
- * shuffle-read batch feeding a native operator being the common shape. Its allocation site cannot
- * know, so it stays charged here while native may also reserve it. Coordinating reservation
- * ownership across the FFI boundary needs changes on both sides; see
- * [[https://github.com/apache/datafusion-comet/issues/5997]].
+ * '''This is for buffers the JVM owns.''' Buffers allocated to be handed straight to native -- by
+ * `NativeUtil`, for the JVM UDF result, and in `CometArrowStream.stream` -- come from the
+ * listener-less root instead, and imports from native come from `CometArrowImportAllocator`, a
+ * listener-less child of it. Native memory is native's to account for: a native operator that
+ * retains a batch reserves its buffers through Comet's pool, which charges the same Spark task,
+ * so charging them here as well would reserve the same memory twice. A batch the JVM read into
+ * this allocator and then streams to native is covered by the same rule without its allocation
+ * site having to know: the stream's reader retains each batch into its own root-backed allocator
+ * and closes the source, which moves the charge off this one, and then calls [[reconcile]] so
+ * that the reservation follows straight away. The root is also what callers with no task get,
+ * such as broadcast coalescing on the driver or a reader built while native pulls a stream, and
+ * what Comet's on-heap mode gets, where charging an off-heap consumer would be wrong.
  *
  * '''Lifetime.''' The task allocator cannot simply be closed when the task ends. The process-wide
  * allocator exists precisely because Arrow buffers can outlive the task that created them, and
@@ -93,8 +91,8 @@ object CometTaskArrowAllocator extends Logging {
    * an allocator belonging to a task other than the current one.
    */
   def forCurrentTask(): BufferAllocator = {
-    // Cheapest check first, and the one that eliminates the most callers: the driver, broadcast
-    // coalescing and the cached batch serializer all allocate with no task in scope.
+    // Cheapest check first, and the one that eliminates the most callers: the driver, and native
+    // threads pulling a stream, allocate with no task in scope.
     val taskContext = TaskContext.get()
     if (taskContext == null) {
       CometArrowAllocator
@@ -106,10 +104,9 @@ object CometTaskArrowAllocator extends Logging {
 
   private def create(taskContext: TaskContext): BufferAllocator = {
     val taskMemoryManager = taskContext.taskMemoryManager()
-    // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory
-    // configured, and charging an off-heap consumer there would be wrong. Note this differs from
-    // CometUnifiedShuffleMemoryAllocator, which throws in that situation; throwing here would
-    // break those tests.
+    // Comet's on-heap mode accounts for nothing and exists so the Spark SQL suite can run without
+    // off-heap memory configured, so it gets the root: the same switch that
+    // `CometShuffleMemoryAllocator.getInstance` makes for shuffle pages.
     if (!accountingEnabled || taskMemoryManager == null ||
       taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) {
       return CometArrowAllocator
@@ -120,6 +117,7 @@ object CometTaskArrowAllocator extends Logging {
     val listener = new CometArrowAllocationListener(taskMemoryManager)
     val allocator = CometArrowAllocator
       .newChildAllocator(s"comet-task-$taskAttemptId", listener, 0L, Long.MaxValue)
+    listener.bind(() => allocator.getAllocatedMemory)
     val created = new TaskAllocator(allocator, listener)
     val previous = perTask.putIfAbsent(taskAttemptId, created)
     if (previous != null) {
@@ -185,6 +183,18 @@ object CometTaskArrowAllocator extends Logging {
     }
   }
 
+  /**
+   * Brings a task allocator's reservation up to date with what it owns, if `allocator` is a task
+   * allocator or a child of one, and does nothing otherwise. For a caller that has just moved
+   * buffers out of it by retaining them in another allocator and closing the source, which Arrow
+   * does without calling any listener, so that the task stops paying for them now rather than at
+   * its next allocation.
+   */
+  private[spark] def reconcile(allocator: BufferAllocator): Unit = allocator.getListener match {
+    case listener: CometArrowAllocationListener => listener.reconcile()
+    case _ =>
+  }
+
   /** Number of tasks currently holding an accounted allocator. Visible for testing. */
   private[comet] def trackedTaskCount: Int = perTask.size()
 
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala
index 4a15355d269..827657bf7f3 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala
@@ -25,6 +25,7 @@ import org.apache.arrow.memory.BufferAllocator
 import org.apache.arrow.vector.{VectorSchemaRoot, VectorUnloader}
 import org.apache.arrow.vector.ipc.ArrowReader
 import org.apache.arrow.vector.types.pojo.Schema
+import org.apache.spark.comet.CometTaskArrowAllocator
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
 import org.apache.comet.vector.{CometVector, CometVectorUtils}
@@ -69,7 +70,17 @@ private[comet] class ColumnarBatchArrowReader(
         // Do not close this borrowed root: src and the helper own its vectors.
       }
     } finally {
+      // Closing the source leaves this reader's allocator owning the buffers it retained, which
+      // Arrow does without telling the listener of the allocator they came from. A task allocator
+      // is told here instead, so the task stops paying for the batch once native has it.
+      val sourceAllocators = (0 until src.numCols())
+        .map(src.column)
+        .collect { case v: CometVector => v.getValueVector.getAllocator }
+        // A NullVector owns no buffers and has no allocator.
+        .filter(_ != null)
+        .distinct
       src.close()
+      sourceAllocators.foreach(CometTaskArrowAllocator.reconcile)
     }
     true
   }
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
index d5bc9c41a33..6ecd54f8f87 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
@@ -231,10 +231,10 @@ object CometArrowStream extends Logging {
       name: String,
       readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = {
     val context = TaskContext.get()
-    // A child of the unaccounted root, not of the task allocator: every byte allocated here is
-    // exported to native and exists only for native to consume. Whichever native operator retains
-    // the batch reserves those same buffers in Comet's pool, which charges the same Spark task, so
-    // reporting them on this side too would reserve the same memory twice.
+    // A child of the listener-less root, not of the task allocator: everything loaded here is
+    // exported to native, which accounts for what it retains. A reader that loads a batch from the
+    // task allocator and then closes the source moves the batch's charge here, so the task stops
+    // paying for it as well. See CometTaskArrowAllocator.
     val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue)
     var reader: ArrowReader = null
     var arrowStream: ArrowArrayStream = null
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index c5bbe055b3d..61b292ef994 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -25,18 +25,25 @@ import java.util.concurrent.atomic.{AtomicLong, AtomicReference}
 
 import org.scalatest.funsuite.AnyFunSuite
 
+import org.apache.arrow.c.Data
 import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.{FieldVector, IntVector, VectorSchemaRoot}
 import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.CometArrowAllocator
+import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator}
+import org.apache.comet.vector.NativeUtil
 
 /**
  * Tests that JVM Arrow allocations are reported to Spark, that they are reported against the task
- * that made them rather than whichever task happens to be on the releasing thread, and, just as
- * importantly, that the paths where they cannot be reported fail quietly rather than throwing.
- * Arrow allocation on these paths cannot fail today and this listener must not change that.
+ * that made them rather than whichever task happens to be on the releasing thread, that what is
+ * charged follows Arrow's ownership as buffers move between allocators, and, just as importantly,
+ * that the paths where they cannot be reported fail quietly rather than throwing. Arrow
+ * allocation on these paths cannot fail today and this listener must not change that.
  */
 class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
@@ -47,20 +54,22 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   private val nextTaskAttemptId = new AtomicLong(1000L)
 
   // ---------------------------------------------------------------------------------------------
-  // Reservation arithmetic. Driven through the listener directly, since Arrow's rounding policy
-  // would otherwise decide the sizes under test.
+  // Reservation arithmetic. Driven through the listener directly, with a stand-in for the
+  // allocator's accountant, since Arrow's rounding policy would otherwise decide the sizes under
+  // test.
   // ---------------------------------------------------------------------------------------------
 
   test("allocations are charged to the current task in whole blocks") {
     withTask() { task =>
       val listener = new CometArrowAllocationListener(task.taskMemoryManager)
+      val owner = new StubOwner(listener)
 
       // Far smaller than a block, so the reservation should round up to exactly one block.
-      listener.onAllocation(128L)
+      owner.allocate(128L)
       assert(listener.reservedBytes == blockSize)
 
       // Still inside the first block, so Spark is not asked again.
-      listener.onAllocation(1024L)
+      owner.allocate(1024L)
       assert(listener.reservedBytes == blockSize)
 
       listener.taskCompleted()
@@ -70,7 +79,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   test("a request larger than a block rounds up to a block multiple") {
     withTask() { task =>
       val listener = new CometArrowAllocationListener(task.taskMemoryManager)
-      listener.onAllocation(blockSize * 3 + 7L)
+      new StubOwner(listener).allocate(blockSize * 3 + 7L)
       // Rounded up rather than sized to the exact deficit, so growth leaves headroom and the next
       // small allocation does not go straight back into Spark.
       assert(listener.reservedBytes == blockSize * 4)
@@ -81,10 +90,11 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   test("releasing returns whole blocks to Spark") {
     withTask() { task =>
       val listener = new CometArrowAllocationListener(task.taskMemoryManager)
-      listener.onAllocation(blockSize * 2)
+      val owner = new StubOwner(listener)
+      owner.allocate(blockSize * 2)
       assert(listener.reservedBytes == blockSize * 2)
 
-      listener.onRelease(blockSize * 2)
+      owner.release(blockSize * 2)
       assert(listener.reservedBytes == 0L)
       listener.taskCompleted()
     }
@@ -129,25 +139,28 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
-  test("the process-wide root is not accounted, which is what the FFI paths rely on") {
+  test("the allocators on either side of the FFI boundary are not accounted") {
     withTask() { task =>
       // Establish the task allocator first, so this asserts "not charged" rather than "no task".
       CometTaskArrowAllocator.forCurrentTask()
-      // Both directions across the C Data Interface use the listener-less root: imported buffers
-      // wrap memory the native side owns, and buffers allocated for export are reserved again by
-      // whichever native operator retains the batch, through a pool that charges the same task.
-      val buf = CometArrowAllocator.buffer(blockSize)
-      try {
-        assert(reservedFor(task) == 0L)
-      } finally {
-        buf.close()
+      // Buffers allocated for export come from the listener-less root, and imports from a
+      // listener-less child of it, because native accounts for what it retains. Arrow reports an
+      // imported buffer to the importing allocator's listener at full capacity, so giving the
+      // import allocator a listener would charge Spark for native memory.
+      for (allocator <- Seq(CometArrowAllocator, CometArrowImportAllocator)) {
+        val buf = allocator.buffer(blockSize)
+        try {
+          assert(reservedFor(task) == 0L, s"${allocator.getName} was charged to the task")
+        } finally {
+          buf.close()
+        }
       }
     }
   }
 
   test("no active task uses the unaccounted root allocator") {
     TaskContext.unset()
-    // Broadcast coalescing and the cached batch serializer can allocate off a task thread.
+    // Broadcast coalescing on the driver, and native threads pulling a stream, have no task.
     assert(CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator)
   }
 
@@ -197,6 +210,125 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
+  // ---------------------------------------------------------------------------------------------
+  // What is charged follows ownership. Arrow moves a buffer's charge between allocators with
+  // `transferBalance`, which calls no listener, so a tally of the callbacks drifts from what the
+  // task's allocator actually owns.
+  // ---------------------------------------------------------------------------------------------
+
+  test("a buffer whose ownership moves to another allocator stops being charged to the task") {
+    withTask() { task =>
+      val stream = CometArrowAllocator.newChildAllocator("stream", 0L, Long.MaxValue)
+      try {
+        val allocator = CometTaskArrowAllocator.forCurrentTask()
+        val source = allocator.buffer(blockSize)
+        assert(reservedFor(task) == blockSize)
+
+        // What ColumnarBatchArrowReader does to every batch it streams to native: retain the
+        // buffers into the stream's allocator, then close the source. Arrow makes the stream's
+        // allocator the owner, and neither listener is told.
+        val retained = source.getReferenceManager.retain(source, stream)
+        try {
+          source.close()
+          assert(allocator.getAllocatedMemory == 0L)
+          assert(listenerFor(task).getUsed == 0L)
+        } finally {
+          // Native drops it later, and the release goes to the stream allocator's listener, which
+          // is the root's no-op. Nothing on the task's allocator hears about it.
+          retained.close()
+        }
+
+        // The reservation catches up at the next callback on the task's allocator: one block for
+        // the new buffer, not one on top of a block for a buffer the task no longer owns.
+        val next = allocator.buffer(128L)
+        assert(reservedFor(task) == blockSize)
+        next.close()
+        assert(reservedFor(task) == 0L)
+      } finally {
+        stream.close()
+      }
+    }
+  }
+
+  test("batches streamed to native stop being charged to the task that read them") {
+    withTask() { task =>
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val importer = CometArrowAllocator.newChildAllocator("importer", 0L, Long.MaxValue)
+      val numBatches = 16
+      // Half a block of values per batch, which Arrow allocates as one block-sized buffer.
+      val rowsPerBatch = (blockSize / 2 / IntVector.TYPE_WIDTH).toInt
+      val batches = Iterator.tabulate(numBatches) { b =>
+        val vector = new IntVector("i", allocator)
+        vector.allocateNew(rowsPerBatch)
+        (0 until rowsPerBatch).foreach(i => vector.set(i, b + i))
+        vector.setValueCount(rowsPerBatch)
+        val root = new VectorSchemaRoot(java.util.Arrays.asList[FieldVector](vector))
+        root.setRowCount(rowsPerBatch)
+        NativeUtil.rootAsBatch(root): ColumnarBatch
+      }
+      try {
+        // The same export a native operator's JVM input takes, drained here by a JVM importer
+        // standing in for native's ScanExec.
+        val stream = CometArrowStream.fromColumnarBatchIter(
+          batches,
+          StructType(Seq(StructField("i", IntegerType))),
+          CometArrowStream.NATIVE_TIMEZONE,
+          "listener-suite")
+        val reader = Data.importArrayStream(importer, stream)
+        var imported = 0
+        try {
+          while (reader.loadNextBatch()) {
+            imported += 1
+            // Native holds this batch now, and the task no longer owns any of it, so it should not
+            // be paying for it either. Summing the callbacks would have kept a block for every
+            // batch read so far, and waiting for the task's next allocation would have kept one
+            // for this batch, which for a broadcast coalesced into one batch is all of it.
+            assert(allocator.getAllocatedMemory == 0L)
+            assert(
+              reservedFor(task) == 0L,
+              s"reserved ${reservedFor(task)} bytes after handing over $imported batches")
+          }
+        } finally {
+          reader.close()
+        }
+        assert(imported == numBatches)
+      } finally {
+        importer.close()
+      }
+    }
+  }
+
+  test("a buffer adopted from another allocator is charged while the task owns it") {
+    withTask() { task =>
+      val other = CometArrowAllocator.newChildAllocator("other", 0L, Long.MaxValue)
+      try {
+        val allocator = CometTaskArrowAllocator.forCurrentTask()
+        val foreign = other.buffer(blockSize)
+        // The mirror image: the task allocator takes a reference and the original owner lets go,
+        // so Arrow makes the task allocator the owner without telling either listener.
+        val adopted = foreign.getReferenceManager.retain(foreign, allocator)
+        try {
+          foreign.close()
+          assert(allocator.getAllocatedMemory == blockSize)
+          assert(listenerFor(task).getUsed == blockSize)
+        } finally {
+          // The final release then reaches this task's listener for bytes it never saw allocated.
+          adopted.close()
+        }
+        // That must not leave the task's figure short, so the next allocation is charged in full.
+        val buf = allocator.buffer(blockSize)
+        try {
+          assert(reservedFor(task) == blockSize)
+        } finally {
+          buf.close()
+        }
+        assert(reservedFor(task) == 0L)
+      } finally {
+        other.close()
+      }
+    }
+  }
+
   // ---------------------------------------------------------------------------------------------
   // Task completion, and buffers that outlive their task.
   // ---------------------------------------------------------------------------------------------
@@ -442,6 +574,24 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   // Fixtures.
   // ---------------------------------------------------------------------------------------------
 
+  /**
+   * Stands in for the task allocator's accountant, and calls the listener the way Arrow would.
+   */
+  private class StubOwner(listener: CometArrowAllocationListener) {
+    private val owned = new AtomicLong(0L)
+    listener.bind(() => owned.get)
+
+    def allocate(bytes: Long): Unit = {
+      owned.addAndGet(bytes)
+      listener.onAllocation(bytes)
+    }
+
+    def release(bytes: Long): Unit = {
+      owned.addAndGet(-bytes)
+      listener.onRelease(bytes)
+    }
+  }
+
   /** Holds memory and refuses to give it back, so `trySpillAndAcquire` throws on its behalf. */
   private class FailingSpillConsumer(tmm: TaskMemoryManager, failure: Exception)
       extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {

From e1da8dec683bfcb4d64dd3e23ead5c77bd23b40a Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sat, 26 Sep 2026 01:49:44 -0600
Subject: [PATCH 14/18] docs: describe per-task Arrow accounting in the guides,
 skills and config

The tuning guide listed JVM Arrow buffers among the memory no budget
covers, and counted them against spark.executor.memoryOverhead. Memory a
task's Arrow allocator owns is now charged to spark.memory.offHeap.size,
so the list keeps only buffers passed between the JVM and native, and a
paragraph says what the charge does to spilling and how to turn it off,
the way the fair_unified change in #6205 was described.

In the memory management guide, the inventory, the allocation-site
diagram and the cgroup diagram showed every JVM Arrow site going to an
allocator nobody accounts for. They now show which sites use the task
allocator. The paragraphs on what is charged describe the accountant
and the reader's reconcile rather than a sum of callbacks, and why
ownership transfers make the difference. The claim that a JVM-read
batch handed to native stays charged on both sides is replaced by what
actually remains, a per-column export through
NativeColumnarToRowConverter, which native does not reserve. Its
shuffle-read example was wrong too: shuffle reads decode natively. The
open-problem entry that described that double charge is removed. The
constraints on a Comet memory consumer now name both consumers, and say
how the listener adopts a stranded grant.

The memory and FFI review skills described CometArrowAllocator as seen
by no budget, which is no longer true of a task's allocations. They now
say where each kind of JVM Arrow memory is charged, and what a reviewer
should check in the listener.

The config description is rewritten in the tuning guide's terms, and now
says it is an executor setting that has to be set at submit time, as
spark.comet.memory.logInterval's does, because it is read once per JVM
from the SparkConf.
---
 .ai/skills/review-comet-ffi-pr/SKILL.md       |  10 +-
 .ai/skills/review-comet-memory-pr/SKILL.md    |  19 +-
 .../contributor-guide/memory_management.md    | 181 ++++++++++--------
 docs/source/user-guide/latest/tuning.md       |  12 +-
 .../scala/org/apache/comet/CometConf.scala    |  16 +-
 5 files changed, 141 insertions(+), 97 deletions(-)

diff --git a/.ai/skills/review-comet-ffi-pr/SKILL.md b/.ai/skills/review-comet-ffi-pr/SKILL.md
index f57ef627367..62f718af556 100644
--- a/.ai/skills/review-comet-ffi-pr/SKILL.md
+++ b/.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, 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.
 
diff --git a/.ai/skills/review-comet-memory-pr/SKILL.md b/.ai/skills/review-comet-memory-pr/SKILL.md
index fbd4610b20a..d99871c78d9 100644
--- a/.ai/skills/review-comet-memory-pr/SKILL.md
+++ b/.ai/skills/review-comet-memory-pr/SKILL.md
@@ -158,15 +158,24 @@ 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`. 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. Check that the listener
+  stays without a lock in `getUsed` and `spill`, 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
diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 13428cc0a0b..dbbe7d7efcc 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -76,14 +76,14 @@ Enabling Comet does not add one new memory consumer, it adds several, and they a
 accounted by the same party. This inventory is worth internalizing before reading the rest of the
 page:
 
-| Allocator                                 | Lives in    | Bounded by                                                    | Visible to Spark? |
-| ----------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
-| Spark execution + storage (on-heap)       | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
-| Spark Tungsten (off-heap)                 | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
-| Comet native (Rust global allocator)      | Native heap | `memory_limit` (see below), enforced only via the memory pool | Reservations only |
-| Comet JVM Arrow that stays in the JVM     | Off-heap    | **Nothing**, but reported to `TaskMemoryManager`              | Yes               |
-| Comet JVM Arrow crossing the FFI boundary | Off-heap    | **Nothing**: a `RootAllocator(Long.MaxValue)`                 | No                |
-| Comet JVM shuffle pages                   | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+| Allocator                                | Lives in    | Bounded by                                                    | Visible to Spark? |
+| ---------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- |
+| Spark execution + storage (on-heap)      | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
+| Spark Tungsten (off-heap)                | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
+| Comet native (Rust global allocator)     | Native heap | `memory_limit` (see below), enforced only via the memory pool | Reservations only |
+| Comet JVM Arrow owned by a task          | Off-heap    | **Nothing**, but charged to `TaskMemoryManager`               | Yes               |
+| Comet JVM Arrow passed to or from native | Off-heap    | **Nothing**: a `RootAllocator(Long.MaxValue)`                 | No                |
+| Comet JVM shuffle pages                  | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
 
 Several observations follow.
 
@@ -108,52 +108,64 @@ pool over JNI. A native allocation therefore consumes the same accounting budget
 while occupying entirely different memory. Raising `spark.memory.offHeap.size` raises the ceiling
 for both at once, and raises the pod's memory request by the same amount.
 
-**Comet's JVM-side Arrow allocations inside a task are reported to Spark, but still unbounded.**
+**Comet's JVM-side Arrow memory is charged to the task that owns it, but it is still unbounded.**
 `CometArrowAllocator` (`spark/src/main/scala/org/apache/comet/package.scala`) is a single
-process-wide `new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream
-export (`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. JVM-owned
-allocations do not use the root directly: `CometTaskArrowAllocator.forCurrentTask()` hands each
-task a child carrying its own `CometArrowAllocationListener`, which charges the bytes to a
-`MemoryConsumer` for that task in whole blocks, so they appear in `showMemoryUsage` and are
-arbitrated against Spark's other off-heap consumers. The limit is still `Long.MaxValue`: the
-listener reports without enforcing, so an allocation here cannot fail, though the bytes do consume
-the off-heap pool and other consumers see correspondingly less headroom. Set
-`spark.comet.arrowAllocator.accounting.enabled=false` to stop reporting them.
+process-wide `new RootAllocator(Long.MaxValue)` with no allocation listener. JVM-owned allocations
+do not use it directly: `CometTaskArrowAllocator.forCurrentTask()` hands each task a child of it
+carrying its own `CometArrowAllocationListener`, a Spark `MemoryConsumer` for that task, which
+reserves what the child owns from the off-heap execution pool in whole 1 MiB blocks. Those bytes
+therefore appear in `showMemoryUsage` and are arbitrated against Spark's other off-heap consumers,
+Comet's native pool among them. The limit is still `Long.MaxValue`, because the listener reports
+without enforcing, so an allocation here cannot fail. Set
+`spark.comet.arrowAllocator.accounting.enabled=false` to stop charging them, in which case
+`forCurrentTask()` returns the root. It also returns the root in on-heap mode, and off a task, as
+for broadcast coalescing on the driver or for a reader built on a native thread that is pulling a
+stream.
 
 **The owner is the allocator, not the calling thread.** Arrow hands `AllocationListener` nothing
 but a size, and it reports a release on whichever thread drops the last reference, which for
 anything exported over the C Data Interface is a Tokio worker with no task context installed.
 Attributing by `TaskContext` would therefore drop those releases and leave a task charged for
 memory it had already freed. Binding one listener to one task's allocator is what makes the release
-land on the task that allocated. Children cut from a task allocator inherit its listener, so the
-paths that make their own children are covered without knowing about any of this.
-
-**Anything crossing the FFI boundary uses a listener-less allocator, and is not accounted.** Both
-directions, for the same reason: native's pool is the authority for bytes native holds. Coming in,
-imports go through `CometArrowImportAllocator`, a listener-less child of the root, because Arrow's
-`wrapForeignAllocation` reports an imported buffer to the allocator's listener at full capacity
-even though no JVM-side allocation happened. Going out, whichever DataFusion operator retains the
-batch reserves those buffers itself -- `ExternalSorter` through `get_reserved_bytes_for_record_batch`,
-the hash join build side through `get_record_batch_memory_size`, both of which read sizes off the
-`ArrayData` and so count imported buffers -- and through a unified pool that charges the same Spark
-task. Reporting either direction on the JVM side would reserve the same memory twice and could
-reject an allocation that fits. So `NativeUtil`, the JVM UDF result in `CometUdfBridge`
-together with the codegen output vector it exports, and `CometNativeArrowSource.stream` all use
-the listener-less root; IPC reads, the cached batch serializer and
-`CometNativeArrowSource.readerBatchIter` use the task allocator. The
-driver, broadcast coalescing and the cached batch serializer also fall back to the root when there
-is no task to charge, as does Comet's on-heap mode.
-
-**The split is by allocation site, so it is not exact.** A buffer that is used in the JVM and only
-later handed to native -- a shuffle-read batch feeding a native operator is the common shape -- is
-allocated somewhere that cannot know its future, so it stays charged on the JVM side while native
-may also reserve it. Closing that needs reservation ownership to be handed over at the boundary,
-which is a change on both sides of it.
+land on the task that allocated. Children cut from a task allocator inherit its listener and roll
+their bytes up into it, so the paths that make their own children are covered without knowing
+about any of this.
+
+**What is charged is what the allocator owns, not a running total of the callbacks.** When two
+allocators hold references to one buffer and the owner releases first, Arrow moves the charge to
+the other with `transferBalance`, which calls no listener, and the final `onRelease` goes to the new
+owner. `ColumnarBatchArrowReader`, which feeds every JVM batch that reaches native through
+`CometArrowStream.wrapColumnarBatchRDD`, does exactly this to each batch: it retains the buffers
+into the stream's allocator and closes the source. A listener that summed its callbacks would
+charge the task for every batch it ever streamed, until the task ended. The listener therefore
+treats a callback only as a prompt to re-read what its allocator owns, from Arrow's own accountant,
+which Arrow keeps correct across transfers. `ColumnarBatchArrowReader` also has the source's task
+allocator reconcile as soon as it closes the source, so the task stops paying for a batch the moment
+native has it, even when that is the last batch the task reads, as it is for a broadcast coalesced
+into one batch. Any other transfer is picked up at the next allocation or release on the task's
+allocator, and whatever is still reserved is returned when the task completes.
+
+**Memory passed to or from native is not charged on the JVM side.** Native memory is native's to
+account for: a native operator that retains a batch reserves its buffers through Comet's pool,
+`ExternalSorter` through `get_reserved_bytes_for_record_batch` and the hash join build side through
+`get_record_batch_memory_size`, and both read sizes off the `ArrayData`, so they count imported
+buffers too. Charging those buffers on the JVM side as well would reserve the same memory twice.
+So `NativeUtil`, `CometUdfBridge` together with the codegen output vector it exports, and
+`CometArrowStream.stream` allocate from the root, and imports go through
+`CometArrowImportAllocator`, which must stay listener-less as well, because Arrow's
+`wrapForeignAllocation` reports an imported buffer to the importing allocator's listener at full
+capacity. Allocations that stay in the JVM, which are `StreamReader`, the cached batch serializer,
+`CometArrowStream.readerBatchIter` and the Python runner, use the task allocator. A batch read that
+way and then streamed to native needs no special handling, because the stream's reader moves its
+charge off the task as described above. What remains is a buffer exported column by column with
+`Data.exportVector`, which stays owned by, and charged to, the task allocator while native reads it.
+`NativeColumnarToRowConverter` is the one caller that does this with a batch the JVM may own, and
+native's columnar-to-row conversion reserves nothing, so the bytes are not charged twice.
 
 **A task allocator is closed only once it has been drained.** The process-wide allocator exists
 because Arrow buffers can outlive the task that created them, and Arrow treats closing an allocator
 that still owns bytes as a leak. At task completion the Spark reservation is dropped and the
-allocator is closed if it is empty; otherwise it is parked and closed by a later task once the
+allocator is closed if it is empty. Otherwise it is parked and closed by a later task once the
 stragglers are released, since `BaseAllocator` keeps every child in a map until it closes.
 
 One further child, `CometArrowImportAllocator` (`comet-ffi-imports`), is what the Arrow C Data
@@ -177,36 +189,43 @@ flowchart LR
     NU["NativeUtil<br>FFI structs, imports, exports"]
     UDF["CometUdfBridge<br>JVM UDF inputs and result"]
     CGO["CometBatchKernelCodegenOutput<br>codegen UDF output"]
-    SR["StreamReader<br>shuffle and IPC reads"]
-    NAS["CometNativeArrowSource<br>stream and readerBatchIter"]
+    NASS["CometArrowStream.stream<br>batches exported to native"]
+    BC["Utils broadcast-coalesce<br>on the driver"]
+    SR["StreamReader<br>broadcast and IPC reads"]
+    NASR["CometArrowStream.readerBatchIter<br>batches read in the JVM"]
     CACHE["ArrowCachedBatchSerializer"]
     PY["CometArrowPythonRunnerBase"]
-    BC["Utils broadcast-coalesce"]
   end
 
   ROOT["CometArrowAllocator<br>RootAllocator, no limit<br>no allocation listener"]
-  SHUF["Comet JVM shuffle pages<br>CometUnifiedShuffleMemoryAllocator"]
-  NPOOL["Comet native memory pool<br>declared reservations only"]
+  TASK["CometTaskArrowAllocator<br>a child of the root per task<br>CometArrowAllocationListener"]
+  subgraph OTHER["other consumers of Spark's off-heap pool"]
+    SHUF["Comet JVM shuffle pages<br>CometUnifiedShuffleMemoryAllocator"]
+    NPOOL["Comet native memory pool<br>declared reservations only"]
+  end
   TMM["Spark off-heap execution pool<br>TaskMemoryManager"]
   NOBODY["accounted by nobody"]
 
   NU --> ROOT
   UDF --> ROOT
   CGO --> ROOT
-  SR --> ROOT
-  NAS --> ROOT
-  CACHE --> ROOT
-  PY --> ROOT
+  NASS --> ROOT
   BC --> ROOT
+  SR --> TASK
+  NASR --> TASK
+  CACHE --> TASK
+  PY --> TASK
   ROOT --> NOBODY
+  TASK -->|"what it owns, in 1 MiB blocks"| TMM
   SHUF --> TMM
   NPOOL -->|"CometTaskMemoryManager over JNI"| TMM
 ```
 
 ### Constraints on a Comet memory consumer
 
-`CometTaskMemoryManager` is the one place where Comet code acts as a Spark `MemoryConsumer`, and the
-rules below are why it is written the way it is. Each is easy to break by accident.
+`CometTaskMemoryManager` and `CometArrowAllocationListener` are the two places where Comet code acts
+as a Spark `MemoryConsumer`, and the rules below are why they are written the way they are. Each is
+easy to break by accident.
 
 **`getUsed` and `spill` must stay lock-free.** Spark calls both while already holding the
 `TaskMemoryManager` monitor, which is why `NativeMemoryConsumer.getUsed` reads an `AtomicLong` and
@@ -227,19 +246,24 @@ reaches this method from native over JNI, so whatever escapes crosses the JNI bo
 **Spark exposes no per-consumer usage figure.** `TaskMemoryManager.getMemoryConsumptionForThisTask`
 is task-wide, and a consumer that reaches `acquireExecutionMemory` directly rather than through
 `MemoryConsumer.acquireMemory` keeps an inherited `used` of zero. That is why `NativeMemoryConsumer`
-overrides `getUsed` to report Comet's own tally instead: without it, Spark's spill-victim ordering
-and `showMemoryUsage` would believe the consumer held nothing.
+overrides `getUsed` to report Comet's own tally instead, and `CometArrowAllocationListener` reports
+what its allocator owns: without that, Spark's spill-victim ordering and `showMemoryUsage` would
+believe the consumer held nothing.
 
 **A partial grant can be stranded.** `acquireExecutionMemory` takes its first grant from the pool
 before asking other consumers to spill, so when a spill throws, the task has been charged for bytes
 the call never returns. Nothing releases them until Spark's final task cleanup, so they are headroom
 nobody can use for the rest of the task. Any caller that swallows the exception has to reconcile
 that grant, and the only figure available for doing so is the task-wide one above.
+`CometArrowAllocationListener` adopts it as the change in that figure across the call, with both
+reads and the call inside the `TaskMemoryManager` monitor so that no other consumer in the task can
+acquire in between.
 
 **A consumer whose `spill` returns zero takes budget it can never give back.**
-`NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim
-nothing from it; it is only ever a spill trigger. The bytes it holds are real, so other consumers in
-the same task see correspondingly less headroom and can spill earlier than they otherwise would.
+`NativeMemoryConsumer.spill` returns `0`, as does `CometArrowAllocationListener.spill`, so Spark can
+select either as a spill victim and reclaim nothing from it; it is only ever a spill trigger. The
+bytes it holds are real, so other consumers in the same task see correspondingly less headroom and
+can spill earlier than they otherwise would.
 
 ## Where Comet's budget comes from
 
@@ -397,14 +421,17 @@ accounting (if any) saw the bytes appear. Whether the buffer is _reserved_ in Co
 separate decision, made by whichever operator holds it, and that operator neither knows nor cares
 which side of the boundary the bytes came from.
 
-**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of the
-unaccounted root and exports the whole per-partition iterator once as an `ArrowArrayStream`.
+**JVM → native (`ScanExec`).** The JVM exports the whole per-partition iterator once as an
+`ArrowArrayStream`, and every batch it hands over belongs to the stream's allocator, a child of the
+listener-less root. A batch built for the stream, such as a row-to-Arrow conversion, is allocated
+there. A batch the JVM had already read into a task allocator is retained there by
+`ColumnarBatchArrowReader` and its source closed, which moves the charge off the task allocator.
 `ScanExec` imports each batch through `AlignedArrowStreamReader` with `CopyMode::UnpackOrClone`:
 dictionary columns are unpacked into new native arrays, everything else is an `Arc` clone of the
 imported buffers. Those bytes stay where Java Arrow put them and are pinned for as long as any native
-reference survives. They are invisible to Spark's `TaskMemoryManager` at allocation time, deliberately:
-this is the export path, so nobody charged for them yet. Whether they are charged _later_ depends
-on who holds them. DataFusion's `ExternalSorter` reserves `get_reserved_bytes_for_record_batch` for
+reference survives. Spark's `TaskMemoryManager` does not see them on the JVM side, deliberately,
+because native decides whether to charge them. Whether they are charged _later_ depends on who
+holds them. DataFusion's `ExternalSorter` reserves `get_reserved_bytes_for_record_batch` for
 every batch it retains, and the hash join build side reserves `get_record_batch_memory_size` for
 each incoming batch; both read buffer sizes off the `ArrayData` and apply equally to imported
 buffers. So an imported batch that a sort or join has buffered _is_ reserved in Comet's pool, and
@@ -510,10 +537,11 @@ flowchart TB
       TUNG["Spark Tungsten off-heap<br>TaskMemoryManager"]
       SHUFP["Comet JVM shuffle pages<br>CometUnifiedShuffleMemoryAllocator"]
       NATRES["Comet native heap, reserved<br>operators that call try_grow<br>declared to Spark over JNI, never measured"]
+      ARROWT["Comet JVM Arrow owned by a task<br>CometTaskArrowAllocator, unbounded"]
     end
     subgraph NONE["accounted by nobody"]
       NATUND["Comet native heap, undeclared<br>kernels, array builders, decompression<br>Parquet metadata, object_store, tokio"]
-      ARROWR["Comet JVM Arrow<br>CometArrowAllocator, unbounded"]
+      ARROWR["Comet JVM Arrow passed to or from native<br>CometArrowAllocator, unbounded"]
       NONHEAP["JVM non-heap<br>metaspace, code cache, thread stacks<br>GC structures, Netty direct buffers"]
       PAGEC["page cache charged to the cgroup<br>file I/O, including spill files"]
       FRAG["allocator overhead<br>fragmentation, padding<br>jemalloc retained and dirty pages"]
@@ -523,10 +551,11 @@ flowchart TB
 
 Spark's accounting covers the first group, though not in the same sense throughout it. The JVM
 heap, Tungsten pages and Comet's shuffle pages are allocated by JVM code that reports what it
-allocated. A native reservation is a number an operator declared before allocating: `try_grow`
-succeeds only once `CometTaskMemoryManager` has charged Spark's off-heap execution pool over JNI, so
-the budget really is spent, but nothing measured the bytes and the reservation is only a lower bound
-on them. The second group is outside every accounting layer.
+allocated, and a task's Arrow allocator reports what it owns after the fact, without being able to
+refuse an allocation. A native reservation is a number an operator declared before allocating:
+`try_grow` succeeds only once `CometTaskMemoryManager` has charged Spark's off-heap execution pool
+over JNI, so the budget really is spent, but nothing measured the bytes and the reservation is only
+a lower bound on them. The second group is outside every accounting layer.
 
 When the total crosses `memory.max`, the kernel OOM killer kills the process. The failure mode is
 significantly worse than a task-level OOM: every task running on that executor dies, every cached
@@ -560,16 +589,10 @@ much they matter:
   its container is still stopped only by the kill.
 - **The memory overhead is sized by hand.** The gap has to fit in `spark.executor.memoryOverhead`,
   and the memory usage log measures it, but nothing sizes the overhead from it.
-- **`CometArrowAllocator` is unbounded.** Allocations that live and die inside a task are now
-  reported to Spark's memory manager, so they are no longer invisible, but nothing caps them: the
-  listener reports without enforcing, and enforcing would mean failing allocations on paths that
-  cannot fail today. Allocations made off a task, and anything crossing the FFI boundary in either
-  direction, are not reported at all.
-- **Reservation ownership is not handed over at the FFI boundary.** Which side accounts for a
-  buffer is decided by where it was allocated, not by who holds it, so a buffer allocated for JVM
-  use and later handed to native is charged on the JVM side while a native operator that retains it
-  charges the same task again. Buffers allocated for export dodge this only because their
-  allocation site knows where they are going.
+- **`CometArrowAllocator` is unbounded.** What a task's allocator owns is charged to Spark, but
+  nothing caps it: the listener reports without enforcing, and enforcing would mean failing
+  allocations on paths that cannot fail today. Arrow memory passed to or from native, and anything
+  allocated off a task, is not charged on the JVM side at all.
 - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be
   resident on either side with no reservation covering it, because reservations are made and
   withdrawn by individual operators while the bytes outlive them.
diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md
index 9322b0398c3..ccfe8adcbda 100644
--- a/docs/source/user-guide/latest/tuning.md
+++ b/docs/source/user-guide/latest/tuning.md
@@ -102,7 +102,7 @@ there is. That includes:
 - per-batch working memory in expression kernels and Arrow array builders,
 - decompression buffers and Parquet reader structures,
 - object store request buffers and the async runtime's own machinery,
-- Arrow buffers allocated on the JVM side, which no budget covers at all,
+- Arrow buffers passed between the JVM and native code, unless a native operator reserves them,
 - allocator overhead: buffer padding, size-class rounding, fragmentation, and pages the allocator retains after a
   free rather than returning to the operating system.
 
@@ -116,6 +116,13 @@ while Comet runs; see [Sizing the Overhead from the Memory Usage Log].
 fraction to each task separately, where Spark's own limit of an even share of the pool per running task is tighter
 whenever more than one task is running, and the `greedy_unified` pool ignores it.
 
+Arrow memory that Comet allocates on the JVM side for a task, for example to read a cached table or to exchange
+batches with a Python worker, is charged to the same pool as a memory consumer of that task. Charging it never fails an allocation, but it
+leaves Comet's native operators and Spark's own consumers less of the pool, so they can spill sooner. Comet 1.1.0 and
+earlier did not charge this memory to the pool, so if you sized `spark.memory.offHeap.size` against one of those
+releases, check how often queries spill. Setting `spark.comet.arrowAllocator.accounting.enabled=false` stops charging
+it, in which case it has to fit in `spark.executor.memoryOverhead` instead.
+
 For more details about Spark off-heap memory mode, please refer to [Spark documentation].
 
 [Spark documentation]: https://spark.apache.org/docs/latest/configuration.html
@@ -159,7 +166,8 @@ container to include it, so the memory that Comet's operators explicitly reserve
 not have room is everything Comet allocates without reserving it — the untracked categories listed
 under [Configuring Comet Memory]. Those allocations are made by the Rust global allocator and live
 in the native heap, outside the JVM heap and outside Spark's off-heap allocations, and nothing in
-the container sizing accounts for them. The same applies to Comet's JVM-side Arrow buffers.
+the container sizing accounts for them. The same applies to Comet's JVM-side Arrow buffers on their
+way to or from native code.
 
 `spark.executor.memoryOverhead` is the only slack the container has for this, and the JVM's own
 non-heap usage — metaspace, code cache, thread stacks, GC structures — is already drawing on it.
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 417d037d23b..5758f2821d7 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -930,14 +930,14 @@ object CometConf extends ShimCometConf {
     conf("spark.comet.arrowAllocator.accounting.enabled")
       .category(CATEGORY_TUNING)
       .doc(
-        "When enabled, JVM-side Arrow allocations made inside a Spark task are reported to " +
-          "Spark's memory manager, so they are visible to Spark's off-heap accounting rather " +
-          "than to no budget at all. Reporting only: this setting never fails an allocation, " +
-          "but the bytes do consume the off-heap pool, so other consumers see correspondingly " +
-          "less headroom. Buffers imported over the Arrow C Data Interface are never reported, " +
-          "because the memory belongs to the native side. Disable to restore the previous " +
-          "behaviour of not accounting for these allocations. " +
-          s"$TUNING_GUIDE.")
+        "Whether to charge the Arrow memory that Comet allocates on the JVM side for a task to " +
+          "Spark's off-heap memory pool, as a memory consumer of that task. This never fails an " +
+          "allocation, but it does leave Comet's native operators and Spark's other consumers " +
+          "correspondingly less of the pool, so they may spill sooner. Arrow memory that is " +
+          "handed to Comet's native code, or received from it, is not charged here, because it " +
+          "is native's to account for. This is an executor setting, read when an executor " +
+          "first allocates Arrow memory for a task, so it must be set when the application is " +
+          s"submitted. Set to false to stop charging these allocations. $TUNING_GUIDE.")
       .booleanConf
       .createWithDefault(true)
 

From 1e13a9ab7d435cd2cba77ac6a4cc2adb87aedc09 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sat, 26 Sep 2026 08:03:19 -0600
Subject: [PATCH 15/18] refactor: rename the accounting config to
 spark.comet.memory.jvmArrowAccounting.enabled

The key's first segment, arrowAllocator, is not one of the category
segments the config conventions list, and main already has a
spark.comet.memory family in spark.comet.memory.logInterval, another
executor setting in the tuning category that is read once per JVM. The
key has not been released, so it is renamed outright rather than kept as
a deprecated alternative, and the Scala symbol follows the same pattern
as COMET_MEMORY_LOG_INTERVAL. The declaration moves next to that one.
---
 .../contributor-guide/memory_management.md    |  2 +-
 docs/source/user-guide/latest/tuning.md       |  2 +-
 .../scala/org/apache/comet/CometConf.scala    | 30 +++++++++----------
 .../comet/CometArrowAllocationListener.scala  |  4 +--
 .../spark/comet/CometTaskArrowAllocator.scala |  4 +--
 ...ometArrowAllocationListenerBenchmark.scala |  2 +-
 6 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index dbbe7d7efcc..515e57833ed 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -117,7 +117,7 @@ reserves what the child owns from the off-heap execution pool in whole 1 MiB blo
 therefore appear in `showMemoryUsage` and are arbitrated against Spark's other off-heap consumers,
 Comet's native pool among them. The limit is still `Long.MaxValue`, because the listener reports
 without enforcing, so an allocation here cannot fail. Set
-`spark.comet.arrowAllocator.accounting.enabled=false` to stop charging them, in which case
+`spark.comet.memory.jvmArrowAccounting.enabled=false` to stop charging them, in which case
 `forCurrentTask()` returns the root. It also returns the root in on-heap mode, and off a task, as
 for broadcast coalescing on the driver or for a reader built on a native thread that is pulling a
 stream.
diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md
index ccfe8adcbda..7eae248271a 100644
--- a/docs/source/user-guide/latest/tuning.md
+++ b/docs/source/user-guide/latest/tuning.md
@@ -120,7 +120,7 @@ Arrow memory that Comet allocates on the JVM side for a task, for example to rea
 batches with a Python worker, is charged to the same pool as a memory consumer of that task. Charging it never fails an allocation, but it
 leaves Comet's native operators and Spark's own consumers less of the pool, so they can spill sooner. Comet 1.1.0 and
 earlier did not charge this memory to the pool, so if you sized `spark.memory.offHeap.size` against one of those
-releases, check how often queries spill. Setting `spark.comet.arrowAllocator.accounting.enabled=false` stops charging
+releases, check how often queries spill. Setting `spark.comet.memory.jvmArrowAccounting.enabled=false` stops charging
 it, in which case it has to fit in `spark.executor.memoryOverhead` instead.
 
 For more details about Spark off-heap memory mode, please refer to [Spark documentation].
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 5758f2821d7..e779e8f6cad 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -374,6 +374,21 @@ object CometConf extends ShimCometConf {
     .checkValue(_ >= 0, "The memory usage log interval must not be negative")
     .createWithDefault(TimeUnit.SECONDS.toMillis(10))
 
+  val COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.memory.jvmArrowAccounting.enabled")
+      .category(CATEGORY_TUNING)
+      .doc(
+        "Whether to charge the Arrow memory that Comet allocates on the JVM side for a task to " +
+          "Spark's off-heap memory pool, as a memory consumer of that task. This never fails an " +
+          "allocation, but it does leave Comet's native operators and Spark's other consumers " +
+          "correspondingly less of the pool, so they may spill sooner. Arrow memory that is " +
+          "handed to Comet's native code, or received from it, is not charged here, because it " +
+          "is native's to account for. This is an executor setting, read when an executor " +
+          "first allocates Arrow memory for a task, so it must be set when the application is " +
+          s"submitted. Set to false to stop charging these allocations. $TUNING_GUIDE.")
+      .booleanConf
+      .createWithDefault(true)
+
   val COMET_SHUFFLE_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.shuffle.enabled")
       .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.enabled")
@@ -926,21 +941,6 @@ object CometConf extends ShimCometConf {
       .doubleConf
       .createWithDefault(1.0)
 
-  val COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED: ConfigEntry[Boolean] =
-    conf("spark.comet.arrowAllocator.accounting.enabled")
-      .category(CATEGORY_TUNING)
-      .doc(
-        "Whether to charge the Arrow memory that Comet allocates on the JVM side for a task to " +
-          "Spark's off-heap memory pool, as a memory consumer of that task. This never fails an " +
-          "allocation, but it does leave Comet's native operators and Spark's other consumers " +
-          "correspondingly less of the pool, so they may spill sooner. Arrow memory that is " +
-          "handed to Comet's native code, or received from it, is not charged here, because it " +
-          "is native's to account for. This is an executor setting, read when an executor " +
-          "first allocates Arrow memory for a task, so it must be set when the application is " +
-          s"submitted. Set to false to stop charging these allocations. $TUNING_GUIDE.")
-      .booleanConf
-      .createWithDefault(true)
-
   val COMET_NATIVE_LOAD_REQUIRED: ConfigEntry[Boolean] = conf("spark.comet.nativeLoadRequired")
     .category(CATEGORY_EXEC)
     .doc(
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index c5cfbe5cdb0..5ba80dc05b6 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -284,7 +284,7 @@ object CometArrowAllocationListener extends Logging {
       logWarning(
         s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " +
           "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " +
-          s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " +
+          s"Set ${CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key}=false to stop " +
           "reporting these allocations to Spark.")
     }
   }
@@ -295,7 +295,7 @@ object CometArrowAllocationListener extends Logging {
         "Failed to report a JVM Arrow allocation to Spark's memory manager. The allocation " +
           "itself is unaffected, so this is a reporting gap rather than a failure, but Spark's " +
           "view of these bytes will be short until the task ends. " +
-          s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " +
+          s"Set ${CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key}=false to stop " +
           "reporting these allocations to Spark.",
         e)
     }
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
index 4d120bc6119..13578f96ae1 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
@@ -82,8 +82,8 @@ object CometTaskArrowAllocator extends Logging {
    */
   private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env =>
     env.conf.getBoolean(
-      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key,
-      CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.defaultValue.get)
+      CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key,
+      CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.defaultValue.get)
   }
 
   /**
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
index 07e2896f2f5..6e79b6ea47d 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
@@ -33,7 +33,7 @@ import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, TaskM
  *
  * Each case runs the same allocate/release loop twice: once against a plain `RootAllocator`,
  * which is what Comet did before [[CometArrowAllocationListener]] existed and what
- * `spark.comet.arrowAllocator.accounting.enabled=false` restores, and once against a task
+ * `spark.comet.memory.jvmArrowAccounting.enabled=false` restores, and once against a task
  * allocator carrying the listener. Reservation call counts are printed under each table, because
  * the interesting variable is not the per-buffer bookkeeping but how often a buffer size crosses
  * a block boundary and has to go into `TaskMemoryManager` at all.

From 6514a06b479ce4ac4b2b26d4c367e481f986d603 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sat, 26 Sep 2026 10:10:26 -0600
Subject: [PATCH 16/18] feat: reserve JVM Arrow memory before allocating, and
 refuse what Spark cannot cover

The listener reserved from Spark in onAllocation, after Arrow had already
made the buffer. That callback may not throw, so a short grant could
only be kept and the allocation let stand: grow semantics, with the
shortfall as overcommit that Spark did not know about and no ledger to
repay it.

Reservation now happens in onPreAllocation, which Arrow calls before it
allocates and documents as the one callback that may throw to terminate
the allocation. The listener asks Spark for what the allocation needs,
which is what the allocator owns plus allocations already admitted but
not yet reported plus this one, rounded up to a block. If Spark cannot
cover it, the listener returns whatever Spark did grant and throws
Arrow's OutOfMemoryException, which Arrow's vectors already clean up
after. SparkOutOfMemoryError is an Error and would skip those NonFatal
paths. This is try_grow semantics: nothing is allocated that is not
covered. Nothing can spill these buffers, so a refusal fails the task,
the way a Spark consumer fails when it cannot acquire memory.

- A grant short of the rounded block is accepted when it covers the
  actual need, so the rounding never causes a refusal.
- Allocations admitted but not yet reported count against the
  reservation, so two threads cannot both be admitted against the same
  headroom.
- When the acquisition throws, any grant Spark took first is measured
  under the TaskMemoryManager monitor as before, then released and not
  kept, and the failure becomes the refusal's cause. An interrupt
  re-arms the thread's flag.
- onAllocation, onFailedAllocation, onRelease and reconcile only settle
  or shrink the reservation. Shrinking now keeps what is owned and
  pending rounded up to a block, so a sub-block grant is returned once
  nothing needs it.

The failure-containment tests become refusal tests: the allocation is
refused, no buffer is left, nothing stays reserved, and nothing is
charged to the task. New tests cover a refusal followed by an allocation
that fits, a sub-block grant being accepted, and an in-flight allocation
counting against the reservation. Each rule was checked by removing it,
and each removal fails its test. The benchmark's native-pressure pool
now fits the Arrow peak, since a smaller one would refuse rather than
slow down.

The memory management and tuning guides, the review skills and the
config description now say allocations are bounded and can fail.
---
 .ai/skills/review-comet-ffi-pr/SKILL.md       |   8 +-
 .ai/skills/review-comet-memory-pr/SKILL.md    |  13 +-
 .../contributor-guide/memory_management.md    |  39 +--
 docs/source/user-guide/latest/tuning.md       |   6 +-
 .../scala/org/apache/comet/CometConf.scala    |   7 +-
 .../main/scala/org/apache/comet/package.scala |   4 +-
 .../comet/CometArrowAllocationListener.scala  | 293 ++++++++++--------
 .../spark/comet/CometTaskArrowAllocator.scala |  15 +-
 ...ometArrowAllocationListenerBenchmark.scala |  34 +-
 .../CometArrowAllocationListenerSuite.scala   | 185 +++++++----
 10 files changed, 351 insertions(+), 253 deletions(-)

diff --git a/.ai/skills/review-comet-ffi-pr/SKILL.md b/.ai/skills/review-comet-ffi-pr/SKILL.md
index 62f718af556..facf9ba926a 100644
--- a/.ai/skills/review-comet-ffi-pr/SKILL.md
+++ b/.ai/skills/review-comet-ffi-pr/SKILL.md
@@ -138,10 +138,10 @@ An FFI change is usually also a memory change, and the two are easy to review se
 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, 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.
+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.
 
diff --git a/.ai/skills/review-comet-memory-pr/SKILL.md b/.ai/skills/review-comet-memory-pr/SKILL.md
index d99871c78d9..9ba6d2ec6f1 100644
--- a/.ai/skills/review-comet-memory-pr/SKILL.md
+++ b/.ai/skills/review-comet-memory-pr/SKILL.md
@@ -168,14 +168,17 @@ Two allocators sit outside Comet's pool entirely, and they are easy to confuse:
   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`. A PR that
+  `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. Check that the listener
-  stays without a lock in `getUsed` and `spill`, 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.
+  `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
diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 515e57833ed..cbed07ad918 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -81,7 +81,7 @@ page:
 | Spark execution + storage (on-heap)      | JVM heap    | `spark.executor.memory` and the unified memory manager        | Yes               |
 | Spark Tungsten (off-heap)                | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
 | Comet native (Rust global allocator)     | Native heap | `memory_limit` (see below), enforced only via the memory pool | Reservations only |
-| Comet JVM Arrow owned by a task          | Off-heap    | **Nothing**, but charged to `TaskMemoryManager`               | Yes               |
+| Comet JVM Arrow owned by a task          | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
 | Comet JVM Arrow passed to or from native | Off-heap    | **Nothing**: a `RootAllocator(Long.MaxValue)`                 | No                |
 | Comet JVM shuffle pages                  | Off-heap    | `spark.memory.offHeap.size` via `TaskMemoryManager`           | Yes               |
 
@@ -108,15 +108,19 @@ pool over JNI. A native allocation therefore consumes the same accounting budget
 while occupying entirely different memory. Raising `spark.memory.offHeap.size` raises the ceiling
 for both at once, and raises the pod's memory request by the same amount.
 
-**Comet's JVM-side Arrow memory is charged to the task that owns it, but it is still unbounded.**
+**Comet's JVM-side Arrow memory is charged to the task that owns it, and bounded by the pool.**
 `CometArrowAllocator` (`spark/src/main/scala/org/apache/comet/package.scala`) is a single
 process-wide `new RootAllocator(Long.MaxValue)` with no allocation listener. JVM-owned allocations
 do not use it directly: `CometTaskArrowAllocator.forCurrentTask()` hands each task a child of it
 carrying its own `CometArrowAllocationListener`, a Spark `MemoryConsumer` for that task, which
 reserves what the child owns from the off-heap execution pool in whole 1 MiB blocks. Those bytes
 therefore appear in `showMemoryUsage` and are arbitrated against Spark's other off-heap consumers,
-Comet's native pool among them. The limit is still `Long.MaxValue`, because the listener reports
-without enforcing, so an allocation here cannot fail. Set
+Comet's native pool among them. The reservation is made before Arrow allocates, in
+`onPreAllocation`, the one `AllocationListener` callback Arrow allows to throw. If Spark cannot
+cover the allocation, any partial grant is handed back and the allocation is refused with Arrow's
+`OutOfMemoryException`. That is `try_grow` rather than `grow`: nothing is allocated that is not
+covered, and there is no overcommit. Nothing can spill these buffers, so a refusal fails the
+allocating task, the way a Spark consumer fails when it cannot acquire memory. Set
 `spark.comet.memory.jvmArrowAccounting.enabled=false` to stop charging them, in which case
 `forCurrentTask()` returns the root. It also returns the root in on-heap mode, and off a task, as
 for broadcast coalescing on the driver or for a reader built on a native thread that is pulling a
@@ -255,9 +259,9 @@ before asking other consumers to spill, so when a spill throws, the task has bee
 the call never returns. Nothing releases them until Spark's final task cleanup, so they are headroom
 nobody can use for the rest of the task. Any caller that swallows the exception has to reconcile
 that grant, and the only figure available for doing so is the task-wide one above.
-`CometArrowAllocationListener` adopts it as the change in that figure across the call, with both
+`CometArrowAllocationListener` measures it as the change in that figure across the call, with both
 reads and the call inside the `TaskMemoryManager` monitor so that no other consumer in the task can
-acquire in between.
+acquire in between, and releases it before refusing the allocation.
 
 **A consumer whose `spill` returns zero takes budget it can never give back.**
 `NativeMemoryConsumer.spill` returns `0`, as does `CometArrowAllocationListener.spill`, so Spark can
@@ -537,7 +541,7 @@ flowchart TB
       TUNG["Spark Tungsten off-heap<br>TaskMemoryManager"]
       SHUFP["Comet JVM shuffle pages<br>CometUnifiedShuffleMemoryAllocator"]
       NATRES["Comet native heap, reserved<br>operators that call try_grow<br>declared to Spark over JNI, never measured"]
-      ARROWT["Comet JVM Arrow owned by a task<br>CometTaskArrowAllocator, unbounded"]
+      ARROWT["Comet JVM Arrow owned by a task<br>CometTaskArrowAllocator, refused when the pool is short"]
     end
     subgraph NONE["accounted by nobody"]
       NATUND["Comet native heap, undeclared<br>kernels, array builders, decompression<br>Parquet metadata, object_store, tokio"]
@@ -549,13 +553,13 @@ flowchart TB
   end
 ```
 
-Spark's accounting covers the first group, though not in the same sense throughout it. The JVM
-heap, Tungsten pages and Comet's shuffle pages are allocated by JVM code that reports what it
-allocated, and a task's Arrow allocator reports what it owns after the fact, without being able to
-refuse an allocation. A native reservation is a number an operator declared before allocating:
-`try_grow` succeeds only once `CometTaskMemoryManager` has charged Spark's off-heap execution pool
-over JNI, so the budget really is spent, but nothing measured the bytes and the reservation is only
-a lower bound on them. The second group is outside every accounting layer.
+Spark's accounting covers the first group, though not in the same sense throughout it. The JVM heap,
+Tungsten pages and Comet's shuffle pages are allocated by JVM code that reports what it allocated,
+and a task's Arrow allocator reserves what an allocation needs before Arrow makes it, refusing the
+allocation when Spark cannot cover it. A native reservation is a number an operator declared before
+allocating: `try_grow` succeeds only once `CometTaskMemoryManager` has charged Spark's off-heap
+execution pool over JNI, so the budget really is spent, but nothing measured the bytes and the
+reservation is only a lower bound on them. The second group is outside every accounting layer.
 
 When the total crosses `memory.max`, the kernel OOM killer kills the process. The failure mode is
 significantly worse than a task-level OOM: every task running on that executor dies, every cached
@@ -589,10 +593,9 @@ much they matter:
   its container is still stopped only by the kill.
 - **The memory overhead is sized by hand.** The gap has to fit in `spark.executor.memoryOverhead`,
   and the memory usage log measures it, but nothing sizes the overhead from it.
-- **`CometArrowAllocator` is unbounded.** What a task's allocator owns is charged to Spark, but
-  nothing caps it: the listener reports without enforcing, and enforcing would mean failing
-  allocations on paths that cannot fail today. Arrow memory passed to or from native, and anything
-  allocated off a task, is not charged on the JVM side at all.
+- **`CometArrowAllocator` is bounded only for what a task owns.** A task's allocator is refused
+  what Spark cannot cover, but Arrow memory passed to or from native, and anything allocated off a
+  task, is charged nowhere on the JVM side, and the root itself has no limit.
 - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be
   resident on either side with no reservation covering it, because reservations are made and
   withdrawn by individual operators while the bytes outlive them.
diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md
index 7eae248271a..f6b88d2f425 100644
--- a/docs/source/user-guide/latest/tuning.md
+++ b/docs/source/user-guide/latest/tuning.md
@@ -117,10 +117,12 @@ fraction to each task separately, where Spark's own limit of an even share of th
 whenever more than one task is running, and the `greedy_unified` pool ignores it.
 
 Arrow memory that Comet allocates on the JVM side for a task, for example to read a cached table or to exchange
-batches with a Python worker, is charged to the same pool as a memory consumer of that task. Charging it never fails an allocation, but it
+batches with a Python worker, is charged to the same pool as a memory consumer of that task. An allocation the pool
+cannot cover fails the task that makes it, the way a Spark operator fails when it cannot acquire memory, and the charge
 leaves Comet's native operators and Spark's own consumers less of the pool, so they can spill sooner. Comet 1.1.0 and
 earlier did not charge this memory to the pool, so if you sized `spark.memory.offHeap.size` against one of those
-releases, check how often queries spill. Setting `spark.comet.memory.jvmArrowAccounting.enabled=false` stops charging
+releases, check how often queries spill, and raise it if tasks fail with `Unable to reserve ... for a JVM Arrow
+allocation`. Setting `spark.comet.memory.jvmArrowAccounting.enabled=false` stops charging
 it, in which case it has to fit in `spark.executor.memoryOverhead` instead.
 
 For more details about Spark off-heap memory mode, please refer to [Spark documentation].
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index e779e8f6cad..e51282c3c08 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -379,9 +379,10 @@ object CometConf extends ShimCometConf {
       .category(CATEGORY_TUNING)
       .doc(
         "Whether to charge the Arrow memory that Comet allocates on the JVM side for a task to " +
-          "Spark's off-heap memory pool, as a memory consumer of that task. This never fails an " +
-          "allocation, but it does leave Comet's native operators and Spark's other consumers " +
-          "correspondingly less of the pool, so they may spill sooner. Arrow memory that is " +
+          "Spark's off-heap memory pool, as a memory consumer of that task. An allocation the " +
+          "pool cannot cover fails the task that makes it, and the charge leaves Comet's native " +
+          "operators and Spark's other consumers correspondingly less of the pool, so they may " +
+          "spill sooner. Arrow memory that is " +
           "handed to Comet's native code, or received from it, is not charged here, because it " +
           "is native's to account for. This is an executor setting, read when an executor " +
           "first allocates Arrow memory for a task, so it must be set when the application is " +
diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala
index e2c407843d9..8c050cec756 100644
--- a/spark/src/main/scala/org/apache/comet/package.scala
+++ b/spark/src/main/scala/org/apache/comet/package.scala
@@ -39,8 +39,8 @@ package object comet {
    * buffers through Comet's pool, which charges the same Spark task, so reporting them here as
    * well would reserve the same memory twice. JVM-owned allocations should go through
    * `CometTaskArrowAllocator.forCurrentTask()` instead, which cuts a per-task child whose
-   * listener reports what that child owns, and which hands back this allocator when there is no
-   * task.
+   * listener charges what that child owns to Spark and refuses an allocation Spark cannot cover,
+   * and which hands back this allocator when there is no task.
    */
   val CometArrowAllocator = new RootAllocator(Long.MaxValue)
 
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index 5ba80dc05b6..5b826a5ab0b 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -24,30 +24,41 @@ import java.util.function.LongSupplier
 
 import scala.util.control.NonFatal
 
-import org.apache.arrow.memory.AllocationListener
+import org.apache.arrow.memory.{AllocationListener, AllocationOutcome, OutOfMemoryException}
 import org.apache.spark.internal.Logging
 import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager}
 
 import org.apache.comet.CometConf
 
 /**
- * Accounts one task's JVM-side Arrow allocations against Spark's off-heap execution pool.
+ * Charges one task's JVM-side Arrow allocations to Spark's off-heap execution pool, and refuses
+ * an allocation the pool cannot cover.
  *
  * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, so until now the
  * off-heap bytes it hands out were counted by nobody: not Spark's `TaskMemoryManager`, and not
  * Comet's native memory pool. They are still resident in the container, which makes them a blind
- * spot when an executor is killed for exceeding its memory limit. This closes the reporting half
- * of that gap: the bytes appear in `TaskMemoryManager.showMemoryUsage` and are arbitrated against
- * Spark's other off-heap consumers.
+ * spot when an executor is killed for exceeding its memory limit. Here they appear in
+ * `TaskMemoryManager.showMemoryUsage`, are arbitrated against Spark's other off-heap consumers,
+ * and are bounded by the same pool.
+ *
+ * '''Reserved before it is allocated.''' Arrow calls [[onPreAllocation]] before it allocates a
+ * buffer, and documents it as the one callback that may throw to terminate the allocation. The
+ * reservation is taken there, in whole blocks, and if Spark cannot cover the allocation whatever
+ * it did grant is handed straight back and the allocation is refused with Arrow's own
+ * `OutOfMemoryException`, the exception Arrow's vectors already clean up after. That is
+ * `try_grow` rather than `grow`: nothing is allocated that is not covered, and there is no
+ * overcommit. Nothing can spill these buffers either, since the JVM holds them until the batch is
+ * done with, so a refusal fails the allocating task, the way a Spark consumer fails when it
+ * cannot acquire memory.
  *
  * '''Ownership.''' One instance is created per task and attached to that task's Arrow allocator
- * by [[CometTaskArrowAllocator]]. Arrow calls `onAllocation` on the listener of the allocator
- * that made a buffer, and `onRelease` on the listener of the allocator that '''owns''' it once
- * the last reference drops, on whichever thread that happens to be. `AllocationListener` is
- * handed nothing but a size, so binding the listener to an allocator is the only way to attribute
- * a release. Reading `TaskContext` inside the callbacks would get it wrong: a buffer exported
- * over the C Data Interface is released by whichever native thread drops it last, and that thread
- * has no task context installed.
+ * by [[CometTaskArrowAllocator]]. Arrow calls the listener of the allocator that makes a buffer
+ * on allocation, and the listener of the allocator that '''owns''' it once the last reference
+ * drops, on whichever thread that happens to be. `AllocationListener` is handed nothing but a
+ * size, so binding the listener to an allocator is the only way to attribute a release. Reading
+ * `TaskContext` inside the callbacks would get it wrong: a buffer exported over the C Data
+ * Interface is released by whichever native thread drops it last, and that thread has no task
+ * context installed.
  *
  * '''What is charged is what the allocator owns, not a running total of the callbacks.''' The two
  * come apart because ownership moves between allocators without either listener hearing about it.
@@ -59,36 +70,30 @@ import org.apache.comet.CometConf
  * the task for each such batch until the task ended. So the callbacks only say when to look, and
  * the size comes from the task allocator's accountant, which Arrow keeps right across transfers.
  * `ColumnarBatchArrowReader` calls [[reconcile]] as soon as it closes a source, so a task stops
- * paying for a batch once native has it. Any other transfer is reflected at the next allocation
- * or release on the task's allocator, and whatever is still reserved when the task ends goes back
- * to Spark then.
- *
- * '''Reporting only.''' A short grant is logged and the allocation proceeds, because Arrow
- * allocation on these paths cannot fail today and making it fail is a behavioural change that
- * belongs in its own commit. Enforcement belongs in `onPreAllocation`, the only callback
- * permitted to throw, and in `onFailedAllocation`, not here. See
- * [[https://github.com/apache/datafusion-comet/issues/5997]].
+ * paying for a batch once native has it. Any other transfer out is reflected at the next
+ * allocation or release on the task's allocator. A transfer in is charged at the next allocation,
+ * which is refused if the pool cannot cover both.
  *
- * '''Neither callback may throw.''' Arrow's `AllocationListener` documents that, and
+ * '''Only [[onPreAllocation]] may throw.''' Arrow documents that the other callbacks cannot, and
  * `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so
- * throwing from here loses the buffer Arrow has already created and never hands back. Spark's
- * acquisition is fallible in three ways, and only the first is caught by `NonFatal`: it runs
- * other consumers' `spill`, which turns a task interrupt into a `RuntimeException` and an I/O
- * failure into a `SparkOutOfMemoryError`, and the execution pool itself parks in `lock.wait()`,
- * so killing a task can raise a plain `InterruptedException` here. Every call into the memory
- * manager is wrapped and reported rather than propagated, and an interrupt additionally re-arms
- * the thread's flag so the cancellation is not swallowed. A failed acquisition can also leave the
- * task charged for bytes Spark never reported back; see [[acquire]].
+ * throwing from any of them loses a buffer Arrow has already created. They only ever shrink the
+ * reservation, and contain anything the memory manager throws. The acquisition in
+ * [[onPreAllocation]] is fallible in three ways, and only the first is caught by `NonFatal`: it
+ * runs other consumers' `spill`, which turns a task interrupt into a `RuntimeException` and an
+ * I/O failure into a `SparkOutOfMemoryError`, and the execution pool itself parks in
+ * `lock.wait()`, so killing a task can raise a plain `InterruptedException`. Each becomes the
+ * refusal's cause, with the thread's interrupt flag re-armed for an interrupt, and any grant
+ * Spark took before failing is handed back first; see [[acquire]].
  *
  * '''Lock order.''' This listener's monitor is taken before Spark's and never the other way
- * round: [[adjust]] holds ours across `acquireExecutionMemory`, and [[acquire]] additionally
- * takes the `TaskMemoryManager` monitor itself, so that the two usage snapshots either side of
- * that call cannot be split by another consumer. [[getUsed]] and [[spill]] must therefore stay
- * lock-free, because Spark calls both while holding its own monitor: were either to take ours, a
- * native reservation arriving through `CometTaskMemoryManager` on a Comet Tokio thread could hold
- * Spark's monitor and wait for ours while an Arrow allocation on the same task held ours and
- * waited for Spark's. For the same reason, nothing reachable from a `spill` callback may allocate
- * JVM Arrow memory.
+ * round: [[onPreAllocation]] holds ours across `acquireExecutionMemory`, and [[acquire]]
+ * additionally takes the `TaskMemoryManager` monitor itself, so that the two usage snapshots
+ * either side of that call cannot be split by another consumer. [[getUsed]] and [[spill]] must
+ * therefore stay lock-free, because Spark calls both while holding its own monitor: were either
+ * to take ours, a native reservation arriving through `CometTaskMemoryManager` on a Comet Tokio
+ * thread could hold Spark's monitor and wait for ours while an Arrow allocation on the same task
+ * held ours and waited for Spark's. For the same reason, nothing reachable from a `spill`
+ * callback may allocate JVM Arrow memory.
  */
 private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryManager)
     extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP)
@@ -107,31 +112,59 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
   /** Bytes currently reserved with Spark. Guarded by this listener's monitor. */
   private var reserved = 0L
 
+  /**
+   * Bytes admitted by [[onPreAllocation]] whose allocation Arrow has not reported yet. They count
+   * against the reservation, so that two allocations on different threads cannot both be admitted
+   * against the same headroom. Guarded by this listener's monitor.
+   */
+  private var pending = 0L
+
   /** Set once the owning task has finished. Volatile so [[getUsed]] can read it lock-free. */
   @volatile private var completed = false
 
   /** Attaches the listener to what it accounts for. Called once, before the allocator is used. */
   private[comet] def bind(owned: LongSupplier): Unit = this.owned = owned
 
-  override def onAllocation(size: Long): Unit = adjustQuietly()
+  /**
+   * Reserves what the allocation needs before Arrow makes it, or refuses it. Once the task has
+   * finished there is nothing left to charge, so a straggling allocation goes through.
+   */
+  override def onPreAllocation(size: Long): Unit = synchronized {
+    if (!completed) {
+      val needed = owned.getAsLong + pending + size
+      if (reserved < needed) {
+        reserved += acquireAtLeast(needed - reserved, size)
+      }
+      pending += size
+    }
+  }
+
+  override def onAllocation(size: Long): Unit = settle(size)
+
+  /** Arrow's own limit refused the allocation after it was admitted here, so hand it back. */
+  override def onFailedAllocation(size: Long, outcome: AllocationOutcome): Boolean = {
+    settle(size)
+    shrinkQuietly()
+    false
+  }
 
-  override def onRelease(size: Long): Unit = adjustQuietly()
+  override def onRelease(size: Long): Unit = shrinkQuietly()
 
   /** Brings the reservation up to date after buffers moved out without a callback. */
-  private[comet] def reconcile(): Unit = adjustQuietly()
+  private[comet] def reconcile(): Unit = shrinkQuietly()
 
   /**
    * Reports what the task's allocator owns. Spark reads this for spill-victim ordering,
    * `showMemoryUsage` and end-of-task leak reporting. The inherited `used` counter stays at zero
-   * because this consumer never calls `acquireMemory` or `allocatePage`; Arrow has already
-   * obtained the memory and we are only accounting for it.
+   * because this consumer never calls `acquireMemory` or `allocatePage`; the reservation is made
+   * on Arrow's behalf, not through the consumer's own page accounting.
    *
    * Reports zero once the task has finished, so that buffers deliberately allowed to outlive
    * their task are not reported by `cleanUpAllAllocatedMemory` as a Spark memory leak.
    */
   override def getUsed: Long = if (completed) 0L else owned.getAsLong
 
-  /** Comet's native operators cannot be made to spill from here. See issue #5997. */
+  /** The JVM holds these buffers until the batch is done with, so there is nothing to spill. */
   override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
 
   /**
@@ -151,66 +184,71 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
         }
       }
     } catch {
-      case e: InterruptedException => reportAndReinterrupt(e)
-      case NonFatal(e) => warnOnMemoryManagerFailure(e)
-      case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
+      case NonFatal(e) => warnOnReleaseFailure(e)
     }
   }
 
   /** Bytes currently reserved with Spark on this task's behalf. Visible for testing. */
   private[comet] def reservedBytes: Long = synchronized(reserved)
 
-  private def adjustQuietly(): Unit = {
+  /** The allocation is now in the allocator's accountant, so it is no longer pending. */
+  private def settle(size: Long): Unit = synchronized {
+    if (!completed) {
+      pending = math.max(0L, pending - size)
+    }
+  }
+
+  /**
+   * Returns what the allocator no longer needs, keeping what it owns and has pending rounded up
+   * to a block so that allocating and releasing around a boundary does not reach Spark every
+   * time. Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes
+   * on the executor-wide pool, so a per-block loop would take that lock once per megabyte freed.
+   * Releasing cannot fail in practice, but this is reached from callbacks that may not throw, so
+   * anything unforeseen is contained.
+   */
+  private def shrinkQuietly(): Unit = {
     try {
-      adjust()
+      synchronized {
+        if (!completed) {
+          val keep = roundUpToBlock(owned.getAsLong + pending)
+          if (reserved > keep) {
+            taskMemoryManager.releaseExecutionMemory(reserved - keep, this)
+            reserved = keep
+          }
+        }
+      }
     } catch {
-      // Growth handles its own failures in `acquire`, so this is the net for the release path and
-      // for anything unforeseen. All three are reachable from the memory manager:
-      // `acquireExecutionMemory` runs other consumers' `spill`, `TaskMemoryManager` turns an
-      // interrupted spill into a RuntimeException and an IOException into a SparkOutOfMemoryError,
-      // and the execution pool itself parks in `lock.wait()`. The last two slip past NonFatal,
-      // which excludes Errors and InterruptedException.
-      case e: InterruptedException => reportAndReinterrupt(e)
-      case NonFatal(e) => warnOnMemoryManagerFailure(e)
-      case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
+      case NonFatal(e) => warnOnReleaseFailure(e)
     }
   }
 
-  private def adjust(): Unit = synchronized {
-    if (!completed) {
-      val ownedBytes = owned.getAsLong
-      if (reserved < ownedBytes) {
-        // Round up so `reserved` stays a block multiple and growth always leaves headroom.
-        // Requesting the bare deficit would land exactly on `ownedBytes` for any buffer at or
-        // above the block size, sending the very next allocation straight back into Spark's lock.
-        val request = roundUpToBlock(ownedBytes - reserved)
-        val granted = acquire(request)
-        reserved += granted
-        if (granted < request) {
-          warnOnShortGrant(request, granted)
-        }
-      } else {
-        // Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes on
-        // the executor-wide pool, so a per-block loop would take that lock once per megabyte freed.
-        val excess = ((reserved - ownedBytes) / BLOCK_SIZE) * BLOCK_SIZE
-        if (excess > 0L) {
-          taskMemoryManager.releaseExecutionMemory(excess, this)
-          reserved -= excess
-        }
+  /**
+   * Asks Spark for `shortfall` bytes, rounded up to a block so that `reserved` leaves headroom
+   * and the next small allocation does not go straight back into Spark's lock, and returns what
+   * it granted. A grant short of the block but covering `shortfall` is accepted: the rounding is
+   * there to save lock traffic, and must not be the reason an allocation that fits is refused.
+   * Anything less is handed back and the allocation refused.
+   */
+  private def acquireAtLeast(shortfall: Long, size: Long): Long = {
+    val granted = acquire(roundUpToBlock(shortfall), size)
+    if (granted < shortfall) {
+      if (granted > 0L) {
+        taskMemoryManager.releaseExecutionMemory(granted, this)
       }
+      throw refusal(size, shortfall, granted, null)
     }
+    granted
   }
 
   /**
-   * Asks Spark for `request` bytes and returns what this consumer ends up holding, which is not
-   * always what Spark returns.
+   * Asks Spark for `request` bytes and returns what it granted, or refuses the allocation if the
+   * acquisition fails.
    *
    * `acquireExecutionMemory` takes its first grant from the pool and only then asks other
    * consumers to spill, so when a spill throws it has already charged the task for bytes it never
-   * reports back. Nothing would release them: [[taskCompleted]] only knows about `reserved`, and
-   * Spark itself only reclaims them in `cleanUpAllAllocatedMemory` at the very end of the task,
-   * so until then they are headroom nobody can use. They are adopted here instead, measured as
-   * the change in what the pool says this task holds.
+   * reports back, and nothing would release them before `cleanUpAllAllocatedMemory` at the very
+   * end of the task. They are measured as the change in what the pool says this task holds, and
+   * released before the refusal is thrown.
    *
    * Spark reports that figure per task rather than per consumer, so it only measures '''our'''
    * grant if nothing else in the task can move it while we are looking. Both snapshots and the
@@ -218,46 +256,55 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
    * the same monitor `acquireExecutionMemory` takes and holds for its whole duration, spills
    * included, and it is reentrant, so taking it here only widens that window to cover the two
    * reads. Every acquisition in the task funnels through that method, so with it held no other
-   * consumer can take memory between a snapshot and the call and have it adopted here.
+   * consumer can take memory between a snapshot and the call and have it released here.
    *
    * What the monitor does not cover is a release, which reaches the pool without it. Another
    * consumer returning memory, or a spill that frees some bytes before throwing, makes the figure
-   * too small, which is the safe direction: too small degrades to what would have happened
-   * anyway. `request` bounds it from above.
+   * too small, which is the safe direction: too small leaves bytes charged to the task until it
+   * ends, which is what would have happened anyway. `request` bounds it from above.
    */
-  private def acquire(request: Long): Long = taskMemoryManager.synchronized {
+  private def acquire(request: Long, size: Long): Long = taskMemoryManager.synchronized {
     val heldBefore = taskMemoryManager.getMemoryConsumptionForThisTask
     try {
       taskMemoryManager.acquireExecutionMemory(request, this)
     } catch {
       case e: InterruptedException =>
-        reportAndReinterrupt(e)
-        adoptOrphanedGrant(heldBefore, request)
-      case NonFatal(e) =>
-        warnOnMemoryManagerFailure(e)
-        adoptOrphanedGrant(heldBefore, request)
-      case e: SparkOutOfMemoryError =>
-        warnOnMemoryManagerFailure(e)
-        adoptOrphanedGrant(heldBefore, request)
+        // Refusing is the only way out of an Arrow callback, but the cancellation must not be
+        // lost with it: re-arm the flag so the task still observes it.
+        Thread.currentThread().interrupt()
+        throw refusedAfterFailure(heldBefore, request, size, e)
+      case NonFatal(e) => throw refusedAfterFailure(heldBefore, request, size, e)
+      case e: SparkOutOfMemoryError => throw refusedAfterFailure(heldBefore, request, size, e)
     }
   }
 
-  private def adoptOrphanedGrant(heldBefore: Long, request: Long): Long = {
-    val orphaned = taskMemoryManager.getMemoryConsumptionForThisTask - heldBefore
-    math.max(0L, math.min(orphaned, request))
+  private def refusedAfterFailure(
+      heldBefore: Long,
+      request: Long,
+      size: Long,
+      cause: Throwable): OutOfMemoryException = {
+    val orphaned = math.max(
+      0L,
+      math.min(taskMemoryManager.getMemoryConsumptionForThisTask - heldBefore, request))
+    if (orphaned > 0L) {
+      taskMemoryManager.releaseExecutionMemory(orphaned, this)
+    }
+    refusal(size, request, 0L, cause)
   }
 
-  /**
-   * An interrupt cannot be allowed out of an Arrow callback any more than anything else can, but
-   * swallowing the cancellation would be wrong too. Spark's execution pool parks in `lock.wait()`
-   * when a task is below its fair share, so killing a task lands here, and `NonFatal`
-   * deliberately excludes `InterruptedException`. Re-arming the flag leaves the cancellation for
-   * the task to observe at its next interruptible point, which is the only place it can act on it
-   * anyway.
-   */
-  private def reportAndReinterrupt(e: InterruptedException): Unit = {
-    Thread.currentThread().interrupt()
-    warnOnMemoryManagerFailure(e)
+  private def refusal(
+      size: Long,
+      needed: Long,
+      granted: Long,
+      cause: Throwable): OutOfMemoryException = {
+    // As Spark's own consumers do before failing, so the task's consumers are in the log.
+    taskMemoryManager.showMemoryUsage()
+    new OutOfMemoryException(
+      s"Unable to reserve $needed bytes of Spark off-heap execution memory for a JVM Arrow " +
+        s"allocation of $size bytes, got $granted. Increase spark.memory.offHeap.size, or set " +
+        s"${CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key}=false to stop charging " +
+        "JVM Arrow allocations to Spark.",
+      cause)
   }
 }
 
@@ -273,30 +320,16 @@ object CometArrowAllocationListener extends Logging {
 
   private val NothingOwned: LongSupplier = () => 0L
 
-  private val shortGrantLogged = new AtomicBoolean(false)
-  private val memoryManagerFailureLogged = new AtomicBoolean(false)
+  private val releaseFailureLogged = new AtomicBoolean(false)
 
   private def roundUpToBlock(bytes: Long): Long =
     ((bytes + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE
 
-  private def warnOnShortGrant(requested: Long, granted: Long): Unit = {
-    if (shortGrantLogged.compareAndSet(false, true)) {
-      logWarning(
-        s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " +
-          "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " +
-          s"Set ${CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key}=false to stop " +
-          "reporting these allocations to Spark.")
-    }
-  }
-
-  private def warnOnMemoryManagerFailure(e: Throwable): Unit = {
-    if (memoryManagerFailureLogged.compareAndSet(false, true)) {
+  private def warnOnReleaseFailure(e: Throwable): Unit = {
+    if (releaseFailureLogged.compareAndSet(false, true)) {
       logWarning(
-        "Failed to report a JVM Arrow allocation to Spark's memory manager. The allocation " +
-          "itself is unaffected, so this is a reporting gap rather than a failure, but Spark's " +
-          "view of these bytes will be short until the task ends. " +
-          s"Set ${CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key}=false to stop " +
-          "reporting these allocations to Spark.",
+        "Failed to return JVM Arrow memory to Spark's memory manager. The task keeps it " +
+          "reserved until it ends, when Spark releases everything the task still holds.",
         e)
     }
   }
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
index 13578f96ae1..03af31b8298 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
@@ -34,13 +34,14 @@ import org.apache.comet.{CometArrowAllocator, CometConf}
  * Hands out the Arrow allocator that JVM-owned allocations should use, one per Spark task.
  *
  * Each task gets a child of `CometArrowAllocator` carrying its own
- * [[CometArrowAllocationListener]], so what that allocator owns is reported to the task's
- * `TaskMemoryManager`. The allocator, not the calling thread, is what identifies the owner:
- * Arrow's `AllocationListener` is given only a size, and a buffer is released on whichever thread
- * drops the last reference, which for anything that reaches native over the C Data Interface is a
- * Comet Tokio worker with no task context installed. Child allocators cut from the returned
- * allocator inherit its listener and roll their bytes up into it, so the paths that make their
- * own children are covered too.
+ * [[CometArrowAllocationListener]], so what that allocator owns is charged to the task's
+ * `TaskMemoryManager`, and an allocation Spark cannot cover fails with an `OutOfMemoryException`
+ * from `buffer`. The allocator, not the calling thread, is what identifies the owner: Arrow's
+ * `AllocationListener` is given only a size, and a buffer is released on whichever thread drops
+ * the last reference, which for anything that reaches native over the C Data Interface is a Comet
+ * Tokio worker with no task context installed. Child allocators cut from the returned allocator
+ * inherit its listener and roll their bytes up into it, so the paths that make their own children
+ * are covered too.
  *
  * '''This is for buffers the JVM owns.''' Buffers allocated to be handed straight to native -- by
  * `NativeUtil`, for the JVM UDF result, and in `CometArrowStream.stream` -- come from the
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
index 6e79b6ea47d..831ca6a6250 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
@@ -29,7 +29,7 @@ import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, TaskMemoryManager, TestMemoryManager}
 
 /**
- * Measures what reporting JVM Arrow allocations to Spark costs, since it is on by default.
+ * Measures what charging JVM Arrow allocations to Spark costs, since it is on by default.
  *
  * Each case runs the same allocate/release loop twice: once against a plain `RootAllocator`,
  * which is what Comet did before [[CometArrowAllocationListener]] existed and what
@@ -116,23 +116,25 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
 
     val root = new RootAllocator(Long.MaxValue)
     try {
-      // Two blocks of pool for both arms, with the pressure thread running throughout. In the
-      // unaccounted arm only that thread touches the pool, which is the point: the delta is what
-      // the Arrow side adds once it competes for the same budget.
-      withTaskAllocator(poolBytes = blockSize * 2) { (accounted, memory) =>
-        withNativePressure(memory) {
-          benchmark.addCase("not accounted") { _ =>
-            churn(root, blockSize, buffersPerIteration)
-          }
-          benchmark.addCase("accounted") { _ =>
+      // Room for the whole set plus the pressure thread's block, which runs throughout. A smaller
+      // pool would refuse the allocations rather than slow them down, so what this measures is
+      // contention for the pool's lock rather than for its capacity. In the unaccounted arm only
+      // the pressure thread touches the pool: the delta is what the Arrow side adds by competing.
+      withTaskAllocator(poolBytes = blockSize * (buffersPerIteration + 1)) {
+        (accounted, memory) =>
+          withNativePressure(memory) {
+            benchmark.addCase("not accounted") { _ =>
+              churn(root, blockSize, buffersPerIteration)
+            }
+            benchmark.addCase("accounted") { _ =>
+              churn(accounted, blockSize, buffersPerIteration)
+            }
+            benchmark.run()
+
+            memory.reset()
             churn(accounted, blockSize, buffersPerIteration)
+            writeLine(s"  accounted: ${memory.summary(buffersPerIteration)} per iteration")
           }
-          benchmark.run()
-
-          memory.reset()
-          churn(accounted, blockSize, buffersPerIteration)
-          writeLine(s"  accounted: ${memory.summary(buffersPerIteration)} per iteration")
-        }
       }
     } finally {
       root.close()
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index 61b292ef994..9e513c06f2c 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -26,7 +26,7 @@ import java.util.concurrent.atomic.{AtomicLong, AtomicReference}
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.arrow.c.Data
-import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.memory.{BufferAllocator, OutOfMemoryException}
 import org.apache.arrow.vector.{FieldVector, IntVector, VectorSchemaRoot}
 import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
 import org.apache.spark.executor.TaskMetrics
@@ -35,15 +35,15 @@ import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
 import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator}
+import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator, CometConf}
 import org.apache.comet.vector.NativeUtil
 
 /**
- * Tests that JVM Arrow allocations are reported to Spark, that they are reported against the task
- * that made them rather than whichever task happens to be on the releasing thread, that what is
- * charged follows Arrow's ownership as buffers move between allocators, and, just as importantly,
- * that the paths where they cannot be reported fail quietly rather than throwing. Arrow
- * allocation on these paths cannot fail today and this listener must not change that.
+ * Tests that JVM Arrow allocations are charged to Spark, that they are charged to the task that
+ * made them rather than whichever task happens to be on the releasing thread, that what is
+ * charged follows Arrow's ownership as buffers move between allocators, and that an allocation
+ * Spark cannot cover is refused before Arrow makes it, leaving no buffer, reservation or grant
+ * behind.
  */
 class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
@@ -377,63 +377,119 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   }
 
   // ---------------------------------------------------------------------------------------------
-  // Failure containment: Spark's acquisition is fallible, Arrow's callbacks are not allowed to be.
+  // Refusal. An allocation Spark cannot cover is refused in onPreAllocation, before Arrow has made
+  // it, and nothing is left behind: no buffer, no reservation, and no grant charged to the task.
   // ---------------------------------------------------------------------------------------------
 
+  test("an allocation the pool cannot cover is refused, and one that fits still succeeds") {
+    // Two blocks of budget, one of them held by a consumer that cannot give anything back.
+    withTask(pool = blockSize * 2) { task =>
+      val other = new PlainConsumer(task.taskMemoryManager)
+      assert(other.take(blockSize) == blockSize)
+
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val refused = intercept[OutOfMemoryException](allocator.buffer(blockSize * 2))
+      assert(refused.getMessage.contains(CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key))
+      // Spark granted the block it had before coming up short. That partial grant is handed back
+      // rather than kept, and Arrow never allocated anything.
+      assert(allocator.getAllocatedMemory == 0L)
+      assert(reservedFor(task) == 0L)
+      assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize)
+
+      val buf = allocator.buffer(blockSize)
+      try {
+        assert(reservedFor(task) == blockSize)
+      } finally {
+        buf.close()
+      }
+    }
+  }
+
+  test("a grant short of a block is accepted when it covers the allocation") {
+    // Reservations are requested in whole blocks to save lock traffic, but a pool with less than a
+    // block left must still admit an allocation that fits in what is left.
+    val left = 64L * 1024
+    withTask(pool = blockSize) { task =>
+      val other = new PlainConsumer(task.taskMemoryManager)
+      assert(other.take(blockSize - left) == blockSize - left)
+
+      val allocator = CometTaskArrowAllocator.forCurrentTask()
+      val buf = allocator.buffer(1024L)
+      try {
+        assert(reservedFor(task) == left)
+        // What does not fit in what is left is still refused.
+        intercept[OutOfMemoryException](allocator.buffer(left))
+        assert(reservedFor(task) == left)
+      } finally {
+        buf.close()
+      }
+      // A reservation that is not a whole block is returned once nothing needs it.
+      assert(reservedFor(task) == 0L)
+    }
+  }
+
+  test("an allocation still in flight counts against the reservation") {
+    withTask(pool = blockSize) { task =>
+      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
+      val owner = new StubOwner(listener)
+
+      // Admitted, but Arrow has not reported it yet...
+      listener.onPreAllocation(blockSize)
+      assert(listener.reservedBytes == blockSize)
+      // ...so a second allocation racing it on another thread cannot be admitted against the same
+      // block.
+      intercept[OutOfMemoryException](listener.onPreAllocation(blockSize))
+      assert(listener.reservedBytes == blockSize)
+
+      owner.made(blockSize)
+      owner.release(blockSize)
+      assert(listener.reservedBytes == 0L)
+      listener.taskCompleted()
+    }
+  }
+
   for ((label, failure) <- Seq(
       "an I/O failure" -> new IOException("spill failed"),
-      "an interrupt" -> new InterruptedIOException("task killed"))) {
-    test(s"$label while spilling does not fail or leak the Arrow allocation") {
-      // Exactly one block of budget, already taken by a consumer that refuses to spill, so the
-      // listener's acquisition has to go through Spark's spill path and comes back throwing.
+      "an interrupted spill" -> new InterruptedIOException("task killed"))) {
+    test(s"$label while acquiring refuses the allocation and leaks nothing") {
+      // Exactly one block of budget, already taken by a consumer whose spill throws, so the
+      // acquisition has to go through Spark's spill path and comes back throwing.
       withTask(pool = blockSize) { task =>
         val hostile = new FailingSpillConsumer(task.taskMemoryManager, failure)
         assert(hostile.take(blockSize) == blockSize)
 
         val allocator = CometTaskArrowAllocator.forCurrentTask()
-        // `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so
-        // throwing from the listener would lose this buffer: Arrow neither returns nor frees it.
-        val buf = allocator.buffer(blockSize)
-        try {
-          assert(allocator.getAllocatedMemory == blockSize)
-          // Nothing was reserved, which is what says the acquisition really did go down the spill
-          // path and throw rather than quietly succeeding and making this test vacuous.
-          assert(reservedFor(task) == 0L)
-        } finally {
-          buf.close()
-        }
-        // Zero here is the leak check: a buffer Arrow created but never handed back would still
-        // be counted.
+        val refused = intercept[OutOfMemoryException](allocator.buffer(blockSize))
+        // The memory manager's failure is the cause, which is what says the acquisition really
+        // went down the spill path rather than being refused for want of memory alone.
+        assert(refused.getCause != null)
         assert(allocator.getAllocatedMemory == 0L)
+        assert(reservedFor(task) == 0L)
+        assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize)
       }
     }
   }
 
-  test("an interrupt while spilling is re-armed rather than thrown or swallowed") {
+  test("an interrupt while acquiring refuses the allocation and leaves the flag set") {
     // Spark's execution pool parks in `lock.wait()` when a task is below its fair share, so a task
-    // kill raises a plain InterruptedException out of `acquireExecutionMemory`. `NonFatal` excludes
-    // it, so before this it escaped `onAllocation` and Arrow lost the buffer it had just created.
-    // TestMemoryManager never parks, so the interrupt is injected through a failing spill instead.
+    // kill raises a plain InterruptedException out of `acquireExecutionMemory`. TestMemoryManager
+    // never parks, so the interrupt is injected through a failing spill instead.
     withTask(pool = blockSize) { task =>
       val hostile =
         new FailingSpillConsumer(task.taskMemoryManager, new InterruptedException("task killed"))
       assert(hostile.take(blockSize) == blockSize)
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
-      val buf = allocator.buffer(blockSize)
-      try {
-        assert(allocator.getAllocatedMemory == blockSize)
-        assert(reservedFor(task) == 0L)
-      } finally {
-        buf.close()
-      }
+      val refused = intercept[OutOfMemoryException](allocator.buffer(blockSize))
+      assert(refused.getCause.isInstanceOf[InterruptedException])
       assert(allocator.getAllocatedMemory == 0L)
+      assert(reservedFor(task) == 0L)
       // Cleared here as well as asserted, so the flag does not leak into the next test.
       assert(Thread.interrupted(), "the interrupt was swallowed instead of being re-armed")
     }
   }
 
-  test("a partial grant lost to a failing spill is adopted rather than stranded") {
+  test("a partial grant lost to a failing spill is released rather than stranded") {
     // One block already taken, one still in the pool, and a two-block request: Spark hands over the
     // block it has and only then asks the other consumer to spill, which throws. It never reports
     // the block it already took, so nothing would release it before the task ended.
@@ -443,26 +499,20 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       assert(hostile.take(blockSize) == blockSize)
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
-      val buf = allocator.buffer(blockSize * 2)
-      try {
-        assert(reservedFor(task) == blockSize)
-        assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize * 2)
-      } finally {
-        buf.close()
-      }
+      intercept[OutOfMemoryException](allocator.buffer(blockSize * 2))
       assert(reservedFor(task) == 0L)
-      // Only the other consumer's block is left. Without adopting the orphan this would still be
-      // two blocks, with one of them charged to the task and owned by nobody.
+      // Only the other consumer's block is left. Without releasing the orphan this would be two
+      // blocks, one of them charged to the task and owned by nobody.
       assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize)
     }
   }
 
-  test("a concurrent acquisition is not adopted as this listener's lost grant") {
+  test("a concurrent acquisition is not released as this listener's lost grant") {
     // The orphan is measured as a change in the task's total consumption, which Spark reports per
     // task rather than per consumer. If another consumer could acquire between the snapshot taken
-    // before the acquisition and the acquisition itself, its bytes would be adopted here and handed
-    // back when this listener next shrank, leaving two consumers holding the same bytes between
-    // them. Both snapshots and the call therefore run as one transaction under the
+    // before the acquisition and the acquisition itself, its bytes would be counted as the orphan
+    // and released here, leaving that consumer holding bytes the pool no longer charges for. Both
+    // snapshots and the call therefore run as one transaction under the
     // TaskMemoryManager monitor. This forces the interleaving that transaction exists to exclude.
     val spare = 1024L
     val task = newTask(pool = blockSize * 2 + spare)
@@ -489,22 +539,18 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       })
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
-      val buf = allocator.buffer(blockSize * 2)
-      try {
-        interloperThread.join(30000L)
-        assert(
-          !interloperThread.isAlive,
-          "the interloper never finished; the transaction deadlocked")
-        // Nothing is claimed twice: what the listener adopted has to fit alongside what the other
-        // two consumers hold. Without the transaction the listener adopts the interloper's bytes on
-        // top of its own grant, and this sum comes out over what the task actually holds.
-        assert(
-          listenerFor(task).reservedBytes + hostile.getUsed + interloper.getUsed ==
-            task.taskMemoryManager.getMemoryConsumptionForThisTask,
-          "the listener adopted bytes belonging to another consumer")
-      } finally {
-        buf.close()
-      }
+      intercept[OutOfMemoryException](allocator.buffer(blockSize * 2))
+      interloperThread.join(30000L)
+      assert(
+        !interloperThread.isAlive,
+        "the interloper never finished; the transaction deadlocked")
+      // Every byte the task is charged for belongs to someone. Without the transaction the
+      // listener counts the interloper's bytes as its own orphan and releases them, and the task
+      // comes out charged for less than the other two consumers hold.
+      assert(
+        listenerFor(task).reservedBytes + hostile.getUsed + interloper.getUsed ==
+          task.taskMemoryManager.getMemoryConsumptionForThisTask,
+        "the listener released bytes belonging to another consumer")
     }
   }
 
@@ -581,7 +627,14 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     private val owned = new AtomicLong(0L)
     listener.bind(() => owned.get)
 
+    /** What `BaseAllocator.buffer` does around a successful allocation. */
     def allocate(bytes: Long): Unit = {
+      listener.onPreAllocation(bytes)
+      made(bytes)
+    }
+
+    /** The second half of an allocation that `onPreAllocation` has already admitted. */
+    def made(bytes: Long): Unit = {
       owned.addAndGet(bytes)
       listener.onAllocation(bytes)
     }

From 1f1e24eb3e58ce6ac2fc647c5ee364af98aff2f7 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sat, 26 Sep 2026 11:00:52 -0600
Subject: [PATCH 17/18] feat: replace the accounting switch with
 spark.comet.legacy.unboundedJvmArrowMemory

Charging JVM Arrow memory to Spark, and refusing an allocation the pool
cannot cover, is a behavior change under the versioning policy: a query
that ran on 1.1.0 can now fail with an error it did not raise before. The
policy's escape hatch for that is a spark.comet.legacy.* key that
defaults to false and restores the previous behavior, so that key
replaces spark.comet.memory.jvmArrowAccounting.enabled rather than
shipping next to it. The old key was never released, so it needs no
alias, and there is no permanent off switch for the next major release
to carry.

This is the first legacy key, so it also adds CATEGORY_LEGACY and a
Legacy Settings section to configs.md, as the config conventions ask,
and a section to the upgrade guide describing the change for 1.2.0.

The key is now read from the SparkConf when a task first needs its
allocator, not once per JVM. It is still an executor setting fixed at
submit time. But the value now belongs to whichever SparkContext is
running, so a test can exercise both branches. The new test does that
by running a task in a local context with and without the key, and it
fails if the allocator ignores the key.
---
 .../contributor-guide/memory_management.md    |  6 ++---
 docs/source/user-guide/latest/configs.md      | 10 +++++++
 .../user-guide/latest/migration-guide.md      | 19 +++++++++++++
 docs/source/user-guide/latest/tuning.md       |  5 ++--
 .../scala/org/apache/comet/CometConf.scala    | 26 +++++++++---------
 .../comet/CometArrowAllocationListener.scala  |  4 +--
 .../spark/comet/CometTaskArrowAllocator.scala | 15 ++++++-----
 ...ometArrowAllocationListenerBenchmark.scala |  8 +++---
 .../CometArrowAllocationListenerSuite.scala   | 27 +++++++++++++++++--
 9 files changed, 88 insertions(+), 32 deletions(-)

diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index cbed07ad918..2a20f2db3b3 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -120,9 +120,9 @@ Comet's native pool among them. The reservation is made before Arrow allocates,
 cover the allocation, any partial grant is handed back and the allocation is refused with Arrow's
 `OutOfMemoryException`. That is `try_grow` rather than `grow`: nothing is allocated that is not
 covered, and there is no overcommit. Nothing can spill these buffers, so a refusal fails the
-allocating task, the way a Spark consumer fails when it cannot acquire memory. Set
-`spark.comet.memory.jvmArrowAccounting.enabled=false` to stop charging them, in which case
-`forCurrentTask()` returns the root. It also returns the root in on-heap mode, and off a task, as
+allocating task, the way a Spark consumer fails when it cannot acquire memory. The deprecated
+`spark.comet.legacy.unboundedJvmArrowMemory=true` restores the earlier behavior of charging
+nothing, in which case `forCurrentTask()` returns the root. It also returns the root in on-heap mode, and off a task, as
 for broadcast coalescing on the driver or for a reader built on a native thread that is pulling a
 stream.
 
diff --git a/docs/source/user-guide/latest/configs.md b/docs/source/user-guide/latest/configs.md
index 59d5c711faf..16267761d60 100644
--- a/docs/source/user-guide/latest/configs.md
+++ b/docs/source/user-guide/latest/configs.md
@@ -72,6 +72,16 @@ shuffle, including applications that rely only on remote reliable storage or dec
 <!--BEGIN:CONFIG_TABLE[tuning]-->
 <!--END:CONFIG_TABLE-->
 
+## Legacy Settings
+
+Each of these settings restores a behavior that an earlier release had, after a release changed it.
+They are deprecated from the moment they are added, and may be removed in any future major release,
+at which point the newer behavior becomes unconditional. The
+[Upgrade Guide](migration-guide.md) describes each change and the setting that reverts it.
+
+<!--BEGIN:CONFIG_TABLE[legacy]-->
+<!--END:CONFIG_TABLE-->
+
 ## Development & Testing Settings
 
 These settings exist for Comet's own test suites and for debugging. They are **not covered by the
diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md
index 1de5f884672..8248b61586b 100644
--- a/docs/source/user-guide/latest/migration-guide.md
+++ b/docs/source/user-guide/latest/migration-guide.md
@@ -53,6 +53,25 @@ Treat setting one of these keys as a temporary measure. If you find you cannot s
 legacy behavior, please open an issue describing your use case so it can be considered before the
 key is removed.
 
+## Upgrading to Comet 1.2.0
+
+### JVM-side Arrow memory is charged to Spark's off-heap memory pool
+
+The Arrow memory that Comet allocates on the JVM side for a task, for example to read a cached table
+or to exchange batches with a Python worker, is now charged to `spark.memory.offHeap.size` as a
+memory consumer of that task. An allocation the pool cannot cover fails the task with an
+`OutOfMemoryException` whose message begins `Unable to reserve`, the way a Spark operator fails when
+it cannot acquire memory. The charge also leaves Comet's native operators and Spark's own consumers
+less of the pool, so queries can spill sooner than before. Earlier releases charged this memory to no
+pool, so it never failed an allocation and had to fit in `spark.executor.memoryOverhead` instead.
+
+Arrow memory that is handed to Comet's native code, or received from it, is not charged on the JVM
+side, because native accounts for what it retains.
+
+To restore the previous behavior, set `spark.comet.legacy.unboundedJvmArrowMemory=true` when the
+application is submitted. It is read on the executors, so setting it in a running session has no
+effect.
+
 ## Upgrading to Comet 1.1.0
 
 Comet `1.1.0` makes no behavior changes that need a `spark.comet.legacy.*` key.
diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md
index f6b88d2f425..2af36f9569b 100644
--- a/docs/source/user-guide/latest/tuning.md
+++ b/docs/source/user-guide/latest/tuning.md
@@ -122,8 +122,9 @@ cannot cover fails the task that makes it, the way a Spark operator fails when i
 leaves Comet's native operators and Spark's own consumers less of the pool, so they can spill sooner. Comet 1.1.0 and
 earlier did not charge this memory to the pool, so if you sized `spark.memory.offHeap.size` against one of those
 releases, check how often queries spill, and raise it if tasks fail with `Unable to reserve ... for a JVM Arrow
-allocation`. Setting `spark.comet.memory.jvmArrowAccounting.enabled=false` stops charging
-it, in which case it has to fit in `spark.executor.memoryOverhead` instead.
+allocation`. Setting the deprecated `spark.comet.legacy.unboundedJvmArrowMemory=true` restores the earlier behavior
+while you adjust, in which case this memory has to fit in `spark.executor.memoryOverhead` instead; see the
+[Upgrade Guide](migration-guide.md).
 
 For more details about Spark off-heap memory mode, please refer to [Spark documentation].
 
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index e51282c3c08..13e846491ca 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -68,6 +68,7 @@ object CometConf extends ShimCometConf {
   private val CATEGORY_SHUFFLE = "shuffle"
   private val CATEGORY_TUNING = "tuning"
   private val CATEGORY_TESTING = "testing"
+  private val CATEGORY_LEGACY = "legacy"
 
   def register(conf: ConfigEntry[_]): Unit = {
     assert(conf.category.nonEmpty, s"${conf.key} does not have a category defined")
@@ -374,21 +375,20 @@ object CometConf extends ShimCometConf {
     .checkValue(_ >= 0, "The memory usage log interval must not be negative")
     .createWithDefault(TimeUnit.SECONDS.toMillis(10))
 
-  val COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED: ConfigEntry[Boolean] =
-    conf("spark.comet.memory.jvmArrowAccounting.enabled")
-      .category(CATEGORY_TUNING)
+  val COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY: ConfigEntry[Boolean] =
+    conf("spark.comet.legacy.unboundedJvmArrowMemory")
+      .category(CATEGORY_LEGACY)
       .doc(
-        "Whether to charge the Arrow memory that Comet allocates on the JVM side for a task to " +
-          "Spark's off-heap memory pool, as a memory consumer of that task. An allocation the " +
-          "pool cannot cover fails the task that makes it, and the charge leaves Comet's native " +
-          "operators and Spark's other consumers correspondingly less of the pool, so they may " +
-          "spill sooner. Arrow memory that is " +
-          "handed to Comet's native code, or received from it, is not charged here, because it " +
-          "is native's to account for. This is an executor setting, read when an executor " +
-          "first allocates Arrow memory for a task, so it must be set when the application is " +
-          s"submitted. Set to false to stop charging these allocations. $TUNING_GUIDE.")
+        "When true, restores the pre-1.2.0 behavior of the Arrow memory that Comet " +
+          "allocates on the JVM side for a task, which was charged to no memory pool and so " +
+          "could never fail an allocation. Since 1.2.0 it is charged to Spark's off-heap " +
+          "memory pool as a memory consumer of the task, and an allocation the pool cannot " +
+          "cover fails the task. Arrow memory handed to Comet's native code, or received from " +
+          "it, is not affected. This is an executor setting, read when an executor first " +
+          "allocates Arrow memory for a task, so it must be set when the application is " +
+          "submitted. This config is deprecated and will be removed in a future major release.")
       .booleanConf
-      .createWithDefault(true)
+      .createWithDefault(false)
 
   val COMET_SHUFFLE_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.shuffle.enabled")
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index 5b826a5ab0b..38426601382 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -302,8 +302,8 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
     new OutOfMemoryException(
       s"Unable to reserve $needed bytes of Spark off-heap execution memory for a JVM Arrow " +
         s"allocation of $size bytes, got $granted. Increase spark.memory.offHeap.size, or set " +
-        s"${CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key}=false to stop charging " +
-        "JVM Arrow allocations to Spark.",
+        s"${CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key}=true to restore the previous " +
+        "behavior of not charging JVM Arrow allocations to Spark.",
       cause)
   }
 }
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
index 03af31b8298..4378e864d3a 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
@@ -77,14 +77,17 @@ object CometTaskArrowAllocator extends Logging {
   private val closeLock = new Object
 
   /**
-   * Resolved once per JVM. Read from the `SparkConf` rather than `SQLConf`, because this is
-   * reached from executor threads where `SQLConf` does not carry Comet's settings. The `Option`
-   * guard covers tests that install a task context without a `SparkEnv`.
+   * Whether tasks get an accounted allocator, which the deprecated
+   * `spark.comet.legacy.unboundedJvmArrowMemory` turns off. Read from the `SparkConf` rather than
+   * `SQLConf`, because this is reached from executor threads where `SQLConf` does not carry
+   * Comet's settings, and when a task first needs its allocator rather than once per JVM, so that
+   * the value belongs to whichever `SparkContext` is running. Tests that install a task context
+   * by hand have no `SparkEnv`, and account as the default does.
    */
-  private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env =>
+  private def accountingEnabled: Boolean = !Option(SparkEnv.get).exists { env =>
     env.conf.getBoolean(
-      CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key,
-      CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.defaultValue.get)
+      CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key,
+      CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.defaultValue.get)
   }
 
   /**
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
index 831ca6a6250..187cb9c6060 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
@@ -33,10 +33,10 @@ import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, TaskM
  *
  * Each case runs the same allocate/release loop twice: once against a plain `RootAllocator`,
  * which is what Comet did before [[CometArrowAllocationListener]] existed and what
- * `spark.comet.memory.jvmArrowAccounting.enabled=false` restores, and once against a task
- * allocator carrying the listener. Reservation call counts are printed under each table, because
- * the interesting variable is not the per-buffer bookkeeping but how often a buffer size crosses
- * a block boundary and has to go into `TaskMemoryManager` at all.
+ * `spark.comet.legacy.unboundedJvmArrowMemory=true` restores, and once against a task allocator
+ * carrying the listener. Reservation call counts are printed under each table, because the
+ * interesting variable is not the per-buffer bookkeeping but how often a buffer size crosses a
+ * block boundary and has to go into `TaskMemoryManager` at all.
  *
  * To run this benchmark:
  * {{{
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index 9e513c06f2c..dc69aab39e5 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -28,7 +28,7 @@ import org.scalatest.funsuite.AnyFunSuite
 import org.apache.arrow.c.Data
 import org.apache.arrow.memory.{BufferAllocator, OutOfMemoryException}
 import org.apache.arrow.vector.{FieldVector, IntVector, VectorSchemaRoot}
-import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
+import org.apache.spark.{SparkConf, SparkContext, TaskContext, TaskContextImpl}
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager}
 import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
@@ -173,6 +173,29 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
+  test("the legacy setting hands tasks the unaccounted root, and is off by default") {
+    // Read from the executors' SparkConf, so each arm runs a task in a context of its own.
+    def taskGetsRoot(legacy: Option[String]): Boolean = {
+      val conf = new SparkConf()
+        .setMaster("local[1]")
+        .setAppName(getClass.getSimpleName)
+        .set("spark.memory.offHeap.enabled", "true")
+        .set("spark.memory.offHeap.size", poolBytes.toString)
+      legacy.foreach(conf.set(CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key, _))
+      val sc = new SparkContext(conf)
+      try {
+        sc.parallelize(Seq(0), 1)
+          .map(_ => CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator)
+          .collect()
+          .head
+      } finally {
+        sc.stop()
+      }
+    }
+    assert(!taskGetsRoot(None), "a task was not accounted by default")
+    assert(taskGetsRoot(Some("true")), "the legacy setting did not restore the unaccounted root")
+  }
+
   // ---------------------------------------------------------------------------------------------
   // Ownership: a release is attributed to the task that allocated, not to the releasing thread.
   // ---------------------------------------------------------------------------------------------
@@ -389,7 +412,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
       val refused = intercept[OutOfMemoryException](allocator.buffer(blockSize * 2))
-      assert(refused.getMessage.contains(CometConf.COMET_MEMORY_JVM_ARROW_ACCOUNTING_ENABLED.key))
+      assert(refused.getMessage.contains(CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key))
       // Spark granted the block it had before coming up short. That partial grant is handed back
       // rather than kept, and Arrow never allocated anything.
       assert(allocator.getAllocatedMemory == 0L)

From 62b1806b57df27191cb2127737d0d7715b8cbf65 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sat, 26 Sep 2026 11:48:34 -0600
Subject: [PATCH 18/18] refactor: simplify the JVM Arrow listener and its tests

A cleanup pass over the reuse, simplification, efficiency and altitude
of this change.

In the listener, onAllocation now trims the reservation as well as
settling it, as every other callback does. Since the switch to
onPreAllocation it only settled, so a transfer out of a task allocator
was corrected only by an explicit reconcile or a later release, although
both scaladocs said the next allocation too. Now any such transfer
corrects itself at the task's next allocation or release, and the
reader's reconcile only makes it immediate. The ownership test uses a
four-block source, which fails at four blocks held if onAllocation stops
trimming.

Other listener changes:
- Growth rounds the target, owned plus pending plus the new allocation,
  up to a block, rather than rounding the shortfall. A reservation that
  is not block aligned then no longer acquires a block and trims part of
  it on the same allocation.
- settle, task completion and the trim share one releaseDownTo, and the
  guard in settle that did nothing is gone.
- The three catch branches in acquire, which were identical, are one,
  and the helper that only they called is inlined.

In CometTaskArrowAllocator:
- The per-batch reconcile moves from ColumnarBatchArrowReader into
  closeAndReconcile. It walks the columns without allocating and
  reconciles the batch's task allocator once, rather than once per
  distinct allocator.
- A finished allocator is parked and then closed if drained, instead of
  being tried twice.
- The cheaper TaskMemoryManager checks now run before the SparkConf
  read.
- The two visible-for-testing counters no test used are deleted, and so
  is the reservedBytesForTask wrapper.
- The rationale for keeping native-bound memory on the root is kept in
  one place, not repeated at each call site.

The suite and the benchmark share TestTasks, which builds and installs a
hand-made task. They replace their copied native-style consumer with the
real CometTaskMemoryManager, and the suite uses Spark's
ThreadUtils.runInNewThread for its detached release. The concurrency
test now asks for the whole pool on the native side, so that Spark does
walk its consumers during short grants, as its comment always said. A
refused Arrow allocation there is a legitimate outcome. The benchmark's
native-pressure case is the same method as the others, with two extra
parameters.
---
 .../main/scala/org/apache/comet/package.scala |  11 +-
 .../comet/CometArrowAllocationListener.scala  | 133 +++++-----
 .../spark/comet/CometTaskArrowAllocator.scala |  85 +++---
 .../arrow/ColumnarBatchArrowReader.scala      |  14 +-
 .../arrow/CometNativeArrowSource.scala        |   6 +-
 ...ometArrowAllocationListenerBenchmark.scala | 194 ++++----------
 .../CometArrowAllocationListenerSuite.scala   | 250 ++++++------------
 .../org/apache/spark/comet/TestTasks.scala    |  89 +++++++
 8 files changed, 337 insertions(+), 445 deletions(-)
 create mode 100644 spark/src/test/scala/org/apache/spark/comet/TestTasks.scala

diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala
index 8c050cec756..0007f581cb7 100644
--- a/spark/src/main/scala/org/apache/comet/package.scala
+++ b/spark/src/main/scala/org/apache/comet/package.scala
@@ -34,13 +34,10 @@ package object comet {
    * leaked. To avoid this, we use a single allocator for the whole execution process.
    *
    * It carries no allocation listener, so neither it nor a child cut with the three-argument
-   * `newChildAllocator` reports anything to Spark's memory manager. That is what buffers
-   * allocated to be handed to native want: a native operator that retains a batch reserves its
-   * buffers through Comet's pool, which charges the same Spark task, so reporting them here as
-   * well would reserve the same memory twice. JVM-owned allocations should go through
-   * `CometTaskArrowAllocator.forCurrentTask()` instead, which cuts a per-task child whose
-   * listener charges what that child owns to Spark and refuses an allocation Spark cannot cover,
-   * and which hands back this allocator when there is no task.
+   * `newChildAllocator` charges anything to Spark. That suits buffers handed to native, which
+   * accounts for what it retains. JVM-owned allocations within a task go through
+   * `CometTaskArrowAllocator.forCurrentTask()` instead, which describes which allocations go
+   * where.
    */
   val CometArrowAllocator = new RootAllocator(Long.MaxValue)
 
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
index 38426601382..abc44af7d24 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala
@@ -34,10 +34,8 @@ import org.apache.comet.CometConf
  * Charges one task's JVM-side Arrow allocations to Spark's off-heap execution pool, and refuses
  * an allocation the pool cannot cover.
  *
- * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, so until now the
- * off-heap bytes it hands out were counted by nobody: not Spark's `TaskMemoryManager`, and not
- * Comet's native memory pool. They are still resident in the container, which makes them a blind
- * spot when an executor is killed for exceeding its memory limit. Here they appear in
+ * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, but the bytes a task's
+ * allocator hands out are resident in the container all the same. Charged here, they appear in
  * `TaskMemoryManager.showMemoryUsage`, are arbitrated against Spark's other off-heap consumers,
  * and are bounded by the same pool.
  *
@@ -69,21 +67,17 @@ import org.apache.comet.CometConf
  * child of the unaccounted root, and then closes the source. Summing the callbacks would charge
  * the task for each such batch until the task ended. So the callbacks only say when to look, and
  * the size comes from the task allocator's accountant, which Arrow keeps right across transfers.
- * `ColumnarBatchArrowReader` calls [[reconcile]] as soon as it closes a source, so a task stops
- * paying for a batch once native has it. Any other transfer out is reflected at the next
- * allocation or release on the task's allocator. A transfer in is charged at the next allocation,
- * which is refused if the pool cannot cover both.
+ * Every callback trims the reservation to what is owned and pending, so a transfer out is
+ * reflected at the next allocation or release on the task's allocator, and the reader calls
+ * [[reconcile]] as soon as it closes a source, so that a task stops paying for a batch once
+ * native has it. A transfer in is charged at the next allocation, which is refused if the pool
+ * cannot cover both.
  *
  * '''Only [[onPreAllocation]] may throw.''' Arrow documents that the other callbacks cannot, and
  * `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so
- * throwing from any of them loses a buffer Arrow has already created. They only ever shrink the
- * reservation, and contain anything the memory manager throws. The acquisition in
- * [[onPreAllocation]] is fallible in three ways, and only the first is caught by `NonFatal`: it
- * runs other consumers' `spill`, which turns a task interrupt into a `RuntimeException` and an
- * I/O failure into a `SparkOutOfMemoryError`, and the execution pool itself parks in
- * `lock.wait()`, so killing a task can raise a plain `InterruptedException`. Each becomes the
- * refusal's cause, with the thread's interrupt flag re-armed for an interrupt, and any grant
- * Spark took before failing is handed back first; see [[acquire]].
+ * throwing from any of them loses a buffer Arrow has already created. They only ever settle or
+ * trim the reservation, and contain anything the memory manager throws. How the acquisition in
+ * [[onPreAllocation]] can fail is described at [[acquire]].
  *
  * '''Lock order.''' This listener's monitor is taken before Spark's and never the other way
  * round: [[onPreAllocation]] holds ours across `acquireExecutionMemory`, and [[acquire]]
@@ -133,7 +127,7 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
     if (!completed) {
       val needed = owned.getAsLong + pending + size
       if (reserved < needed) {
-        reserved += acquireAtLeast(needed - reserved, size)
+        reserved += acquireAtLeast(needed - reserved, roundUpToBlock(needed) - reserved, size)
       }
       pending += size
     }
@@ -144,14 +138,13 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
   /** Arrow's own limit refused the allocation after it was admitted here, so hand it back. */
   override def onFailedAllocation(size: Long, outcome: AllocationOutcome): Boolean = {
     settle(size)
-    shrinkQuietly()
     false
   }
 
-  override def onRelease(size: Long): Unit = shrinkQuietly()
+  override def onRelease(size: Long): Unit = reconcile()
 
-  /** Brings the reservation up to date after buffers moved out without a callback. */
-  private[comet] def reconcile(): Unit = shrinkQuietly()
+  /** Trims the reservation after buffers moved out of the allocator without a callback. */
+  private[comet] def reconcile(): Unit = settle(0L)
 
   /**
    * Reports what the task's allocator owns. Spark reads this for spill-victim ordering,
@@ -178,10 +171,7 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
     try {
       synchronized {
         completed = true
-        if (reserved > 0L) {
-          taskMemoryManager.releaseExecutionMemory(reserved, this)
-          reserved = 0L
-        }
+        releaseDownTo(0L)
       }
     } catch {
       case NonFatal(e) => warnOnReleaseFailure(e)
@@ -191,30 +181,19 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
   /** Bytes currently reserved with Spark on this task's behalf. Visible for testing. */
   private[comet] def reservedBytes: Long = synchronized(reserved)
 
-  /** The allocation is now in the allocator's accountant, so it is no longer pending. */
-  private def settle(size: Long): Unit = synchronized {
-    if (!completed) {
-      pending = math.max(0L, pending - size)
-    }
-  }
-
   /**
-   * Returns what the allocator no longer needs, keeping what it owns and has pending rounded up
-   * to a block so that allocating and releasing around a boundary does not reach Spark every
-   * time. Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes
-   * on the executor-wide pool, so a per-block loop would take that lock once per megabyte freed.
-   * Releasing cannot fail in practice, but this is reached from callbacks that may not throw, so
-   * anything unforeseen is contained.
+   * Moves `size` bytes out of pending, since Arrow has put them in the allocator's accountant by
+   * the time it calls back, and trims the reservation to what is owned and pending, rounded up to
+   * a block. The rounding keeps a change within a block from reaching Spark, although a change
+   * that crosses a block boundary still does every time. Releasing cannot fail in practice, but
+   * this is reached from callbacks that may not throw, so anything unforeseen is contained.
    */
-  private def shrinkQuietly(): Unit = {
+  private def settle(size: Long): Unit = {
     try {
       synchronized {
+        pending = math.max(0L, pending - size)
         if (!completed) {
-          val keep = roundUpToBlock(owned.getAsLong + pending)
-          if (reserved > keep) {
-            taskMemoryManager.releaseExecutionMemory(reserved - keep, this)
-            reserved = keep
-          }
+          releaseDownTo(roundUpToBlock(owned.getAsLong + pending))
         }
       }
     } catch {
@@ -223,14 +202,26 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
   }
 
   /**
-   * Asks Spark for `shortfall` bytes, rounded up to a block so that `reserved` leaves headroom
-   * and the next small allocation does not go straight back into Spark's lock, and returns what
-   * it granted. A grant short of the block but covering `shortfall` is accepted: the rounding is
-   * there to save lock traffic, and must not be the reason an allocation that fits is refused.
-   * Anything less is handed back and the allocation refused.
+   * Returns everything reserved above `keep`, in one call rather than one per block, because
+   * `releaseExecutionMemory` synchronizes on the executor-wide pool. Called with this listener's
+   * monitor held.
+   */
+  private def releaseDownTo(keep: Long): Unit = {
+    if (reserved > keep) {
+      taskMemoryManager.releaseExecutionMemory(reserved - keep, this)
+      reserved = keep
+    }
+  }
+
+  /**
+   * Asks Spark for `request` bytes, which takes the reservation up to a block boundary so that
+   * the next small allocation does not go straight back into Spark's lock, and returns what it
+   * granted. A grant short of `request` is accepted as long as it covers `shortfall`, because the
+   * rounding is there to save lock traffic and must not be the reason an allocation that fits is
+   * refused. Anything less is handed back and the allocation refused.
    */
-  private def acquireAtLeast(shortfall: Long, size: Long): Long = {
-    val granted = acquire(roundUpToBlock(shortfall), size)
+  private def acquireAtLeast(shortfall: Long, request: Long, size: Long): Long = {
+    val granted = acquire(request, size)
     if (granted < shortfall) {
       if (granted > 0L) {
         taskMemoryManager.releaseExecutionMemory(granted, this)
@@ -244,6 +235,12 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
    * Asks Spark for `request` bytes and returns what it granted, or refuses the allocation if the
    * acquisition fails.
    *
+   * It is fallible in three ways, and only the first is caught by `NonFatal`: it runs other
+   * consumers' `spill`, which turns a task interrupt into a `RuntimeException` and an I/O failure
+   * into a `SparkOutOfMemoryError`, and the execution pool itself parks in `lock.wait()`, so
+   * killing a task can raise a plain `InterruptedException`. Each becomes the refusal's cause,
+   * and an interrupt also re-arms the thread's flag.
+   *
    * `acquireExecutionMemory` takes its first grant from the pool and only then asks other
    * consumers to spill, so when a spill throws it has already charged the task for bytes it never
    * reports back, and nothing would release them before `cleanUpAllAllocatedMemory` at the very
@@ -268,28 +265,20 @@ private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryM
     try {
       taskMemoryManager.acquireExecutionMemory(request, this)
     } catch {
-      case e: InterruptedException =>
-        // Refusing is the only way out of an Arrow callback, but the cancellation must not be
-        // lost with it: re-arm the flag so the task still observes it.
-        Thread.currentThread().interrupt()
-        throw refusedAfterFailure(heldBefore, request, size, e)
-      case NonFatal(e) => throw refusedAfterFailure(heldBefore, request, size, e)
-      case e: SparkOutOfMemoryError => throw refusedAfterFailure(heldBefore, request, size, e)
-    }
-  }
-
-  private def refusedAfterFailure(
-      heldBefore: Long,
-      request: Long,
-      size: Long,
-      cause: Throwable): OutOfMemoryException = {
-    val orphaned = math.max(
-      0L,
-      math.min(taskMemoryManager.getMemoryConsumptionForThisTask - heldBefore, request))
-    if (orphaned > 0L) {
-      taskMemoryManager.releaseExecutionMemory(orphaned, this)
+      case e @ (_: InterruptedException | _: SparkOutOfMemoryError | NonFatal(_)) =>
+        if (e.isInstanceOf[InterruptedException]) {
+          // Refusing is the only way out of an Arrow callback, but the cancellation must not be
+          // lost with it: re-arm the flag so the task still observes it.
+          Thread.currentThread().interrupt()
+        }
+        val orphaned = math.max(
+          0L,
+          math.min(taskMemoryManager.getMemoryConsumptionForThisTask - heldBefore, request))
+        if (orphaned > 0L) {
+          taskMemoryManager.releaseExecutionMemory(orphaned, this)
+        }
+        throw refusal(size, request, 0L, e)
     }
-    refusal(size, request, 0L, cause)
   }
 
   private def refusal(
diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
index 4378e864d3a..7405872d76b 100644
--- a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
+++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala
@@ -27,8 +27,10 @@ import org.apache.arrow.memory.BufferAllocator
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.memory.MemoryMode
+import org.apache.spark.sql.vectorized.ColumnarBatch
 
 import org.apache.comet.{CometArrowAllocator, CometConf}
+import org.apache.comet.vector.CometVector
 
 /**
  * Hands out the Arrow allocator that JVM-owned allocations should use, one per Spark task.
@@ -48,13 +50,11 @@ import org.apache.comet.{CometArrowAllocator, CometConf}
  * listener-less root instead, and imports from native come from `CometArrowImportAllocator`, a
  * listener-less child of it. Native memory is native's to account for: a native operator that
  * retains a batch reserves its buffers through Comet's pool, which charges the same Spark task,
- * so charging them here as well would reserve the same memory twice. A batch the JVM read into
- * this allocator and then streams to native is covered by the same rule without its allocation
- * site having to know: the stream's reader retains each batch into its own root-backed allocator
- * and closes the source, which moves the charge off this one, and then calls [[reconcile]] so
- * that the reservation follows straight away. The root is also what callers with no task get,
- * such as broadcast coalescing on the driver or a reader built while native pulls a stream, and
- * what Comet's on-heap mode gets, where charging an off-heap consumer would be wrong.
+ * so charging them here as well would reserve the same memory twice. A batch read into this
+ * allocator and then streamed to native stops being charged here once the stream has it; see
+ * [[closeAndReconcile]]. The root is also what callers with no task get, such as broadcast
+ * coalescing on the driver or a reader built while native pulls a stream, and what Comet's
+ * on-heap mode gets, where charging an off-heap consumer would be wrong.
  *
  * '''Lifetime.''' The task allocator cannot simply be closed when the task ends. The process-wide
  * allocator exists precisely because Arrow buffers can outlive the task that created them, and
@@ -111,8 +111,8 @@ object CometTaskArrowAllocator extends Logging {
     // Comet's on-heap mode accounts for nothing and exists so the Spark SQL suite can run without
     // off-heap memory configured, so it gets the root: the same switch that
     // `CometShuffleMemoryAllocator.getInstance` makes for shuffle pages.
-    if (!accountingEnabled || taskMemoryManager == null ||
-      taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) {
+    if (taskMemoryManager == null ||
+      taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP || !accountingEnabled) {
       return CometArrowAllocator
     }
 
@@ -143,22 +143,20 @@ object CometTaskArrowAllocator extends Logging {
     val finished = perTask.remove(taskAttemptId)
     if (finished != null) {
       finished.listener.taskCompleted()
-      closeLock.synchronized {
-        if (!tryClose(finished.allocator)) {
-          awaitingClose.add(finished.allocator)
-        }
-      }
+      closeLock.synchronized(awaitingClose.add(finished.allocator))
     }
     closeDrained()
   }
 
-  /** Closes any parked allocator whose stragglers have since been released. */
+  /** Closes every parked allocator that has been drained, including one parked just now. */
   private def closeDrained(): Unit = {
     if (!awaitingClose.isEmpty) {
       closeLock.synchronized {
         val parked = awaitingClose.iterator()
         while (parked.hasNext) {
-          if (tryClose(parked.next())) {
+          val allocator = parked.next()
+          if (allocator.getAllocatedMemory == 0L) {
+            closeQuietly(allocator)
             parked.remove()
           }
         }
@@ -166,16 +164,6 @@ object CometTaskArrowAllocator extends Logging {
     }
   }
 
-  /** Closes the allocator if it has been drained. Returns false if it must stay open. */
-  private def tryClose(allocator: BufferAllocator): Boolean = {
-    if (allocator.getAllocatedMemory != 0L) {
-      false
-    } else {
-      closeQuietly(allocator)
-      true
-    }
-  }
-
   private def closeQuietly(allocator: BufferAllocator): Unit = {
     try {
       allocator.close()
@@ -188,31 +176,38 @@ object CometTaskArrowAllocator extends Logging {
   }
 
   /**
-   * Brings a task allocator's reservation up to date with what it owns, if `allocator` is a task
-   * allocator or a child of one, and does nothing otherwise. For a caller that has just moved
-   * buffers out of it by retaining them in another allocator and closing the source, which Arrow
-   * does without calling any listener, so that the task stops paying for them now rather than at
-   * its next allocation.
+   * Closes a batch whose buffers another allocator has just retained, then trims the reservation
+   * of the task allocator they came from. Closing hands ownership to the retaining allocator
+   * without calling any listener, so without this the task would go on paying for them until its
+   * next allocation or release. A batch is built within one task, so it has at most one task
+   * allocator, and a `NullVector` owns no buffers and has no allocator at all.
    */
-  private[spark] def reconcile(allocator: BufferAllocator): Unit = allocator.getListener match {
-    case listener: CometArrowAllocationListener => listener.reconcile()
-    case _ =>
+  private[spark] def closeAndReconcile(batch: ColumnarBatch): Unit = {
+    var listener: CometArrowAllocationListener = null
+    var i = 0
+    while (i < batch.numCols()) {
+      batch.column(i) match {
+        case vector: CometVector =>
+          val allocator = vector.getValueVector.getAllocator
+          if (allocator != null) {
+            allocator.getListener match {
+              case l: CometArrowAllocationListener => listener = l
+              case _ =>
+            }
+          }
+        case _ =>
+      }
+      i += 1
+    }
+    batch.close()
+    if (listener != null) {
+      listener.reconcile()
+    }
   }
 
-  /** Number of tasks currently holding an accounted allocator. Visible for testing. */
-  private[comet] def trackedTaskCount: Int = perTask.size()
-
-  /** Number of finished tasks whose allocator is still draining. Visible for testing. */
-  private[comet] def awaitingCloseCount: Int = awaitingClose.size()
-
   /** The listener accounting for the given task, if it has one. Visible for testing. */
   private[comet] def listenerForTask(
       taskAttemptId: Long): Option[CometArrowAllocationListener] = {
     Option(perTask.get(taskAttemptId)).map(_.listener)
   }
-
-  /** Bytes currently reserved with Spark on behalf of the given task. Visible for testing. */
-  private[comet] def reservedBytesForTask(taskAttemptId: Long): Long = {
-    listenerForTask(taskAttemptId).map(_.reservedBytes).getOrElse(0L)
-  }
 }
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala
index 827657bf7f3..58f1b2835a8 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala
@@ -70,17 +70,9 @@ private[comet] class ColumnarBatchArrowReader(
         // Do not close this borrowed root: src and the helper own its vectors.
       }
     } finally {
-      // Closing the source leaves this reader's allocator owning the buffers it retained, which
-      // Arrow does without telling the listener of the allocator they came from. A task allocator
-      // is told here instead, so the task stops paying for the batch once native has it.
-      val sourceAllocators = (0 until src.numCols())
-        .map(src.column)
-        .collect { case v: CometVector => v.getValueVector.getAllocator }
-        // A NullVector owns no buffers and has no allocator.
-        .filter(_ != null)
-        .distinct
-      src.close()
-      sourceAllocators.foreach(CometTaskArrowAllocator.reconcile)
+      // Closing leaves this reader's allocator owning the buffers it retained, so the task
+      // allocator they came from stops paying for them here.
+      CometTaskArrowAllocator.closeAndReconcile(src)
     }
     true
   }
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
index 6ecd54f8f87..f709cb30739 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
@@ -231,10 +231,8 @@ object CometArrowStream extends Logging {
       name: String,
       readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = {
     val context = TaskContext.get()
-    // A child of the listener-less root, not of the task allocator: everything loaded here is
-    // exported to native, which accounts for what it retains. A reader that loads a batch from the
-    // task allocator and then closes the source moves the batch's charge here, so the task stops
-    // paying for it as well. See CometTaskArrowAllocator.
+    // A child of the listener-less root rather than of the task allocator, since everything loaded
+    // here is exported to native, which accounts for what it retains. See CometTaskArrowAllocator.
     val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue)
     var reader: ArrowReader = null
     var arrowStream: ArrowArrayStream = null
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
index 187cb9c6060..9043d7f8430 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala
@@ -19,14 +19,12 @@
 
 package org.apache.spark.comet
 
-import java.util.Properties
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
 
 import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, RootAllocator}
-import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
+import org.apache.spark.{CometTaskMemoryManager, TaskContext}
 import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
-import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.memory.{MemoryConsumer, MemoryManager, TaskMemoryManager}
 
 /**
  * Measures what charging JVM Arrow allocations to Spark costs, since it is on by default.
@@ -45,11 +43,11 @@ import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, TaskM
  */
 object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
 
-  private val blockSize = 1024L * 1024L
+  private val blockSize = CometArrowAllocationListener.BLOCK_SIZE
   private val poolBytes = 1024L * 1024L * 1024L
 
   override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
-    runBenchmark("JVM Arrow allocations reported to Spark") {
+    runBenchmark("JVM Arrow allocations charged to Spark") {
       // Many sub-block buffers, the shape a codegen output vector's validity and offset buffers
       // take. None of them comes close to a block, so the listener should never reach Spark after
       // the first one.
@@ -63,9 +61,17 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
       // Worst case for the block batching: every allocation crosses a boundary, so every one of
       // them takes the executor-wide lock in `acquireExecutionMemory`.
       allocateAndRelease("block-sized buffers", bufferSize = blockSize, buffersPerIteration = 8)
-      // Same, with a second thread reserving from the same constrained pool the way Comet's native
-      // side does through CometTaskMemoryManager.
-      allocateUnderNativePressure()
+      // Same, with a second thread reserving from the same pool the way Comet's native side does.
+      // The pool has room for the whole set plus that thread's block, because a smaller one would
+      // refuse the allocations rather than slow them down, so what this measures is contention for
+      // the pool's lock rather than for its capacity. In the unaccounted arm only the pressure
+      // thread touches the pool: the delta is what the Arrow side adds by competing.
+      allocateAndRelease(
+        "block-sized buffers under native pressure",
+        bufferSize = blockSize,
+        buffersPerIteration = 8,
+        pool = blockSize * 9,
+        nativePressure = true)
       // Per call site rather than per buffer, but it is the cost the root allocator `val` did not
       // have: a TaskContext lookup and a concurrent map read.
       allocatorLookup()
@@ -75,7 +81,9 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
   private def allocateAndRelease(
       name: String,
       bufferSize: Long,
-      buffersPerIteration: Int): Unit = {
+      buffersPerIteration: Int,
+      pool: Long = poolBytes,
+      nativePressure: Boolean = false): Unit = {
     val benchmark =
       new Benchmark(
         s"$name (${buffersPerIteration}x$bufferSize)",
@@ -86,55 +94,23 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
     // steady-state cost of allocating and releasing rather than the cost of standing a task up.
     val root = new RootAllocator(Long.MaxValue)
     try {
-      withTaskAllocator() { (accounted, memory) =>
-        benchmark.addCase("not accounted") { _ =>
-          churn(root, bufferSize, buffersPerIteration)
-        }
-        benchmark.addCase("accounted") { _ =>
+      withTaskAllocator(pool) { (accounted, memory) =>
+        withNativePressure(nativePressure) {
+          benchmark.addCase("not accounted") { _ =>
+            churn(root, bufferSize, buffersPerIteration)
+          }
+          benchmark.addCase("accounted") { _ =>
+            churn(accounted, bufferSize, buffersPerIteration)
+          }
+          benchmark.run()
+
+          // One more round with the counters zeroed, to report how often a single iteration
+          // reaches the memory manager. That, rather than the per-buffer bookkeeping, is the cost
+          // that scales with buffer size.
+          memory.reset()
           churn(accounted, bufferSize, buffersPerIteration)
+          writeLine(s"  accounted: ${memory.summary(buffersPerIteration)} per iteration")
         }
-        benchmark.run()
-
-        // One more round with the counters zeroed, to report how often a single iteration reaches
-        // the memory manager. That, rather than the per-buffer bookkeeping, is the cost that
-        // scales with buffer size.
-        memory.reset()
-        churn(accounted, bufferSize, buffersPerIteration)
-        writeLine(s"  accounted: ${memory.summary(buffersPerIteration)} per iteration")
-      }
-    } finally {
-      root.close()
-    }
-  }
-
-  private def allocateUnderNativePressure(): Unit = {
-    val buffersPerIteration = 8
-    val benchmark = new Benchmark(
-      s"block-sized buffers under native pressure (${buffersPerIteration}x$blockSize)",
-      buffersPerIteration,
-      output = output)
-
-    val root = new RootAllocator(Long.MaxValue)
-    try {
-      // Room for the whole set plus the pressure thread's block, which runs throughout. A smaller
-      // pool would refuse the allocations rather than slow them down, so what this measures is
-      // contention for the pool's lock rather than for its capacity. In the unaccounted arm only
-      // the pressure thread touches the pool: the delta is what the Arrow side adds by competing.
-      withTaskAllocator(poolBytes = blockSize * (buffersPerIteration + 1)) {
-        (accounted, memory) =>
-          withNativePressure(memory) {
-            benchmark.addCase("not accounted") { _ =>
-              churn(root, blockSize, buffersPerIteration)
-            }
-            benchmark.addCase("accounted") { _ =>
-              churn(accounted, blockSize, buffersPerIteration)
-            }
-            benchmark.run()
-
-            memory.reset()
-            churn(accounted, blockSize, buffersPerIteration)
-            writeLine(s"  accounted: ${memory.summary(buffersPerIteration)} per iteration")
-          }
       }
     } finally {
       root.close()
@@ -193,71 +169,38 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
    * Runs the body against a task allocator, torn down afterwards, so that repeated iterations do
    * not accumulate reservations or allocators.
    */
-  private def withTaskAllocator[T](poolBytes: Long = poolBytes)(
+  private def withTaskAllocator[T](pool: Long = poolBytes)(
       f: (BufferAllocator, CountingTaskMemoryManager) => T): T = {
-    val memory = newTaskMemoryManager(poolBytes)
-    val context = newTaskContext(memory)
-    val previous = TaskContext.get()
-    TaskContext.setTaskContext(context)
-    try {
-      f(CometTaskArrowAllocator.forCurrentTask(), memory)
-    } finally {
-      try {
-        context.markTaskCompleted(None)
-        memory.cleanUpAllAllocatedMemory()
-      } finally {
-        if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous)
-      }
-    }
+    val (context, memory) = TestTasks.newTask(pool)(new CountingTaskMemoryManager(_, _))
+    TestTasks.withInstalled(context)(f(CometTaskArrowAllocator.forCurrentTask(), memory))
   }
 
-  /** Hammers the same pool from another thread, the way native reservations do. */
-  private def withNativePressure[T](memory: TaskMemoryManager)(f: => T): T = {
-    val stop = new AtomicBoolean(false)
-    val consumer = new NativeLikeConsumer(memory)
-    val thread = new Thread(() => {
-      while (!stop.get()) {
-        val granted = consumer.reserve(blockSize)
-        consumer.release(granted)
+  /**
+   * Hammers the task's pool from another thread for the duration of the body, when asked to, the
+   * way native reservations do. Built on the task's thread, which `CometTaskMemoryManager` needs.
+   */
+  private def withNativePressure[T](enabled: Boolean)(f: => T): T = {
+    if (!enabled) {
+      f
+    } else {
+      val stop = new AtomicBoolean(false)
+      val native = new CometTaskMemoryManager(0L, TaskContext.get().taskAttemptId())
+      val thread = new Thread(() => {
+        while (!stop.get()) {
+          native.releaseMemory(native.acquireMemory(blockSize))
+        }
+      })
+      thread.setDaemon(true)
+      thread.setName("native-reservations")
+      thread.start()
+      try f
+      finally {
+        stop.set(true)
+        thread.join()
       }
-    })
-    thread.setDaemon(true)
-    thread.setName("native-reservations")
-    thread.start()
-    try f
-    finally {
-      stop.set(true)
-      thread.join()
     }
   }
 
-  private val nextTaskAttemptId = new AtomicLong(1L)
-
-  private def newTaskMemoryManager(poolBytes: Long): CountingTaskMemoryManager = {
-    val conf = new SparkConf(false)
-      .set("spark.memory.offHeap.enabled", "true")
-      .set("spark.memory.offHeap.size", poolBytes.toString)
-    val memoryManager = new TestMemoryManager(conf)
-    memoryManager.limit(poolBytes)
-    new CountingTaskMemoryManager(memoryManager, nextTaskAttemptId.getAndIncrement())
-  }
-
-  private def newTaskContext(memory: CountingTaskMemoryManager): TaskContextImpl = {
-    new TaskContextImpl(
-      stageId = 0,
-      stageAttemptNumber = 0,
-      partitionId = 0,
-      numPartitions = 1,
-      taskAttemptId = memory.getTaskAttemptId,
-      attemptNumber = 0,
-      taskMemoryManager = memory,
-      localProperties = new Properties,
-      metricsSystem = null,
-      taskMetrics = TaskMetrics.empty,
-      cpus = 1,
-      resources = Map.empty)
-  }
-
   private def writeLine(line: String): Unit = {
     // scalastyle:off println
     println(line)
@@ -270,8 +213,6 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
     private val acquires = new AtomicLong(0L)
     private val releases = new AtomicLong(0L)
 
-    def getTaskAttemptId: Long = taskAttemptId
-
     // Only the Arrow listener's own calls are counted, so the pressure thread's traffic does not
     // land in the reported figure.
     override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = {
@@ -292,23 +233,4 @@ object CometArrowAllocationListenerBenchmark extends BenchmarkBase {
     def summary(buffers: Int): String =
       s"${acquires.get()} acquire and ${releases.get()} release calls for $buffers buffers"
   }
-
-  /** Stands in for `CometTaskMemoryManager`: reserves from Spark directly and never spills. */
-  private class NativeLikeConsumer(memory: TaskMemoryManager)
-      extends MemoryConsumer(memory, 0L, MemoryMode.OFF_HEAP) {
-    private val reserved = new AtomicLong(0L)
-    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
-    override def getUsed: Long = reserved.get()
-    def reserve(bytes: Long): Long = {
-      val granted = memory.acquireExecutionMemory(bytes, this)
-      reserved.addAndGet(granted)
-      granted
-    }
-    def release(bytes: Long): Unit = {
-      if (bytes > 0L) {
-        reserved.addAndGet(-bytes)
-        memory.releaseExecutionMemory(bytes, this)
-      }
-    }
-  }
 }
diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
index dc69aab39e5..0f6f1670cb8 100644
--- a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala
@@ -20,20 +20,21 @@
 package org.apache.spark.comet
 
 import java.io.{InterruptedIOException, IOException}
-import java.util.Properties
 import java.util.concurrent.atomic.{AtomicLong, AtomicReference}
 
+import scala.jdk.CollectionConverters._
+
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.arrow.c.Data
-import org.apache.arrow.memory.{BufferAllocator, OutOfMemoryException}
+import org.apache.arrow.memory.OutOfMemoryException
 import org.apache.arrow.vector.{FieldVector, IntVector, VectorSchemaRoot}
-import org.apache.spark.{SparkConf, SparkContext, TaskContext, TaskContextImpl}
-import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.{CometTaskMemoryManager, SparkConf, SparkContext, TaskContext, TaskContextImpl}
+import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager}
 import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
 import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.ThreadUtils
 
 import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator, CometConf}
 import org.apache.comet.vector.NativeUtil
@@ -50,9 +51,6 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   private val blockSize = CometArrowAllocationListener.BLOCK_SIZE
   private val poolBytes = 64L * 1024 * 1024
 
-  /** Task attempt ids are keys in a process-wide map, so no two tests may share one. */
-  private val nextTaskAttemptId = new AtomicLong(1000L)
-
   // ---------------------------------------------------------------------------------------------
   // Reservation arithmetic. Driven through the listener directly, with a stand-in for the
   // allocator's accountant, since Arrow's rounding policy would otherwise decide the sizes under
@@ -61,8 +59,8 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
   test("allocations are charged to the current task in whole blocks") {
     withTask() { task =>
-      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
-      val owner = new StubOwner(listener)
+      val owner = new StubOwner(task.taskMemoryManager)
+      val listener = owner.listener
 
       // Far smaller than a block, so the reservation should round up to exactly one block.
       owner.allocate(128L)
@@ -78,8 +76,9 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
   test("a request larger than a block rounds up to a block multiple") {
     withTask() { task =>
-      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
-      new StubOwner(listener).allocate(blockSize * 3 + 7L)
+      val owner = new StubOwner(task.taskMemoryManager)
+      owner.allocate(blockSize * 3 + 7L)
+      val listener = owner.listener
       // Rounded up rather than sized to the exact deficit, so growth leaves headroom and the next
       // small allocation does not go straight back into Spark.
       assert(listener.reservedBytes == blockSize * 4)
@@ -89,8 +88,8 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
   test("releasing returns whole blocks to Spark") {
     withTask() { task =>
-      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
-      val owner = new StubOwner(listener)
+      val owner = new StubOwner(task.taskMemoryManager)
+      val listener = owner.listener
       owner.allocate(blockSize * 2)
       assert(listener.reservedBytes == blockSize * 2)
 
@@ -209,7 +208,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       // pins it and drops it later from a Tokio worker with no task context installed. Reading
       // TaskContext in onRelease would ignore this release and leave the task charged for memory
       // it had already freed.
-      onDetachedThread(buf.close())
+      ThreadUtils.runInNewThread("detached-release")(buf.close())
 
       assert(reservedFor(task) == 0L)
       assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
@@ -244,8 +243,8 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       val stream = CometArrowAllocator.newChildAllocator("stream", 0L, Long.MaxValue)
       try {
         val allocator = CometTaskArrowAllocator.forCurrentTask()
-        val source = allocator.buffer(blockSize)
-        assert(reservedFor(task) == blockSize)
+        val source = allocator.buffer(blockSize * 4)
+        assert(reservedFor(task) == blockSize * 4)
 
         // What ColumnarBatchArrowReader does to every batch it streams to native: retain the
         // buffers into the stream's allocator, then close the source. Arrow makes the stream's
@@ -261,8 +260,8 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
           retained.close()
         }
 
-        // The reservation catches up at the next callback on the task's allocator: one block for
-        // the new buffer, not one on top of a block for a buffer the task no longer owns.
+        // With no reconcile, the reservation catches up at the next callback on the task's
+        // allocator: one block for the new buffer, not four for a buffer the task no longer owns.
         val next = allocator.buffer(128L)
         assert(reservedFor(task) == blockSize)
         next.close()
@@ -357,8 +356,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   // ---------------------------------------------------------------------------------------------
 
   test("task completion releases the whole reservation and closes the allocator") {
-    val task = newTask()
-    val allocatorName = withInstalledTask(task) {
+    val allocatorName = withTask() { task =>
       val allocator = CometTaskArrowAllocator.forCurrentTask()
       val buf = allocator.buffer(blockSize)
       assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize)
@@ -374,8 +372,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   }
 
   test("a buffer outliving its task parks the allocator until it is released") {
-    val task = newTask()
-    withInstalledTask(task) {
+    withTask() { task =>
       val allocator = CometTaskArrowAllocator.forCurrentTask()
       val buf = allocator.buffer(blockSize)
 
@@ -407,7 +404,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   test("an allocation the pool cannot cover is refused, and one that fits still succeeds") {
     // Two blocks of budget, one of them held by a consumer that cannot give anything back.
     withTask(pool = blockSize * 2) { task =>
-      val other = new PlainConsumer(task.taskMemoryManager)
+      val other = new OtherConsumer(task.taskMemoryManager)
       assert(other.take(blockSize) == blockSize)
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
@@ -433,7 +430,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     // block left must still admit an allocation that fits in what is left.
     val left = 64L * 1024
     withTask(pool = blockSize) { task =>
-      val other = new PlainConsumer(task.taskMemoryManager)
+      val other = new OtherConsumer(task.taskMemoryManager)
       assert(other.take(blockSize - left) == blockSize - left)
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
@@ -453,8 +450,8 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
   test("an allocation still in flight counts against the reservation") {
     withTask(pool = blockSize) { task =>
-      val listener = new CometArrowAllocationListener(task.taskMemoryManager)
-      val owner = new StubOwner(listener)
+      val owner = new StubOwner(task.taskMemoryManager)
+      val listener = owner.listener
 
       // Admitted, but Arrow has not reported it yet...
       listener.onPreAllocation(blockSize)
@@ -478,7 +475,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
       // Exactly one block of budget, already taken by a consumer whose spill throws, so the
       // acquisition has to go through Spark's spill path and comes back throwing.
       withTask(pool = blockSize) { task =>
-        val hostile = new FailingSpillConsumer(task.taskMemoryManager, failure)
+        val hostile = new OtherConsumer(task.taskMemoryManager, Some(failure))
         assert(hostile.take(blockSize) == blockSize)
 
         val allocator = CometTaskArrowAllocator.forCurrentTask()
@@ -499,7 +496,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     // never parks, so the interrupt is injected through a failing spill instead.
     withTask(pool = blockSize) { task =>
       val hostile =
-        new FailingSpillConsumer(task.taskMemoryManager, new InterruptedException("task killed"))
+        new OtherConsumer(task.taskMemoryManager, Some(new InterruptedException("task killed")))
       assert(hostile.take(blockSize) == blockSize)
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
@@ -518,7 +515,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     // the block it already took, so nothing would release it before the task ended.
     withTask(pool = blockSize * 2) { task =>
       val hostile =
-        new FailingSpillConsumer(task.taskMemoryManager, new IOException("spill failed"))
+        new OtherConsumer(task.taskMemoryManager, Some(new IOException("spill failed")))
       assert(hostile.take(blockSize) == blockSize)
 
       val allocator = CometTaskArrowAllocator.forCurrentTask()
@@ -538,13 +535,12 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     // snapshots and the call therefore run as one transaction under the
     // TaskMemoryManager monitor. This forces the interleaving that transaction exists to exclude.
     val spare = 1024L
-    val task = newTask(pool = blockSize * 2 + spare)
-    withInstalledTask(task) {
+    withTask(pool = blockSize * 2 + spare) { task =>
       val hostile =
-        new FailingSpillConsumer(task.taskMemoryManager, new IOException("spill failed"))
+        new OtherConsumer(task.taskMemoryManager, Some(new IOException("spill failed")))
       assert(hostile.take(blockSize) == blockSize)
 
-      val interloper = new PlainConsumer(task.taskMemoryManager)
+      val interloper = new OtherConsumer(task.taskMemoryManager)
       // Once the acquisition has taken what was left, the pool is empty and Spark answers the
       // interloper by asking the hostile consumer to spill, which throws. That is a legitimate
       // outcome for the interloper and not what is under test here; what it ends up holding is.
@@ -586,19 +582,17 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     withTask() { task =>
       val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize)
       try {
-        val listener = CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).get
+        val listener = listenerFor(task)
         val used = new AtomicLong(-1L)
         val spilled = new AtomicLong(-1L)
         listener.synchronized {
           // A native reservation arriving through CometTaskMemoryManager on a Tokio thread holds
           // Spark's monitor here. If either call waited on this one, it would deadlock against an
           // Arrow allocation on the same task that already holds this monitor and wants Spark's.
-          val probe = new Thread(() => {
+          val probe = daemonThread("lock-order-probe") {
             used.set(listener.getUsed)
             spilled.set(listener.spill(blockSize, listener))
-          })
-          probe.setDaemon(true)
-          probe.setName("lock-order-probe")
+          }
           probe.start()
           probe.join(30000L)
           assert(!probe.isAlive, "getUsed or spill blocked on the reservation monitor")
@@ -612,19 +606,21 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   }
 
   test("concurrent Arrow and native reservations make progress") {
-    // Two blocks of budget shared by both consumers, so most requests are short and Spark walks
-    // its consumer list, calling getUsed on the Arrow listener while holding its own monitor.
+    // Two blocks of budget, and the native side asks for both, so whenever the Arrow side holds
+    // one the native grant comes up short and Spark walks its consumer list, calling getUsed and
+    // spill on the Arrow listener while holding its own monitor. Either side can come up short,
+    // and an Arrow allocation refused for it is a legitimate outcome: finishing is what is tested.
     withTask(pool = blockSize * 2) { task =>
       val allocator = CometTaskArrowAllocator.forCurrentTask()
-      val native = new NativeLikeConsumer(task.taskMemoryManager)
+      val native = new CometTaskMemoryManager(0L, task.taskAttemptId)
       val failure = new AtomicReference[Throwable]()
 
       val arrowThread = loopingThread("arrow-allocations", failure) {
-        val buf = allocator.buffer(blockSize)
-        buf.close()
+        try allocator.buffer(blockSize).close()
+        catch { case _: OutOfMemoryException => }
       }
       val nativeThread = loopingThread("native-reservations", failure) {
-        native.release(native.reserve(blockSize))
+        native.releaseMemory(native.acquireMemory(blockSize * 2))
       }
 
       Seq(arrowThread, nativeThread).foreach(_.start())
@@ -635,7 +631,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
 
       Option(failure.get).foreach(e => fail("a worker failed", e))
       assert(allocator.getAllocatedMemory == 0L)
-      native.release(native.used())
+      assert(native.getUsed == 0L)
     }
   }
 
@@ -644,9 +640,11 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   // ---------------------------------------------------------------------------------------------
 
   /**
-   * Stands in for the task allocator's accountant, and calls the listener the way Arrow would.
+   * A listener with a stand-in for the task allocator's accountant, calling the listener the way
+   * Arrow would.
    */
-  private class StubOwner(listener: CometArrowAllocationListener) {
+  private class StubOwner(tmm: TaskMemoryManager) {
+    val listener = new CometArrowAllocationListener(tmm)
     private val owned = new AtomicLong(0L)
     listener.bind(() => owned.get)
 
@@ -668,18 +666,17 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
-  /** Holds memory and refuses to give it back, so `trySpillAndAcquire` throws on its behalf. */
-  private class FailingSpillConsumer(tmm: TaskMemoryManager, failure: Exception)
-      extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
-    def take(bytes: Long): Long = acquireMemory(bytes)
-    override def spill(size: Long, trigger: MemoryConsumer): Long = throw failure
-  }
-
-  /** Any other consumer in the task: takes memory, and cannot give it back. */
-  private class PlainConsumer(tmm: TaskMemoryManager)
+  /**
+   * Any other consumer in the task: holds memory and cannot give it back. With a `spillFailure`,
+   * its spill throws that instead, so `trySpillAndAcquire` throws on its behalf.
+   */
+  private class OtherConsumer(tmm: TaskMemoryManager, spillFailure: Option[Exception] = None)
       extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
     def take(bytes: Long): Long = acquireMemory(bytes)
-    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+    override def spill(size: Long, trigger: MemoryConsumer): Long = spillFailure match {
+      case Some(failure) => throw failure
+      case None => 0L
+    }
   }
 
   /**
@@ -691,7 +688,7 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
   private class HookedTaskMemoryManager(
       memoryManager: MemoryManager,
       taskAttemptId: Long,
-      hook: AtomicReference[Runnable])
+      val hook: AtomicReference[Runnable] = new AtomicReference[Runnable]())
       extends TaskMemoryManager(memoryManager, taskAttemptId) {
     override def getMemoryConsumptionForThisTask(): Long = {
       val held = super.getMemoryConsumptionForThisTask()
@@ -701,103 +698,26 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
-  /** Stands in for `CometTaskMemoryManager`: reserves from Spark directly and never spills. */
-  private class NativeLikeConsumer(tmm: TaskMemoryManager)
-      extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) {
-    private val reserved = new AtomicLong(0L)
-    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
-    override def getUsed: Long = reserved.get()
-    def used(): Long = reserved.get()
-    def reserve(bytes: Long): Long = {
-      val granted = tmm.acquireExecutionMemory(bytes, this)
-      reserved.addAndGet(granted)
-      granted
-    }
-    def release(bytes: Long): Unit = {
-      if (bytes > 0L) {
-        reserved.addAndGet(-bytes)
-        tmm.releaseExecutionMemory(bytes, this)
-      }
-    }
-  }
-
   private case class TaskFixture(
-      taskAttemptId: Long,
       context: TaskContextImpl,
-      taskMemoryManager: TaskMemoryManager,
-      snapshotHook: AtomicReference[Runnable])
+      taskMemoryManager: HookedTaskMemoryManager) {
+    def taskAttemptId: Long = context.taskAttemptId
+    def snapshotHook: AtomicReference[Runnable] = taskMemoryManager.hook
+  }
 
   private def reservedFor(task: TaskFixture): Long =
-    CometTaskArrowAllocator.reservedBytesForTask(task.taskAttemptId)
+    CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).map(_.reservedBytes).getOrElse(0L)
 
   private def listenerFor(task: TaskFixture): CometArrowAllocationListener =
     CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).get
 
-  private def newTask(offHeap: Boolean = true, pool: Long = poolBytes): TaskFixture = {
-    val conf = new SparkConf(false)
-    if (offHeap) {
-      conf
-        .set("spark.memory.offHeap.enabled", "true")
-        .set("spark.memory.offHeap.size", pool.toString)
-    }
-    val memoryManager = new TestMemoryManager(conf)
-    memoryManager.limit(pool)
-    val taskAttemptId = nextTaskAttemptId.getAndIncrement()
-    val snapshotHook = new AtomicReference[Runnable]()
-    val taskMemoryManager =
-      new HookedTaskMemoryManager(memoryManager, taskAttemptId, snapshotHook)
-    val context = new TaskContextImpl(
-      stageId = 0,
-      stageAttemptNumber = 0,
-      partitionId = 0,
-      numPartitions = 1,
-      taskAttemptId = taskAttemptId,
-      attemptNumber = 0,
-      taskMemoryManager = taskMemoryManager,
-      localProperties = new Properties,
-      metricsSystem = null,
-      taskMetrics = TaskMetrics.empty,
-      cpus = 1,
-      resources = Map.empty)
-    TaskFixture(taskAttemptId, context, taskMemoryManager, snapshotHook)
-  }
-
-  /** Installs the task on this thread, restoring whatever was there before. */
-  private def withInstalledTask[T](task: TaskFixture)(f: => T): T = {
-    val previous = TaskContext.get()
-    TaskContext.setTaskContext(task.context)
-    try {
-      f
-    } finally {
-      try {
-        // Fires the completion listener that drops the reservation; harmless if already run.
-        task.context.markTaskCompleted(None)
-        task.taskMemoryManager.cleanUpAllAllocatedMemory()
-      } finally {
-        if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous)
-      }
-    }
-  }
-
-  private def withTask(offHeap: Boolean = true, pool: Long = poolBytes)(
-      f: TaskFixture => Unit): Unit = {
-    val task = newTask(offHeap, pool)
-    withInstalledTask(task)(f(task))
-  }
-
-  /** Runs the body on a fresh thread, which by construction carries no task context. */
-  private def onDetachedThread(body: => Unit): Unit = {
-    val failure = new AtomicReference[Throwable]()
-    val thread = new Thread(() => {
-      try body
-      catch { case t: Throwable => failure.set(t) }
-    })
-    thread.setDaemon(true)
-    thread.setName("detached-release")
-    thread.start()
-    thread.join(30000L)
-    assert(!thread.isAlive, "the detached release did not finish")
-    Option(failure.get).foreach(t => throw t)
+  /** Runs the body in a task of its own, completed afterwards, and returns what the body does. */
+  private def withTask[T](offHeap: Boolean = true, pool: Long = poolBytes)(
+      f: TaskFixture => T): T = {
+    val (context, taskMemoryManager) =
+      TestTasks.newTask(pool, offHeap)(new HookedTaskMemoryManager(_, _))
+    val task = TaskFixture(context, taskMemoryManager)
+    TestTasks.withInstalled(context)(f(task))
   }
 
   /** An unstarted daemon thread, so a test can choose the moment it runs. */
@@ -822,30 +742,20 @@ class CometArrowAllocationListenerSuite extends AnyFunSuite {
     }
   }
 
+  /** Runs the body 500 times on an unstarted daemon thread, recording the first failure. */
   private def loopingThread(name: String, failure: AtomicReference[Throwable])(
-      body: => Unit): Thread = {
-    val thread = new Thread(() => {
-      try {
-        var i = 0
-        while (i < 500) {
-          body
-          i += 1
-        }
-      } catch {
-        case t: Throwable => failure.compareAndSet(null, t)
+      body: => Unit): Thread = daemonThread(name) {
+    try {
+      var i = 0
+      while (i < 500) {
+        body
+        i += 1
       }
-    })
-    thread.setDaemon(true)
-    thread.setName(name)
-    thread
-  }
-
-  private def rootChildNames(): Set[String] = {
-    val names = Set.newBuilder[String]
-    val children = CometArrowAllocator.getChildAllocators.iterator()
-    while (children.hasNext) {
-      names += children.next().asInstanceOf[BufferAllocator].getName
+    } catch {
+      case t: Throwable => failure.compareAndSet(null, t)
     }
-    names.result()
   }
+
+  private def rootChildNames(): Set[String] =
+    CometArrowAllocator.getChildAllocators.asScala.map(_.getName).toSet
 }
diff --git a/spark/src/test/scala/org/apache/spark/comet/TestTasks.scala b/spark/src/test/scala/org/apache/spark/comet/TestTasks.scala
new file mode 100644
index 00000000000..b29f9dadb1d
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/comet/TestTasks.scala
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.comet
+
+import java.util.Properties
+import java.util.concurrent.atomic.AtomicLong
+
+import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl}
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.memory.{MemoryManager, TaskMemoryManager, TestMemoryManager}
+
+/**
+ * Spark tasks built by hand, each over a memory pool of its own, for driving the JVM Arrow
+ * accounting without a `SparkContext`.
+ */
+private[comet] object TestTasks {
+
+  /** Task attempt ids are keys in a process-wide map, so no two tasks may share one. */
+  private val nextTaskAttemptId = new AtomicLong(1000L)
+
+  /**
+   * A task over a pool of `pool` bytes, off-heap unless `offHeap` is false, whose
+   * `TaskMemoryManager` comes from `newTaskMemoryManager` so that a caller can subclass it.
+   */
+  def newTask[M <: TaskMemoryManager](pool: Long, offHeap: Boolean = true)(
+      newTaskMemoryManager: (MemoryManager, Long) => M): (TaskContextImpl, M) = {
+    val conf = new SparkConf(false)
+    if (offHeap) {
+      conf
+        .set("spark.memory.offHeap.enabled", "true")
+        .set("spark.memory.offHeap.size", pool.toString)
+    }
+    val memoryManager = new TestMemoryManager(conf)
+    memoryManager.limit(pool)
+    val taskAttemptId = nextTaskAttemptId.getAndIncrement()
+    val taskMemoryManager = newTaskMemoryManager(memoryManager, taskAttemptId)
+    val context = new TaskContextImpl(
+      stageId = 0,
+      stageAttemptNumber = 0,
+      partitionId = 0,
+      numPartitions = 1,
+      taskAttemptId = taskAttemptId,
+      attemptNumber = 0,
+      taskMemoryManager = taskMemoryManager,
+      localProperties = new Properties,
+      metricsSystem = null,
+      taskMetrics = TaskMetrics.empty,
+      cpus = 1,
+      resources = Map.empty)
+    (context, taskMemoryManager)
+  }
+
+  /**
+   * Runs the body with the task installed on this thread, then completes the task, which fires
+   * the completion listener that drops the Arrow reservation and is harmless if it already ran,
+   * and restores whatever was installed before.
+   */
+  def withInstalled[T](context: TaskContextImpl)(f: => T): T = {
+    val previous = TaskContext.get()
+    TaskContext.setTaskContext(context)
+    try {
+      f
+    } finally {
+      try {
+        context.markTaskCompleted(None)
+        context.taskMemoryManager.cleanUpAllAllocatedMemory()
+      } finally {
+        if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous)
+      }
+    }
+  }
+}