Resolve lost wakeup in SpillPoolReader with multiple concurrent SpillPoolWriters by pepijnve · Pull Request #23522 · apache/datafusion

pepijnve · 2026-07-13T00:37:02Z

Which issue does this PR close?

Closes SMJ + aggregate finish with one partition but stuck forever when number of partitions is big #23447

Rationale for this change

When multiple SpillPoolWriter clones concurrently push batches to the same channel, more than one non-finished SpillFile can be in flight. This happens because each SpillPoolWriter clone takes the current_write_file at the start of push_batch and puts it back when it's done. When multiple push_batch calls happen concurrently, only the first one will be able to take the current_write_file and the others will all create their own new spill file. Which one gets put back for subsequent use is a race condition.

If this occurred and the writers are all dropped before rotation happens in, multiple files in the files deque will be have writer_finished == false. The last writer drop logic in SpillPoolWriter::drop only finishes whatever file is the current_write_file as finished.

This can lead to a stalled situation when SpillPoolFile::poll_next catches up with the writer and returns Pending because writer_finished == false. A waker for the file is registered, but since the last writer drop logic only finishes and wakes whatever happens to be current_write_file, which may not be the current read file, the waker may end up never being notified.

There is a secondary waker that is registered on the spill pool itself, but due to fine grained locking, it is possible for the wake call in the last writer drop logic to be called before the waker registration.

What changes are included in this PR?

Add support for tracking multiple unfinished write files
Close all unfinished write files when the last writer is dropped
Removed writer_dropped field which was an unnecessary denormalisation of active_writer_count == 0

An additional benefit of tracking all unfinished write files is that excessive creation of tiny spill files is avoided when many writers are pushing batches concurrently.

Are these changes tested?

Reproduction case from linked issue was used to confirm fix

Are there any user-facing changes?

No

SemyonSinchenko · 2026-07-13T08:54:46Z

An additional benefit of tracking all unfinished write files is that excessive creation of tiny spill files is avoided when many writers are pushing batches concurrently.

That is interesting because I constantly facing the "Too many open files" when there is a heavy spill. I did not know it is a problem and I was thinking it is OK, just used ulimit -n 524288.

jayzhan211 · 2026-07-14T12:35:05Z

