DataFusion's Spillable Shuffle: Writing Batches to Disk Without Reordering Them

TL;DR: RepartitionExec gives every output partition one shared memory reservation. When try_grow fails for a batch, the batch is written to that partition’s spill pool on disk and a zero-byte marker travels through the in-memory channel instead. The reader sees the marker and switches to reading the spill file, blocking the channel poll until it has the batch — that is how FIFO order survives a disk round trip. Contrast the aggregation spill, which drains the entire hash table: the shuffle spills selectively, one batch at a time.

The shuffle is the worst operator in the plan

A RepartitionExec is an exchange: N input tasks each produce batches for M output partitions. That makes it simultaneously the most parallel and the most dangerous operator in the plan.

The danger is a deadlock that has nothing to do with locks. Suppose you give each of the M output channels a bounded buffer of size K, and each of the N input tasks writes to all M channels in a fixed order. Input task 1 fills channel 0 and blocks; input task 2 fills channel 1 and blocks; the consumer is waiting on channel 0 but input task 0 is blocked behind input task 1 on channel 1. Classic cyclic wait. Every bounded-buffer exchange has this shape.

DataFusion’s answer is to make the channels unbounded and to put a global gate in front of them:

// datafusion/physical-plan/src/repartition/distributor_channels.rs:37-40
//! There are `N` virtual MPSC (multi-producer, single-consumer) channels with unbounded capacity. However, if all
//! buffers/channels are non-empty, than a global gate will be closed preventing new data from being written (the
//! sender futures will be [pending](Poll::Pending)) until at least one channel is empty (and not closed).

The gate is a single atomic counter of empty channels:

// datafusion/physical-plan/src/repartition/distributor_channels.rs:412-425
struct Gate {
    /// Number of currently empty (and still open) channels.
    empty_channels: AtomicUsize,
    /// Wakers for the sender side, including their channel IDs.
    ///
    /// This is `None` if the there are non-empty channels.
    send_wakers: Mutex<Option<Vec<(Waker, usize)>>>,
}

And the sender checks it before enqueuing:

// datafusion/physical-plan/src/repartition/distributor_channels.rs:218-225
// does ANY receiver need data?
// if so, allow sender to create another
if this.gate.empty_channels.load(Ordering::SeqCst) == 0 {
    let mut guard = this.gate.send_wakers.lock();
    if let Some(send_wakers) = &mut *guard {
        send_wakers.push((cx.waker().clone(), this.channel.id));
        return Poll::Pending;
    }
}

Read the gate’s guarantee precisely, because it is narrower than it looks: at least one channel is empty at any moment. That is a liveness guarantee, not a memory bound. A sender is allowed to push into an already-non-empty channel as long as some other channel is empty, so a single channel can still accumulate a long run of batches while a sibling sits empty.

What the gate buys is the absence of a cyclic wait. When the gate closes, every output channel has data, so the consumer is guaranteed to be able to drain one, which empties it and reopens the gate (distributor_channels.rs:305-320). No sender can be waiting on a resource that only a blocked peer can release, because no peer is blocked — the consumer is always running.

The memory bound comes from somewhere else: the output partition’s MemoryReservation. The gate is what keeps the exchange deadlock-free; the memory pool is what keeps it from eating the machine. Conflating the two is easy, and the code keeps them strictly separate — Gate knows nothing about bytes.

The lock ordering is documented and strict — “channel lock is always acquired prior to gate lock” (distributor_channels.rs:121-123) — because the reverse order would let a sender and a receiver deadlock on the two mutexes.

Where the memory actually goes

Each output partition gets one MemoryConsumer shared by all N input tasks:

// datafusion/physical-plan/src/repartition/mod.rs:506-511
let reservation = Arc::new(
    MemoryConsumer::new(format!("{name}[{partition}]"))
        .with_can_spill(true)
        .register(context.memory_pool()),
);

The reservation is per output partition, not per (input, output) pair — because the memory that matters is the memory the consumer of that partition is holding. When the consumer reads a batch out of the channel, it shrinks the reservation (repartition/mod.rs:2559-2562).

The try_grow happens on the producer side, at send time:

