docs: explain allocator hazards and diagram where memory is allocated

[andygrove]

Which issue does this PR close?

No separate issue. This is a documentation-only improvement to the memory management contributor
guide, describing behavior as it exists on main.

Rationale for this change

The memory management guide explains the layout of Comet's allocators thoroughly, but almost none
of the hazards. The constraints that keep Comet's memory accounting correct are currently implicit
in CometTaskMemoryManager, so a reasonable-looking change can break them silently. Making
NativeMemoryConsumer.getUsed() take a lock, for example, would risk a deadlock against a native
reservation arriving on a Comet Tokio worker, and nothing in the guide explains why it reads an
AtomicLong instead.

Three smaller gaps came up while reading the page:

• It uses "off-heap" and "native heap" throughout without ever distinguishing them, even though the
  difference is the reason Comet's accounting is hard: both sit outside the JVM heap and both count
  toward container RSS, but only one is allocated by JVM code and can be reported to Spark.
• The accounting-gap section mentions memory allocated by "C dependencies" without naming any, which
  leaves the reader unable to act on it. Several of Comet's compression crates look like C bindings
  and are in fact pure Rust, so the guess a reader would make is wrong.
• The page had no diagrams, and "which allocator does this call site use, and who gets charged" is
  much easier to show than to describe.

What changes are included in this PR?

Most changes are to docs/source/contributor-guide/memory_management.md.

• A new "Constraints on a Comet memory consumer" subsection, documenting the rules that
  CometTaskMemoryManager already follows on main: why getUsed and spill must stay lock-free,
  why scala.util.control.NonFatal does not contain an acquireExecutionMemory call (a spilling
  consumer can raise SparkOutOfMemoryError, which is an Error, and ExecutionMemoryPool
  parks in lock.wait() so a task kill raises InterruptedException), why Spark exposes no
  per-consumer usage figure and a consumer that bypasses acquireMemory keeps an inherited used
  of zero, how a partial grant can be stranded when a later spill throws, and why a consumer whose
  spill() returns 0 takes budget it can never give back.

• Two mermaid diagrams. One maps each JVM allocation site to its allocator and on to whoever is
  charged, which on main means every Arrow site reaching the same unaccounted root. The other
  shows everything the pod cgroup counts toward RSS, grouped into what Spark can see, what only
  Comet's native pool has declared, and what nobody accounts for.

• An explanation of off-heap versus native heap, placed directly under the allocator table where
  both terms appear, including the point that spark.memory.offHeap.size budgets both even though
  the bytes live in different places.

• Concrete C dependencies in the non-Rust allocation bullet: libzstd, libhdfs and the object
  store TLS stack in a default build, plus the allocator itself under the jemalloc or mimalloc
  feature. It also names the codecs that are pure Rust in this build (snap, lz4_flex,
  zlib-rs, libbz2-rs-sys) and therefore do pass through GlobalAlloc.

• Plainer wording for spark.comet.exec.memoryPool.fraction, which was described twice as a
  "haircut".

• Terminology consistency in two other pages. Auditing the user and contributor guides for
  off-heap versus native heap usage turned up two genuine mix-ups. ffi.md had a native-to-JVM
  lifecycle table whose column header read Off-heap/Native even though nothing in that direction
  is off-heap, and an ArrowBuf annotated (off-heap) inside a box labelled JVM Heap. tuning.md
  said only that Comet "shares an off-heap memory pool with Spark", which reads as though Comet's
  own allocations are off-heap; it now adds that the pool is a shared budget rather than a shared
  allocator. Every other hit across both guides was checked and left alone: JVM shuffle pages
  really are off-heap, tracing.md correctly says native memory, and the spark-submit examples are
  just config names.

