andygrove opened a new pull request, #6222: URL: https://github.com/apache/datafusion-comet/pull/6222
Backport of #6128 to `branch-1.0`. Cherry-picked from `d3406c6d2196a03f5fbfd781091ac9c0f520b943`. The changes to both pools and the new `spark_memory.rs` are byte-identical to upstream. One two-line hunk in `memory_pools/mod.rs` was applied by hand, and the docs changes were dropped. Both adaptations are described under "What changes are included" below. ## Which issue does this PR close? Closes #6127 and #1733 on `branch-1.0`. Listed in #6201. ## Rationale for this change The bug ships in 1.0.0 through the same code as on `main` before #6128. `fair_pool.rs` and `unified_pool.rs` on `branch-1.0` are identical to `main`'s copies just before #6128. Both implement `MemoryPool::grow` as `try_grow(..).unwrap()`, so the task panics whenever Spark grants less than a `grow` call asks for. `grow` must not fail, because DataFusion calls it for memory that already exists. For example, DataFusion 54.1's sort-merge join calls it when it reads a spilled buffered batch back from disk. `fair_unified`, one of the two affected pools, is the default off-heap pool on `branch-1.0`. ## What changes are included in this PR? The fix is the original one, so see #6128 for the details. The adaptations: - In `memory_pools/mod.rs`, the `mod spark_memory;` declaration and the new `task_attempt_id` argument to `CometFairMemoryPool::new` were added by hand. On `main`, #5494 rewrote `create_memory_pool` before #6128, so upstream's hunk had nothing to apply to here. If #5494 is backported later, its `mod.rs` conflict needs to keep that argument. - The changes to `docs/source/contributor-guide/memory_management.md` and `.ai/skills/review-comet-memory-pr/SKILL.md` are dropped, because neither file exists on `branch-1.0`. They come from #5933 and #6018. ## How are these changes tested? Same tests as the original PR, run locally on `branch-1.0`: - All nine new tests pass: six for `SparkMemory`, one for `CometFairMemoryPool` and two for `CometUnifiedMemoryPool`. - The bug is present on `branch-1.0`, and the tests catch it. With `grow` in both pools reverted to `try_grow(..).unwrap()` and the tests kept, the three pool-level tests fail with a panic in `grow`, while the six `SparkMemory` tests still pass. - `cargo fmt --all -- --check` and `cargo clippy --all-targets --workspace -- -D warnings` pass. The JNI path has no Rust unit tests. The TPC-H SF100 runs described in #6128 were on `main`, and I did not repeat them on `branch-1.0`. ## Are there any user-facing changes? The same as #6128, with no config or API changes. A task whose `grow` call would have panicked now keeps running. The bytes Spark did not grant are carried as overcommit until they are released. While any overcommit is outstanding, `try_grow` in either pool asks Spark for it on top of the request, all or nothing. So `try_grow` is refused until Spark can cover both, and the operator spills. Overcommit is reported in the pools' `Display` output and in their `try_grow` errors. -- 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]
