Skip to content

feat: remove memory accounting from Comet's on-heap mode - #6066

Open
andygrove wants to merge 4 commits into
apache:mainfrom
andygrove:remove-onheap-memory-accounting
Open

andygrove wants to merge 4 commits into
apache:mainfrom
andygrove:remove-onheap-memory-accounting

Conversation

@andygrove

@andygrove andygrove commented Sep 20, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #6063.

Rationale for this change

On-heap mode is not a production configuration. spark.comet.exec.onHeap.enabled defaults to
false, lives in CATEGORY_TESTING, and CometDriverPlugin.init disables Comet outright when
off-heap is off and the flag is not set. Comet's own suites run off-heap (CometTestBase sets
spark.memory.offHeap.enabled=true with 2 GiB). The only things that reach the on-heap path are
Spark's own SQL suite and the Iceberg suites, which set ENABLE_COMET_ONHEAP=true because Spark's
test harness does not configure off-heap memory.

That path nevertheless carried a complete second memory-accounting implementation, and the
accounting it performed did not protect anything.

Comet cannot honestly join Spark's ledger in on-heap mode. NativeMemoryConsumer is hardcoded to
MemoryMode.OFF_HEAP, and Spark sizes the off-heap execution pool from spark.memory.offHeap.size
alone (MemoryManager.scala:61-66, byte-identical on 3.4.3, 3.5.9, 4.0.4 and 4.1.3), which is 0
when off-heap is disabled. Registering as an ON_HEAP consumer instead would be a category error:
Comet's bytes are in the Rust heap, so Spark would evict cached blocks and spill its own sorters to
make room for memory that is not on the heap, while doing nothing about the native RSS that
actually gets an executor OOM-killed.

