diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index f9b908861b84d..8da20d3d23d90 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, SpillPoolSink, SpillPoolWriter}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ @@ -56,7 +56,7 @@ 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, + assert_or_internal_err, internal_datafusion_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; @@ -164,10 +164,40 @@ type InputPartitionsToCurrentPartitionReceiver = Vec, reservation: SharedMemoryReservation, - spill_writer: SpillPoolWriter, + spill_writer: SpillPoolSink, shared_coalescer: Option, } +/// 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(SpillPoolWriter), +} + +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) -> Result { + match self { + 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_sink()), + } + } +} + impl OutputChannel { fn coalesce(&mut self, batch: RecordBatch) -> Result> { match &self.shared_coalescer { @@ -293,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, @@ -305,9 +335,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, @@ -465,15 +497,29 @@ impl RepartitionExecState { .execution .max_spill_file_size_bytes .get(); - 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::spsc_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::mpsc_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. @@ -506,23 +552,22 @@ 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`]. + Ok(( *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(), }, - ) + )) }) - .collect(); + .collect::>>()?; // Extract senders for wait_for_task before moving txs let senders: HashMap<_, _> = txs @@ -3386,14 +3431,14 @@ mod tests { #[cfg(test)] mod test { - use arrow::array::record_batch; - use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::assert_batches_eq; - 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 datafusion_physical_expr::expressions::col; @@ -3572,6 +3617,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::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 + /// 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 = ConfigNonZeroUsize::try_new(1).unwrap(); + // 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 75da18315fb7b..6e964d7a6497b 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; @@ -47,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>>, @@ -55,15 +56,14 @@ 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. - active_writer_count: usize, + /// 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. + 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 { @@ -73,9 +73,8 @@ impl SpillPoolShared { files: VecDeque::new(), spill_manager, waker: None, - writer_dropped: false, - current_write_file: None, - active_writer_count: 1, + open_write_files: VecDeque::new(), + remaining_writer_count: 1, } } @@ -92,69 +91,112 @@ impl SpillPoolShared { } } -/// Writer for a spill pool. Provides coordinated write access with FIFO semantics. +/// Writer for a spill pool that can be cloned to produce additional 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. -/// All clones share the same current write file 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. +/// Created by [`mpsc_channel`]. See that function for architecture diagrams and usage +/// examples. 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, - /// Shared state with readers (includes current_write_file for coordination) - shared: Arc>, + /// The underlying shared writer. Kept private and never cloned, so this pool always has + /// exactly one writer. + inner: SpillPoolSink, +} + +impl SpillPoolWriter { + /// Spills a batch to the pool, rotating files when necessary. + /// + /// 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) + } +} + +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`. + self.inner.shared.lock().remaining_writer_count += 1; + SpillPoolSink { + max_file_size_bytes: self.inner.max_file_size_bytes, + shared: Arc::clone(&self.inner.shared), + } + } } 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), + inner: self.new_sink(), } } } -impl SpillPoolWriter { +impl Drop for SpillPoolSink { + 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. +/// +/// 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, + /// Shared state with readers (includes current_write_file for coordination) + shared: Arc>, +} + +impl SpillPoolSink { /// 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 - /// ``` + /// See [`spsc_channel`] for overall architecture and examples. /// /// # Errors /// @@ -170,8 +212,10 @@ 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 - if shared.current_write_file.is_none() { + // 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) drop(shared); @@ -194,107 +238,62 @@ 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; - } - - // 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; - - 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); - } - } - - Ok(()) - } -} - -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; + let mut file_shared = 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; } - // 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 - drop(shared); + // Wake reader waiting on this specific file + file_shared.wake(); - let mut file_shared = current_file.lock(); + let max_file_size_reached = file_shared.estimated_size > self.max_file_size_bytes; - // Finish the current writer if it exists + if max_file_size_reached { + // Finish the IPC writer if let Some(mut writer) = file_shared.writer.take() { - // Ignore errors on drop - we're in destructor - let _ = writer.finish(); + 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 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); - shared = self.shared.lock(); + // Put back the current file for further writing + let mut shared = self.shared.lock(); + shared.open_write_files.push_back(write_file); } - // Mark writer as dropped and wake pool-level readers - shared.writer_dropped = true; - shared.wake(); + Ok(()) } } -/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, single-consumer) -/// semantics. +/// 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 [`mpsc_channel`] instead. /// -/// 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 -/// to the spill file, without waiting for the file to be sealed, while writers continue to +/// 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 @@ -321,18 +320,18 @@ impl Drop for SpillPoolWriter { /// │ Writer Side Shared State Reader Side │ /// │ ─────────── ──────────── ─────────── │ /// │ │ -/// │ SpillPoolWriter ┌────────────────────┐ SpillPoolReader │ +/// │ SpillPoolSink ┌────────────────────┐ RecordBatchStream │ /// │ │ │ VecDeque │ │ │ /// │ │ │ ┌────┐┌────┐ │ │ │ /// │ 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? │ /// │ │ │ │ @@ -340,7 +339,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 @@ -383,7 +382,7 @@ impl Drop for SpillPoolWriter { /// /// # 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. /// @@ -410,7 +409,7 @@ impl Drop for 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 { @@ -459,14 +458,14 @@ impl Drop for 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), }; @@ -476,6 +475,51 @@ 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 [`spsc_channel`] for the general architecture description +/// of the spill pool. +/// +/// Additional writers can be created by cloning the returned [`SpillPoolWriter`]. +/// +/// 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 [`spsc_channel`] instead. +/// +/// # File Management +/// +/// 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 +/// 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 `(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( + max_file_size_bytes: usize, + spill_manager: Arc, +) -> (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. /// Protected by a Mutex to coordinate between concurrent readers and the writer. struct ActiveSpillFileShared { @@ -608,9 +652,9 @@ 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. +/// 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 @@ -637,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 /// @@ -723,7 +767,7 @@ impl Stream for SpillPoolReader { } // No files in queue - check if writer is done - if shared.writer_dropped { + if shared.remaining_writer_count == 0 { // Writer is done and no more files will be added - EOF return Poll::Ready(None); } @@ -747,7 +791,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 { @@ -764,24 +808,35 @@ mod tests { fn create_spill_channel( max_file_size: usize, + ) -> (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)); + + spsc_channel(max_file_size, spill_manager) + } + + fn create_shared_spill_channel( + max_file_size: usize, ) -> (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)); - 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) } @@ -1205,6 +1260,57 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "multi_thread", worker_threads = 10)] + async fn test_concurrent_writers() -> Result<()> { + let (writer, mut reader) = create_shared_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); @@ -1323,7 +1429,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 { @@ -1385,7 +1491,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. @@ -1464,7 +1570,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 {