feat: add an async API for spill file writing

[Phoenix500526]

Which issue does this PR close?

• Closes Add an async API for spill file writing #23247.

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.

[Phoenix500526]

cc @alamb

[codecov-commenter]

Codecov Report

❌ Patch coverage is 74.39863% with 149 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.90%. Comparing base (0da2151) to head (a979e98).

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 ⚠️
...atafusion/physical-plan/src/spill/spill_manager.rs	70.27%	17 Missing and 5 partials ⚠️
datafusion/physical-plan/src/sorts/sort.rs	70.00%	17 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24923      +/-   ##
==========================================
- Coverage   81.92%   81.90%   -0.02%     
==========================================
  Files        1132     1132              
  Lines      421192   421682     +490     
  Branches   421192   421682     +490     
==========================================
+ Hits       345041   345384     +343     
- Misses      55756    55864     +108     
- Partials    20395    20434      +39     

☔ 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]

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_store → datafusion-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]

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

[Phoenix500526]

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.

[jayzhan211]

@Phoenix500526

The spill loop turns a memory shortfall into a hard failure, where the old code
deliberately did not:

// before
let reservation_failed = self.reservation.try_grow(sorted_size).is_err();
// Even if the reservation is not enough, the batch is already in
// memory, so it's okay to combine it with previously sorted batches
globally_sorted_batches.push(batch);
if reservation_failed { self.consume_and_spill_append(...)?; }

// after (sort.rs:517)
let workspace = self.merge_pool.borrow(sorted_size);
self.reservation.try_grow(sorted_size - workspace.size())?;   // <-- new failure path

That old comment is still true and the new code contradicts it. At this point the
batch was just produced by sorted_stream.next(), so it is already resident.
Returning ResourcesExhausted reclaims nothing — it only kills a query that would
otherwise have completed. Refusing an allocation is the right call when refusing
prevents it; here it doesn't.

This is a real regression, not a hypothetical: the PR relaxes an existing test to
keep it green —

-  .with_sort_spill_reservation_bytes(1),
+  .with_sort_spill_reservation_bytes(spill_workspace),
-  .with_memory_limit(batches_memory, 1.0)
+  .with_memory_limit(batches_memory + spill_workspace, 1.0)

— and test_spill_output_respects_memory_limit then asserts that the old
configuration now yields ResourcesExhausted. Users running a tight
sort_spill_reservation_bytes, or sorting wide Utf8View/string batches where one
output batch exceeds the remaining merge workspace, get a new OOM on a workload
that spilled fine before.

Please keep the workspace-borrow optimization but stop propagating the shortfall:

-            let spill_workspace = match self.reservation.try_grow(sorted_size) {
-                Ok(()) => None,
-                Err(_) => {
-                    let workspace = self.merge_pool.borrow(sorted_size);
-                    self.reservation.try_grow(sorted_size - workspace.size())?;
-                    Some(workspace)
-                }
-            };
+            let spill_workspace = match self.reservation.try_grow(sorted_size) {
+                Ok(()) => None,
+                Err(_) => {
+                    // Reuse already-reserved workspace where we can.
+                    let workspace = self.merge_pool.borrow(sorted_size);
+                    let remainder = sorted_size - workspace.size();
+                    // The batch is already in memory: failing here would not
+                    // reclaim it, so record the usage and spill it immediately.
+                    if self.reservation.try_grow(remainder).is_err() {
+                        self.reservation.grow(remainder);
+                    }
+                    Some(workspace)
+                }
+            };

With that, test_spill_output_respects_memory_limit should be dropped and the two
config changes in
should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_to_spill
reverted. If you believe the strict behaviour is genuinely the one we want, it needs
to be called out in the PR description and in the upgrade guide as a breaking
behavioural change — it shouldn't land as a side effect of adding an async writer API.

[Phoenix500526]

@Phoenix500526

The spill loop turns a memory shortfall into a hard failure, where the old code deliberately did not:

// before
let reservation_failed = self.reservation.try_grow(sorted_size).is_err();
// Even if the reservation is not enough, the batch is already in
// memory, so it's okay to combine it with previously sorted batches
globally_sorted_batches.push(batch);
if reservation_failed { self.consume_and_spill_append(...)?; }

// after (sort.rs:517)
let workspace = self.merge_pool.borrow(sorted_size);
self.reservation.try_grow(sorted_size - workspace.size())?;   // <-- new failure path

That old comment is still true and the new code contradicts it. At this point the batch was just produced by sorted_stream.next(), so it is already resident. Returning ResourcesExhausted reclaims nothing — it only kills a query that would otherwise have completed. Refusing an allocation is the right call when refusing prevents it; here it doesn't.

This is a real regression, not a hypothetical: the PR relaxes an existing test to keep it green —

-  .with_sort_spill_reservation_bytes(1),
+  .with_sort_spill_reservation_bytes(spill_workspace),
-  .with_memory_limit(batches_memory, 1.0)
+  .with_memory_limit(batches_memory + spill_workspace, 1.0)

— and test_spill_output_respects_memory_limit then asserts that the old configuration now yields ResourcesExhausted. Users running a tight sort_spill_reservation_bytes, or sorting wide Utf8View/string batches where one output batch exceeds the remaining merge workspace, get a new OOM on a workload that spilled fine before.

Please keep the workspace-borrow optimization but stop propagating the shortfall:

-            let spill_workspace = match self.reservation.try_grow(sorted_size) {
-                Ok(()) => None,
-                Err(_) => {
-                    let workspace = self.merge_pool.borrow(sorted_size);
-                    self.reservation.try_grow(sorted_size - workspace.size())?;
-                    Some(workspace)
-                }
-            };
+            let spill_workspace = match self.reservation.try_grow(sorted_size) {
+                Ok(()) => None,
+                Err(_) => {
+                    // Reuse already-reserved workspace where we can.
+                    let workspace = self.merge_pool.borrow(sorted_size);
+                    let remainder = sorted_size - workspace.size();
+                    // The batch is already in memory: failing here would not
+                    // reclaim it, so record the usage and spill it immediately.
+                    if self.reservation.try_grow(remainder).is_err() {
+                        self.reservation.grow(remainder);
+                    }
+                    Some(workspace)
+                }
+            };

With that, test_spill_output_respects_memory_limit should be dropped and the two config changes in should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_to_spill reverted. If you believe the strict behaviour is genuinely the one we want, it needs to be called out in the PR description and in the upgrade guide as a breaking behavioural change — it shouldn't land as a side effect of adding an async writer API.

Hi, @jayzhan211 thanks for pointing this out. In a979e98fc, I restored the existing spill fallback behavior while retaining workspace reuse and memory accounting across async writes. I also removed the strict-failure test and restored the original memory settings in the existing spill test. I opened #25183 to track the memory-limit guarantees and potential compatibility changes separately. We can discuss stricter enforcement there if needed, keeping this PR focused on async spill writing.

[jayzhan211]

Thanks @Phoenix500526 🚀