Memory Management#

This page describes how memory is budgeted, accounted, and enforced across the JVM/native boundary. It is aimed at contributors working on memory pools, operators that reserve memory, or anyone debugging an out-of-memory report. For user-facing tuning advice, see the Tuning Guide.

This page covers off-heap mode (spark.memory.offHeap.enabled=true) only. Comet also has an on-heap mode, but it exists so that the Spark SQL test suite can run against Comet without changing Spark’s memory configuration. It must not be used in production, and it is not described here. The pool types that only on-heap mode exposes belong to the CATEGORY_TESTING config group for the same reason.

Overview#

A Comet executor has to satisfy three separate memory budgets at once, and they are enforced by three different parties:

Budget	Enforced by	What happens when it is exceeded
JVM heap	The JVM	`java.lang.OutOfMemoryError`; Spark treats it as fatal and exits the executor
Spark’s memory pools	Spark bookkeeping	A consumer is asked to spill, or the task fails with `SparkOutOfMemoryError`
Container RSS	The OS / cgroup	`SIGKILL` of the whole executor process (exit 137, `OOMKilled`)

The first two are accounting: a running total of bytes that consumers have voluntarily declared. The third is physical: the kernel measures resident pages and does not care what any accounting layer believes.

The first row is easy to get wrong. Spark’s Executor.isFatalError exempts SparkOutOfMemoryError, which is the error Spark’s own memory manager throws when a consumer cannot get the bytes it asked for, so that failure is confined to the task and the executor keeps running. An ordinary JVM OutOfMemoryError is not exempt: the executor hands it to SparkUncaughtExceptionHandler, which calls System.exit(SparkExitCode.OOM) (exit code 52). Real heap exhaustion therefore usually costs the whole executor, not one task. An executor loss on its own does not tell you which budget was exceeded; the exit code does.

Comet’s difficulty is that its allocations are made by Rust code, so they are invisible to the JVM heap and to Spark’s own off-heap accounting, yet they land squarely in container RSS. Comet therefore maintains its own budget that is meant to shadow the physical one, and the accuracy of that shadow is the central problem this page is about.

Who allocates what#

Enabling Comet does not add one new memory consumer, it adds several, and they are not all 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.

The JVM shuffle allocator is an ordinary Spark consumer. CometShuffleMemoryAllocator.getInstance returns CometUnifiedShuffleMemoryAllocator, a Spark MemoryConsumer drawing from spark.memory.offHeap.size, so shuffle pages are arbitrated against Spark’s other consumers in the same task like any other allocation.

Where Comet’s budget comes from#

CometExecIterator.getMemoryConfig computes the budget once per executor and passes it across JNI to Java_org_apache_comet_Native_createPlan as memory_limit. Comet shares Spark’s off-heap pool rather than asking for a separate allocation:

memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction

spark.comet.exec.memoryPool.fraction defaults to 1.0. Lowering it is the current workaround for Comet’s under-accounting (see The accounting gap). It holds back a slice of the off-heap pool that Comet is not allowed to reserve, on the assumption that Comet’s real usage overshoots its reservations by roughly that slice.

A second value, memory_limit_per_task, is computed and passed alongside it, but only the on-heap pool types read it.

Resolving the pool type#

parse_memory_pool_config (native/core/src/execution/memory_pools/config.rs) turns the pool-type string and the limit into a MemoryPoolConfig. Two pool types are valid in off-heap mode:

Pool type	Sized from	Notes
`fair_unified` (default)	`memory_limit`	Delegates to Spark’s `TaskMemoryManager`; task-shared
`greedy_unified`	n/a (pool size `0`)	Spark owns the limit entirely; task-shared

Any other pool type is rejected with a configuration error.

The pool stack#

create_memory_pool builds a base pool and createPlan then wraps it in decorators. Reading from the inside out, a Comet plan in the default configuration sees:

[LoggingMemoryPool]        <- only when spark.comet.debug.memory=true
  [TaskSharedMemoryPool]   <- RAII handle for the per-task registry
    [TrackConsumersPool]   <- DataFusion; names the top 10 consumers in error messages
      [CometFairMemoryPool]  <- delegates acquire/release to Spark over JNI

Each decorator forwards every MemoryPool method to its inner pool, so reserved() at any level reports the base pool’s number.

The unified pools#

CometUnifiedMemoryPool and CometFairMemoryPool (unified_pool.rs, fair_pool.rs) are the bridge to Spark. Their try_grow calls CometTaskMemoryManager.acquireMemory over JNI, which goes through Spark’s ordinary TaskMemoryManager. That means:

