andygrove opened a new pull request, #6713: URL: https://github.com/apache/datafusion-comet/pull/6713
## 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`: #6544's Rust tests build the pools over a fake Spark through `create_pool`, a seam that #6271 added. #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 #6271's log reads, is not on this branch, and #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 #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 `ignore`s on this branch. That includes the two new #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 #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. -- 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]
