feat: report JVM Arrow allocations to Spark's memory manager

.github/workflows/pr_build_linux.yml

@@ -537,6 +537,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

.github/workflows/pr_build_macos.yml

@@ -185,6 +185,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

docs/source/contributor-guide/memory_management.md

@@ -64,23 +64,63 @@ 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               |
-
-Two observations follow.
-
-**Comet's JVM-side Arrow allocator is unbounded and accounted by nobody.** `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.
+| 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.
+
+**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.
+
+**The owner is the allocator, not the calling thread.** Arrow hands `AllocationListener` nothing
+but a size, and it reports a release on whichever thread drops the last reference, which for
+anything exported over the C Data Interface is a Tokio worker with no task context installed.
+Attributing by `TaskContext` would therefore drop those releases and leave a task charged for
+memory it had already freed. Binding one listener to one task's allocator is what makes the release
+land on the task that allocated. Children cut from a task allocator inherit its listener, so the
+paths that make their own children are covered without knowing about any of this.
+
+**Anything crossing the FFI boundary uses 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`
+together with the codegen output vector it exports, and `CometNativeArrowSource.stream` all use
+the listener-less root; IPC reads, the cached batch serializer and
+`CometNativeArrowSource.readerBatchIter` use the task allocator. The
+driver, broadcast coalescing and the cached batch serializer also fall back to the root when there
+is no task to charge, as does Comet's on-heap mode.
+
+**The split is by allocation site, so it is not exact.** A buffer that is used in the JVM and only
+later handed to native -- a shuffle-read batch feeding a native operator is the common shape -- is
+allocated somewhere that cannot know its future, so it stays charged on the JVM side while native
+may also reserve it. Closing that needs reservation ownership to be handed over at the boundary,
+which is a change on both sides of it.
+
+**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
@@ -205,13 +245,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
+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 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 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
@@ -286,7 +326,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
@@ -321,7 +361,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** and participates in no budget.
+- **`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.

spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java

@@ -209,6 +209,10 @@ private static void evaluateInternal(
             });
     assert udf != null : "reflective instantiation returned null for " + udfClassName;
 
+    // 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];

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -852,6 +852,21 @@ 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 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)
+
   val COMET_NATIVE_LOAD_REQUIRED: ConfigEntry[Boolean] = conf("spark.comet.nativeLoadRequired")
     .category(CATEGORY_EXEC)
     .doc(

spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala

@@ -87,21 +87,27 @@ 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 = {
+    // 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, 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)

spark/src/main/scala/org/apache/comet/package.scala

@@ -32,6 +32,19 @@ 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.
+   *
+   * It carries no allocation listener, so allocating from it directly is not reported to Spark's
+   * 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.
+   *
+   * 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)
 

spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala

@@ -47,7 +47,12 @@ import org.apache.comet.CometArrowAllocator
 class NativeUtil extends AutoCloseable {
   import Utils._
 
-  /** Use the global allocator */
+  // 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 */

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.