The fixed-size DataFusion pool that stood in for one was sized from spark.comet.memoryOverhead,
which was in the initial commit (#1, February 2024), a month before CometTaskMemoryManager and
the unified pool existed (#83). It survived because deleting it broke the tests: #1062 required
off-heap and removed the on-heap branch, and #1177 restored it with the rationale "after #1062 we
have not been running Spark tests for native execution". Per #6054, the driver plugin could not
even fold that figure into the container on three of the five supported Spark versions.

What changes are included in this PR?

Native. parse_memory_pool_config returns UnboundedMemoryPool whenever off-heap is disabled
and ignores the pool-type string there. Six of the nine MemoryPoolType variants (greedy,
fair_spill, and their _task_shared and _global pairings) and their arms in mod.rs are gone,
as is memory_limit_per_task, which no off-heap pool read. createPlan loses the corresponding
JNI parameter.

JVM shuffle. CometBoundedShuffleMemoryAllocator becomes
CometUnboundedShuffleMemoryAllocator: a page table over UnsafeMemoryAllocator with no budget.
The pages still have to be Unsafe-allocated in either memory mode, because SpillWriter hands
their addresses to writeSortedFileNative for Rust to dereference and
TaskMemoryManager.allocatePage would return long[] heap pages here. With no shared budget there
is nothing to wait for, so the blocking-allocation protocol added in #5493 goes away — per-thread
retention accounting, the two fail-fast liveness predicates, the timeout, the progress logging and
the task-kill polling — and allocateBlocking collapses into allocate and leaves
CometShuffleMemoryAllocatorTrait. The allocator is now created per task like the off-heap one
rather than being an executor-wide singleton, and getUsed reads an AtomicLong instead of taking
the allocator's monitor, since TaskMemoryManager calls it while holding its own. Its page table
is the one limit it still has, and exhausting it now reports SparkOutOfMemoryError rather than
IllegalStateException, so the callers that already answer a refused acquisition by spilling and
retrying can make progress instead of failing the task. A page number indexes the allocator's own
table, so an address encoded by one instance cannot be resolved by another, and each caller has to
create one allocator and share it for the whole task. Both classes now say so.

Configs. spark.comet.memoryOverhead, spark.comet.exec.onHeap.memoryPool,
spark.comet.shuffle.jvm.memoryFactor and spark.comet.shuffle.jvm.memoryWaitTimeout are removed,
along with getCometMemoryOverhead*, getCometShuffleMemorySize and shouldOverrideMemoryConf.
spark.comet.exec.onHeap.enabled stays: it is still the switch that keeps Comet off in on-heap
mode unless a test opts in, and its doc now says the mode accounts for nothing.

Driver plugin. CometDriverPlugin.init no longer adjusts spark.executor.memoryOverhead;
there is nothing left to add. Both ShimCometDriverPlugin files are deleted, since
getMemoryOverheadMinMib was only needed by the removed calculation.

CI and docs. Both pr_build_linux.yml and pr_build_macos.yml name the renamed allocator
suite. Without that check-suites fails Preflight and the suite runs in neither job. The JVM
shuffle contributor guide's Memory Management section described a single allocator that spills when
an allocation fails, which is now true of the off-heap path only, so it describes both paths and
the row-count triggers instead.

Diffs. The four dev/diffs patches drop .set("spark.comet.memoryOverhead", ...). They were
regenerated through the documented flow (clone at the tag, apply, edit, git diff <tag>), not
hand-edited, and each was verified to apply cleanly to a pristine tree afterwards. Regenerating
4.0.4 and 4.1.3 also corrects a stale index line for ParquetRowIndexSuite.scala, whose recorded
post-image hash did not match what git apply actually produces. That line is inert — CI applies
with plain git apply — but it is now consistent with the hunks.

Behavior change

Nothing bounds Comet's allocations in on-heap mode any more. The on-heap pool was the only bound
the Spark SQL suite ran under, so any memory-pressure-driven native spill it triggers today stops
happening, and nothing caps Comet's RSS in those jobs. Given the suite's data sizes the spill
coverage is probably near zero already, but that is an assumption rather than a measurement.
Row-count spilling still triggers independently of any pool, but only the bypass writer really
keeps it. CometDiskBlockWriter spills at
min(spark.comet.shuffle.jvm.spillThreshold, spark.comet.shuffle.jvm.batchSize), and the batch
size defaults to 8192, so that path stays bounded per writer. CometShuffleExternalSorter compares
against spark.comet.shuffle.jvm.spillThreshold alone, which defaults to Int.MaxValue, and its
remaining triggers were allocation failures that the unbounded allocator no longer raises. So on
the sort-based path a task in on-heap mode now buffers its whole map output until it closes. That
writer is selected when the partition count exceeds spark.shuffle.sort.bypassMergeThreshold or
partitions times cores exceeds spark.comet.shuffle.jvm.maxWritersPerExecutor.
spark.comet.shuffle.native.maxBufferBytes still triggers independently of any pool. If CI
memory regresses, the cheap recovery is a single executor-wide GreedyMemoryPool with a fixed cap,
which is one arm in parse_memory_pool_config rather than the whole subsystem.

Off-heap mode is untouched.

Overlap with #6054

#6054 is in the merge queue and also removes the driver plugin's spark.executor.memoryOverhead
mutation, shouldOverrideMemoryConf and the shim files. This PR has to remove them too, because it
deletes the config they consumed. I will rebase onto main once #6054 lands and keep its
warnIfExecutorMemoryOverheadUnset addition.

How are these changes tested?

Existing tests, adjusted:

  • CometUnboundedShuffleMemoryAllocatorSuite (renamed) keeps the getUsed accounting and
    page-table-exhaustion tests, drops the two that asserted budget exhaustion, and adds one
    asserting that a 64 MiB request now succeeds. The page-table test asserts the
    SparkOutOfMemoryError above, and that freeing a page lets the next allocation through, which is
    what lets a spilling writer recover.
  • CometDiskBlockWriterSuite loses the five tests that covered the deleted blocking protocol. The
    three that cover behavior this PR keeps — write() reclaiming buffered pages on a fatal error, the
    SpillSorter constructor not leaking on a failed allocation, and the unsafe writer allocating
    nothing before write() — are ported to off-heap, where TestMemoryManager.limit supplies the
    same pressure and the assertions become getUsed == 0 /
    getMemoryConsumptionForThisTask == 0. That also moves them onto the production allocator.
  • SpillSorterSuite now builds one allocator per test in beforeEach instead of relying on the
    executor-wide singleton to hand the same instance to two call sites.
  • CometPluginsSuite's two override assertions become "the configured value stands" and "nothing
    is added when none was configured".
  • The two CometExecSuite "spill sort with (multiple) dictionaries" tests lose a withSQLConf
    wrapper for spark.comet.memoryOverhead that was already inert twice over: the suite runs
    off-heap, and the config is read from SparkConf rather than SQLConf.

Local verification so far: the four regenerated diffs apply cleanly to pristine v3.4.3, v3.5.9,
v4.0.4 and v4.1.3 trees; CometDiskBlockWriterSuite and SpillSorterSuite pass. Broader suite
runs are still in progress and I will report them here. This change touches the serde-adjacent JNI
signature, a native operator path and dev/diffs, so it needs the Spark SQL and Iceberg suites
before it is queued.

@github-actions github-actions Bot added enhancement New feature or request area:shuffle Shuffle (JVM and native) area:memory Memory pools, reservations, OOM handling labels Sep 20, 2026
@andygrove andygrove added run-spark-4.1-tests Run the Spark 4.1 SQL tests on this pull request instead of waiting for the merge queue run-iceberg-tests run-spark-3.4-tests Run the Spark 3.4 SQL tests on this PR run-spark-3.5-tests Run the Spark 3.5 SQL tests on this PR run-spark-4.0-tests labels Sep 20, 2026
@andygrove
andygrove marked this pull request as ready for review September 21, 2026 05:04
@rich7420

Copy link
Copy Markdown
Contributor

@andygrove please resolve the conflicts, overall lgtm

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correctness

This removes the separate accounting used by Comet's opt-in on-heap testing mode. Native plans now use UnboundedMemoryPool, and JVM shuffle uses a task-owned Unsafe page table. The production off-heap pool selection, budget formula and Spark allocation path remain unchanged.

I checked the maintained Spark 3.5 and 4.0 sources. Their on-heap Tungsten allocator returns heap-backed pages, so keeping Unsafe pages for rows dereferenced by Rust is necessary. Both JVM writer entry points share one allocator with every component that decodes its page addresses. The JNI caller, declaration and native entry point remove the same argument. Page-table exhaustion is reported before allocating or changing usage, and the existing managed-OOM handlers can spill and retry. Constructor and writer cleanup remain responsible for these pages, since Spark's task cleanup does not own them.

[P2] Preserve the compatibility path for released on-heap settings. The inline comment covers the unconditional removal of configuration behavior shipped in 1.0.0 while this branch targets 1.1.0. The repository's configuration compatibility policy requires a deprecation/migration path. I found no additional P1/P2 issue in the allocator or JNI changes.

Validation

The head has 99 successful checks and 70 skips. Spark SQL jobs for 3.4.3, 3.5.9, 4.0.4 and 4.1.3 and Iceberg 1.8–1.11 workflows completed successfully. Sampled logs explicitly enable on-heap Comet. The shuffle job ran the renamed allocator suite, SpillSorterSuite and CometDiskBlockWriterSuite, with 491 Scala tests passing.

Those jobs checked out historical merge b2c2942440f8, combining this head f1c4e61045f7 with d62d3382ac5d. Every authored file matches this head at that tested merge. The current API merge is unavailable because of the conflict already discussed on the PR. This evidence does not qualify a later conflict resolution. Local validation was source and caller analysis plus git diff --check. The maintained Spark 3.4 and 4.1 branches were unavailable, so their runtime CI coverage is separate from the canonical-source comparison.

Performance

Removing the executor-wide on-heap allocator also removes its shared monitor contention, retention maps and blocking allocation protocol. Off-heap allocation behavior is unchanged. The tradeoff is higher possible memory use in the testing mode: native memory pressure no longer causes spilling, and the JVM sort path can retain map output until its row, pointer-array or page-table boundary. The bypass writer retains its row-count spill trigger. The completed CI runs establish test completion, not an RSS bound or a measured performance improvement. No benchmark was run.

Could you add a small before/after on-heap memory-pressure microbenchmark, recording peak RSS, spill counts and elapsed time for a multi-batch sort or aggregation? That would quantify the effect of removing the old bound; the green CI runs do not measure it.

Design

The two remaining allocation backends have a clear purpose: Spark arbitrates off-heap execution, while the testing mode supplies native-addressable pages without a second budget. Task ownership also makes the page namespace easier to reason about. The documentation explains the intentionally unbounded mode and loss of memory-pressure spill coverage. The compatibility transition for existing explicit settings needs to accompany that design change.

Abstraction & complexity

Deleting the six on-heap pool variants, the unused off-heap per-task limit, and allocateBlocking removes substantial special-case code. The remaining allocation trait still serves the writers' common allocate/free/address operations. Tests that depended on the removed blocking protocol are deleted, and the retained failure-cleanup tests use a deliberately constrained Spark memory manager. I found no additional abstraction or complexity issue worth raising.

// changing Spark's memory configuration, and native memory cannot be charged to Spark's
// on-heap pool, so nothing is accounted. See the memory management contributor guide.
logDebug("on-heap mode: native memory is unbounded and unaccounted")
MemoryConfig(offHeapMode, memoryPoolType = "unbounded", memoryLimit = 0)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correctness

[P2] Keep a compatibility path for the released on-heap memory controls

Could we preserve a deprecated path for the existing settings before making this unconditional in 1.1? For example, an on-heap test invocation with spark.comet.exec.onHeap.enabled=true, spark.comet.memoryOverhead=128m and spark.comet.exec.onHeap.memoryPool=greedy previously selected a finite native reservation pool. It now silently selects an unbounded pool, regardless of those explicit settings. The shuffle budget settings disappear too.

memoryOverhead, onHeap.memoryPool and shuffle.jvm.memoryFactor all shipped in 1.0.0. The versioning policy covers spark.comet.* configuration semantics, including these testing-category keys, and requires deprecation before removal in a major release. Keeping the old behavior available through the documented compatibility mechanism, with an upgrade-guide entry, would let this simplification land without silently removing an existing test harness's configured bound.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@sunchao Hmm .. technically, this is correct, but this is a test-only config intended for running Spark SQL test suite and it is disabled by default. Our docs already say that no user should use on-heap mode. I'd like to push back against this review issue. WDYT?

On-heap mode exists so that Spark's own SQL suite and the Iceberg suites can
run against Comet without changing Spark's memory configuration. It is off by
default, sits in CATEGORY_TESTING, and the driver plugin disables Comet
entirely unless a test opts in. Comet's own suites run off-heap.

The accounting that mode performed did not protect anything. Native memory is
not on the JVM heap, so there is no Spark pool it can honestly be charged to:
NativeMemoryConsumer is hardcoded to MemoryMode.OFF_HEAP and Spark sizes that
pool from spark.memory.offHeap.size alone, which is 0 when off-heap is
disabled. The fixed-size DataFusion pool that stood in for one was sized from
spark.comet.memoryOverhead, a config that predates the unified pool and that
the driver plugin could not fold into the container on three of five supported
Spark versions.

On-heap mode now gets UnboundedMemoryPool on the native side and an unbounded
per-task Unsafe page allocator on the JVM side, which removes:

- six of the nine MemoryPoolType variants and memory_limit_per_task
- spark.comet.memoryOverhead, spark.comet.exec.onHeap.memoryPool,
  spark.comet.shuffle.jvm.memoryFactor and
  spark.comet.shuffle.jvm.memoryWaitTimeout
- the shared-pool blocking-allocation protocol in what is now
  CometUnboundedShuffleMemoryAllocator, along with allocateBlocking
- the driver plugin's spark.executor.memoryOverhead mutation and the
  ShimCometDriverPlugin files that only it needed

spark.comet.exec.onHeap.enabled stays: it is still the switch that keeps Comet
off in on-heap mode unless a test opts in.

Closes apache#6063
Four follow-ups from a self review of this PR.

The renamed `CometUnboundedShuffleMemoryAllocatorSuite` was never renamed in
`pr_build_linux.yml` or `pr_build_macos.yml`, so `check-suites` failed Preflight
and the suite ran in neither job.

`CometUnboundedShuffleMemoryAllocator` never refuses a page, which leaves its
page table as the only limit it has. Exhausting it threw `IllegalStateException`,
which nothing catches, so the task failed. Report `SparkOutOfMemoryError`
instead, which `SpillWriter.acquireNewPageIfNecessary` and
`CometShuffleExternalSorter.growPointerArrayIfNecessary` already answer by
spilling and retrying.

A page number indexes one allocator's own table, so an address encoded by one
instance cannot be resolved by another. That was previously guaranteed by the
executor-wide singleton and is now a requirement on callers, so say so on both
the allocator and the factory.

The JVM shuffle guide described a single allocator that spills when an
allocation fails, which is true of the off-heap path only.
Scalafix `RemoveUnused` fails the Lint Java jobs on the 3.4, 3.5 and 4.0
profiles:

    -  private implicit val signaler: Signaler = ThreadSignaler
    +  ThreadSignaler

This PR removed the five tests that covered the blocking-allocation protocol,
and they were the only callers of `failAfter`. With those gone the `TimeLimits`
mixin and the `Signaler` that served it are unreferenced, so remove both rather
than the val alone.

Lint did not catch this earlier in the PR because Preflight was failing on the
renamed allocator suite and gated every downstream job. Spark 4.1 is excluded
from the lint matrix, so formatting locally on the default profile does not run
scalafix at all.
…pool

apache#6054 landed on main and its comment above warnIfExecutorMemoryOverheadUnset
named spark.comet.memoryOverhead, which this branch removes, and said the share
operators reserve is charged against a memory pool, which is now true in
off-heap mode only.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@andygrove I agree with the test-only scope in your reply. The flag is disabled by default, the removed controls are explicitly testing or internal, and the existing guide already excludes on-heap mode from production. I would drop my request to retain the old pools or add a legacy flag solely for this testing mode.

[P2] I would narrow the remaining compatibility concern to policy and merge order. The currently merged policy still covers these keys. The explicit testing/internal exemption in #6089 is a reasonable remedy. Could we land that exemption first and rebase this PR onto it? My review of #6089 objects to treating absence from configs.md as proof that a key is internal, not to the explicit testing/internal carve-out. Once that narrower exemption lands, this compatibility concern is resolved without restoring the old accounting machinery.

Re-reviewed 97217c9cbbf9 against 09b44ad6fa17, including the rebase. The production off-heap allocator and native pool implementations remain unchanged, the JNI parameter removal matches on both sides, and task-local allocator sharing and cleanup remain consistent. The rebase keeps #6054’s warnIfExecutorMemoryOverheadUnset. I found no new P1/P2 implementation issue.

Current CI has 85 successes, 10 skips and 2 failures. The shuffle job passed 495 tests, including the allocator, sorter and failure-cleanup cases. However, the Spark 4.0 SQL shard failed SQLAppStatusListenerMemoryLeakSuite’s “no memory leak” test at noLiveData(), which checks retained listener state. Required Checks consequently failed. This log does not establish an allocator leak, and I have not established that the failure is unrelated to the PR. Both sampled jobs checked out 33d8ceec9ee8, whose tree exactly matches this head.

There is still no new before/after RSS or spill measurement, so the passing jobs establish completion for their workloads, not a memory bound or performance improvement. Local checks were source/ABI analysis and git diff --check. Maintained Spark 3.5/4.0 sources were checked; maintained 3.4/4.1 sources remain unavailable. No local full suite or benchmark was run.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:memory Memory pools, reservations, OOM handling area:shuffle Shuffle (JVM and native) enhancement New feature or request run-iceberg-tests run-spark-3.4-tests Run the Spark 3.4 SQL tests on this PR run-spark-3.5-tests Run the Spark 3.5 SQL tests on this PR run-spark-4.0-tests run-spark-4.1-tests Run the Spark 4.1 SQL tests on this pull request instead of waiting for the merge queue

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Remove memory accounting from Comet's on-heap mode

3 participants