+        if !shared.current_write_files.is_empty() {
+            // Copy and clear `current_write_files` so we can release shared lock before locking files
+            let files = shared.current_write_files.clone();
+            shared.current_write_files.clear();


Do we really need to release shared lock before locking files

Does it make sense to take file out from current_write_files consecutively then we could avoid clone?

Do we really need to release shared lock before locking files

I'm not entirely sure. My assumption was that this was already carefully tuned, so I didn't want to change the lock nesting. If I understood it correctly, the intention is to avoid IO while the shared lock is held. InProgressSpilFile::finish can do IO so the shared lock is released before calling finish.

Does it make sense to take file out from current_write_files consecutively then we could avoid clone?

The clone/clear is actually kind of pointless. Might as well just mem::take the entire VecDequeue here. I've made that change.

I'm not entirely sure. My assumption was that this was already carefully tuned, so I didn't want to change the lock nesting. If I understood it correctly, the intention is to avoid IO while the shared lock is held. InProgressSpilFile::finish can do IO so the shared lock is released before calling finish.

we also need to be sure we don't have a deadlock due to lock inversion (aka try to take the file/shared locks in different orders across different code paths)

alamb · 2026-07-15T12:58:44Z

+            file_shared.writer_finished = true;
+            // Wake reader waiting on this file (it's now finished)
+            file_shared.wake();
+            // Don't put back current_write_file - let it rotate


Suggested change

// Don't put back current_write_file - let it rotate

// Don't put back current_write_files - let it rotate

Singular is actually correct here, but I'll change the wording a bit to clarify what's being done.

alamb

This looks good to me -- thank you @pepijnve and @jayzhan211 (BTW 👋 it is great to see you again). I think adding some more comments explaining the expected state of the various queues would help make this code easier to understand in the future, but doesn't have to happen as part of this PR

cc @adriangb who added I think added SpillPoolShared in the first place

#18207

cc @xanderbailey who added some of this share pool machnery

#20672

alamb · 2026-07-15T13:04:01Z

 use super::spill_manager::SpillManager;

 /// Shared state between the writer and readers of a spill pool.
 /// This contains the queue of files and coordination state.


I this these comments would be nice to update to reflect the new design

Specifically note that the current_write_files is a list of files that are currently open but not being acively written to. When a file is being actively written to, active writer count is incremented by one and an entry is removed from current_write_files

I'll do another pass over the code and update the comments. I didn't pay sufficient attention to those.

alamb · 2026-07-15T13:05:32Z


            // No files in queue - check if writer is done
-            if shared.writer_dropped {
+            if shared.active_writer_count == 0 {


I wonder (maybe as a follow on PR) if adding some sort of RAAI guard for updating active_writer_count would make the code les error prone.

Or maybe even having a separate list of files being actively written (rather than just a count 🤔 )

The SpillPoolWriter plays the role of RAAI guard itself, doesn't it? It's incremented when a SpillPoolWriter is cloned, and decremented on drop.

The name active_writer_count is a bit misleading. It does not reflect the number of writers currently writing a batch. It reflects the number of writers that still exist.

alamb · 2026-07-15T13:07:08Z

+        if !shared.current_write_files.is_empty() {
+            // Copy and clear `current_write_files` so we can release shared lock before locking files
+            let files = shared.current_write_files.clone();
+            shared.current_write_files.clear();


I'm not entirely sure. My assumption was that this was already carefully tuned, so I didn't want to change the lock nesting. If I understood it correctly, the intention is to avoid IO while the shared lock is held. InProgressSpilFile::finish can do IO so the shared lock is released before calling finish.

we also need to be sure we don't have a deadlock due to lock inversion (aka try to take the file/shared locks in different orders across different code paths)

pepijnve · 2026-07-15T15:13:22Z

I've updated the documentation and code comments for correctness. I've deemphasised FIFO semantics quite a bit. The existing documentation stated that this was always guaranteed, but that was not actually the case. Prior to this PR the following chain of events could happen:

Writer A starts a batch push of B1, takes current_write_file F1
Writer B starts a batch push B2, current_write_file is None so it starts a new file F2
Writer B finishes its batch push, and sets current_write_file to F2
Writer A finished its batch push, and sets current_write_file to F1
Writer B pushes another batch B3 which is written current_write_file which is F1

The reader will observe batches in the order B1, B3, B2.

This MR assumes that we can take this even further, since the FIFO guarantee already doesn't hold when multiple writers are present. Would be good to get some feedback from @adriangb indeed since the non-FIFO behaviour was already present in the first commit of this code unless I'm completely misreading it.

I'm going to try to write a test that demonstrates the behaviour.

adriangb · 2026-07-15T15:31:26Z

I'm worried about dropping the FIFO semantics (even if they were broken).
Consider a plan like Projection(inner=Repartition(inner=Sort, preserve_order=true))): if RepartitionExec does not actually preserve order under spilling then the sort order will be lost and the results will be wrong.

pepijnve · 2026-07-15T15:50:16Z

I'm worried about dropping the FIFO semantics (even if they were broken). Consider a plan like Projection(inner=Repartition(inner=Sort, preserve_order=true))): if RepartitionExec does not actually preserve order under spilling then the sort order will be lost and the results will be wrong.

I'm probably missing a bit too much context. Could you help me define the exact order guarantee that's being promised by the spill pool? Single writer remains FIFO, but in the presence of multiple concurrent writers, what order can you actually guarantee for the reader? The best you can do I think is guarantee that the relative order of the batches per writer is retained. Is that sufficient?

pepijnve · 2026-07-15T16:35:27Z

BTW, I had based myself a bit on this comment in RepartitionExecState

// Create spill channels based on mode:
// - preserve_order: one spill channel per (input, output) pair for proper FIFO ordering
// - non-preserve-order: one shared spill channel per output partition since all inputs
//   share the same receiver

When preserve_order is requested, you get an SPSC channel per input and FIFO is guaranteed. When it is not requested, you get a single MPSC channel for all inputs and order is undefined. That aspect hasn't changed. What I've tried to clarify in the documentation is that the FIFO guarantee is only in the SPSC case, not the MPSC one.

adriangb · 2026-07-15T17:14:53Z

That makes sense, I was just checking that and came to the same conclusion. As long as SPSC preserves ordering / FIFO we are good 👍🏻

The spill pool's FIFO guarantee only holds for a single writer: with multiple concurrent `SpillPoolWriter` clones the reader can observe batches out of write order. That is fine for non-preserve-order `RepartitionExec` (the output is an unordered multiset), but for `preserve_order = true` the per-(input, output) stream feeds an order-sensitive `StreamingMerge`, so losing FIFO would silently produce wrong (unsorted) results. Previously the "single writer per ordered pool" invariant was upheld only by convention: one `preserve_order` bool drove two independent decisions (channel count vs. writer cloning) in two places, coupled only by comments. A future edit could break one without the other. Encode the invariant in the type system instead: - `channel()` now returns `SpillPoolWriter`, which is **not** `Clone`, so an ordered pool can only ever have one writer (enforced at compile time). It wraps the shared implementation and delegates `push_batch`. - `shared_channel()` returns the `Clone` `SharedSpillPoolWriter` for the multi-producer, per-writer-FIFO case. - `RepartitionExec` selects the topology in one place: `preserve_order` builds one dedicated ordered writer per input (moved, never cloned) via `PartitionSpillWriters::PerInput`; non-preserve builds one shared writer cloned across inputs via `PartitionSpillWriters::Shared`. Now feeding an ordered pool with a shared multi-producer writer simply does not compile. Adds `test_preserve_order_with_spill_file_rotation`, which forces a spill file per batch (`max_spill_file_size_bytes = 1`) and asserts each output partition stays sorted — exercising FIFO across file rotation. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WsjHVeKZtrSHugLwNFURjL

pepijnve · 2026-07-15T18:20:42Z

I've added a test_concurrent_writers test case that uses a 10 thread worker pool with 10 writers so that there's actual concurrency in the test. In contrast to test_concurrent_reader_writer there is no FIFO order assertion. Instead I'm asserting that no batches were lost.

Enforce single-writer spill pools for preserve_order via the type system

pepijnve · 2026-07-15T18:26:35Z

Thanks for the suggestion @adriangb I was thinking of doing something similar as well. I'll handle the merge conflict with main

# Conflicts: # datafusion/physical-plan/src/repartition/mod.rs

…e SpillWriter wrapper type

pepijnve · 2026-07-15T19:25:44Z

I flipped around the nesting of SpillPoolWriter and SharedSpillPoolWriter. You can then create a SpillPoolWriter via SharedSpillPoolWriter::new_writer. This has the benefit of not needing the SpillWriter wrapper in RepartitionExec. The tradeoff is that the writer type itself no longer strictly communicates the spsc vs mpsc and ordering characteristic of the channel. We still retain the guarantee from the compiler that channel is single producer and shared_channel can be multi producer.

I've edited the documentation a bit to be less verbose. Claude seemed kind of happy to repeat the fact that indeed SharedSpillPoolWriter is Clone over and over again.

adriangb

This is super nice!

adriangb · 2026-07-15T22:04:08Z

+        match self {
+            PartitionSpillWriters::PerInput(writers) => writers[input]
+                .take()
+                .expect("spill writer for input partition requested more than once"),


It looks like this gets called from a fallible function, we could (not saying we should) make this error instead of panic.

adriangb · 2026-07-15T22:05:52Z

+/// A tuple of `(SharedSpillPoolWriter, SendableRecordBatchStream)` that share the same
+/// underlying pool. The reader is returned as a stream for immediate use with
+/// async stream combinators. The writer can be cloned to create additional writers.
+pub fn shared_channel(


could this be pub(crate)?

So... I completely forgot about backwards compatibility for a moment. The spill_pool module, channel function, and SpillPoolWriter type were all public. The changes we're making here are breaking for existing code. Should I reshuffle things a bit to maintain the MPSC capable variant under the original names and introduce new ones for the SPSC only variant instead?

edit: I've gone ahead and done this already. Could use some help on naming these things. Best I could come up with is SpillPoolWriter for the MP version and SpillPoolSink for the non-clonable one. A bit too generic to my liking, but I didn't have anything better right away.

Technically, yes. If it's not a huge burden lets use new names and deprecate the old methods / names according to the API health guide (https://datafusion.apache.org/contributor-guide/api-health.html).

I'm fine with renaming and having deprecated aliases. What should the new names be though?

For the factory functions I went for spsc_channel and mpsc_channel. The existing channel is marked deprecated and simply calls mpsc_channel which is what the effective behaviour was before.

The writers is a bit trickier. We originally repurposed SpillPoolWriter to be the SP one and introduced SharedSpillPoolWriter as the MP one. That's not an option since we end up changing the definition of an existing type.

jayzhan211 · 2026-07-16T14:15:24Z

+                let mut file_shared = file.lock();
+
+                // Finish the current writer if it exists
+                if let Some(mut writer) = file_shared.writer.take() {


alamb · 2026-07-16T18:13:55Z

@jayzhan211 is back! 🎉

alamb

This PR is looking really nice to me @pepijnve @adriangb and @jayzhan211

However, since the scope has significantly grown since I first reviewed the PR I don't think it would be a good idea to backport this change as is without some "bake time" on main first...

#22547

Perhaps we can apply the earlier (smaller) fix for 54?

alamb · 2026-07-16T18:16:57Z

+/// The set of spill-pool writers for a single output partition, before they are handed to the
+/// per-input tasks. The variant encodes the repartition mode so the wrong writer topology cannot
+/// be constructed for a given mode.
+enum PartitionSpillWriters {


this is much clearer

github-actions · 2026-07-16T18:28:30Z

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details

     Cloning apache/main
    Building datafusion-physical-plan v54.0.0 (current)
       Built [  35.668s] (current)
     Parsing datafusion-physical-plan v54.0.0 (current)
      Parsed [   0.147s] (current)
    Building datafusion-physical-plan v54.0.0 (baseline)
       Built [  35.298s] (baseline)
     Parsing datafusion-physical-plan v54.0.0 (baseline)
      Parsed [   0.150s] (baseline)
    Checking datafusion-physical-plan v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   1.031s] 223 checks: 222 pass, 1 fail, 0 warn, 30 skip

--- failure function_marked_deprecated: function #[deprecated] added ---

Description:
A function is now #[deprecated]. Downstream crates will get a compiler warning when using this function.
        ref: https://doc.rust-lang.org/reference/attributes/diagnostics.html#the-deprecated-attribute
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.48.0/src/lints/function_marked_deprecated.ron

Failed in:
  function datafusion_physical_plan::spill::spill_pool::channel in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/spill/spill_pool.rs:480

     Summary semver requires new minor version: 0 major and 1 minor checks failed
    Finished [  74.105s] datafusion-physical-plan

pepijnve · 2026-07-16T18:36:15Z

This PR is looking really nice to me @pepijnve @adriangb and @jayzhan211

However, since the scope has significantly grown since I first reviewed the PR I don't think it would be a good idea to backport this change as is without some "bake time" on main first...

Release DataFusion 54.1.0 (minor/patch) Release #22547

Perhaps we can apply the earlier (smaller) fix for 54?

I can go back and redo just the code changes and test additions on a back port branch. Makes sense to keep it as small as possible for a patch.

codecov-commenter · 2026-07-16T18:52:02Z

Codecov Report

❌ Patch coverage is 89.55224% with 21 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@95de385). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/repartition/mod.rs	86.41%	2 Missing and 9 partials ⚠️
datafusion/physical-plan/src/spill/spill_pool.rs	91.66%	6 Missing and 4 partials ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #23522   +/-   ##
=======================================
  Coverage        ?   80.65%           
=======================================
  Files           ?     1086           
  Lines           ?   366430           
  Branches        ?   366430           
=======================================
  Hits            ?   295561           
  Misses          ?    53244           
  Partials        ?    17625

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

alamb · 2026-07-16T18:54:01Z

(I am resolving the CI failures now)

pepijnve · 2026-07-16T19:08:56Z

❌ Patch coverage is 89.55224% with 21 lines in your changes missing coverage. Please review.

Does every patch now require 100% coverage of changed lines?

alamb · 2026-07-16T19:37:35Z

Does every patch now require 100% coverage of changed lines?

No!

We are sorting through the coverage reports (and they are somewhat useless now b/c there isn't a base to compare against):

Code coverage report not working #23647

pepijnve · 2026-07-16T19:43:32Z

No!

That's a relief. In the meantime I've created the minimal backport PR version of this.

xanderbailey · 2026-07-16T21:52:01Z

This is a very nice improvement, thanks for working on this!

#23654) ## Which issue does this PR close? - part of #22547 54.x branch backport of fix for #23447 ## Rationale for this change See main PR #23522 ## What changes are included in this PR? See main PR #23522 ## Are these changes tested? Additional test cases added to verify existing behaviour Fix manually tested with reproduction code from #23447 ## Are there any user-facing changes? No

hareshkh · 2026-07-17T12:22:44Z

+    /// Maximum size in bytes before rotating to a new file.
+    /// Typically set from configuration `datafusion.execution.max_spill_file_size_bytes`.
+    max_file_size_bytes: usize,
+    /// Shared state with readers (includes current_write_file for coordination)


nit:

(includes current_write_file for coordination)

is stale

Since this branch has been merged already it would be best to make a PR with a fix

github-actions Bot added the physical-plan Changes to the physical-plan crate label Jul 13, 2026

pepijnve changed the title ~~Improve robustness of SpillPoolReader with multiple concurrent writers~~ Resolve lost wakeup in SpillPoolReader with multiple concurrent SpillPoolWriters Jul 13, 2026

pepijnve mentioned this pull request Jul 13, 2026

Release DataFusion 54.1.0 (minor/patch) Release #22547

Open

31 tasks

jayzhan211 reviewed Jul 14, 2026

View reviewed changes

pepijnve force-pushed the issue_23447 branch from 518c07f to 1042dea Compare July 14, 2026 14:04

Improve robustness of SpillPoolReader with multiple concurrent writers

191cceb

pepijnve force-pushed the issue_23447 branch from 1042dea to 191cceb Compare July 14, 2026 14:09

alamb reviewed Jul 15, 2026

View reviewed changes

alamb approved these changes Jul 15, 2026

View reviewed changes

Update documentation

c50bc33

adriangb and others added 2 commits July 15, 2026 12:43

Add MPSC usage test

7cb38fb

adriangb mentioned this pull request Jul 15, 2026

Enforce single-writer spill pools for preserve_order via the type system pepijnve/datafusion#3

Merged

Merge pull request #3 from pydantic/spill-pool-typed-writers

73db692

Enforce single-writer spill pools for preserve_order via the type system

pepijnve added 5 commits July 15, 2026 20:32

Fix test compile error

363b280

Merge remote-tracking branch 'upstream/main' into issue_23447

b8f9719

# Conflicts: # datafusion/physical-plan/src/repartition/mod.rs

Invert SharedSpillPoolWriter/SpillPoolWriter relationship to eliminat…

2f2c52d

…e SpillWriter wrapper type

Formatting

ac6dcb3

Rename shared_writer to new_writer

1e5d680

pepijnve added 2 commits July 15, 2026 21:36

More documentation tweaks

37075ad

Fix typo in docstring

bb67abd

adriangb approved these changes Jul 15, 2026

View reviewed changes

pepijnve added 5 commits July 16, 2026 09:40

Eliminate panic code path

79d4648

Fix compile error

64e6125

Revert name changes to avoid backwards incompatible change

f5032f1

Update docstrings

8972148

Documentation updates

fc84765

jayzhan211 reviewed Jul 16, 2026

View reviewed changes

jayzhan211 approved these changes Jul 16, 2026

View reviewed changes

alamb approved these changes Jul 16, 2026

View reviewed changes

Merge remote-tracking branch 'apache/main' into issue_23447

536055c

github-actions Bot added the auto detected api change Auto detected API change label Jul 16, 2026

fmt

1d46d57

fix links

c1c3d88

pepijnve mentioned this pull request Jul 16, 2026

[branch-54] Resolve lost wakeup in SpillPoolReader with multiple conc… #23654

Merged

alamb added this pull request to the merge queue Jul 16, 2026

Merged via the queue into apache:main with commit 12fa0ce Jul 16, 2026
40 checks passed

hareshkh reviewed Jul 17, 2026

View reviewed changes

	// Don't put back current_write_file - let it rotate
	// Don't put back current_write_files - let it rotate

Uh oh!

Conversation

pepijnve commented Jul 13, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

SemyonSinchenko commented Jul 13, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pepijnve Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pepijnve commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

adriangb commented Jul 15, 2026

Uh oh!

pepijnve commented Jul 15, 2026

Uh oh!

pepijnve commented Jul 15, 2026

Uh oh!

adriangb commented Jul 15, 2026

Uh oh!

pepijnve commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pepijnve commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pepijnve commented Jul 15, 2026

Uh oh!

adriangb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pepijnve Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Jul 16, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

pepijnve Jul 14, 2026 •

edited

Loading

pepijnve commented Jul 15, 2026 •

edited

Loading

pepijnve commented Jul 15, 2026 •

edited

Loading

pepijnve commented Jul 15, 2026 •

edited

Loading

pepijnve Jul 16, 2026 •

edited

Loading

github-actions Bot commented Jul 16, 2026 •

edited

Loading

codecov-commenter commented Jul 16, 2026 •

edited

Loading