Resolve lost wakeup in SpillPoolReader with multiple concurrent SpillPoolWriters (apache#23522)

## Which issue does this PR close?

- Closes apache#23447

## Rationale for this change

When multiple `SpillPoolWriter` clones concurrently push batches to the
same channel, more than one non-finished `SpillFile` can be in flight.
This happens because each `SpillPoolWriter` clone takes the
`current_write_file` at the start of `push_batch` and puts it back when
it's done. When multiple `push_batch` calls happen concurrently, only
the first one will be able to take the `current_write_file` and the
others will all create their own new spill file. Which one gets put back
for subsequent use is a race condition.

If this occurred and the writers are all dropped before rotation happens
in, multiple files in the `files` deque will be have `writer_finished ==
false`. The last writer drop logic in `SpillPoolWriter::drop` only
finishes whatever file is the `current_write_file` as finished.

This can lead to a stalled situation when `SpillPoolFile::poll_next`
catches up with the writer and returns `Pending` because
`writer_finished == false`. A waker for the file is registered, but
since the last writer drop logic only finishes and wakes whatever
happens to be `current_write_file`, which may not be the current read
file, the waker may end up never being notified.

There is a secondary waker that is registered on the spill pool itself,
but due to fine grained locking, it is possible for the wake call in the
last writer drop logic to be called before the waker registration.

## What changes are included in this PR?

- Add support for tracking multiple unfinished write files
- Close all unfinished write files when the last writer is dropped
- Removed `writer_dropped` field which was an unnecessary
denormalisation of `active_writer_count == 0`

An additional benefit of tracking all unfinished write files is that
excessive creation of tiny spill files is avoided when many writers are
pushing batches concurrently.

## Are these changes tested?

Reproduction case from linked issue was used to confirm fix

## Are there any user-facing changes?

No

---------

Co-authored-by: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: 4 people

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<DistributionReceiver<MaybeB
 struct OutputChannel {
     sender: DistributionSender<MaybeBatch>,
     reservation: SharedMemoryReservation,
-    spill_writer: SpillPoolWriter,
+    spill_writer: SpillPoolSink,
     shared_coalescer: Option<SharedCoalescer>,
 }
 
+/// 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<Option<SpillPoolSink>>),
+    /// 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<SpillPoolSink> {
+        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<Vec<RecordBatch>> {
         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<SharedCoalescer>,
-    /// Spill writers for writing spilled data.
-    /// SpillPoolWriter is Clone, so multiple writers can share state in non-preserve-order mode.
-    spill_writers: Vec<SpillPoolWriter>,
+    /// 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<SendableRecordBatchStream>,
@@ -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::<Result<HashMap<_, _>>>()?;
 
             // 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<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(())
+    }
+
     #[tokio::test]
     async fn test_hash_partitioning_with_spilling() -> Result<()> {
         use datafusion_execution::runtime_env::RuntimeEnvBuilder;