JonSnow1807 opened a new pull request, #7607: URL: https://github.com/apache/datafusion-comet/pull/7607
## Which issue does this PR close? Closes #7244. ## Rationale for this change Two native memory paths treat an exception from Spark during a memory request as a refusal: - The multi-partition shuffle repartitioner spills on any `try_grow` error, so an exception (for example a failed spill of another consumer, or a killed task) makes it spill and carry on. - An infallible `grow` on the Spark-backed pools logs the failed call and books the size as overcommit, so the exception only reaches the log. ## What changes are included in this PR? - `multi_partition.rs`: spill only when the error is a refusal (`ResourcesExhausted`, found with `find_root`, the same check as the Iceberg writer). Any other error fails the shuffle. - `SparkMemory` keeps the first request to Spark that fails, in `grow` or `try_grow`, and every later `try_grow` of the task then fails without asking Spark. The first one after a failed `grow` gets the original error, so the task fails with the exception Spark threw. - The fair pool checks this under its lock, before its own limits, so its own refusal can’t hide the failure. - Overcommit accounting is unchanged, and `grow` stays infallible. - Notes in `memory_management.md` and `native_shuffle.md`. The failure is kept rather than handed out once because DataFusion’s sort and multi-level merge take any `try_grow` error as a refusal and ask again, which would drop a one-time error. One gap is left: a failed `grow` with no later `try_grow` in the task is still only logged. Closing it would mean checking for the failure when a plan’s stream ends in `jni_api.rs`, which changes the JNI plan lifecycle, so I left it out. I can add it here or in a follow-up. ## How are these changes tested? New unit tests, each failing on main and passing with this change: - `pool_errors_other_than_refusal_fail_the_shuffle` - `a_failed_acquire_fails_every_later_try_acquire` - `a_failed_try_acquire_fails_every_later_one` - `a_failed_grow_fails_every_later_try_grow_of_the_task` (greedy and fair unified pools) `a_refused_reservation_spills` checks that a refusal, also when wrapped in `Context` or `Shared`, still spills. `cargo test -p datafusion-comet-shuffle` and `cargo test -p datafusion-comet --lib` pass, and clippy is clean. -- 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]
