fix: [branch-1.1] let a spilled final aggregate read its spill files back past its memory share (#6544)

[andygrove]

Which issue does this PR close?

Closes #6254 on branch-1.1. #6544 already closed it on main.

Rationale for this change

This is the branch-1.1 backport of #6544, which has the backport-1.1 label. #6254 is a 1.1.0 regression (it carries the regression label). DataFusion 55 moved Comet's final hash aggregates onto FinalHashAggregateStream. Once one has spilled, it merges its sorted spill files and replays them through an OrderedFinalAggregateStream that has no way to spill, so a refused memory request during the replay fails the task. The merge reserves read buffers for as many spill files as fit, in a sibling reservation of the same consumer, so the replay often finds the consumer's share already taken. A final aggregate over input sorted on some of its grouping keys runs as an OrderedFinalAggregateStream and fails the same way.

The upstream fix, apache/datafusion#25383, is only on DataFusion main, and this branch is on DataFusion 55.1. #6544 works around the failure in Comet's memory pools, and this backport takes that workaround as it is.

Not proposed for branch-1.0: it is on DataFusion 54, which doesn't fail here, and #6544 has no backport-1.0 label.

What changes are included in this PR?

A cherry-pick (-x) of #6544's squash commit; see #6544 for what the change does. In short, a new SpillReplayPool wraps both Comet pools, between TaskSharedMemoryPool and TrackConsumersPool. When the pool refuses a try_grow from a final aggregate reading its spill files back, the wrapper records the request with the pool's grow instead. Every other call and every other refusal passes through unchanged. The change comes with seven Rust tests, two CometAggregateSuite tests and a section in the memory management guide.

The wrapper's logic in spill_replay.rs, its wiring in create_pool, both CometAggregateSuite test bodies and the guide's new text apply as on main. Three files conflicted (mod.rs, memory_management.md and CometAggregateSuite.scala), and the two pool files needed a small port. The adaptations:

• mod.rs, fair_pool.rs, unified_pool.rs: fix: let a spilled final aggregate read its spill files back past its memory share #6544's Rust tests build the pools over a fake Spark through create_pool, a seam that fix: count pool overcommit as untracked memory in the native memory usage log #6271 added. fix: count pool overcommit as untracked memory in the native memory usage log #6271 (the native memory usage log counting overcommit) is not on branch-1.1, so only the seam is ported, as it is on main: create_memory_pool calls create_pool with the Spark connection as a closure, with_spark is pub(super), and the pools' new constructors, which nothing calls any more, are removed. Nothing else in the two pool files changes, and the wrapper is wired in exactly as on main.
• spill_replay.rs: main's overcommit(), which fix: count pool overcommit as untracked memory in the native memory usage log #6271's log reads, is not on this branch, and fix: let a spilled final aggregate read its spill files back past its memory share #6544 only made it look through the wrapper. With no such reader, unwrap_spill_replay would be dead code and is dropped. Three tests drop their overcommit(&pool) assertions, which restate pool.reserved() minus what the fake Spark holds, and the one that shrinks the replay asserts pool.reserved() instead. The rest of the file is identical to main's.
• memory_management.md: the new section follows the task-shared pools section directly, because the paragraph on the shared CometTaskMemoryManager that precedes it on main comes from fix: return a native plan's memory before its Spark task ends #6261, which is not on this branch. The pool stack diagram and the section text are the same as on main.
• CometAggregateSuite.scala: only two import lines differ. This branch doesn't import Column, CometBaseAggregateExec or CometSortAggregateExec, so the merged lines add QueryTest and CometSortExec to this branch's imports.

How are these changes tested?

Run locally on this branch with JDK 17:

• Rust: cargo test -p datafusion-comet --lib passes (580 tests), including the seven in spill_replay.rs. cargo fmt --check and workspace clippy with --all-targets -D warnings are clean. My local toolchain is Rust 1.97, and CI uses stable.
• CometAggregateSuite, CometTaskMetricsSuite and CometExecIteratorLifecycleSuite on the default Spark 4.1 profile: 153 pass, with 2 that CometAggregateSuite already ignores on this branch. That includes the two new A native final aggregate that has spilled can fail the task during its replay #6254 tests, which each check that a native final aggregate ran and spilled.
• On Spark 3.5 and Scala 2.12 the sources compile, scalastyle, spotless and scalafix (RemoveUnused, in CHECK mode) pass, and the two new tests pass.
• Without the fix, both new tests fail. I took the SpillReplayPool wrapper out of create_pool, which leaves this branch's own pool stack, rebuilt the native library and ran the two tests. They fail with Failed to acquire 115952 bytes where this consumer already holds 3031120 bytes and the fair limit is 3145728 bytes and Failed to acquire 112528 bytes where this consumer already holds 2555968 bytes and the fair limit is 2621440 bytes, the same errors that fix: let a spilled final aggregate read its spill files back past its memory share #6544 reports. So this branch has the bug and the tests catch it. The fix changes no JVM code, so the JVM side was the same in both runs.
• prettier --check on the guide passes.

By CI's path routing, which I checked with dev/ci/compute-changes.py, this PR runs the all-profiles Linux build, macOS, the Spark SQL jobs for 3.5, 4.0 and 4.1, and the four Iceberg versions. The Spark SQL tests run Comet on-heap, which uses an unbounded pool, so they don't exercise this change, and no run-* label is needed.

[sunchao]

Summary

• Prior state and problem: DataFusion 55.1 can fill a final aggregate’s memory share with spill-merge buffers, leaving its non-spillable replay unable to proceed.
• Design approach: SpillReplayPool identifies replay through the consumer name and an occupied sibling reservation, then records refused growth through the existing grow mechanism.
• Correctness / compatibility analysis: Verified DataFusion’s reservation lifetimes, merge admission, replay and cleanup paths. Ordinary refusals remain intact, and releases repay overcommit before returning bytes to Spark. Compared relevant Spark aggregation and memory-manager sources across 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. No introduced P1/P2 issues found within this review.
• Key design decisions: The wrapper isolates the DataFusion-specific workaround from both pool implementations. Its assumptions are documented, and issue #6583 tracks removal. The backport preserves the upstream wrapper logic.
• Implementation sketch: Both off-heap pools receive the wrapper through create_pool. Seven Rust tests cover replay, rejection and accounting. Two Scala regression tests check results and native spilling. The memory guide describes the exception.
• Behavioral changes worth calling out: Replay may exceed the fair share, pool limit or Spark grant while remaining accounted for. The wrapper adds dispatch and final-aggregate map/mutex bookkeeping. Other consumers avoid its tracking mutex. Performance assessment was limited to code paths.
• Suggested improvements: None meeting the reproducible P1/P2 threshold.

Reviewed all six changed files at 8bcacb60688e0b3d03e7551fd24cd2eca7f2d56a against 7b7eec69282a7abc33a0e8e687ca594c72f140ec. The PR remains non-draft. The snapshot and refreshed discussion contained no reviews, issue comments or review threads.

Routed skills: review-comet-pr and review-comet-memory-pr.

Exact-head CI: The primary run has 7 successful checks, 18 in progress and 4 skipped, with no failures reported. Rust, JVM, macOS, Spark SQL and Iceberg validation remains pending. Completed label runs skipped substantive suites and do not establish integration success.

Local validation: cargo test -p datafusion-comet --lib memory_pools --offline --no-default-features --locked passed all 26 selected tests, including all seven new tests. No local JVM, Spark SQL or Iceberg suites were run. Default features and throughput/RSS measurements were not validated locally. Project source remains unchanged.