• A precise account of what the memory pool does and does not track, replacing the tuning
  guide's claim that Comet's memory accounting "isn't 100% accurate", which told the reader nothing
  actionable. It now states that the pool tracks only memory an operator explicitly reserves, names
  the cases where that happens (the sort buffer, the build side of a hash join, hash aggregation
  state, and the shuffle writer's buffered partitions), and lists what goes uncounted: per-batch
  working memory in kernels and array builders, decompression and Parquet reader buffers, object
  store and async-runtime machinery, JVM-side Arrow buffers, and allocator overhead. That makes the
  reason for spark.comet.exec.memoryPool.fraction concrete rather than a vague accuracy caveat.

How are these changes tested?

Documentation only, so there are no code tests. Verified by building the docs with the project's
Sphinx environment:

• sphinx-build -b html docs/source succeeds, with the warning count unchanged from the same build
  before these edits and no warning naming memory_management.
• Both mermaid fences render as diagrams rather than literal code blocks: the generated
  memory_management.html contains two class="mermaid" blocks and zero highlight-mermaid code
  blocks. This uses the sphinxcontrib-mermaid support and myst_fence_as_directive already on
  main.
• prettier --write reports the page unchanged, so it satisfies the prettier --check "**/*.md"
  preflight.

[andygrove]

diagrams:

[sunchao]

Correctness

Reviewed 0f9cbaa225b8 against 67168ca1f494. The existing guide described Comet's memory budgets and the gap between reservations and resident memory. This PR adds allocator call-site and container-accounting diagrams, explains the hazards around Spark memory consumers, and clarifies the FFI wrapper labels and tuning advice. All three changed files are documentation, so it changes no SQL evaluation, types, null handling, overflow, fallback or runtime allocation behavior.

There is one P2 correctness issue in the new container diagram: native reservations are shown as declared to Comet's pool only, outside Spark-visible accounting. Both supported off-heap pools forward these reservations to Spark through JNI. Their physical allocation remains native, but their reservation consumes Spark's execution budget. The inline comment asks for the diagram and accompanying visibility wording to preserve that distinction.

I checked the acquisition, release, consumer usage, spill and interruption descriptions against Comet's current implementation and the maintained Spark 3.5 and 4.0 sources. Those sources support the documented monitor/callback and partial-grant hazards. I also traced the eight JVM Arrow allocation sites and checked the locked codec versions and available allocator source. The required maintained Spark 3.4 and 4.1 branches were unavailable, so this does not qualify those versions.

Validation and CI

Both Mermaid diagrams render on GitHub at the reviewed commit. The local Sphinx toolchain is unavailable, so I did not run a Sphinx build. Comet CI is green at the reviewed head, but its build, Spark SQL, PyArrow, Iceberg and benchmark jobs were skipped. This is source and documentation validation, with no local Spark/JVM execution or performance measurement.

Performance

The change adds no executable work, allocations or hot-path overhead. The expanded tuning explanation usefully separates explicit reservations from kernel, codec, Arrow and allocator overhead. No benchmark is needed for the documentation-only changes, and the review does not establish a measured margin for spark.comet.exec.memoryPool.fraction. That margin remains workload-dependent.

Design

The two diagrams provide useful complementary views: where allocations originate, and how container usage relates to accounting. The consumer constraints also explain why the bridge uses its own usage counter and why a failed acquisition cannot be treated as a clean zero-byte grant. The actionable improvement is to keep allocation location and budget accounting distinct throughout, including the native reservation group called out above.

Abstraction & complexity

The PR introduces no runtime abstraction. Using existing allocator, pool and consumer names keeps the explanation close to the implementation, and changing ArrowBuf to a handle in the FFI diagram clarifies the wrapper's role. The diagrams earn their space by connecting call sites to those existing concepts. I found no additional complexity issue beyond the inconsistent accounting category.

[sunchao]

Correctness

[P2] Show native reservations as charged to Spark's accounting

Could we place NATRES under Spark-visible accounting, or label it as visible to both Spark and Comet? This page covers off-heap mode, where both fair_unified and greedy_unified forward successful try_grow reservations through JNI to CometTaskMemoryManager.acquireMemory. That calls Spark's acquireExecutionMemory and updates the native consumer's reported usage. For example, a successful 128 MiB native reservation consumes 128 MiB of Spark's off-heap execution budget, even though Rust allocates the actual buffers. The current separate “Comet's native pool only” group tells readers that these reservations are invisible to Spark, contradicting the earlier shared-budget explanation and allocator diagram. The unreserved native allocations belong outside Spark's accounting, but the declared reservations do not. Please make the new “no JVM metric counts it” paragraph consistent with this distinction too.

[andygrove]

You're right, and the diagram contradicted this page's own "the budget is shared even though the memory is not" paragraph a few sections up. I've moved the reservation node into the Spark-visible group and labelled it declared to Spark over JNI, never measured, which keeps allocation location and budget accounting distinct without needing a group of its own.

The caption underneath now says Spark's accounting covers that group in two different senses: the JVM heap, Tungsten pages and shuffle pages are reported by the JVM code that allocated them, whereas a native reservation is a number an operator declared before allocating, so the budget really is spent but the reservation is only a lower bound on the bytes behind it.

Three other places made the same claim, so I fixed those too. The visibility cell for the native allocator in the "Who allocates what" table now reads "Reservations only" rather than "No", the native-heap bullet points forward to the shared-budget paragraph instead of stopping at "no JVM metric counts it", and the overview no longer says native allocations are invisible to Spark's off-heap accounting.

Rebuilt with Sphinx: still 61 warnings, unchanged from before the edit, both fences still render as diagrams rather than code blocks, and prettier is clean.

[sunchao]

Re-reviewed 22d3078e57d5. The prior P2 is fixed: the container diagram now places declared native reservations under Spark-visible accounting, and the table, overview and visibility wording agree. Both off-heap pools charge those reservations through JNI. The reservation figure remains distinct from measured physical allocation. I found no remaining or new P1/P2 issue in the incremental change.

Both Mermaid diagrams render at this commit. Source checks cover the maintained Spark 3.5/4.0 branches. Current Comet CI is green, with build/test jobs skipped. Its checkout contains the same three changed documentation files as the reviewed head. The author reports a successful Sphinx build with unchanged warnings. I did not rerun Sphinx locally.

[andygrove]

Thanks for the review @sunchao

[comphead]

@andygrove for native heap do you mean Spark exec memory overhead area?