fix: let a spilled final aggregate read its spill files back past its memory share

docs/source/contributor-guide/memory_management.md

@@ -253,14 +253,17 @@ ignored and the pool is always `UnboundedMemoryPool`.
 the inside out, a Comet plan in the default configuration sees:
 
 ```text
-[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
+[LoggingMemoryPool]          <- only when spark.comet.debug.memory=true
+  [TaskSharedMemoryPool]     <- RAII handle for the per-task registry
+    [SpillReplayPool]        <- lets a spilled final aggregate read its spill files back
+      [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.
+reports the base pool's number. `SpillReplayPool` also turns one kind of refused `try_grow` into a
+`grow`; see
+[Final aggregates reading their spill files back](#final-aggregates-reading-their-spill-files-back).
 
 ### The unified pools
 
@@ -346,6 +349,18 @@ hands every native plan in a task the same one and drops it when the task comple
 covers the whole task, so `CometExecIterator.close()` warns about memory still in use only when
 the task's last open native plan closes.
 
+### Final aggregates reading their spill files back
+
+`SpillReplayPool` (`spill_replay.rs`) wraps both Comet pools to work around
+[issue #6254](https://github.com/apache/datafusion-comet/issues/6254). Once one of DataFusion 55's
+final aggregates has spilled, it reads its spill files back through an aggregate that cannot spill,
+so a refused `try_grow` there fails the task. When the pool refuses such a request,
+`SpillReplayPool` records it with the pool's `grow` instead, which skips `CometFairMemoryPool`'s
+local checks and carries what Spark does not grant as overcommit. Every other refusal is passed on
+unchanged. `spill_replay.rs` describes how the wrapper recognizes these requests and what that relies
+on in DataFusion. [Issue #6583](https://github.com/apache/datafusion-comet/issues/6583) tracks
+removing it.
+
 ## How DataFusion consumes the pool
 
 Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservation` API:

native/core/src/execution/memory_pools/mod.rs

@@ -20,13 +20,15 @@ mod fair_pool;
 pub mod logging_pool;
 mod plan_pool;
 mod spark_memory;
+mod spill_replay;
 mod task_shared;
 mod unified_pool;
 
 use datafusion::execution::memory_pool::{MemoryPool, TrackConsumersPool, UnboundedMemoryPool};
 use fair_pool::CometFairMemoryPool;
 use jni::objects::{Global, JObject};
 use spark_memory::SparkMemory;
+use spill_replay::SpillReplayPool;
 use std::num::NonZeroUsize;
 use std::sync::Arc;
 use unified_pool::CometUnifiedMemoryPool;
@@ -70,10 +72,16 @@ fn create_pool(
 
     match pool_type {
         MemoryPoolType::GreedyUnified => acquire_task_shared_pool(task_attempt_id, || {
-            tracked(CometUnifiedMemoryPool::with_spark(spark()))
+            Arc::new(SpillReplayPool::new(
+                task_attempt_id,
+                tracked(CometUnifiedMemoryPool::with_spark(spark())),
+            ))
         }),
         MemoryPoolType::FairUnified => acquire_task_shared_pool(task_attempt_id, || {
-            tracked(CometFairMemoryPool::with_spark(spark(), pool_size))
+            Arc::new(SpillReplayPool::new(
+                task_attempt_id,
+                tracked(CometFairMemoryPool::with_spark(spark(), pool_size)),
+            ))
         }),
         MemoryPoolType::Unbounded => Arc::new(UnboundedMemoryPool::default()),
     }
@@ -85,6 +93,7 @@ fn create_pool(
 /// takes nothing from Spark or that the function did not create.
 pub(crate) fn overcommit(pool: &Arc<dyn MemoryPool>) -> usize {
     let pool = task_shared::unwrap_task_shared(pool).unwrap_or(pool);
+    let pool = spill_replay::unwrap_spill_replay(pool).unwrap_or(pool);
     if let Some(tracked) = pool.downcast_ref::<TrackConsumersPool<CometUnifiedMemoryPool>>() {
         tracked.inner().overcommit()
     } else if let Some(tracked) = pool.downcast_ref::<TrackConsumersPool<CometFairMemoryPool>>() {