diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 465ca4a99e961..14e31b0cef12e 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2913,7 +2913,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; @@ -3214,6 +3214,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(()) + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) } diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 2639188a2609d..14e579faef573 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; @@ -56,11 +57,9 @@ 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>>, + open_write_files: VecDeque>>, /// 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. @@ -74,8 +73,7 @@ impl SpillPoolShared { files: VecDeque::new(), spill_manager, waker: None, - writer_dropped: false, - current_write_file: None, + open_write_files: VecDeque::new(), active_writer_count: 1, } } @@ -171,8 +169,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); @@ -193,52 +193,49 @@ 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 = 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; + let max_file_size_reached = 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 max_file_size_reached { + // 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 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.open_write_files.push_back(write_file); } Ok(()) @@ -258,31 +255,32 @@ 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.open_write_files.is_empty() { + let files = mem::take(&mut shared.open_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(); } } @@ -722,7 +720,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); } @@ -746,7 +744,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 { @@ -1204,6 +1202,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);