Skip to content

feat: add an async API for spill file writing - #24923

Open
Phoenix500526 wants to merge 4 commits into
apache:mainfrom
Phoenix500526:issue/23247
Open

feat: add an async API for spill file writing #24923
Phoenix500526 wants to merge 4 commits into
apache:mainfrom
Phoenix500526:issue/23247

Conversation

@Phoenix500526

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

SpillFile supports asynchronous reads, but spill writes currently use std::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?

  • Add the public AsyncSpillWriter trait and SpillFile::open_async_writer.
  • Preserve compatibility with existing SpillFile implementations through a default adapter backed by open_writer.
  • Separate Arrow IPC stream encoding from I/O so encoded buffers can be sent to synchronous or asynchronous writers.
  • Migrate external sort and stream-based spill herlpers to the asychronous writing path.
  • Keep memory reservations active while batches are waiting for asynchronous writes to complete.
  • And best-effort cleanup for incomplete uploads after errors or query cancellation.
  • Update the object store spill example to stream data through multipart uploads instread of buffering the complete spill file in memory.

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:

  • Arrow IPC round trips through an asynchronous spill writer.
  • Cleanup after write errors, failed aborts, and task cancellation.
  • Memory accounting while an asynchronous spill write is pending.
  • Multipart part uploads before finish is called.
  • Releasing the source Arrow allocation when a small multipart tail is buffered.

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 AsyncSpillWriter and override SpillFile::open_async_writer to use asynchronous storage APIs directly.

This is an additive API change. Existing synchronous SpillFile implementations continue to work through the default adapter and do not need to be updated.

@github-actions github-actions Bot added execution Related to the execution crate physical-plan Changes to the physical-plan crate labels Sep 3, 2026
@Phoenix500526

Copy link
Copy Markdown
Contributor Author

cc @alamb

@codecov-commenter

codecov-commenter commented Sep 3, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 76.44509% with 163 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.68%. Comparing base (2c31327) to head (33f4ab2).
⚠️ Report is 2 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-plan/src/spill/mod.rs 79.57% 36 Missing and 32 partials ⚠️
.../physical-plan/src/spill/in_progress_spill_file.rs 59.13% 34 Missing and 4 partials ⚠️
datafusion/physical-plan/src/sorts/sort.rs 80.37% 22 Missing and 9 partials ⚠️
...atafusion/physical-plan/src/spill/spill_manager.rs 70.27% 17 Missing and 5 partials ⚠️
datafusion/physical-plan/src/sorts/builder.rs 81.81% 1 Missing and 3 partials ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. Undeclared feature. datafusion/physical-plan/Cargo.toml:91 is tokio = { workspace = true }["macros", "rt", "sync"]. cargo tree -p datafusion-physical-plan -e features -i tokio shows time arriving only through object_storedatafusion-execution. This compiles by accident of workspace feature unification — the same fragility the codec comment ~30 lines below warns about (#21917). If object_store drops time, the crate stops building.

  2. Runtime panic. tokio::time::timeout panics 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 whose JoinHandle is dropped, unwind builds swallow it silently: the abort never runs, nothing is logged, and the documented best-effort cleanup quietly becomes no cleanup. Under panic = "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);
         });

@Phoenix500526

Copy link
Copy Markdown
Contributor Author

@jayzhan211 thanks for catching this — addressed in 14cf7f5: the drop-time timeout is removed, the semaphore limit is retained, and tokio/time is explicitly enabled. I also documented that backends should bound their own abort latency; the update is ready for another look.

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
@Phoenix500526

Copy link
Copy Markdown
Contributor Author

Hi, @jayzhan211 , after rebasing this PR, I've pushed three follow-up commits:

  • 29bab9cb0a applies your suggestion: removes the timeout from drop cleanup, keeps the concurrency limit, and docs that backends should bound their own abort latency. It also declares the tokio::time feature.
  • 7b7a396343 is a very simple modification that fixes a clippy warning in the PostgreSQL test helper by throwing BigDecimal instead of taking ownership.
  • 33f4ab2c8e8 adapts the async spill path to the retained workspace introduced in PR fix: preserve external sort workspace across spilling #24740. It replaces unchecked reservation growth with borrowing from available workspace, while keeping the output batch accounted for throughout the async write. It also writes each batch before requesting the next one and releases fully consumed merge inputs so their budget can be reused. The tests cover insufficient memory, reservation retention during pending writes, and cleanup after cancellation.

The remaining CI check, cargo test hash collisions(amd64), was canceled after reaching the six-hour limit. This appears unrelated to the spill changes: I found the same count-distinct tests taking hours and hitting the timeout on main and other branches. I've opened #25011 with the logs and reproduction steps to track this and ask about the intended test coverage before proposing any changes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

execution Related to the execution crate physical-plan Changes to the physical-plan crate sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add an async API for spill file writing

3 participants