// datafusion/physical-plan/src/repartition/mod.rs:232-252
async fn send(&mut self, batch: RecordBatch) -> Result<(), SendError<MaybeBatch>> {
    let size = batch.get_array_memory_size();

    // Decide the payload outside of any await: never hold a MutexGuard
    // across an await point.
    let (payload, is_memory_batch) = {
        match self.reservation.try_grow(size) {
            Ok(_) => (Ok(RepartitionBatch::Memory(batch)), true),
            Err(_) => match self.spill_writer.push_batch(&batch) {
                Ok(()) => (Ok(RepartitionBatch::Spilled), false),
                Err(err) => (Err(err), false),
            },
        }
    };

    let result = self.sender.send(Some(payload)).await;
    if result.is_err() && is_memory_batch {
        self.reservation.shrink(size);
    }
    result
}

Four things to notice:

1. The fallback is inline. try_grow failing is not an error; it selects the spill path.
2. The batch is not cloned. RepartitionBatch::Spilled is a marker — it carries no data. The batch itself already went to disk.
3. The payload decision happens outside any await. The comment is explicit: “never hold a MutexGuard across an await point”. Computing the payload requires the reservation’s mutex; doing that before the send().await keeps the critical section short and avoids the classic SendFuture-holds-lock-across-await bug.
4. A failed send refunds the memory. If the receiver hung up (a LIMIT dropped the stream), the reservation is shrunk — otherwise a cancelled query would leak its reservation until the whole task was dropped.

The marker and the two-state reader

The reader side is a small state machine:

// datafusion/physical-plan/src/repartition/mod.rs:2473-2482
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StreamState {
    /// Reading from the memory channel (normal operation)
    ReadingMemory,
    /// Waiting for a spilled batch from the spill stream.
    /// Must not poll channel until spilled batch is received to preserve ordering.
    ReadingSpilled,
}

// datafusion/physical-plan/src/repartition/mod.rs:2546-2590
StreamState::ReadingMemory => {
    let value = match self.receiver.recv().poll_unpin(cx) {
        Poll::Ready(v) => v,
        Poll::Pending => return Poll::Pending,
    };
    match value {
        Some(Some(v)) => match v {
            Ok(RepartitionBatch::Memory(batch)) => {
                // Release memory and return batch
                self.reservation.shrink(batch.get_array_memory_size());
                return Poll::Ready(Some(Ok(batch)));
            }
            Ok(RepartitionBatch::Spilled) => {
                // Batch was spilled, transition to reading from spill stream
                // We must block on spill stream until we get the batch
                // to preserve ordering
                self.state = StreamState::ReadingSpilled;
            }
            Err(e) => return Poll::Ready(Some(Err(e))),
        },
        ...

and

// datafusion/physical-plan/src/repartition/mod.rs:2591-2612
StreamState::ReadingSpilled => {
    match self.spill_stream.poll_next_unpin(cx) {
        Poll::Ready(Some(Ok(batch))) => {
            self.state = StreamState::ReadingMemory;
            return Poll::Ready(Some(Ok(batch)));
        }
        ...
        Poll::Ready(None) => {
            // Spill stream ended — release its resources before
            // we go back to draining the memory channel.
            ...
            self.state = StreamState::ReadingMemory;
        }
        Poll::Pending => {
            // Spilled batch not ready yet, must wait
            // This preserves ordering by blocking until spill data arrives
            return Poll::Pending;
        }
    }
}

This is the crux. The channel is a FIFO of slots, and a slot may be either a real batch or a “there is a batch on disk” marker. When the reader hits a marker it must not go back to polling the channel — otherwise it would return the next batch while the spilled one is still in flight, reordering the output. So it blocks on the spill stream until it has the batch, then resumes the channel.

This is why the marker approach beats the alternative (writing the batch to disk and sending a promise that resolves later): the promise version needs a second queue and a join, and the join is where the ordering bugs live. The marker version reuses the channel’s own FIFO as the ordering guarantee.

The spill pool

The spill pool is a small, carefully designed concurrent structure: datafusion/physical-plan/src/spill/spill_pool.rs.

// datafusion/physical-plan/src/spill/spill_pool.rs:34-64
/// Shared state between the writer and readers of a spill pool.
/// This contains the queue of files and coordination state.
///
/// # Locking Design
///
/// This struct uses **fine-grained locking** with nested `Arc<Mutex<>>`:
/// - `SpillPoolShared` is wrapped in `Arc<Mutex<>>` (outer lock)
/// - Each `ActiveSpillFileShared` is wrapped in `Arc<Mutex<>>` (inner lock)
///
/// This enables:
/// 1. **Short critical sections**: The outer lock is held only for queue operations
/// 2. **I/O outside locks**: Disk I/O happens while holding only the file-specific lock
/// 3. **Concurrent operations**: Reader can access the queue while writer does I/O
///
/// **Lock ordering discipline**: Never hold both locks simultaneously to prevent deadlock.

The reader can start reading a file before the writer has sealed it. The writer appends a batch, flushes, and wakes the reader; the reader reads whatever has been flushed. Files rotate when they exceed max_spill_file_size_bytes (default 128 MiB, datafusion/common/src/config.rs:1037):

// datafusion/physical-plan/src/spill/spill_pool.rs:262-276
let max_file_size_reached = file_shared.estimated_size > self.max_file_size_bytes;

if max_file_size_reached {
    // Finish the IPC writer
    if let Some(mut writer) = file_shared.writer.take() {
        writer.finish()?;
    }
    file_shared.writer_finished = true;
    file_shared.wake();
} else {
    drop(file_shared);
    let mut shared = self.shared.lock();
    shared.open_write_files.push_back(write_file);
}

SPSC vs MPSC: the ordering contract dictates the topology

// datafusion/physical-plan/src/repartition/mod.rs:513-546
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.
    ...
    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 {
    // 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])
};

This is the “spill everything vs spill selectively” question, and for the shuffle the answer is dictated by a semantic requirement rather than by preference:

• preserve_order: the output must be globally sorted, and the downstream StreamingMerge merges the per-input streams. Each input’s data must stay in order relative to itself. So each (input, output) pair gets its own single-producer pool — spsc_channel — and the writer is moved into the input task, never cloned. The type system enforces it:

// datafusion/physical-plan/src/repartition/mod.rs:185-213
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),
}

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()),
    }
}

