From 191cceb9ec1e2789bffaee866c0cdcead5413d7a Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Mon, 13 Jul 2026 02:14:43 +0200 Subject: [PATCH 01/17] Improve robustness of SpillPoolReader with multiple concurrent writers --- .../physical-plan/src/spill/spill_pool.rs | 124 +++++++++--------- 1 file changed, 61 insertions(+), 63 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 75da18315fb7..d83e86c4fed9 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -17,6 +17,7 @@ use futures::{Stream, StreamExt}; use std::collections::VecDeque; +use std::mem; use std::sync::Arc; use std::task::Waker; @@ -55,14 +56,12 @@ struct SpillPoolShared { spill_manager: Arc, /// Pool-level waker to notify when new files are available (single reader) waker: Option, - /// Whether the writer has been dropped (no more files will be added) - writer_dropped: bool, - /// Writer's reference to the current file (shared by all cloned writers). - /// Has its own lock to allow I/O without blocking queue access. - current_write_file: Option>>, - /// Number of active writer clones. Only when this reaches zero should - /// `writer_dropped` be set to true. This prevents premature EOF signaling - /// when one writer clone is dropped while others are still active. + /// FIFO queue of open write files. The queue may contain multiple items when multiple + /// writers concurrently write to the pool. + /// Each write file has its own lock to allow I/O without blocking queue access. + current_write_files: VecDeque>>, + /// Number of active writer clones. As long as this value is greater than zero, readers should + /// assume batches may still be pushed. This prevents premature EOF signaling. active_writer_count: usize, } @@ -73,8 +72,7 @@ impl SpillPoolShared { files: VecDeque::new(), spill_manager, waker: None, - writer_dropped: false, - current_write_file: None, + current_write_files: VecDeque::new(), active_writer_count: 1, } } @@ -171,7 +169,9 @@ impl SpillPoolWriter { let mut shared = self.shared.lock(); // Create new file if we don't have one yet - if shared.current_write_file.is_none() { + let current_write_file = if !shared.current_write_files.is_empty() { + shared.current_write_files.pop_front().unwrap() + } else { let spill_manager = Arc::clone(&shared.spill_manager); // Release shared lock before disk I/O (fine-grained locking) drop(shared); @@ -194,52 +194,48 @@ impl SpillPoolWriter { // Re-acquire lock and push to shared queue shared = self.shared.lock(); shared.files.push_back(Arc::clone(&file_shared)); - shared.current_write_file = Some(file_shared); shared.wake(); // Wake readers waiting for new files - } + file_shared + }; - let current_write_file = shared.current_write_file.take(); // Release shared lock before file I/O (fine-grained locking) // This allows readers to access the queue while we do disk I/O drop(shared); // Write batch to current file - lock only the specific file - if let Some(current_file) = current_write_file { - // Now lock just this file for I/O (separate from shared lock) - let mut file_shared = current_file.lock(); - - // Append the batch - if let Some(ref mut writer) = file_shared.writer { - writer.append_batch(batch)?; - // make sure we flush the writer for readers - writer.flush()?; - file_shared.batches_written += 1; - file_shared.estimated_size += batch_size; - } + let mut file_shared = current_write_file.lock(); + + // Append the batch + if let Some(ref mut writer) = file_shared.writer { + writer.append_batch(batch)?; + // make sure we flush the writer for readers + writer.flush()?; + file_shared.batches_written += 1; + file_shared.estimated_size += batch_size; + } - // Wake reader waiting on this specific file - file_shared.wake(); + // Wake reader waiting on this specific file + file_shared.wake(); - // Check if we need to rotate - let needs_rotation = file_shared.estimated_size > self.max_file_size_bytes; + // Check if we need to rotate + let needs_rotation = file_shared.estimated_size > self.max_file_size_bytes; - if needs_rotation { - // Finish the IPC writer - if let Some(mut writer) = file_shared.writer.take() { - writer.finish()?; - } - // Mark as finished so readers know not to wait for more data - 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 - } else { - // Release file lock - drop(file_shared); - // Put back the current file for further writing - let mut shared = self.shared.lock(); - shared.current_write_file = Some(current_file); + if needs_rotation { + // Finish the IPC writer + if let Some(mut writer) = file_shared.writer.take() { + writer.finish()?; } + // Mark as finished so readers know not to wait for more data + 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 + } else { + // Release file lock + drop(file_shared); + // Put back the current file for further writing + let mut shared = self.shared.lock(); + shared.current_write_files.push_back(current_write_file); } Ok(()) @@ -259,31 +255,33 @@ impl Drop for SpillPoolWriter { return; } - // Finalize the current file when the last writer is dropped - if let Some(current_file) = shared.current_write_file.take() { - // Release shared lock before locking file + // Finalize any spill files that were not finished yet + + if !shared.current_write_files.is_empty() { + let files = mem::take(&mut shared.current_write_files); drop(shared); - let mut file_shared = current_file.lock(); + for file in files { + let mut file_shared = file.lock(); - // Finish the current writer if it exists - if let Some(mut writer) = file_shared.writer.take() { - // Ignore errors on drop - we're in destructor - let _ = writer.finish(); - } + // Finish the current writer if it exists + if let Some(mut writer) = file_shared.writer.take() { + // Ignore errors on drop - we're in destructor + let _ = writer.finish(); + } - // Mark as finished so readers know not to wait for more data - file_shared.writer_finished = true; + // Mark as finished so readers know not to wait for more data + file_shared.writer_finished = true; - // Wake reader waiting on this file (it's now finished) - file_shared.wake(); + // Wake reader waiting on this file (it's now finished) + file_shared.wake(); + drop(file_shared); + } - drop(file_shared); shared = self.shared.lock(); } - // Mark writer as dropped and wake pool-level readers - shared.writer_dropped = true; + // Wake pool-level readers shared.wake(); } } @@ -723,7 +721,7 @@ impl Stream for SpillPoolReader { } // No files in queue - check if writer is done - if shared.writer_dropped { + if shared.active_writer_count == 0 { // Writer is done and no more files will be added - EOF return Poll::Ready(None); } From c50bc3350090cce3c8ebb2ff82cb145b6d6daeaa Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 17:05:53 +0200 Subject: [PATCH 02/17] Update documentation --- .../physical-plan/src/spill/spill_pool.rs | 94 ++++++++++--------- 1 file changed, 51 insertions(+), 43 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index d83e86c4fed9..125b39e50e00 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -48,7 +48,7 @@ use super::spill_manager::SpillManager; /// **Lock ordering discipline**: Never hold both locks simultaneously to prevent deadlock. /// Always: acquire outer lock → release outer lock → acquire inner lock (if needed). struct SpillPoolShared { - /// Queue of ALL files (including the current write file if it exists). + /// Queue of ALL files (including the current write files if any exist). /// Readers always read from the front of this queue (FIFO). /// Each file has its own lock to enable concurrent reader/writer access. files: VecDeque>>, @@ -59,10 +59,11 @@ struct SpillPoolShared { /// FIFO queue of open write files. The queue may contain multiple items when multiple /// writers concurrently write to the pool. /// Each write file has its own lock to allow I/O without blocking queue access. - current_write_files: VecDeque>>, - /// Number of active writer clones. As long as this value is greater than zero, readers should - /// assume batches may still be pushed. This prevents premature EOF signaling. - active_writer_count: usize, + open_write_files: VecDeque>>, + /// Number of `SpillPoolWriter` instances that have not been dropped yet. As long as this value + /// is greater than zero, readers should assume batches may still be pushed. This prevents + /// premature EOF signaling. + remaining_writer_count: usize, } impl SpillPoolShared { @@ -72,8 +73,8 @@ impl SpillPoolShared { files: VecDeque::new(), spill_manager, waker: None, - current_write_files: VecDeque::new(), - active_writer_count: 1, + open_write_files: VecDeque::new(), + remaining_writer_count: 1, } } @@ -90,15 +91,21 @@ impl SpillPoolShared { } } -/// Writer for a spill pool. Provides coordinated write access with FIFO semantics. +/// Writer for a spill pool. Provides coordinated write access. /// /// Created by [`channel`]. See that function for architecture diagrams and usage examples. /// /// The writer is `Clone`, allowing multiple writers to coordinate on the same pool. -/// All clones share the same current write file and coordinate file rotation. +/// +/// All clones share the same pool of write files and coordinate file rotation. /// The writer automatically manages file rotation based on the `max_file_size_bytes` -/// configured in [`channel`]. When the last writer clone is dropped, it finalizes the -/// current file so readers can access all written data. +/// configured in [`channel`]. When the last writer clone is dropped, it finalizes any open write +/// files so readers can access all written data. +/// +/// When only a single writer exists, the reader will receive batches in the order the writer +/// writes them (FIFO semantics). +/// When multiple writers are used this guarantee is dropped. To prevent writers from blocking +/// each other, concurrent writes by multiple writers may cause new write files to be created. pub struct SpillPoolWriter { /// Maximum size in bytes before rotating to a new file. /// Typically set from configuration `datafusion.execution.max_spill_file_size_bytes`. @@ -107,18 +114,6 @@ pub struct SpillPoolWriter { shared: Arc>, } -impl Clone for SpillPoolWriter { - fn clone(&self) -> Self { - // Increment the active writer count so that `writer_dropped` is only - // set to true when the *last* clone is dropped. - self.shared.lock().active_writer_count += 1; - Self { - max_file_size_bytes: self.max_file_size_bytes, - shared: Arc::clone(&self.shared), - } - } -} - impl SpillPoolWriter { /// Spills a batch to the pool, rotating files when necessary. /// @@ -168,9 +163,9 @@ impl SpillPoolWriter { // Fine-grained locking: Lock shared state briefly for queue access let mut shared = self.shared.lock(); - // Create new file if we don't have one yet - let current_write_file = if !shared.current_write_files.is_empty() { - shared.current_write_files.pop_front().unwrap() + // Create new file if there is none available to append to + let write_file = if !shared.open_write_files.is_empty() { + shared.open_write_files.pop_front().unwrap() } else { let spill_manager = Arc::clone(&shared.spill_manager); // Release shared lock before disk I/O (fine-grained locking) @@ -203,7 +198,7 @@ impl SpillPoolWriter { drop(shared); // Write batch to current file - lock only the specific file - let mut file_shared = current_write_file.lock(); + let mut file_shared = write_file.lock(); // Append the batch if let Some(ref mut writer) = file_shared.writer { @@ -217,10 +212,9 @@ impl SpillPoolWriter { // Wake reader waiting on this specific file file_shared.wake(); - // Check if we need to rotate - let needs_rotation = file_shared.estimated_size > self.max_file_size_bytes; + let max_file_size_reached = file_shared.estimated_size > self.max_file_size_bytes; - if needs_rotation { + if max_file_size_reached { // Finish the IPC writer if let Some(mut writer) = file_shared.writer.take() { writer.finish()?; @@ -229,25 +223,39 @@ impl SpillPoolWriter { 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 + + // Don't place `write_file` back in the `open_write_files` queue so we don't + // try writing to it again } else { // Release file lock drop(file_shared); // Put back the current file for further writing let mut shared = self.shared.lock(); - shared.current_write_files.push_back(current_write_file); + shared.open_write_files.push_back(write_file); } Ok(()) } } +impl Clone for SpillPoolWriter { + fn clone(&self) -> Self { + // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` + // implementation. + self.shared.lock().remaining_writer_count += 1; + Self { + max_file_size_bytes: self.max_file_size_bytes, + shared: Arc::clone(&self.shared), + } + } +} + 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; + shared.remaining_writer_count -= 1; + let is_last_writer = shared.remaining_writer_count == 0; if !is_last_writer { // Other writer clones are still active; do not finalize or @@ -256,9 +264,8 @@ impl Drop for SpillPoolWriter { } // Finalize any spill files that were not finished yet - - if !shared.current_write_files.is_empty() { - let files = mem::take(&mut shared.current_write_files); + if !shared.open_write_files.is_empty() { + let files = mem::take(&mut shared.open_write_files); drop(shared); for file in files { @@ -291,7 +298,8 @@ impl Drop for SpillPoolWriter { /// /// This is the recommended way to create a spill pool. The writer is `Clone`, allowing /// multiple producers to coordinate writes to the same pool. The reader can consume batches -/// in FIFO order. The reader can start reading immediately after a writer appends a batch +/// in FIFO order if a single writer is used. Otherwise the read order is undefined. +/// The reader can start reading immediately after a writer appends a batch /// to the spill file, without waiting for the file to be sealed, while writers continue to /// write more data. /// @@ -324,13 +332,13 @@ impl Drop for SpillPoolWriter { /// │ │ │ ┌────┐┌────┐ │ │ │ /// │ push_batch() │ │ F1 ││ F2 │ ... │ next().await │ /// │ │ │ └────┘└────┘ │ │ │ -/// │ ▼ │ (FIFO order) │ ▼ │ +/// │ ▼ │ │ ▼ │ /// │ ┌─────────┐ │ │ ┌──────────┐ │ /// │ │Current │───────▶│ Coordination: │◀───│ Current │ │ /// │ │Write │ │ - Wakers │ │ Read │ │ /// │ │File │ │ - Batch counts │ │ File │ │ /// │ └─────────┘ │ - Writer status │ └──────────┘ │ -/// │ │ └────────────────────┘ │ │ +/// │ │ └────────────────────┘ │ │ /// │ │ │ │ /// │ Size > limit? Read all batches? │ /// │ │ │ │ @@ -338,7 +346,7 @@ impl Drop for SpillPoolWriter { /// │ Rotate to new file Pop from queue │ /// └─────────────────────────────────────────────────────────────────────────┘ /// -/// Writer produces → Shared FIFO queue → Reader consumes +/// Writer produces → Shared queue → Reader consumes /// ``` /// /// # File State Machine @@ -606,7 +614,7 @@ impl Stream for SpillPoolFile { } } -/// A stream that reads from a SpillPool in FIFO order. +/// A stream that reads from a SpillPool. The reader guarantees FIFO order if a single writer is used. /// /// Created by [`channel`]. See that function for architecture diagrams and usage examples. /// @@ -721,7 +729,7 @@ impl Stream for SpillPoolReader { } // No files in queue - check if writer is done - if shared.active_writer_count == 0 { + if shared.remaining_writer_count == 0 { // Writer is done and no more files will be added - EOF return Poll::Ready(None); } From 8c3b09fc7c147599d37d04dc7d3e49df5be15214 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 15 Jul 2026 12:43:05 -0500 Subject: [PATCH 03/17] Enforce single-writer spill pools for preserve_order via the type system MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Claude-Session: https://claude.ai/code/session_01WsjHVeKZtrSHugLwNFURjL --- .../physical-plan/src/repartition/mod.rs | 192 ++++++++++++++++-- .../physical-plan/src/spill/spill_pool.rs | 116 ++++++++--- 2 files changed, 265 insertions(+), 43 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index a07d110e6604..6df232536542 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -39,7 +39,7 @@ use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::spill::spill_manager::SpillManager; -use crate::spill::spill_pool::{self, SpillPoolWriter}; +use crate::spill::spill_pool::{self, SharedSpillPoolWriter, SpillPoolWriter}; use crate::statistics::StatisticsArgs; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ @@ -164,10 +164,63 @@ type InputPartitionsToCurrentPartitionReceiver = Vec, reservation: SharedMemoryReservation, - spill_writer: SpillPoolWriter, + spill_writer: SpillWriter, shared_coalescer: Option, } +/// Per-input-task handle to a spill pool. +/// +/// The variant is fixed by the repartition mode and dispatches [`push_batch`] to the matching +/// writer type. Making the mode a type distinction (rather than a runtime flag) is what prevents +/// a `preserve_order` pool from ever being fed by a shared multi-producer writer, which would +/// silently break the ordering the merge downstream relies on. +/// +/// [`push_batch`]: SpillWriter::push_batch +enum SpillWriter { + /// `preserve_order` mode: a dedicated single-producer FIFO pool, moved into exactly one input + /// task and never cloned, so batches are read back in write order. + Ordered(SpillPoolWriter), + /// Non-preserve-order mode: a shared multi-producer pool cloned across all input tasks. + Shared(SharedSpillPoolWriter), +} + +impl SpillWriter { + fn push_batch(&self, batch: &RecordBatch) -> Result<()> { + match self { + SpillWriter::Ordered(writer) => writer.push_batch(batch), + SpillWriter::Shared(writer) => writer.push_batch(batch), + } + } +} + +/// 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 { + /// `preserve_order`: one single-producer FIFO writer per input partition. Each is `take`n + /// exactly once (moved into the matching input task), so the pool always has one writer. + PerInput(Vec>), + /// Non-preserve-order: one shared writer, cloned into every input task. + Shared(SharedSpillPoolWriter), +} + +impl PartitionSpillWriters { + /// Hand out the writer for input partition `input`. + /// + /// In `PerInput` mode this moves the dedicated writer out (it must only be requested once per + /// input); in `Shared` mode it clones the shared writer. + fn take_for_input(&mut self, input: usize) -> SpillWriter { + match self { + PartitionSpillWriters::PerInput(writers) => SpillWriter::Ordered( + writers[input] + .take() + .expect("spill writer for input partition requested more than once"), + ), + PartitionSpillWriters::Shared(writer) => SpillWriter::Shared(writer.clone()), + } + } +} + impl OutputChannel { fn coalesce(&mut self, batch: RecordBatch) -> Result> { match &self.shared_coalescer { @@ -305,9 +358,11 @@ struct PartitionChannels { /// partition. `None` in preserve-order mode (downstream /// `StreamingMergeBuilder` handles batching). shared_coalescer: Option, - /// Spill writers for writing spilled data. - /// SpillPoolWriter is Clone, so multiple writers can share state in non-preserve-order mode. - spill_writers: Vec, + /// Spill writers for writing spilled data, before they are handed to the per-input tasks. + /// The variant is chosen by the repartition mode (see [`PartitionSpillWriters`]): a dedicated + /// single-producer FIFO writer per input in preserve-order mode, or one shared writer in + /// non-preserve-order mode. + spill_writers: PartitionSpillWriters, /// Spill readers for reading spilled data - one per input partition (FIFO semantics). /// Each (input, output) pair gets its own reader to maintain proper ordering. spill_readers: Vec, @@ -464,15 +519,26 @@ impl RepartitionExecState { .options() .execution .max_spill_file_size_bytes; - let num_spill_channels = if preserve_order { - num_input_partitions + let (spill_writers, spill_readers) = if preserve_order { + // preserve_order: one dedicated single-producer FIFO pool per input partition. + // Each writer is moved into exactly one input task (never cloned), so the ordering + // the downstream merge relies on is preserved across the spill boundary. + let mut writers = Vec::with_capacity(num_input_partitions); + let mut readers = Vec::with_capacity(num_input_partitions); + for _ in 0..num_input_partitions { + let (writer, reader) = + spill_pool::channel(max_file_size, Arc::clone(&spill_manager)); + writers.push(Some(writer)); + readers.push(reader); + } + (PartitionSpillWriters::PerInput(writers), readers) } else { - 1 + // non-preserve-order: one shared multi-producer pool per output partition, since + // all inputs share the same receiver and the output is an unordered multiset. + let (writer, reader) = + spill_pool::shared_channel(max_file_size, Arc::clone(&spill_manager)); + (PartitionSpillWriters::Shared(writer), vec![reader]) }; - let (spill_writers, spill_readers): (Vec<_>, Vec<_>) = (0 - ..num_spill_channels) - .map(|_| spill_pool::channel(max_file_size, Arc::clone(&spill_manager))) - .unzip(); // Coalesce on the producer side, before the channel's gate, so // the consumer never sees the per-input-task small batches. @@ -505,18 +571,17 @@ impl RepartitionExecState { std::mem::take(streams_and_metrics).into_iter().enumerate() { let txs: HashMap<_, _> = channels - .iter() + .iter_mut() .map(|(partition, channels)| { - // In preserve_order mode: each input gets its own spill writer (index i) - // In non-preserve-order mode: all inputs share spill writer 0 via clone - let spill_writer_idx = if preserve_order { i } else { 0 }; + // Hand this input task its spill writer: in preserve_order mode this moves + // the input's dedicated FIFO writer out; otherwise it clones the shared + // writer. See [`PartitionSpillWriters::take_for_input`]. ( *partition, OutputChannel { sender: channels.tx[i].clone(), reservation: Arc::clone(&channels.reservation), - spill_writer: channels.spill_writers[spill_writer_idx] - .clone(), + spill_writer: channels.spill_writers.take_for_input(i), shared_coalescer: channels.shared_coalescer.clone(), }, ) @@ -3380,7 +3445,7 @@ mod tests { #[cfg(test)] mod test { - use arrow::array::record_batch; + use arrow::array::{UInt32Array, record_batch}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::assert_batches_eq; @@ -3566,6 +3631,95 @@ mod test { Ok(()) } + /// Regression test for order preservation across spill *file rotation*. + /// + /// A `preserve_order` repartition relies on each per-(input, output) spill pool delivering + /// batches in strict FIFO order (see [`spill_pool::channel`] / [`SpillPoolWriter`]). This uses + /// the same memory profile as [`Self::test_preserve_order_with_spilling`] — which is tuned to + /// force spilling while still completing — but additionally sets `max_spill_file_size_bytes` + /// to 1 so every spilled batch lands in its own file. That exercises the FIFO-across-rotation + /// path: if ordering were lost across rotated files (e.g. by feeding an ordered pool with a + /// shared multi-producer writer), the downstream `StreamingMerge` would emit out-of-order rows + /// and the sortedness assertion below would fail. + #[tokio::test] + async fn test_preserve_order_with_spill_file_rotation() -> Result<()> { + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + // Same sorted input as `test_preserve_order_with_spilling`: + // Partition1: [1,3], [5,7], [9,11]; Partition2: [2,4], [6,8], [10,12] + let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap(); + let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap(); + let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap(); + let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap(); + let batch5 = record_batch!(("c0", UInt32, [9, 11])).unwrap(); + let batch6 = record_batch!(("c0", UInt32, [10, 12])).unwrap(); + let schema = batch1.schema(); + let sort_exprs = LexOrdering::new([PhysicalSortExpr { + expr: col("c0", &schema).unwrap(), + options: SortOptions::default().asc(), + }]) + .unwrap(); + let partition1 = vec![batch1, batch3, batch5]; + let partition2 = vec![batch2, batch4, batch6]; + let input_partitions = vec![partition1, partition2]; + + // Force a new spill file per spilled batch to exercise FIFO across rotation. + let mut session_config = SessionConfig::new(); + session_config + .options_mut() + .execution + .max_spill_file_size_bytes = 1; + // Same tight limit as `test_preserve_order_with_spilling`: forces spilling while leaving + // the merge enough non-spillable headroom to complete. + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(608, 1.0) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))? + .with_preserve_order(); + + // Each output partition merges sorted substreams, so its rows must be non-decreasing. + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + let mut last: Option = None; + while let Some(result) = stream.next().await { + let batch = result?; + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for r in 0..col.len() { + let v = col.value(r); + if let Some(prev) = last { + assert!( + prev <= v, + "output partition {i} not sorted: {prev} came before {v}" + ); + } + last = Some(v); + } + } + } + + let metrics = exec.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur for order-preserving repartition at this \ + memory limit. If this fails, the memory limit may need adjustment." + ); + Ok(()) + } + #[tokio::test] async fn test_hash_partitioning_with_spilling() -> Result<()> { use datafusion_execution::runtime_env::RuntimeEnvBuilder; diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 125b39e50e00..af5028e0568e 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -91,22 +91,29 @@ impl SpillPoolShared { } } -/// Writer for a spill pool. Provides coordinated write access. +/// Multi-producer writer for a spill pool. Provides coordinated write access shared by +/// multiple concurrent writers. /// -/// Created by [`channel`]. See that function for architecture diagrams and usage examples. -/// -/// The writer is `Clone`, allowing multiple writers to coordinate on the same pool. +/// Created by [`shared_channel`]. See that function for architecture diagrams and usage +/// examples. /// +/// This writer is `Clone`, allowing multiple producers to coordinate on the same pool. /// All clones share the same pool of write files and coordinate file rotation. /// The writer automatically manages file rotation based on the `max_file_size_bytes` -/// configured in [`channel`]. When the last writer clone is dropped, it finalizes any open write -/// files so readers can access all written data. +/// configured in [`shared_channel`]. When the last writer clone is dropped, it finalizes any +/// open write files so readers can access all written data. /// -/// When only a single writer exists, the reader will receive batches in the order the writer -/// writes them (FIFO semantics). -/// When multiple writers are used this guarantee is dropped. To prevent writers from blocking -/// each other, concurrent writes by multiple writers may cause new write files to be created. -pub struct SpillPoolWriter { +/// # Ordering guarantee +/// +/// Each individual writer's batches are read back in the order that writer wrote them +/// (**per-writer FIFO**). When multiple writers are used, no *global* order across writers is +/// promised: to prevent writers from blocking each other, concurrent writes may create separate +/// write files that the reader interleaves arbitrarily. +/// +/// If you need strict end-to-end FIFO (a single writer whose batches are read back in exact +/// write order), use [`channel`] and [`SpillPoolWriter`] instead — that type is not `Clone`, +/// so the single-writer invariant is enforced by the compiler. +pub struct SharedSpillPoolWriter { /// 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, @@ -114,7 +121,7 @@ pub struct SpillPoolWriter { shared: Arc>, } -impl SpillPoolWriter { +impl SharedSpillPoolWriter { /// Spills a batch to the pool, rotating files when necessary. /// /// If the current file would exceed `max_file_size_bytes` after adding @@ -238,7 +245,7 @@ impl SpillPoolWriter { } } -impl Clone for SpillPoolWriter { +impl Clone for SharedSpillPoolWriter { fn clone(&self) -> Self { // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` // implementation. @@ -250,7 +257,7 @@ impl Clone for SpillPoolWriter { } } -impl Drop for SpillPoolWriter { +impl Drop for SharedSpillPoolWriter { fn drop(&mut self) { let mut shared = self.shared.lock(); @@ -293,14 +300,45 @@ impl Drop for SpillPoolWriter { } } -/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, single-consumer) -/// semantics. +/// Single-producer writer for a spill pool that guarantees strict FIFO ordering. +/// +/// Created by [`channel`]. Unlike [`SharedSpillPoolWriter`], this type is **not** `Clone`, so +/// exactly one writer can ever exist for the pool. That makes the "single writer" invariant a +/// compile-time property rather than a runtime convention: the reader is guaranteed to observe +/// batches in the exact order they were written. /// -/// This is the recommended way to create a spill pool. The writer is `Clone`, allowing -/// multiple producers to coordinate writes to the same pool. The reader can consume batches -/// in FIFO order if a single writer is used. Otherwise the read order is undefined. -/// The reader can start reading immediately after a writer appends a batch -/// to the spill file, without waiting for the file to be sealed, while writers continue to +/// Use this variant whenever downstream correctness depends on preserving order across the spill +/// boundary — e.g. a `RepartitionExec` with `preserve_order = true`, where each input partition +/// owns its own pool and its output feeds an order-sensitive merge. Reach for [`shared_channel`] +/// / [`SharedSpillPoolWriter`] only when several producers must share one pool and the consumer +/// treats the output as an unordered multiset. +pub struct SpillPoolWriter { + /// The underlying shared writer. Kept private and never cloned, so this pool always has + /// exactly one writer. + inner: SharedSpillPoolWriter, +} + +impl SpillPoolWriter { + /// Spills a batch to the pool, rotating files when necessary. + /// + /// See [`SharedSpillPoolWriter::push_batch`] for the rotation semantics; the only difference + /// is that this writer cannot be cloned, so batches are always read back in write order. + pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { + self.inner.push_batch(batch) + } +} + +/// Creates a paired writer and reader for a spill pool with SPSC (single-producer, +/// single-consumer) semantics and strict FIFO ordering. +/// +/// This is the recommended way to create a spill pool. The returned [`SpillPoolWriter`] is +/// **not** `Clone`, so exactly one writer exists for the pool and the reader is guaranteed to +/// consume batches in the exact order they were written. If you instead need several producers +/// to share one pool, use [`shared_channel`] (which returns a `Clone` [`SharedSpillPoolWriter`] +/// and only promises per-writer FIFO). +/// +/// The reader can start reading immediately after the writer appends a batch +/// to the spill file, without waiting for the file to be sealed, while the writer continues to /// write more data. /// /// Internally this coordinates rotating spill files based on size limits, and @@ -469,10 +507,29 @@ pub fn channel( max_file_size_bytes: usize, spill_manager: Arc, ) -> (SpillPoolWriter, SendableRecordBatchStream) { + let (inner, reader) = shared_channel(max_file_size_bytes, spill_manager); + (SpillPoolWriter { inner }, reader) +} + +/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, +/// single-consumer) semantics. +/// +/// Unlike [`channel`], the returned [`SharedSpillPoolWriter`] is `Clone`, so multiple producers +/// can push to the same pool concurrently. This relaxes the ordering guarantee to **per-writer +/// FIFO**: each writer's batches are read back in that writer's write order, but no global order +/// across writers is promised (see [`SharedSpillPoolWriter`]). +/// +/// Use this only when the consumer treats the spilled output as an unordered multiset. If +/// downstream correctness depends on order (e.g. `preserve_order = true`), use [`channel`] +/// instead so the compiler enforces a single writer. +pub fn shared_channel( + max_file_size_bytes: usize, + spill_manager: Arc, +) -> (SharedSpillPoolWriter, SendableRecordBatchStream) { let schema = Arc::clone(spill_manager.schema()); let shared = Arc::new(Mutex::new(SpillPoolShared::new(spill_manager))); - let writer = SpillPoolWriter { + let writer = SharedSpillPoolWriter { max_file_size_bytes, shared: Arc::clone(&shared), }; @@ -779,6 +836,17 @@ mod tests { channel(max_file_size, spill_manager) } + fn create_shared_spill_channel( + max_file_size: usize, + ) -> (SharedSpillPoolWriter, SendableRecordBatchStream) { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = create_test_schema(); + let spill_manager = Arc::new(SpillManager::new(env, metrics, schema)); + + shared_channel(max_file_size, spill_manager) + } + fn create_spill_channel_with_metrics( max_file_size: usize, ) -> (SpillPoolWriter, SendableRecordBatchStream, SpillMetrics) { @@ -1376,7 +1444,7 @@ mod tests { /// Verifies that the reader stays alive as long as any writer clone exists. /// - /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning + /// `SharedSpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning /// mode multiple input partition tasks share clones of the same writer. /// The reader must not see EOF until **all** clones have been dropped, /// even if the queue is temporarily empty between writes from different @@ -1391,7 +1459,7 @@ mod tests { /// 5. EOF is only signalled after writer2 is also dropped. #[tokio::test] async fn test_clone_drop_does_not_signal_eof_prematurely() -> Result<()> { - let (writer1, mut reader) = create_spill_channel(1024 * 1024); + let (writer1, mut reader) = create_shared_spill_channel(1024 * 1024); let writer2 = writer1.clone(); // Synchronization: tell writer2 when it may proceed. From 7cb38fb481239dd1949f5d1d3f1df43dcd8a9fe4 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 20:17:47 +0200 Subject: [PATCH 04/17] Add MPSC usage test --- .../physical-plan/src/spill/spill_pool.rs | 53 ++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 125b39e50e00..2c7fc4db159e 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -753,7 +753,7 @@ mod tests { use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; use arrow::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common_runtime::SpawnedTask; + use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_execution::runtime_env::RuntimeEnv; fn create_test_schema() -> SchemaRef { @@ -1211,6 +1211,57 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "multi_thread", worker_threads = 10)] + async fn test_concurrent_writers() -> Result<()> { + let (writer, mut reader) = create_spill_channel(1024 * 1024); + + // Spawn writer tasks + let mut writer_join_set = JoinSet::new(); + for w in 0..10 { + let writer = writer.clone(); + writer_join_set.spawn(async move { + for b in 0..10 { + let batch = create_test_batch((w * 100) + (b * 10), 10); + writer.push_batch(&batch).unwrap(); + } + }); + } + drop(writer); + + // Reader task (runs concurrently) + let reader_handle = SpawnedTask::spawn(async move { + let mut batch_order = vec![]; + loop { + match reader.next().await { + None => break, + Some(batch) => { + let batch = batch.unwrap(); + + assert_eq!(batch.num_rows(), 10); + + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + batch_order.push(col.value(0) / 10); + } + } + } + batch_order + }); + + // Wait for both to complete + writer_join_set.join_all().await; + let mut batch_order = reader_handle.await.unwrap(); + + // When used with multiple writers, order is not guaranteed + batch_order.sort(); + assert_eq!(batch_order, (0i32..100i32).collect::>()); + + Ok(()) + } + #[tokio::test] async fn test_reader_catches_up_to_writer() -> Result<()> { let (writer, mut reader) = create_spill_channel(1024 * 1024); From 363b2803e227e04b0374e8907cec88349e9ff824 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 20:32:29 +0200 Subject: [PATCH 05/17] Fix test compile error --- datafusion/physical-plan/src/spill/spill_pool.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 227557b544d9..19bad3ffb7ce 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -1281,7 +1281,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 10)] async fn test_concurrent_writers() -> Result<()> { - let (writer, mut reader) = create_spill_channel(1024 * 1024); + let (writer, mut reader) = create_shared_spill_channel(1024 * 1024); // Spawn writer tasks let mut writer_join_set = JoinSet::new(); From 2f2c52df76e4cb8576cd9d245532d79f88f3af08 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 21:20:01 +0200 Subject: [PATCH 06/17] Invert SharedSpillPoolWriter/SpillPoolWriter relationship to eliminate SpillWriter wrapper type --- .../physical-plan/src/repartition/mod.rs | 41 +-- .../physical-plan/src/spill/spill_pool.rs | 256 +++++++----------- 2 files changed, 111 insertions(+), 186 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 89c521217fa7..4c9862363825 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -39,7 +39,7 @@ use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::spill::spill_manager::SpillManager; -use crate::spill::spill_pool::{self, SharedSpillPoolWriter}; +use crate::spill::spill_pool::{self, SharedSpillPoolWriter, SpillPoolWriter}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ @@ -164,35 +164,10 @@ type InputPartitionsToCurrentPartitionReceiver = Vec, reservation: SharedMemoryReservation, - spill_writer: SpillWriter, + spill_writer: SpillPoolWriter, shared_coalescer: Option, } -/// Per-input-task handle to a spill pool. -/// -/// The variant is fixed by the repartition mode and dispatches [`push_batch`] to the matching -/// writer type. Making the mode a type distinction (rather than a runtime flag) is what prevents -/// a `preserve_order` pool from ever being fed by a shared multi-producer writer, which would -/// silently break the ordering the merge downstream relies on. -/// -/// [`push_batch`]: SpillWriter::push_batch -enum SpillWriter { - /// `preserve_order` mode: a dedicated single-producer FIFO pool, moved into exactly one input - /// task and never cloned, so batches are read back in write order. - Ordered(SpillPoolWriter), - /// Non-preserve-order mode: a shared multi-producer pool cloned across all input tasks. - Shared(SharedSpillPoolWriter), -} - -impl SpillWriter { - fn push_batch(&self, batch: &RecordBatch) -> Result<()> { - match self { - SpillWriter::Ordered(writer) => writer.push_batch(batch), - SpillWriter::Shared(writer) => writer.push_batch(batch), - } - } -} - /// 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. @@ -209,14 +184,12 @@ impl PartitionSpillWriters { /// /// In `PerInput` mode this moves the dedicated writer out (it must only be requested once per /// input); in `Shared` mode it clones the shared writer. - fn take_for_input(&mut self, input: usize) -> SpillWriter { + fn take_for_input(&mut self, input: usize) -> SpillPoolWriter { match self { - PartitionSpillWriters::PerInput(writers) => SpillWriter::Ordered( - writers[input] - .take() - .expect("spill writer for input partition requested more than once"), - ), - PartitionSpillWriters::Shared(writer) => SpillWriter::Shared(writer.clone()), + PartitionSpillWriters::PerInput(writers) => writers[input] + .take() + .expect("spill writer for input partition requested more than once"), + PartitionSpillWriters::Shared(writer) => writer.shared_writer(), } } } diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 19bad3ffb7ce..6f56e4c2c396 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -91,8 +91,7 @@ impl SpillPoolShared { } } -/// Multi-producer writer for a spill pool. Provides coordinated write access shared by -/// multiple concurrent writers. +/// Writer for a spill pool that can be cloned to produce additional writers. /// /// Created by [`shared_channel`]. See that function for architecture diagrams and usage /// examples. @@ -102,18 +101,89 @@ impl SpillPoolShared { /// The writer automatically manages file rotation based on the `max_file_size_bytes` /// configured in [`shared_channel`]. When the last writer clone is dropped, it finalizes any /// open write files so readers can access all written data. -/// -/// # Ordering guarantee -/// -/// Each individual writer's batches are read back in the order that writer wrote them -/// (**per-writer FIFO**). When multiple writers are used, no *global* order across writers is -/// promised: to prevent writers from blocking each other, concurrent writes may create separate -/// write files that the reader interleaves arbitrarily. -/// -/// If you need strict end-to-end FIFO (a single writer whose batches are read back in exact -/// write order), use [`channel`] and [`SpillPoolWriter`] instead — that type is not `Clone`, -/// so the single-writer invariant is enforced by the compiler. pub struct SharedSpillPoolWriter { + /// The underlying shared writer. Kept private and never cloned, so this pool always has + /// exactly one writer. + inner: SpillPoolWriter, +} + +impl SharedSpillPoolWriter { + /// Spills a batch to the pool, rotating files when necessary. + /// + /// See [`SharedSpillPoolWriter::push_batch`] for the rotation semantics. + pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { + self.inner.push_batch(batch) + } +} + +impl SharedSpillPoolWriter { + /// Returns a new shared writer that can be used to spill batches to the pool. + pub fn shared_writer(&self) -> SpillPoolWriter { + // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` + // implementation of `SpillPoolWriter`. + self.inner.shared.lock().remaining_writer_count += 1; + SpillPoolWriter { + max_file_size_bytes: self.inner.max_file_size_bytes, + shared: Arc::clone(&self.inner.shared), + } + } +} + +impl Clone for SharedSpillPoolWriter { + fn clone(&self) -> Self { + Self { + inner: self.shared_writer(), + } + } +} + +impl Drop for SpillPoolWriter { + fn drop(&mut self) { + let mut shared = self.shared.lock(); + + shared.remaining_writer_count -= 1; + let is_last_writer = shared.remaining_writer_count == 0; + + if !is_last_writer { + // Other writer clones are still active; do not finalize or + // signal EOF to readers. + return; + } + + // Finalize any spill files that were not finished yet + if !shared.open_write_files.is_empty() { + let files = mem::take(&mut shared.open_write_files); + drop(shared); + + for file in files { + let mut file_shared = file.lock(); + + // Finish the current writer if it exists + if let Some(mut writer) = file_shared.writer.take() { + // Ignore errors on drop - we're in destructor + let _ = writer.finish(); + } + + // Mark as finished so readers know not to wait for more data + file_shared.writer_finished = true; + + // Wake reader waiting on this file (it's now finished) + file_shared.wake(); + drop(file_shared); + } + + shared = self.shared.lock(); + } + + // Wake pool-level readers + shared.wake(); + } +} + +/// Single writer for a spill pool that cannot be cloned further. +/// +/// Created by [`channel`] and [`SharedSpillPoolWriter::shared_writer`]. +pub struct SpillPoolWriter { /// 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, @@ -121,41 +191,11 @@ pub struct SharedSpillPoolWriter { shared: Arc>, } -impl SharedSpillPoolWriter { +impl SpillPoolWriter { /// Spills a batch to the pool, rotating files when necessary. /// - /// If the current file would exceed `max_file_size_bytes` after adding - /// this batch, the file is finalized and a new one is started. - /// /// See [`channel`] for overall architecture and examples. /// - /// # File Rotation Logic - /// - /// ```text - /// push_batch() - /// │ - /// ▼ - /// Current file exists? - /// │ - /// ├─ No ──▶ Create new file ──▶ Add to shared queue - /// │ Wake readers - /// ▼ - /// Write batch to current file - /// │ - /// ▼ - /// estimated_size > max_file_size_bytes? - /// │ - /// ├─ No ──▶ Keep current file for next batch - /// │ - /// ▼ - /// Yes: finish() current file - /// Mark writer_finished = true - /// Wake readers - /// │ - /// ▼ - /// Next push_batch() creates new file - /// ``` - /// /// # Errors /// /// Returns an error if disk I/O fails or disk quota is exceeded. @@ -245,97 +285,10 @@ impl SharedSpillPoolWriter { } } -impl Clone for SharedSpillPoolWriter { - fn clone(&self) -> Self { - // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` - // implementation. - self.shared.lock().remaining_writer_count += 1; - Self { - max_file_size_bytes: self.max_file_size_bytes, - shared: Arc::clone(&self.shared), - } - } -} - -impl Drop for SharedSpillPoolWriter { - fn drop(&mut self) { - let mut shared = self.shared.lock(); - - shared.remaining_writer_count -= 1; - let is_last_writer = shared.remaining_writer_count == 0; - - if !is_last_writer { - // Other writer clones are still active; do not finalize or - // signal EOF to readers. - return; - } - - // Finalize any spill files that were not finished yet - if !shared.open_write_files.is_empty() { - let files = mem::take(&mut shared.open_write_files); - drop(shared); - - for file in files { - let mut file_shared = file.lock(); - - // Finish the current writer if it exists - if let Some(mut writer) = file_shared.writer.take() { - // Ignore errors on drop - we're in destructor - let _ = writer.finish(); - } - - // Mark as finished so readers know not to wait for more data - file_shared.writer_finished = true; - - // Wake reader waiting on this file (it's now finished) - file_shared.wake(); - drop(file_shared); - } - - shared = self.shared.lock(); - } - - // Wake pool-level readers - shared.wake(); - } -} - -/// Single-producer writer for a spill pool that guarantees strict FIFO ordering. -/// -/// Created by [`channel`]. Unlike [`SharedSpillPoolWriter`], this type is **not** `Clone`, so -/// exactly one writer can ever exist for the pool. That makes the "single writer" invariant a -/// compile-time property rather than a runtime convention: the reader is guaranteed to observe -/// batches in the exact order they were written. -/// -/// Use this variant whenever downstream correctness depends on preserving order across the spill -/// boundary — e.g. a `RepartitionExec` with `preserve_order = true`, where each input partition -/// owns its own pool and its output feeds an order-sensitive merge. Reach for [`shared_channel`] -/// / [`SharedSpillPoolWriter`] only when several producers must share one pool and the consumer -/// treats the output as an unordered multiset. -pub struct SpillPoolWriter { - /// The underlying shared writer. Kept private and never cloned, so this pool always has - /// exactly one writer. - inner: SharedSpillPoolWriter, -} - -impl SpillPoolWriter { - /// Spills a batch to the pool, rotating files when necessary. - /// - /// See [`SharedSpillPoolWriter::push_batch`] for the rotation semantics; the only difference - /// is that this writer cannot be cloned, so batches are always read back in write order. - pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { - self.inner.push_batch(batch) - } -} - /// Creates a paired writer and reader for a spill pool with SPSC (single-producer, /// single-consumer) semantics and strict FIFO ordering. /// -/// This is the recommended way to create a spill pool. The returned [`SpillPoolWriter`] is -/// **not** `Clone`, so exactly one writer exists for the pool and the reader is guaranteed to -/// consume batches in the exact order they were written. If you instead need several producers -/// to share one pool, use [`shared_channel`] (which returns a `Clone` [`SharedSpillPoolWriter`] -/// and only promises per-writer FIFO). +/// If you need a spill pool that supports several producers, use [`shared_channel`] instead. /// /// The reader can start reading immediately after the writer appends a batch /// to the spill file, without waiting for the file to be sealed, while the writer continues to @@ -507,29 +460,10 @@ pub fn channel( max_file_size_bytes: usize, spill_manager: Arc, ) -> (SpillPoolWriter, SendableRecordBatchStream) { - let (inner, reader) = shared_channel(max_file_size_bytes, spill_manager); - (SpillPoolWriter { inner }, reader) -} - -/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, -/// single-consumer) semantics. -/// -/// Unlike [`channel`], the returned [`SharedSpillPoolWriter`] is `Clone`, so multiple producers -/// can push to the same pool concurrently. This relaxes the ordering guarantee to **per-writer -/// FIFO**: each writer's batches are read back in that writer's write order, but no global order -/// across writers is promised (see [`SharedSpillPoolWriter`]). -/// -/// Use this only when the consumer treats the spilled output as an unordered multiset. If -/// downstream correctness depends on order (e.g. `preserve_order = true`), use [`channel`] -/// instead so the compiler enforces a single writer. -pub fn shared_channel( - max_file_size_bytes: usize, - spill_manager: Arc, -) -> (SharedSpillPoolWriter, SendableRecordBatchStream) { let schema = Arc::clone(spill_manager.schema()); let shared = Arc::new(Mutex::new(SpillPoolShared::new(spill_manager))); - let writer = SharedSpillPoolWriter { + let writer = SpillPoolWriter { max_file_size_bytes, shared: Arc::clone(&shared), }; @@ -539,6 +473,24 @@ pub fn shared_channel( (writer, Box::pin(reader)) } +/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, +/// single-consumer) semantics. +/// +/// Additional writers can be created by cloning the returned [`SharedSpillPoolWriter`]. +/// +/// In contrast to [`channel`], this implementation provides no guarantees regarding +/// the read order of the returned [`SendableRecordBatchStream`]. +/// +/// If you need strict end-to-end FIFO (a single writer whose batches are read back in exact +/// write order), use [`channel`] instead. +pub fn shared_channel( + max_file_size_bytes: usize, + spill_manager: Arc, +) -> (SharedSpillPoolWriter, SendableRecordBatchStream) { + let (inner, reader) = channel(max_file_size_bytes, spill_manager); + (SharedSpillPoolWriter { inner }, reader) +} + /// Shared state between writer and readers for an active spill file. /// Protected by a Mutex to coordinate between concurrent readers and the writer. struct ActiveSpillFileShared { From ac6dcb350a4a64e76de46da8e77979386441ba2e Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 21:20:16 +0200 Subject: [PATCH 07/17] Formatting --- datafusion/physical-plan/src/spill/spill_pool.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 6f56e4c2c396..3d436ff381ea 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -123,8 +123,8 @@ impl SharedSpillPoolWriter { // implementation of `SpillPoolWriter`. self.inner.shared.lock().remaining_writer_count += 1; SpillPoolWriter { - max_file_size_bytes: self.inner.max_file_size_bytes, - shared: Arc::clone(&self.inner.shared), + max_file_size_bytes: self.inner.max_file_size_bytes, + shared: Arc::clone(&self.inner.shared), } } } From 1e5d6807a26b88c73c9183f1d50c25f31caf7199 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 21:22:29 +0200 Subject: [PATCH 08/17] Rename shared_writer to new_writer --- datafusion/physical-plan/src/repartition/mod.rs | 2 +- datafusion/physical-plan/src/spill/spill_pool.rs | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 4c9862363825..2579a96c6a38 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -189,7 +189,7 @@ impl PartitionSpillWriters { PartitionSpillWriters::PerInput(writers) => writers[input] .take() .expect("spill writer for input partition requested more than once"), - PartitionSpillWriters::Shared(writer) => writer.shared_writer(), + PartitionSpillWriters::Shared(writer) => writer.new_writer(), } } } diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 3d436ff381ea..a056d1cd2a8f 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -118,7 +118,7 @@ impl SharedSpillPoolWriter { impl SharedSpillPoolWriter { /// Returns a new shared writer that can be used to spill batches to the pool. - pub fn shared_writer(&self) -> SpillPoolWriter { + pub fn new_writer(&self) -> SpillPoolWriter { // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` // implementation of `SpillPoolWriter`. self.inner.shared.lock().remaining_writer_count += 1; @@ -132,7 +132,7 @@ impl SharedSpillPoolWriter { impl Clone for SharedSpillPoolWriter { fn clone(&self) -> Self { Self { - inner: self.shared_writer(), + inner: self.new_writer(), } } } @@ -182,7 +182,7 @@ impl Drop for SpillPoolWriter { /// Single writer for a spill pool that cannot be cloned further. /// -/// Created by [`channel`] and [`SharedSpillPoolWriter::shared_writer`]. +/// Created by [`channel`] and [`SharedSpillPoolWriter::new_writer`]. pub struct SpillPoolWriter { /// Maximum size in bytes before rotating to a new file. /// Typically set from configuration `datafusion.execution.max_spill_file_size_bytes`. From 37075adedfc0d6804545d42138d4cc50ddfdbf55 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 21:36:33 +0200 Subject: [PATCH 09/17] More documentation tweaks --- .../physical-plan/src/spill/spill_pool.rs | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index a056d1cd2a8f..f41c9f5f0858 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -97,10 +97,6 @@ impl SpillPoolShared { /// examples. /// /// This writer is `Clone`, allowing multiple producers to coordinate on the same pool. -/// All clones share the same pool of write files and coordinate file rotation. -/// The writer automatically manages file rotation based on the `max_file_size_bytes` -/// configured in [`shared_channel`]. When the last writer clone is dropped, it finalizes any -/// open write files so readers can access all written data. pub struct SharedSpillPoolWriter { /// The underlying shared writer. Kept private and never cloned, so this pool always has /// exactly one writer. @@ -474,7 +470,8 @@ pub fn channel( } /// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, -/// single-consumer) semantics. +/// single-consumer) semantics. See [`channel`] for the general architecture description +/// of the spill pool. /// /// Additional writers can be created by cloning the returned [`SharedSpillPoolWriter`]. /// @@ -483,6 +480,23 @@ pub fn channel( /// /// If you need strict end-to-end FIFO (a single writer whose batches are read back in exact /// write order), use [`channel`] instead. +/// +/// # File Management +/// +/// The shared channel uses the same size-based rotation trigger as the [single producer channel](channel). +/// All writers share the smae pool of write files and coordinate file rotation. The number of open +/// files is kept as small as possible. When more writes occur concurrently than there are open write +/// files an additional file will be opened to write to. This prevents multiple writers from blocking +/// each other. +/// +/// When the last writer clone is dropped, it finalizes any remaining open write files so that all +/// written data can be accessed by the reader. +/// +/// # Returns +/// +/// 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( max_file_size_bytes: usize, spill_manager: Arc, From bb67abd6968fe974d2304fc745de59007aedde2f Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Wed, 15 Jul 2026 22:25:28 +0200 Subject: [PATCH 10/17] Fix typo in docstring --- datafusion/physical-plan/src/spill/spill_pool.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index f41c9f5f0858..bdcfebc2348c 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -484,7 +484,7 @@ pub fn channel( /// # File Management /// /// The shared channel uses the same size-based rotation trigger as the [single producer channel](channel). -/// All writers share the smae pool of write files and coordinate file rotation. The number of open +/// All writers share the same pool of write files and coordinate file rotation. The number of open /// files is kept as small as possible. When more writes occur concurrently than there are open write /// files an additional file will be opened to write to. This prevents multiple writers from blocking /// each other. From 79d4648727f15fc6a4478783ec43842299bbca23 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Thu, 16 Jul 2026 09:40:22 +0200 Subject: [PATCH 11/17] Eliminate panic code path --- .../physical-plan/src/repartition/mod.rs | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 2579a96c6a38..616fb2aa5567 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -55,8 +55,8 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, - assert_or_internal_err, internal_err, + ColumnStatistics, DataFusionError, HashMap, HashSet, ScalarValue, SplitPoint, + assert_or_internal_err, internal_datafusion_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; @@ -184,12 +184,16 @@ impl PartitionSpillWriters { /// /// In `PerInput` mode this moves the dedicated writer out (it must only be requested once per /// input); in `Shared` mode it clones the shared writer. - fn take_for_input(&mut self, input: usize) -> SpillPoolWriter { + fn take_for_input(&mut self, input: usize) -> Result { match self { - PartitionSpillWriters::PerInput(writers) => writers[input] - .take() - .expect("spill writer for input partition requested more than once"), - PartitionSpillWriters::Shared(writer) => writer.new_writer(), + PartitionSpillWriters::PerInput(writers) => { + writers[input].take().ok_or_else(|| { + internal_datafusion_err!( + "spill writer for input partition requested more than once" + ) + }) + } + PartitionSpillWriters::Shared(writer) => Ok(writer.new_writer()), } } } @@ -551,15 +555,15 @@ impl RepartitionExecState { // writer. See [`PartitionSpillWriters::take_for_input`]. ( *partition, - OutputChannel { + Ok(OutputChannel { sender: channels.tx[i].clone(), reservation: Arc::clone(&channels.reservation), - spill_writer: channels.spill_writers.take_for_input(i), + spill_writer: channels.spill_writers.take_for_input(i)?, shared_coalescer: channels.shared_coalescer.clone(), - }, + }), ) }) - .collect(); + .collect::>>()?; // Extract senders for wait_for_task before moving txs let senders: HashMap<_, _> = txs From 64e612596dca8c60a8532186d4516b3d11539c69 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Thu, 16 Jul 2026 13:48:18 +0200 Subject: [PATCH 12/17] Fix compile error --- datafusion/physical-plan/src/repartition/mod.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 616fb2aa5567..a2faf90ba90c 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -553,15 +553,15 @@ impl RepartitionExecState { // Hand this input task its spill writer: in preserve_order mode this moves // the input's dedicated FIFO writer out; otherwise it clones the shared // writer. See [`PartitionSpillWriters::take_for_input`]. - ( + Ok(( *partition, - Ok(OutputChannel { + OutputChannel { sender: channels.tx[i].clone(), reservation: Arc::clone(&channels.reservation), spill_writer: channels.spill_writers.take_for_input(i)?, shared_coalescer: channels.shared_coalescer.clone(), - }), - ) + }, + )) }) .collect::>>()?; From f5032f1edbf589872765c5daa4ba9b210e1b545a Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Thu, 16 Jul 2026 13:48:37 +0200 Subject: [PATCH 13/17] Revert name changes to avoid backwards incompatible change --- .../physical-plan/src/repartition/mod.rs | 26 +++--- .../physical-plan/src/spill/spill_pool.rs | 89 ++++++++++--------- 2 files changed, 63 insertions(+), 52 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index a2faf90ba90c..9082f734b8bb 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -39,7 +39,7 @@ use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::spill::spill_manager::SpillManager; -use crate::spill::spill_pool::{self, SharedSpillPoolWriter, SpillPoolWriter}; +use crate::spill::spill_pool::{self, SpillPoolSink, SpillPoolWriter}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ @@ -55,7 +55,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, HashSet, ScalarValue, SplitPoint, + ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, assert_or_internal_err, internal_datafusion_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; @@ -164,7 +164,7 @@ type InputPartitionsToCurrentPartitionReceiver = Vec, reservation: SharedMemoryReservation, - spill_writer: SpillPoolWriter, + spill_writer: SpillPoolSink, shared_coalescer: Option, } @@ -174,9 +174,9 @@ struct OutputChannel { enum PartitionSpillWriters { /// `preserve_order`: one single-producer FIFO writer per input partition. Each is `take`n /// exactly once (moved into the matching input task), so the pool always has one writer. - PerInput(Vec>), + PerInput(Vec>), /// Non-preserve-order: one shared writer, cloned into every input task. - Shared(SharedSpillPoolWriter), + Shared(SpillPoolWriter), } impl PartitionSpillWriters { @@ -184,7 +184,7 @@ impl PartitionSpillWriters { /// /// In `PerInput` mode this moves the dedicated writer out (it must only be requested once per /// input); in `Shared` mode it clones the shared writer. - fn take_for_input(&mut self, input: usize) -> Result { + fn take_for_input(&mut self, input: usize) -> Result { match self { PartitionSpillWriters::PerInput(writers) => { writers[input].take().ok_or_else(|| { @@ -193,7 +193,7 @@ impl PartitionSpillWriters { ) }) } - PartitionSpillWriters::Shared(writer) => Ok(writer.new_writer()), + PartitionSpillWriters::Shared(writer) => Ok(writer.new_sink()), } } } @@ -323,7 +323,7 @@ impl SharedCoalescer { /// /// See [`RepartitionExec`] for the overall N×M architecture. /// -/// [`spill_pool::channel`]: crate::spill::spill_pool::channel +/// [`spill_pool::channel`]: crate::spill::spill_pool::spsc_channel struct PartitionChannels { /// Senders for each input partition to send data to this output partition tx: InputPartitionsToCurrentPartitionSender, @@ -503,8 +503,10 @@ impl RepartitionExecState { let mut writers = Vec::with_capacity(num_input_partitions); let mut readers = Vec::with_capacity(num_input_partitions); for _ in 0..num_input_partitions { - let (writer, reader) = - spill_pool::channel(max_file_size, Arc::clone(&spill_manager)); + let (writer, reader) = spill_pool::spsc_channel( + max_file_size, + Arc::clone(&spill_manager), + ); writers.push(Some(writer)); readers.push(reader); } @@ -513,7 +515,7 @@ impl RepartitionExecState { // non-preserve-order: one shared multi-producer pool per output partition, since // all inputs share the same receiver and the output is an unordered multiset. let (writer, reader) = - spill_pool::shared_channel(max_file_size, Arc::clone(&spill_manager)); + spill_pool::mpsc_channel(max_file_size, Arc::clone(&spill_manager)); (PartitionSpillWriters::Shared(writer), vec![reader]) }; @@ -3616,7 +3618,7 @@ mod test { /// Regression test for order preservation across spill *file rotation*. /// /// A `preserve_order` repartition relies on each per-(input, output) spill pool delivering - /// batches in strict FIFO order (see [`spill_pool::channel`] / [`SpillPoolWriter`]). This uses + /// batches in strict FIFO order (see [`spill_pool::spsc_channel`] / [`SpillPoolSink`]). This uses /// the same memory profile as [`Self::test_preserve_order_with_spilling`] — which is tuned to /// force spilling while still completing — but additionally sets `max_spill_file_size_bytes` /// to 1 so every spilled batch lands in its own file. That exercises the FIFO-across-rotation diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index bdcfebc2348c..238e103c2ab0 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -93,47 +93,47 @@ impl SpillPoolShared { /// Writer for a spill pool that can be cloned to produce additional writers. /// -/// Created by [`shared_channel`]. See that function for architecture diagrams and usage +/// Created by [`mspc_channel`]. See that function for architecture diagrams and usage /// examples. /// /// This writer is `Clone`, allowing multiple producers to coordinate on the same pool. -pub struct SharedSpillPoolWriter { +pub struct SpillPoolWriter { /// The underlying shared writer. Kept private and never cloned, so this pool always has /// exactly one writer. - inner: SpillPoolWriter, + inner: SpillPoolSink, } -impl SharedSpillPoolWriter { +impl SpillPoolWriter { /// Spills a batch to the pool, rotating files when necessary. /// - /// See [`SharedSpillPoolWriter::push_batch`] for the rotation semantics. + /// See [`SpillPoolWriter::push_batch`] for the rotation semantics. pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { self.inner.push_batch(batch) } } -impl SharedSpillPoolWriter { - /// Returns a new shared writer that can be used to spill batches to the pool. - pub fn new_writer(&self) -> SpillPoolWriter { +impl SpillPoolWriter { + /// Returns a new sink that can be used to spill batches to the pool. + pub fn new_sink(&self) -> SpillPoolSink { // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` // implementation of `SpillPoolWriter`. self.inner.shared.lock().remaining_writer_count += 1; - SpillPoolWriter { + SpillPoolSink { max_file_size_bytes: self.inner.max_file_size_bytes, shared: Arc::clone(&self.inner.shared), } } } -impl Clone for SharedSpillPoolWriter { +impl Clone for SpillPoolWriter { fn clone(&self) -> Self { Self { - inner: self.new_writer(), + inner: self.new_sink(), } } } -impl Drop for SpillPoolWriter { +impl Drop for SpillPoolSink { fn drop(&mut self) { let mut shared = self.shared.lock(); @@ -176,10 +176,10 @@ impl Drop for SpillPoolWriter { } } -/// Single writer for a spill pool that cannot be cloned further. +/// Single writer for a spill pool that cannot be cloned. /// -/// Created by [`channel`] and [`SharedSpillPoolWriter::new_writer`]. -pub struct SpillPoolWriter { +/// Created by [`spsc_channel`] and [`SpillPoolWriter::new_sink`]. +pub struct SpillPoolSink { /// 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, @@ -187,10 +187,10 @@ pub struct SpillPoolWriter { shared: Arc>, } -impl SpillPoolWriter { +impl SpillPoolSink { /// Spills a batch to the pool, rotating files when necessary. /// - /// See [`channel`] for overall architecture and examples. + /// See [`spsc_channel`] for overall architecture and examples. /// /// # Errors /// @@ -284,7 +284,7 @@ impl SpillPoolWriter { /// Creates a paired writer and reader for a spill pool with SPSC (single-producer, /// single-consumer) semantics and strict FIFO ordering. /// -/// If you need a spill pool that supports several producers, use [`shared_channel`] instead. +/// If you need a spill pool that supports several producers, use [`mspc_channel`] instead. /// /// The reader can start reading immediately after the writer appends a batch /// to the spill file, without waiting for the file to be sealed, while the writer continues to @@ -403,7 +403,7 @@ impl SpillPoolWriter { /// # let spill_manager = Arc::new(SpillManager::new(env, metrics, schema.clone())); /// # /// // Create channel with 1MB file size limit -/// let (writer, mut reader) = spill_pool::channel(1024 * 1024, spill_manager); +/// let (writer, mut reader) = spill_pool::spsc_channel(1024 * 1024, spill_manager); /// /// // Spawn writer and reader concurrently; writer wakes reader via wakers /// let writer_task = tokio::spawn(async move { @@ -452,14 +452,14 @@ impl SpillPoolWriter { /// If instead we use file rotation, and as long as the readers can keep up with the writer, /// then we can ensure that once a file is fully read by all readers it can be deleted, /// thus bounding the maximum disk usage to roughly `max_file_size_bytes`. -pub fn channel( +pub fn spsc_channel( max_file_size_bytes: usize, spill_manager: Arc, -) -> (SpillPoolWriter, SendableRecordBatchStream) { +) -> (SpillPoolSink, SendableRecordBatchStream) { let schema = Arc::clone(spill_manager.schema()); let shared = Arc::new(Mutex::new(SpillPoolShared::new(spill_manager))); - let writer = SpillPoolWriter { + let writer = SpillPoolSink { max_file_size_bytes, shared: Arc::clone(&shared), }; @@ -469,21 +469,30 @@ pub fn channel( (writer, Box::pin(reader)) } +/// Alias for [`mpsc_channel`]. +#[deprecated(note = "Use mpsc_channel instead")] +pub fn channel( + max_file_size_bytes: usize, + spill_manager: Arc, +) -> (SpillPoolWriter, SendableRecordBatchStream) { + mpsc_channel(max_file_size_bytes, spill_manager) +} + /// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, -/// single-consumer) semantics. See [`channel`] for the general architecture description +/// single-consumer) semantics. See [`spsc_channel`] for the general architecture description /// of the spill pool. /// -/// Additional writers can be created by cloning the returned [`SharedSpillPoolWriter`]. +/// Additional writers can be created by cloning the returned [`SpillPoolWriter`]. /// -/// In contrast to [`channel`], this implementation provides no guarantees regarding +/// In contrast to [`spsc_channel`], this implementation provides no guarantees regarding /// the read order of the returned [`SendableRecordBatchStream`]. /// /// If you need strict end-to-end FIFO (a single writer whose batches are read back in exact -/// write order), use [`channel`] instead. +/// write order), use [`spsc_channel`] instead. /// /// # File Management /// -/// The shared channel uses the same size-based rotation trigger as the [single producer channel](channel). +/// The shared channel uses the same size-based rotation trigger as the [single producer channel](spsc_channel). /// All writers share the same pool of write files and coordinate file rotation. The number of open /// files is kept as small as possible. When more writes occur concurrently than there are open write /// files an additional file will be opened to write to. This prevents multiple writers from blocking @@ -497,12 +506,12 @@ pub fn channel( /// 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( +pub fn mpsc_channel( max_file_size_bytes: usize, spill_manager: Arc, -) -> (SharedSpillPoolWriter, SendableRecordBatchStream) { - let (inner, reader) = channel(max_file_size_bytes, spill_manager); - (SharedSpillPoolWriter { inner }, reader) +) -> (SpillPoolWriter, SendableRecordBatchStream) { + let (inner, reader) = spsc_channel(max_file_size_bytes, spill_manager); + (SpillPoolWriter { inner }, reader) } /// Shared state between writer and readers for an active spill file. @@ -639,7 +648,7 @@ impl Stream for SpillPoolFile { /// A stream that reads from a SpillPool. The reader guarantees FIFO order if a single writer is used. /// -/// Created by [`channel`]. See that function for architecture diagrams and usage examples. +/// Created by [`spsc_channel`]. See that function for architecture diagrams and usage examples. /// /// The stream automatically handles file rotation and reads from completed files. /// When no data is available, it returns `Poll::Pending` and registers a waker to @@ -793,35 +802,35 @@ mod tests { fn create_spill_channel( max_file_size: usize, - ) -> (SpillPoolWriter, SendableRecordBatchStream) { + ) -> (SpillPoolSink, SendableRecordBatchStream) { let env = Arc::new(RuntimeEnv::default()); let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); let schema = create_test_schema(); let spill_manager = Arc::new(SpillManager::new(env, metrics, schema)); - channel(max_file_size, spill_manager) + spsc_channel(max_file_size, spill_manager) } fn create_shared_spill_channel( max_file_size: usize, - ) -> (SharedSpillPoolWriter, SendableRecordBatchStream) { + ) -> (SpillPoolWriter, SendableRecordBatchStream) { let env = Arc::new(RuntimeEnv::default()); let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); let schema = create_test_schema(); let spill_manager = Arc::new(SpillManager::new(env, metrics, schema)); - shared_channel(max_file_size, spill_manager) + mpsc_channel(max_file_size, spill_manager) } fn create_spill_channel_with_metrics( max_file_size: usize, - ) -> (SpillPoolWriter, SendableRecordBatchStream, SpillMetrics) { + ) -> (SpillPoolSink, SendableRecordBatchStream, SpillMetrics) { let env = Arc::new(RuntimeEnv::default()); let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); let schema = create_test_schema(); let spill_manager = Arc::new(SpillManager::new(env, metrics.clone(), schema)); - let (writer, reader) = channel(max_file_size, spill_manager); + let (writer, reader) = spsc_channel(max_file_size, spill_manager); (writer, reader, metrics) } @@ -1414,7 +1423,7 @@ mod tests { let spill_manager = Arc::new(SpillManager::new(Arc::clone(&env), metrics.clone(), schema)); - let (writer, mut reader) = channel(1024 * 1024, spill_manager); + let (writer, mut reader) = spsc_channel(1024 * 1024, spill_manager); // Write some batches for i in 0..5 { @@ -1555,7 +1564,7 @@ mod tests { let schema = create_test_schema(); let spill_manager = Arc::new(SpillManager::new(runtime, metrics.clone(), schema)); - let (writer, mut reader) = channel(batch_size - 1, spill_manager); + let (writer, mut reader) = spsc_channel(batch_size - 1, spill_manager); // Step 3: Write NUM_BATCHES batches to create approximately NUM_BATCHES files for i in 0..NUM_BATCHES { From 897214831f380cdd2e112641c0ca9c912d496bd7 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Thu, 16 Jul 2026 13:51:56 +0200 Subject: [PATCH 14/17] Update docstrings --- datafusion/physical-plan/src/spill/spill_pool.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 238e103c2ab0..5464f32111bc 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -314,7 +314,7 @@ impl SpillPoolSink { /// │ Writer Side Shared State Reader Side │ /// │ ─────────── ──────────── ─────────── │ /// │ │ -/// │ SpillPoolWriter ┌────────────────────┐ SpillPoolReader │ +/// │ SpillPoolSink ┌────────────────────┐ RecordBatchStream │ /// │ │ │ VecDeque │ │ │ /// │ │ │ ┌────┐┌────┐ │ │ │ /// │ push_batch() │ │ F1 ││ F2 │ ... │ next().await │ @@ -376,7 +376,7 @@ impl SpillPoolSink { /// /// # Returns /// -/// A tuple of `(SpillPoolWriter, SendableRecordBatchStream)` that share the same +/// A tuple of `(SpillPoolSink, SendableRecordBatchStream)` that share the same /// underlying pool. The reader is returned as a stream for immediate use with /// async stream combinators. /// @@ -503,7 +503,7 @@ pub fn channel( /// /// # Returns /// -/// A tuple of `(SharedSpillPoolWriter, SendableRecordBatchStream)` that share the same +/// A tuple of `(SpillPoolWriter, 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 mpsc_channel( @@ -1470,7 +1470,7 @@ mod tests { /// Verifies that the reader stays alive as long as any writer clone exists. /// - /// `SharedSpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning + /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning /// mode multiple input partition tasks share clones of the same writer. /// The reader must not see EOF until **all** clones have been dropped, /// even if the queue is temporarily empty between writes from different From fc8476510ffc8af6db8c77eafb1d3cbf3daed026 Mon Sep 17 00:00:00 2001 From: Pepijn Van Eeckhoudt Date: Thu, 16 Jul 2026 13:57:58 +0200 Subject: [PATCH 15/17] Documentation updates --- datafusion/physical-plan/src/spill/spill_pool.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 5464f32111bc..2ebeb5922f40 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -95,8 +95,6 @@ impl SpillPoolShared { /// /// Created by [`mspc_channel`]. See that function for architecture diagrams and usage /// examples. -/// -/// This writer is `Clone`, allowing multiple producers to coordinate on the same pool. pub struct SpillPoolWriter { /// The underlying shared writer. Kept private and never cloned, so this pool always has /// exactly one writer. @@ -106,7 +104,11 @@ pub struct SpillPoolWriter { impl SpillPoolWriter { /// Spills a batch to the pool, rotating files when necessary. /// - /// See [`SpillPoolWriter::push_batch`] for the rotation semantics. + /// See [`mpsc_channel`] for the rotation semantics. + /// + /// # Errors + /// + /// Returns an error if disk I/O fails or disk quota is exceeded. pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { self.inner.push_batch(batch) } @@ -114,6 +116,10 @@ impl SpillPoolWriter { impl SpillPoolWriter { /// Returns a new sink that can be used to spill batches to the pool. + /// + /// As an alternative to this function, it is also possible to clone the writer. The benefit + /// of this method is that the output type matches the type used by [`spsc_channel`]. This + /// enables cost-free abstraction for producers over SPSC and MPSC channels. pub fn new_sink(&self) -> SpillPoolSink { // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` // implementation of `SpillPoolWriter`. @@ -675,7 +681,7 @@ pub struct SpillPoolReader { impl SpillPoolReader { /// Creates a new reader from shared pool state. /// - /// This is private - use the `channel()` function to create a reader/writer pair. + /// This is private - use the [`spsc_channel`] function to create a reader/writer pair. /// /// # Arguments /// From 1d46d57464ab3f05cd73ddc7302d25938847161d Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 16 Jul 2026 14:47:12 -0400 Subject: [PATCH 16/17] fmt --- datafusion/physical-plan/src/repartition/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 25309b70fa70..8da20d3d23d9 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -3431,14 +3431,14 @@ mod tests { #[cfg(test)] mod test { + use super::*; + use crate::test::TestMemoryExec; + use crate::union::UnionExec; use arrow::array::{UInt32Array, record_batch}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::assert_batches_eq; use datafusion_common::config::ConfigNonZeroUsize; - use super::*; - use crate::test::TestMemoryExec; - use crate::union::UnionExec; use datafusion_physical_expr::expressions::col; From c1c3d885a569866b585aa382782960a5dd95012c Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 16 Jul 2026 14:54:27 -0400 Subject: [PATCH 17/17] fix links --- datafusion/physical-plan/src/spill/spill_pool.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 2ebeb5922f40..6e964d7a6497 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -93,7 +93,7 @@ impl SpillPoolShared { /// Writer for a spill pool that can be cloned to produce additional writers. /// -/// Created by [`mspc_channel`]. See that function for architecture diagrams and usage +/// Created by [`mpsc_channel`]. See that function for architecture diagrams and usage /// examples. pub struct SpillPoolWriter { /// The underlying shared writer. Kept private and never cloned, so this pool always has @@ -290,7 +290,7 @@ impl SpillPoolSink { /// Creates a paired writer and reader for a spill pool with SPSC (single-producer, /// single-consumer) semantics and strict FIFO ordering. /// -/// If you need a spill pool that supports several producers, use [`mspc_channel`] instead. +/// If you need a spill pool that supports several producers, use [`mpsc_channel`] instead. /// /// The reader can start reading immediately after the writer appends a batch /// to the spill file, without waiting for the file to be sealed, while the writer continues to