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

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<u32> = None;
+            while let Some(result) = stream.next().await {
+                let batch = result?;
+                let col = batch
+                    .column(0)
+                    .as_any()
+                    .downcast_ref::<UInt32Array>()
+                    .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<Schema> {
         Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
     }

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<SpillManager>,
     /// Pool-level waker to notify when new files are available (single reader)
     waker: Option<Waker>,
-    /// 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<Arc<Mutex<ActiveSpillFileShared>>>,
+    open_write_files: VecDeque<Arc<Mutex<ActiveSpillFileShared>>>,
     /// 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::<Int32Array>()
+                            .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::<Vec<_>>());
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_reader_catches_up_to_writer() -> Result<()> {
         let (writer, mut reader) = create_spill_channel(1024 * 1024);