Comet competes with Spark’s own off-heap consumers (Tungsten sorters, BytesToBytesMap, and so on) for the same spark.memory.offHeap.size, and Spark’s unified memory manager arbitrates.
Spark can force Spark’s consumers in the same task to spill to satisfy Comet’s request.
The reverse does not hold. Comet registers a NativeMemoryConsumer with the TaskMemoryManager so that Spark has something to charge, but its spill() always returns 0 (CometTaskMemoryManager.java). A Spark allocation can therefore never make a native sorter or aggregate release its reservations. Native operators spill only when their own try_grow fails, so a JVM consumer that is blocked behind native reservations has no way to reclaim them.
A partial grant (acquired < additional) is released immediately and reported as ResourcesExhausted, which is the signal DataFusion uses to spill.

CometFairMemoryPool additionally applies a local check before it asks Spark. It divides pool_size by the number of consumers currently registered with the pool and rejects the request if the pool’s total reserved bytes plus the request would exceed that quotient. Two details matter for tuning:

The comparison is against the shared total (state.used), not against the requesting consumer’s own reservation. With an 8 GiB pool and two registered consumers, once one of them holds 3 GiB a 2 GiB request from the other is rejected, even though neither would exceed 4 GiB. In effect the usable pool shrinks to pool_size / num_consumers in aggregate as soon as more than one consumer is registered.
The pool is task-shared (see below), so num_consumers counts every registered consumer across every native plan in the task, not just the plan making the request.

This is why fair_unified spills earlier than greedy_unified. It is not a per-consumer quota, and reading it as one overstates the memory a multi-operator task can use.

Task-shared pools and their lifetime#

A single Spark task can run more than one native plan concurrently: a shuffle runs the pre-shuffle operators and the shuffle writer as separate native execution contexts. If each got its own pool, the per-task limit would be enforced once per plan rather than once per task.

acquire_task_shared_pool (task_shared.rs) keeps a process-wide HashMap<task_attempt_id, Weak<TaskSharedMemoryPool>>. Plans in the same task upgrade the existing Weak and share one pool; the returned Arc is the only lifetime handle, so the registry entry disappears when the last plan (and its last reservation) drops. There is no explicit release call to forget, and a createPlan that fails partway through cleans up on unwind.

TaskSharedMemoryPool::drop has to handle one race: an acquire can observe an expired Weak and insert a replacement before the dying pool reaches the registry lock. The drop therefore compares pointers and only removes an entry that is still its own.

How DataFusion consumes the pool#

Native operators reserve through DataFusion’s MemoryConsumer / MemoryReservation API:

try_grow(n) may fail. Spillable operators (ExternalSorter, the grouped hash aggregate, sort-merge join) respond to a ResourcesExhausted error by spilling to disk and retrying. This is the only mechanism that turns memory pressure into progress rather than failure.
grow(n) is infallible and panics if the pool refuses. It is used where the caller cannot spill.
shrink(n) returns bytes to the pool.

An operator that never calls try_grow is invisible to the pool no matter how much memory it uses.

Crossing the FFI boundary#

Batches move between the JVM and native over the Arrow C Data and C Stream interfaces, which are zero-copy. Nothing is copied, so the allocator that produced a batch and the runtime that decides when it dies can be on opposite sides of the boundary. See Arrow FFI for the mechanics; what matters here is who is charged and who controls the lifetime.

Two things are easy to conflate here. The allocator of a buffer determines which process-level accounting (if any) saw the bytes appear. Whether the buffer is reserved in Comet’s pool is a 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. 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 on who holds them. DataFusion’s ExternalSorter reserves get_reserved_bytes_for_record_batch for every batch it retains, and the hash join build side reserves get_record_batch_memory_size for each incoming batch; both read buffer sizes off the ArrayData and apply equally to imported buffers. So an imported batch that a sort or join has buffered is reserved in Comet’s pool, and through a unified pool is charged to Spark. The same batch held by an operator that does not reserve (a projection or filter, or ScanExec itself between polls) is charged to nothing.

Native → JVM (CometExecIterator). DataFusion produces the batch in Rust, and whether it is still reserved when it reaches the boundary depends on the producing operator. The sort’s output is wrapped in a ReservationStream that shrinks the reservation by the batch’s size as each batch is emitted, so by the time prepare_output sees a sorted batch its bytes are no longer reserved. prepare_output then calls move_to_spark, which writes an FFI_ArrowArray whose release callback drops the Rust ArrayData; no reservation travels with it. The JVM wraps the pointers in ArrowBufs and the memory is freed when the JVM calls close(). An exported batch is therefore usually not pool-charged while the JVM holds it, but it is resident until the JVM releases it. A slow or backed-up JVM consumer keeps native memory alive that no accounting layer is counting.

The point is not that one direction is charged and the other is not. Buffer lifetime and reservation lifetime are independent. A reservation is a number an operator chose to declare and later withdraw; a buffer lives until its last reference drops, on whichever side of the boundary that happens. The two are kept in step only inside operators written to do so, and only while that operator holds the batch. Once the batch leaves the operator, in either direction, the reservation stops and the bytes keep going. That is the shape of the gap to hold in mind when reading the next section, and it is what any fix at the allocator level has to close.

The accounting gap#