• non-preserve-order: the output is an unordered multiset, so all N inputs can share one multi-producer pool per output partition. mpsc_channel returns a cloneable SpillPoolWriter; new_sink() hands out an extra sink and bumps remaining_writer_count, which is how the pool knows not to signal EOF while another producer is still alive (spill_pool.rs:105-135).

So the number of spill files is N × M in preserve-order mode and M in non-preserve mode. That is a real cost: 16 inputs × 16 outputs = 256 open files in the worst case.

The joins: what spills and what does not

This is where the “spill everything vs spill selectively” split becomes stark, because DataFusion’s joins are on a spectrum.

HashJoinExec: reserves, but does not spill

The hash join does register a consumer and reserve for its build side:

// datafusion/physical-plan/src/joins/hash_join/exec.rs:1905-1917
(None, PartitionMode::Partitioned) => {
    let left_stream = self.left.execute(partition, Arc::clone(&context))?;
    let reservation =
        MemoryConsumer::new(format!("HashJoinInput[{partition}]"))
            .register(context.memory_pool());
    OnceFut::new(collect_left_input(...))
}

and the hash table allocation is gated:

// datafusion/physical-plan/src/joins/hash_join/exec.rs:3080-3099
let estimated_hashtable_size =
    estimate_memory_size::<(u32, u64)>(num_rows, fixed_size_u32)?
        // Each build row also owns an index in the duplicate-key chain.
        .checked_add(num_rows * size_of::<u32>())
        ...;
reservation.try_grow(estimated_hashtable_size)?;

But there is no spill fallback. When try_grow fails, the ResourcesExhausted error propagates and the query fails. And note that the consumer is not registered with can_spill(true) — because it genuinely cannot spill. A hash map cannot be partially spilled without either (a) serializing it and losing the O(1) probe property, or (b) partitioning it, which is a reshuffle.

The practical consequence: a hash join with a large build side is the hardest query to run under a memory limit in DataFusion. The workaround is to force a sort-merge join instead, which does spill.

SortMergeJoinExec: spills selectively, and only what it must

The sort-merge join streams the probe side and buffers the inner side. For semi/anti/mark joins with a filter, it needs to buffer the inner rows of the current key group so each can be evaluated against the outer group. That buffer is the only thing that can grow without bound, and it is spilled — but only when the pool refuses:

// datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs:686-700
// Reserve memory for the newly buffered slice. If the pool
// is exhausted, spill the entire buffer to disk.
if self.try_resize_reservation().is_err() {
    if self.runtime_env.disk_manager.tmp_files_enabled() {
        self.spill_inner_key_buffer(&mut writer)?;
    } else {
        // Re-attempt to get the error message
        self.try_resize_reservation().map_err(|e| {
            DataFusionError::Execution(format!("{e}. Disk spilling disabled."))
        })?;
    }
}

The spill is per key group, appended to a single open file for that group, and the file is finished when the group is complete (bitwise_stream.rs:447-466).

The “skip spilling for small sides” idea shows up in the design of this operator rather than in a branch: the no-filter path never buffers inner rows at all — it just advances the cursor — so it cannot OOM and never touches the disk. The doc comment spells out the asymmetry:

// datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs:104-112
//! Memory usage is bounded and independent of total input size:
//! - One outer batch at a time (not tracked by reservation — single batch,
//!   cannot be spilled since it's needed for filter evaluation)
//! - One inner batch at a time (streaming)
//! - `matched` bitset: one bit per outer row, re-allocated per batch
//! - Inner key group buffer: only for filtered joins, one key group at a time.
//!   Tracked via `MemoryReservation`; spilled to disk when the memory pool
//!   limit is exceeded.

The degenerate case is documented too: a single hot key means one key group buffers an unbounded number of rows — in memory or on disk (bitwise_stream.rs:114-124).

NestedLoopJoinExec: spills the build side in chunks

The nested loop join broadcasts the left side, so it spills the left side when it does not fit:

// datafusion/physical-plan/src/joins/nested_loop_join.rs:740-741
let can_spill = context.runtime_env().disk_manager.tmp_files_enabled()
    && !coordinated_fallback_disabled;

and notably it pre-creates the spill manager up front, with a comment that explains why:

// datafusion/physical-plan/src/joins/nested_loop_join.rs:743-751
// Built here rather than on demand: by the time the load hits the memory limit the
// stream is partly consumed, and it has to be spillable at that point without going
// back to the left child.

That is a genuinely subtle piece of engineering: the spill decision is deferred (you only know you need to spill when you hit the limit), but the capability to spill must be established before the child stream is consumed, because you cannot re-read a child.

Trade-offs and failure modes

• The gate is a liveness mechanism, not a memory bound. It guarantees at least one output channel is empty, which prevents the cyclic wait; it does not limit how many batches a single channel can hold. The memory bound comes from the output partition’s MemoryReservation — and the channel buffer is not accounted to it until the batch is sent. With 1024 output partitions and 1 MiB batches, a single hot partition can accumulate a lot of channel buffer that the memory pool never sees.
• Spilling a shuffle is not free and not always a win. You pay an IPC encode, a file write, a file read, and a decode — for a batch that might have been delivered in microseconds. The try_grow fallback triggers on every batch once the pool is full, so a saturated shuffle writes everything to disk. This is the “spill everything” pathology, and unlike the aggregation, the shuffle has no way to avoid it: it cannot batch up and spill less often, because the consumer wants that batch.
• The marker scheme requires a strictly FIFO channel. The whole ordering argument collapses if the channel reorders. preserve_order mode pays for this with an SPSC channel per input; non-preserve-order mode gets away with one MPSC pool per output.
• File descriptor pressure. N × M spill files in preserve-order mode. With 32 inputs and 32 outputs that is 1024 files, each with its own InProgressSpillFile, IPC writer and buffered encoder.
• Disk quota is the real limit. max_temp_directory_size is enforced by the DiskManager, and exceeding it is an I/O error, not a memory error. A shuffle that spills more than the quota dies in a way that looks unrelated to the problem.
• Hash joins do not spill. If your plan has a hash join with a big build side and a memory limit, the query fails with ResourcesExhausted and no amount of spilling elsewhere will save it. Force a sort-merge join.
• Cancelled queries leak file handles until the abort tasks run. The AsyncIPCStreamWriter drop handler spawns a cleanup task bounded by a semaphore of 8 concurrent aborts (spill/mod.rs:521-527); beyond that, the writer is simply dropped and the backend’s own lifecycle policy takes over.

The through-line: the shuffle spills selectively and reactively, because its unit of work is a self-contained batch that can be written and read back independently. The aggregation spills wholesale, because its unit of work is a hash table that can only be reconstructed by sorting. The join sits in between: it spills the one structure that can grow without bound, and leaves everything else streaming. Each of those choices follows from what the data structure actually is, and that is the interesting part of the design.