You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "xanderbailey (via GitHub)" <gi...@apache.org> on 2026/03/03 14:43:11 UTC
[PR] Fix repartition from dropping data when spilling [datafusion]
xanderbailey opened a new pull request, #20672:
URL: https://github.com/apache/datafusion/pull/20672
## Which issue does this PR close?
- Closes #.
## Rationale for this change
In non-preserve-order repartitioning mode, all input partition tasks share clones of the same `SpillPoolWriter` for each output partition. `SpillPoolWriter` used `#[derive(Clone)]` but its `Drop` implementation unconditionally set `writer_dropped = true` and finalized the current spill file. This meant that when the **first** input task finished and its clone was dropped, the `SpillPoolReader` would see `writer_dropped = true` on an empty queue and return EOF — silently discarding every batch subsequently written by the still-running input tasks.
This bug requires three conditions to trigger:
1. Non-preserve-order repartitioning (so spill writers are cloned across input tasks)
2. Memory pressure causing batches to spill to disk
3. Input tasks finishing at different times (the common case with varying partition sizes)
## What changes are included in this PR?
**`datafusion/physical-plan/src/spill/spill_pool.rs`:**
- Added `active_writer_count: usize` to `SpillPoolShared` to track the number of live writer clones.
- Replaced `#[derive(Clone)]` on `SpillPoolWriter` with a manual `Clone` impl that increments `active_writer_count` under the shared lock.
- Updated `Drop` to decrement `active_writer_count` and only finalize the current file / set `writer_dropped = true` when the count reaches zero (i.e. the last clone is dropped). Non-last clones now return immediately from `Drop`.
- Added regression test `test_clone_drop_does_not_signal_eof_prematurely` that reproduces the exact failure: writer1 writes and drops, the reader drains the queue, then writer2 (still alive) writes. Without the fix the reader returns premature EOF and the assertion fails; with the fix the reader waits and reads both batches.
## Are these changes tested?
Yes. A new unit test (`test_clone_drop_does_not_signal_eof_prematurely`) directly reproduces the bug. It was verified to **fail without the fix** and **pass with the fix**. All 41 existing repartition tests and all 15 existing spill pool tests continue to pass.
## Are there any user-facing changes?
No API changes. Queries that repartition under memory pressure will no longer silently lose rows.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#issuecomment-4016128922
I made a backport here:
- https://github.com/apache/datafusion/pull/20777
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886047254
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -233,6 +249,15 @@ impl Drop for SpillPoolWriter {
fn drop(&mut self) {
let mut shared = self.shared.lock();
+ shared.active_writer_count -= 1;
Review Comment:
I like that idea, it seems more future proof, tried it here [52a06a4](https://github.com/apache/datafusion/pull/20672/commits/52a06a40889fb20837bbc4ad9083a9b92ba2859f) what do you think?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886623707
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -93,6 +89,37 @@ impl SpillPoolShared {
}
}
+/// Tracks the number of live [`SpillPoolWriter`] clones.
+///
+/// Cloning increments the count; dropping decrements it.
+/// [`WriterCount::is_last`] returns `true` when called from the final clone,
+/// which the writer uses to decide whether to finalize the spill pool.
+struct WriterCount(Arc<AtomicUsize>);
+
+impl WriterCount {
+ fn new() -> Self {
+ Self(Arc::new(AtomicUsize::new(1)))
+ }
+
+ /// Returns `true` if this is the only remaining clone.
+ fn is_last(&self) -> bool {
+ self.0.load(Ordering::Acquire) == 1
+ }
+}
+
+impl Clone for WriterCount {
+ fn clone(&self) -> Self {
+ self.0.fetch_add(1, Ordering::Relaxed);
Review Comment:
So thread A clones (increment), thread B drops and calls `is_last()`, but doesn't see A's increment yet — so B thinks it's the last writer and finalizes prematurely. I think today the control flow prevents this from happening `setup_input_partitions` clones and passed through a task spawn boundary before anything can be dropped.
That being said, this might not always be true so I think `SeqCst` seems safer.
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -93,6 +89,37 @@ impl SpillPoolShared {
}
}
+/// Tracks the number of live [`SpillPoolWriter`] clones.
+///
+/// Cloning increments the count; dropping decrements it.
+/// [`WriterCount::is_last`] returns `true` when called from the final clone,
+/// which the writer uses to decide whether to finalize the spill pool.
+struct WriterCount(Arc<AtomicUsize>);
+
+impl WriterCount {
+ fn new() -> Self {
+ Self(Arc::new(AtomicUsize::new(1)))
+ }
+
+ /// Returns `true` if this is the only remaining clone.
+ fn is_last(&self) -> bool {
+ self.0.load(Ordering::Acquire) == 1
+ }
+}
+
+impl Clone for WriterCount {
+ fn clone(&self) -> Self {
+ self.0.fetch_add(1, Ordering::Relaxed);
Review Comment:
So thread A clones (increment), thread B drops and calls `is_last()`, but doesn't see A's increment yet so B thinks it's the last writer and finalizes prematurely. I think today the control flow prevents this from happening `setup_input_partitions` clones and passed through a task spawn boundary before anything can be dropped.
That being said, this might not always be true so I think `SeqCst` seems safer.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2878707746
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -1343,6 +1368,90 @@ mod tests {
Ok(())
}
+ /// Regression test for data loss when multiple writer clones are used.
+ ///
+ /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning
+ /// mode all input partition tasks share clones of the same writer. Before
+ /// the fix, `Drop` unconditionally set `writer_dropped = true` even when
+ /// other clones were still alive. This caused the `SpillPoolReader` to
+ /// return EOF prematurely, silently losing every batch written by the
+ /// remaining writers.
+ ///
+ /// The test sequence is:
+ ///
+ /// 1. writer1 writes a batch, then is dropped.
+ /// 2. The reader consumes that batch.
+ /// 3. The reader polls again — the queue is now empty.
+ /// - **Bug**: `writer_dropped` is already true → `Ready(None)` (EOF).
+ /// - **Fix**: `active_writer_count > 0` → `Pending` (wait for data).
+ /// 4. writer2 (still alive) writes a batch.
+ /// 5. The reader must see that batch — not silently lose it.
+ #[tokio::test]
+ async fn test_clone_drop_does_not_signal_eof_prematurely() -> Result<()> {
+ let (writer1, mut reader) = create_spill_channel(1024 * 1024);
+ let writer2 = writer1.clone();
+
+ // Synchronization: tell writer2 when it may proceed.
+ let (proceed_tx, proceed_rx) = tokio::sync::oneshot::channel::<()>();
+
+ // Spawn writer2 — it waits for the signal before writing.
+ let writer2_handle = SpawnedTask::spawn(async move {
+ proceed_rx.await.unwrap();
+ writer2.push_batch(&create_test_batch(10, 10)).unwrap();
+ // writer2 is dropped here (last clone → true EOF)
+ });
+
+ // Writer1 writes one batch, then drops.
+ writer1.push_batch(&create_test_batch(0, 10))?;
+ drop(writer1);
+
+ // Read writer1's batch.
+ let batch1 = reader.next().await.unwrap()?;
+ assert_eq!(batch1.num_rows(), 10);
+ let col = batch1
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(col.value(0), 0);
+
+ // Signal writer2 to write its batch. It will execute when the
+ // current task yields (i.e. when reader.next() returns Pending).
+ proceed_tx.send(()).unwrap();
+
+ // With the bug the reader returns None here because it already
+ // saw writer_dropped=true on an empty queue. With the fix it
+ // returns Pending, the runtime schedules writer2, and the batch
+ // becomes available.
+ let batch2 = tokio::time::timeout(
+ std::time::Duration::from_secs(5),
+ reader.next(),
+ )
+ .await
+ .expect("Reader timed out — should not hang");
+
+ assert!(
Review Comment:
Without this fix we fail here.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #20672:
URL: https://github.com/apache/datafusion/pull/20672
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886560181
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -93,6 +89,37 @@ impl SpillPoolShared {
}
}
+/// Tracks the number of live [`SpillPoolWriter`] clones.
+///
+/// Cloning increments the count; dropping decrements it.
+/// [`WriterCount::is_last`] returns `true` when called from the final clone,
+/// which the writer uses to decide whether to finalize the spill pool.
+struct WriterCount(Arc<AtomicUsize>);
+
+impl WriterCount {
+ fn new() -> Self {
+ Self(Arc::new(AtomicUsize::new(1)))
+ }
+
+ /// Returns `true` if this is the only remaining clone.
+ fn is_last(&self) -> bool {
+ self.0.load(Ordering::Acquire) == 1
+ }
+}
+
+impl Clone for WriterCount {
+ fn clone(&self) -> Self {
+ self.0.fetch_add(1, Ordering::Relaxed);
Review Comment:
I think using Ordering::SeqCst or Ordering::Acquire might be safer here. It probably is unlikely to cause problems in practice, but I think it would be important that any wrtiter increment is seen before the `is_last` check is done
https://doc.rust-lang.org/std/sync/atomic/enum.Ordering.html
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -247,17 +266,15 @@ impl SpillPoolWriter {
impl Drop for SpillPoolWriter {
fn drop(&mut self) {
- let mut shared = self.shared.lock();
-
- shared.active_writer_count -= 1;
- let is_last_writer = shared.active_writer_count == 0;
-
- if !is_last_writer {
+ if !self.writer_count.is_last() {
Review Comment:
I think there a race condition now -- what prevents a new writer from being added after `is_last` returns false but before the lock is taken below?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886665306
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -247,17 +266,15 @@ impl SpillPoolWriter {
impl Drop for SpillPoolWriter {
fn drop(&mut self) {
- let mut shared = self.shared.lock();
-
- shared.active_writer_count -= 1;
- let is_last_writer = shared.active_writer_count == 0;
-
- if !is_last_writer {
+ if !self.writer_count.is_last() {
Review Comment:
Gave it another shot here [380da74](https://github.com/apache/datafusion/pull/20672/commits/380da74bb74573bd9f075593ad7d03346361de9c)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2890541242
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -443,6 +480,7 @@ pub fn channel(
let writer = SpillPoolWriter {
max_file_size_bytes,
shared: Arc::clone(&shared),
+ writer_count: WriterCount::new(),
Review Comment:
Have reverted
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886623707
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -93,6 +89,37 @@ impl SpillPoolShared {
}
}
+/// Tracks the number of live [`SpillPoolWriter`] clones.
+///
+/// Cloning increments the count; dropping decrements it.
+/// [`WriterCount::is_last`] returns `true` when called from the final clone,
+/// which the writer uses to decide whether to finalize the spill pool.
+struct WriterCount(Arc<AtomicUsize>);
+
+impl WriterCount {
+ fn new() -> Self {
+ Self(Arc::new(AtomicUsize::new(1)))
+ }
+
+ /// Returns `true` if this is the only remaining clone.
+ fn is_last(&self) -> bool {
+ self.0.load(Ordering::Acquire) == 1
+ }
+}
+
+impl Clone for WriterCount {
+ fn clone(&self) -> Self {
+ self.0.fetch_add(1, Ordering::Relaxed);
Review Comment:
So thread A clones (increment), thread B drops and calls `is_last()`, but doesn't see A's increment yet so B thinks it's the last writer and finalizes prematurely. I think today the control flow prevents this from happening `setup_input_partitions` clones and passed through a task spawn boundary before anything can be dropped.
That being said, this might not always be true so I think `SeqCst` seems safer.
I also think I need to go and read the rust book again...
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2884412856
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -1343,6 +1368,88 @@ mod tests {
Ok(())
}
+ /// Regression test for data loss when multiple writer clones are used.
+ ///
+ /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning
+ /// mode all input partition tasks share clones of the same writer. Before
+ /// the fix, `Drop` unconditionally set `writer_dropped = true` even when
+ /// other clones were still alive. This caused the `SpillPoolReader` to
+ /// return EOF prematurely, silently losing every batch written by the
+ /// remaining writers.
Review Comment:
Yep, that makes sense to me.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886645657
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -247,17 +266,15 @@ impl SpillPoolWriter {
impl Drop for SpillPoolWriter {
fn drop(&mut self) {
- let mut shared = self.shared.lock();
-
- shared.active_writer_count -= 1;
- let is_last_writer = shared.active_writer_count == 0;
-
- if !is_last_writer {
+ if !self.writer_count.is_last() {
Review Comment:
Actually this is also a race...
Thread A: is_last() → count is 2 → false → skip finalization
Thread B: is_last() → count is 2 → false → skip finalization
Thread A: WriterCount::drop → count becomes 1
Thread B: WriterCount::drop → count becomes 0
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886665306
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -247,17 +266,15 @@ impl SpillPoolWriter {
impl Drop for SpillPoolWriter {
fn drop(&mut self) {
- let mut shared = self.shared.lock();
-
- shared.active_writer_count -= 1;
- let is_last_writer = shared.active_writer_count == 0;
-
- if !is_last_writer {
+ if !self.writer_count.is_last() {
Review Comment:
Gave it another shot here [fc9a744](https://github.com/apache/datafusion/pull/20672/commits/fc9a744f5a9c84bb94056276425f3fd97450566c)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "neilconway (via GitHub)" <gi...@apache.org>.
neilconway commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2884372912
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -1343,6 +1368,88 @@ mod tests {
Ok(())
}
+ /// Regression test for data loss when multiple writer clones are used.
+ ///
+ /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning
+ /// mode all input partition tasks share clones of the same writer. Before
+ /// the fix, `Drop` unconditionally set `writer_dropped = true` even when
+ /// other clones were still alive. This caused the `SpillPoolReader` to
+ /// return EOF prematurely, silently losing every batch written by the
+ /// remaining writers.
Review Comment:
[ Nit, optional ] I personally try to avoid comments that are very specific to the previous (buggy) implementation, because that implementation won't exist for future readers, and the current / future state of the code might drift without the comments being updated. So I'd personally remove references to "the bug" or the exact previous behavior, and talk more about what a correct implementation ought to do and the properties of a correct implementation that we're checking for.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2885932273
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -233,6 +249,15 @@ impl Drop for SpillPoolWriter {
fn drop(&mut self) {
let mut shared = self.shared.lock();
+ shared.active_writer_count -= 1;
+ let is_last_writer = shared.active_writer_count == 0;
+
+ if !is_last_writer {
+ // Other writer clones are still active; do not finalize or
+ // signal EOF to readers.
+ return;
+ }
+
// Finalize the current file when the last writer is dropped
if let Some(current_file) = shared.current_write_file.take() {
Review Comment:
this would also be great to move into a function of `shared` rather than have the logic in the drop of the sharedwriter
That being said, it wasn't introduced in this PR so no need to change it
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -233,6 +249,15 @@ impl Drop for SpillPoolWriter {
fn drop(&mut self) {
let mut shared = self.shared.lock();
+ shared.active_writer_count -= 1;
Review Comment:
I think in general this pattern of incrementing / decrementing / checking a counter from some other strcuture under a lock is error prone (e.g. if someone else changes this code in the future but forgets to update the counter it may be a hard to find bug)
I would personally suggest encapsulating the counter into some structure that encpsulated the increment/decrement
Something maybe like putting an `Arc<AtomicUsize>` on the shared pool and then passing out objects like
```rust
struct WriterLease {
inner: Arc<AtomicUsize>
}
impl Drop for WriterLease {
...
// decrement the counter
}
```
And then just get one of those for the SpillPollWriter
```rust
pub struct SpillPoolWriter {
lease: WriterLease
...
}
```
However it is a fair point that this PR is probably more efficient as it uses an already existing lock 🤷
I don't feel strongly about this
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -1343,6 +1368,90 @@ mod tests {
Ok(())
}
+ /// Regression test for data loss when multiple writer clones are used.
+ ///
+ /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning
+ /// mode all input partition tasks share clones of the same writer. Before
+ /// the fix, `Drop` unconditionally set `writer_dropped = true` even when
+ /// other clones were still alive. This caused the `SpillPoolReader` to
+ /// return EOF prematurely, silently losing every batch written by the
+ /// remaining writers.
+ ///
+ /// The test sequence is:
+ ///
+ /// 1. writer1 writes a batch, then is dropped.
+ /// 2. The reader consumes that batch.
+ /// 3. The reader polls again — the queue is now empty.
+ /// - **Bug**: `writer_dropped` is already true → `Ready(None)` (EOF).
+ /// - **Fix**: `active_writer_count > 0` → `Pending` (wait for data).
+ /// 4. writer2 (still alive) writes a batch.
+ /// 5. The reader must see that batch — not silently lose it.
+ #[tokio::test]
+ async fn test_clone_drop_does_not_signal_eof_prematurely() -> Result<()> {
+ let (writer1, mut reader) = create_spill_channel(1024 * 1024);
+ let writer2 = writer1.clone();
+
+ // Synchronization: tell writer2 when it may proceed.
+ let (proceed_tx, proceed_rx) = tokio::sync::oneshot::channel::<()>();
+
+ // Spawn writer2 — it waits for the signal before writing.
+ let writer2_handle = SpawnedTask::spawn(async move {
+ proceed_rx.await.unwrap();
+ writer2.push_batch(&create_test_batch(10, 10)).unwrap();
+ // writer2 is dropped here (last clone → true EOF)
+ });
+
+ // Writer1 writes one batch, then drops.
+ writer1.push_batch(&create_test_batch(0, 10))?;
+ drop(writer1);
+
+ // Read writer1's batch.
+ let batch1 = reader.next().await.unwrap()?;
+ assert_eq!(batch1.num_rows(), 10);
+ let col = batch1
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(col.value(0), 0);
+
+ // Signal writer2 to write its batch. It will execute when the
+ // current task yields (i.e. when reader.next() returns Pending).
+ proceed_tx.send(()).unwrap();
+
+ // With the bug the reader returns None here because it already
+ // saw writer_dropped=true on an empty queue. With the fix it
+ // returns Pending, the runtime schedules writer2, and the batch
+ // becomes available.
+ let batch2 = tokio::time::timeout(
+ std::time::Duration::from_secs(5),
+ reader.next(),
+ )
+ .await
+ .expect("Reader timed out — should not hang");
+
+ assert!(
Review Comment:
I also verified that this test fails without the code fix
```shell
andrewlamb@Andrews-MacBook-Pro-3:~/Software/datafusion2$ cargo test -p datafusion-physical-plan test_clone_drop_does_not_signal_eof_prematurely
Finished `test` profile [unoptimized + debuginfo] target(s) in 0.13s
Running unittests src/lib.rs (target/debug/deps/datafusion_physical_plan-33977765615826e4)
running 1 test
test spill::spill_pool::tests::test_clone_drop_does_not_signal_eof_prematurely ... FAILED
failures:
---- spill::spill_pool::tests::test_clone_drop_does_not_signal_eof_prematurely stdout ----
thread 'spill::spill_pool::tests::test_clone_drop_does_not_signal_eof_prematurely' (2314602) panicked at datafusion/physical-plan/src/spill/spill_pool.rs:1400:9:
Reader must not return EOF while a writer clone is still alive
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
failures:
spill::spill_pool::tests::test_clone_drop_does_not_signal_eof_prematurely
test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 1266 filtered out; finished in 0.00s
error: test failed, to rerun pass `-p datafusion-physical-plan --lib`
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#issuecomment-4016116883
> Are we happy to merge this as it is now?
Yeah, I think it looks good now. Thanks @xanderbailey and sorry for the back and forth. I'll merge this one and then make a backport
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#issuecomment-4012398779
Are we happy to merge this as it is now?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886649160
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -247,17 +266,15 @@ impl SpillPoolWriter {
impl Drop for SpillPoolWriter {
fn drop(&mut self) {
- let mut shared = self.shared.lock();
-
- shared.active_writer_count -= 1;
- let is_last_writer = shared.active_writer_count == 0;
-
- if !is_last_writer {
+ if !self.writer_count.is_last() {
Review Comment:
Do you have to atomically decrement and check if it's the last here? Rather than when you drop the `WriteCount`
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "xanderbailey (via GitHub)" <gi...@apache.org>.
xanderbailey commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2884446908
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -1343,6 +1368,88 @@ mod tests {
Ok(())
}
+ /// Regression test for data loss when multiple writer clones are used.
+ ///
+ /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning
+ /// mode all input partition tasks share clones of the same writer. Before
+ /// the fix, `Drop` unconditionally set `writer_dropped = true` even when
+ /// other clones were still alive. This caused the `SpillPoolReader` to
+ /// return EOF prematurely, silently losing every batch written by the
+ /// remaining writers.
Review Comment:
Fixed
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix repartition from dropping data when spilling [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2890342906
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -443,6 +480,7 @@ pub fn channel(
let writer = SpillPoolWriter {
max_file_size_bytes,
shared: Arc::clone(&shared),
+ writer_count: WriterCount::new(),
Review Comment:
I am confused about how this works -- doesn't it mean that each writer gets its own (new) WriterCount? I would have expected the writer count was using something on shared 🤔
Maybe we should go back to the original pattern of updating counts under a lock 🤔
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -88,6 +89,33 @@ impl SpillPoolShared {
}
}
+/// Tracks the number of live [`SpillPoolWriter`] clones.
+///
+/// Cloning increments the count. [`WriterCount::decrement`] atomically
+/// decrements the count and reports whether the caller was the last clone.
+struct WriterCount(Arc<AtomicUsize>);
+
+impl WriterCount {
+ fn new() -> Self {
+ Self(Arc::new(AtomicUsize::new(1)))
+ }
+
+ /// Decrements the count and returns `true` if this was the last clone.
Review Comment:
👍
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org