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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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)