feat: add an async API for spill file writing - #24923
Conversation
|
cc @alamb |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24923 +/- ##
==========================================
- Coverage 81.69% 81.68% -0.02%
==========================================
Files 1127 1127
Lines 415386 415965 +579
Branches 415386 415965 +579
==========================================
+ Hits 339360 339783 +423
- Misses 56091 56200 +109
- Partials 19935 19982 +47 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
55c2846 to
a1b03e7
Compare
jayzhan211
left a comment
There was a problem hiding this comment.
Thanks @Phoenix500526 , I have a suggestion:
tokio::time::timeout in AsyncIPCStreamWriter::drop panics on runtimes without a time driver, and tokio/time isn't declared
spill/mod.rs:652 is the first production-path tokio::time use in datafusion-physical-plan — every other one in this crate is #[cfg(test)]. Two problems:
-
Undeclared feature.
datafusion/physical-plan/Cargo.toml:91istokio = { workspace = true }→["macros", "rt", "sync"].cargo tree -p datafusion-physical-plan -e features -i tokioshowstimearriving only throughobject_store→datafusion-execution. This compiles by accident of workspace feature unification — the same fragility the codec comment ~30 lines below warns about (#21917). Ifobject_storedropstime, the crate stops building. -
Runtime panic.
tokio::time::timeoutpanics when the embedder's runtime has no time driver (Builder::new_multi_thread().enable_io().build(), for example) — and DataFusion doesn't own the runtime. Since the panic happens in a detached task whoseJoinHandleis dropped, unwind builds swallow it silently: the abort never runs, nothing is logged, and the documented best-effort cleanup quietly becomes no cleanup. Underpanic = "abort"it kills the process.
The semaphore already bounds cleanup concurrency, and AsyncSpillWriter's own docs tell backends to configure lifecycle cleanup for abandoned uploads — which is where a deadline belongs. Suggest dropping the timeout and declaring the feature:
--- a/datafusion/physical-plan/Cargo.toml
+++ b/datafusion/physical-plan/Cargo.toml
-tokio = { workspace = true }
+tokio = { workspace = true, features = ["time"] }--- a/datafusion/physical-plan/src/spill/mod.rs
+++ b/datafusion/physical-plan/src/spill/mod.rs
const MAX_CONCURRENT_SPILL_ABORTS: usize = 8;
-const SPILL_ABORT_TIMEOUT: Duration = Duration::from_secs(30);
static SPILL_ABORT_PERMITS: LazyLock<Arc<tokio::sync::Semaphore>> =
LazyLock::new(|| Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_SPILL_ABORTS)));
@@
let _abort_task = handle.spawn(async move {
- match tokio::time::timeout(SPILL_ABORT_TIMEOUT, writer.abort()).await {
- Ok(Ok(())) => {}
- Ok(Err(error)) => {
- debug!("Failed to abort dropped spill writer: {error}");
- }
- Err(_) => {
- debug!("Timed out aborting dropped spill writer");
- }
+ // Backends are responsible for bounding their own abort latency;
+ // see the `AsyncSpillWriter` lifecycle-cleanup contract.
+ if let Err(error) = writer.abort().await {
+ debug!("Failed to abort dropped spill writer: {error}");
}
// `permit` is released when this bounded cleanup task exits.
drop(permit);
});|
@jayzhan211 thanks for catching this — addressed in 14cf7f5: the drop-time timeout is removed, the semaphore limit is retained, and |
Object store spill backends previously had to block while uploading data. Add an async writer path so external sort can stream multipart uploads without buffering complete spill files in memory. CLOSES apache#23247
Embedders may use Tokio runtimes without a time driver. Let backends bound abort latency while retaining the cleanup limit. Refs apache#23247
Decimal formatting only borrows its input. Preserve that ownership contract to satisfy all-feature Clippy checks.
Async spill writes must retain their input budget without bypassing pool limits. Reuse reserved workspace and release consumed merge inputs so spilling can progress within its budget. Refs apache#23247
e3214ec to
33f4ab2
Compare
|
Hi, @jayzhan211 , after rebasing this PR, I've pushed three follow-up commits:
The remaining CI check, |
Which issue does this PR close?
Rationale for this change
SpillFilesupports asynchronous reads, but spill writes currently usestd::io::Write. Backends built on asynchronous storage APIs must either block while uploading data or buffer an entire spill file in memory before uploading it.This PR adds an asynchronous spill writing path so remote backends can upload spill data as Arrow IPC buffers are produced.
What changes are included in this PR?
AsyncSpillWritertrait andSpillFile::open_async_writer.SpillFileimplementations through a default adapter backed byopen_writer.Spill paths that require a partially written local file to remain readable continue to use the synchronous API.
What is the testing strategy for this PR?
The added tests cover:
finishis called.The object store spill example was also run end to end. The full workspace test suite, extended tests, Clippy, Rustdoc, and repository lint checks pass.
Are there any user-facing changes?
Yes. Custom spill backends can implement
AsyncSpillWriterand overrideSpillFile::open_async_writerto use asynchronous storage APIs directly.This is an additive API change. Existing synchronous
SpillFileimplementations continue to work through the default adapter and do not need to be updated.