andygrove opened a new pull request, #5959: URL: https://github.com/apache/datafusion-comet/pull/5959
## Which issue does this PR close? Relates to #4576. Stacks on #5934 and should merge after it; until then the diff includes that PR's commits. ## Rationale for this change Every Comet memory pool only counts what operators voluntarily reserve. Native memory that bypasses the pool (scratch buffers in kernels, intermediate Arrow arrays, allocations inside DataFusion operators that are not reserved) is invisible to it until the executor exceeds its container limit and is killed. #5934 added a process-wide count of the bytes the Rust allocator has actually handed out, but only as a tracing metric. This PR makes that count available in every build and adds the smallest pool that acts on it: before asking Spark for memory, it refuses a reservation when real native usage plus the request would exceed Comet's off-heap budget. The operators that can spill then spill; the ones that cannot fail the task instead of the executor. The pool is opt-in and the default stays `fair_unified`. ## What changes are included in this PR? - `alloc-accounting` is now a default cargo feature. All `#[cfg(feature)]` guards stay, so a build with `--no-default-features` still compiles without the wrapper. The bench and the tracing guide document the opt-out invocation. - New `CheckedMemoryPool<P: MemoryPool>` (`native/core/src/execution/memory_pools/checked_pool.rs`). `try_grow` returns `ResourcesExhausted` when `native_allocated + additional > budget`, naming the request, the bytes in use, the budget and the reserved total, and otherwise delegates to the inner pool. `grow`, `shrink`, `register`, `unregister` and `reserved` delegate; `memory_limit` reports the budget. The inner pool is generic so the gate is unit-testable without a JVM. - New off-heap pool type `greedy_unified_checked` for `spark.comet.exec.memoryPool`: `CheckedMemoryPool` around `CometUnifiedMemoryPool`, budget = `spark.memory.offHeap.size` x `spark.comet.exec.memoryPool.fraction` (the value `fair_unified` already receives), inside the same task-shared and consumer-tracking wrappers as `greedy_unified`. If the native library was built without the feature, selecting this pool is a config error rather than a silently unchecked pool. - Docs: `spark.comet.exec.memoryPool` description and a `tuning.md` entry describing what the check is and is not. Semantics worth stating up front: the balance and the budget are both process-wide. When any task pushes real usage to the budget, every task's next non-zero reservation is denied. There is no per-task attribution, and allocations themselves are never refused; this is a reservation gate, not a hard limit. Per-task attribution and changing the default pool are follow-ups. ## How are these changes tested? Rust unit tests: - `denies_when_real_bytes_plus_request_exceed_the_budget`: pool over `UnboundedMemoryPool` with budget = current balance + 64 MiB; holding a 256 MiB block the pool never heard about makes a 1-byte `try_grow` fail with `ResourcesExhausted` and nothing reserved; dropping the block makes the same call succeed. - Zero-byte grows never fail, `memory_limit` reports the budget, successful grows and shrinks reach the inner pool. - `parse_memory_pool_config` accepts `greedy_unified_checked` in off-heap mode with the memory limit as budget, rejects it in on-heap mode, and (in a `--no-default-features` build) rejects it with a message naming the feature. JVM tests in `CometExecSuite`: - `greedy_unified_checked pool completes a sort within budget`: a dictionary-heavy sort under the new pool matches Spark and runs natively. - `greedy_unified_checked pool denies reservations once real native usage exceeds budget`: with `memoryPool.fraction=0.000001` the budget is a few KiB, below what the executor's native code already holds, so the sort's first reservation is denied with the checked pool's message. Both feature sets were run locally: default (`alloc-accounting` on) and `--no-default-features --features hdfs-opendal`. -- 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]