The pool tracks declared reservations. Container RSS counts pages the process touched. The two diverge for several structural reasons:

Undeclared allocations. Arrow array builders, expression kernels producing intermediate arrays, decompression buffers, Parquet metadata structures, object_store request buffers, and tokio’s own machinery all allocate without reserving. Only operators that were explicitly written to reserve show up in the pool.
Rounding and padding. Arrow buffers are padded to 64-byte boundaries and builders grow by doubling, so a reservation of exactly n bytes routinely corresponds to more than n bytes of heap.
Allocator behavior. malloc-level fragmentation, size-class rounding, and jemalloc’s retained/dirty page cache all add resident bytes that no layer above the allocator can see. Freeing memory does not necessarily return pages to the OS.
Non-Rust allocations. Memory allocated by C dependencies through libc malloc, and anything mmaped, never passes through Rust’s GlobalAlloc.
Batches in flight across the FFI boundary. Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the JVM closes them (see Crossing the FFI boundary).

The practical consequence is that reserved() is a lower bound on Comet’s real footprint, and the gap is workload-dependent. spark.comet.exec.memoryPool.fraction exists purely so operators can hand-tune a haircut that covers the gap for their workload.

To measure the gap on a real query, enable tracing with the jemalloc feature and compare jemalloc_allocated against the summed thread_NNN_comet_memory_reserved values; see Tracing.

What the container sees#

On Kubernetes, Spark sizes the executor pod from ResourceProfile:

pod memory request = pod memory limit
                   = spark.executor.memory
                   + spark.executor.memoryOverhead   (default max(0.1 * executor.memory, 384 MiB))
                   + spark.memory.offHeap.size
                   + pyspark memory                  (Python applications only)

Both the request and the limit are set to this same value, so the pod’s cgroup memory.max is a hard ceiling on the sum of everything in the container. That cgroup counts, among other things:

the JVM heap (spark.executor.memory),
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),
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 memory.max, the kernel OOM killer kills the process. The failure mode is significantly worse than a task-level OOM: every task running on that executor dies, every cached block it held is lost and must be recomputed, and the shuffle files it produced become unavailable to downstream fetches. Spark’s driver sees only ExecutorLostFailure with exit code 137.

Two facts follow that are easy to get wrong:

spark.memory.offHeap.size is part of the pod limit, not extra headroom on top of it. Raising Comet’s off-heap budget on Kubernetes raises the pod’s memory request by the same amount, so the scheduler will place fewer executors per node rather than silently giving Comet more room.
spark.executor.memoryOverhead is the only slack in the container, and the JVM’s own non-heap usage already consumes a large part of it. Comet’s overshoot beyond its declared reservations eats into the same allowance.

YARN behaves analogously. The container size is the same sum, and the NodeManager kills containers that exceed it, but the kill is done by the NodeManager’s monitor rather than the kernel, so it is somewhat less abrupt.

Open problems#

Nothing described above prevents an executor from being OOM-killed. Every enforcement point Comet has bounds declared reservations, and the sections above describe several structural reasons why declared reservations are a lower bound on physical usage. The known gaps, roughly in order of how much they matter:

No signal for real native usage. The only way to observe the gap today is to enable tracing with the jemalloc feature and compare jemalloc_allocated against summed reservations after 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.
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 cannot trigger native spilling. NativeMemoryConsumer.spill() returns 0, so native reservations are released only when a native operator decides to spill on its own failed try_grow. JVM consumers in the same task can be starved behind them.

Issue #4576 tracks work on the first two.

Debugging memory issues#

Tool	What it gives you
`spark.comet.debug.memory=true`	`LoggingMemoryPool` logs every register/grow/shrink with the consumer name
`spark.comet.explain.native.enabled=true`	Native plan with per-operator metrics, including spill counts
Tracing	`jemalloc_allocated` vs summed pool reservations; the accounting gap
`TrackConsumersPool`	Names the top 10 consumers in `ResourcesExhausted` messages (always on)
`thresher`	Third-party crate that dumps a jemalloc heap profile at a threshold

A checklist for triaging an executor OOM kill:

Confirm which budget was exceeded. Exit code 137 / OOMKilled on the pod is the cgroup. Exit code 52 with java.lang.OutOfMemoryError in the executor log is JVM heap exhaustion; Spark treats it as fatal, so the executor is lost either way and the exit code is what distinguishes them. A failed task with SparkOutOfMemoryError and a surviving executor is Spark’s managed memory pool, which is the only one of the three that is recoverable at task level.
Compare jemalloc_allocated against the summed pool reservations from a trace. A large excess points at undeclared native allocations; a small excess points at the budget simply being too small, or at the JVM side.
Check spark.comet.batchSize against the schema width. Peak memory scales with batch_size * columns, and wide or deeply nested schemas amplify it.
Check whether the operators involved can spill at all. ShuffledHashJoin cannot, so spark.comet.exec.forceShuffledHashJoin=true converts a spillable sort-merge join into one that is not.