andygrove opened a new pull request, #6128: URL: https://github.com/apache/datafusion-comet/pull/6128
## Which issue does this PR close? Closes #6127. ## Rationale for this change `MemoryPool::grow` is documented as "Infallibly grow the provided `reservation` by `additional` bytes. This must always succeed". DataFusion calls it for memory that already exists and only needs recording. Both Comet pools implemented it as `self.try_grow(..).unwrap()`, so whenever Spark granted less than requested the task panicked and Spark had to retry it. This happens today on TPC-H SF100 with `spark.memory.offHeap.size=2g`: the sort-merge join's `restore_spilled_batches` calls `grow` for a spilled batch it has just read back from disk, and the task dies at `fair_pool.rs:122`. The panic was introduced deliberately by #1732, to fix #1733: the old `grow` ignored a partial grant, so `shrink` could hand Spark back more than it had granted. This PR keeps that guarantee without panicking, along the lines of the alternative proposed in #1731. ## What changes are included in this PR? - **New `memory_pools/overcommit.rs`.** An `Overcommit` ledger records bytes a pool has recorded beyond what Spark granted, and a `granted` helper clamps Spark's reply. It makes no JNI calls, so it is unit tested directly. - **`grow` in `CometUnifiedMemoryPool` and `CometFairMemoryPool`.** It acquires what Spark will grant, records the full amount in `used`, and carries the shortfall as overcommit. It never fails. In the fair pool it also ignores the per-consumer fair limit, since the memory is already in use. A failed JNI call counts as a zero grant. - **`shrink` in both pools.** It repays overcommit first and releases only the remainder to Spark, so Spark is never handed back more than it granted. - **`reserved()` includes the overcommit,** so the next `try_grow` is refused and that operator spills instead. `try_grow` is unchanged. It does not try to re-acquire outstanding overcommit from Spark: the debt is cleared only by shrinks. ## How are these changes tested? - Unit tests for the ledger: a shrink repays the debt before releasing, a shrink with no debt releases everything, and grants are clamped to the request. - The pools themselves call Spark over JNI and have no Rust unit tests. They were exercised end to end with TPC-H SF100 at `spark.memory.offHeap.size=2g` (2 executors x 8 cores, Spark 4.1.1, `fair_unified`), alternating runs with and without this change. Results will be added here when the runs finish. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
