Opensource internals

One take, N slices: how DataFusion's RepartitionExec splits a batch across shuffle partitions

12 min read Jay Hubbert

DataFusion's hash shuffle performs a single gather per input batch and hands out zero-copy RecordBatch slices, not per-partition copies.

TL;DR: BatchPartitioner never copies a batch once per partition. It computes one row-index vector per partition, concatenates them, performs a single take_arrays over the whole batch, and then hands out RecordBatch::slice(start, len) views over that one reordered buffer. Round-robin partitioning skips the take entirely — the whole batch goes to one output partition, whole.

The problem: a shuffle is a scatter-gather in disguise

Every parallel plan in DataFusion eventually has to move rows between worker threads. RepartitionExec is that operator: it maps N input partitions onto M output partitions according to a Partitioning scheme (datafusion/physical-expr/src/partitioning.rs:122-132):

rust 
pub enum Partitioning {
    /// Allocate batches using a round-robin algorithm and the specified number of partitions
    RoundRobinBatch(usize),
    /// Allocate rows based on a hash of one of more expressions and the specified number of partitions
    Hash(Vec<Arc<dyn PhysicalExpr>>, usize),
    /// Partition rows by source-declared ranges
    Range(RangePartitioning),
    /// Unknown partitioning scheme with a known number of partitions
    UnknownPartitioning(usize),
}

The naive implementation of “split this 8192-row batch into 16 partitions” is: for each partition, scan the batch and copy the matching rows into a fresh builder. That is 16 passes over the data and 16 output allocations. It is also almost never the right shape, because the expensive part of a shuffle is copying bytes, not deciding where rows go.

DataFusion’s answer splits the problem in two:

  1. Routing (decide which row goes where) is done with plain integer indices — no data movement at all.
  2. Materialization (move the bytes) is done once for the whole batch, and the per-partition results are slices of that one materialized buffer.

The router: BatchPartitioner

rust 
// datafusion/physical-plan/src/repartition/mod.rs:649-673
pub struct BatchPartitioner {
    state: BatchPartitionerState,
    timer: metrics::Time,
}

enum BatchPartitionerState {
    Hash {
        exprs: Vec<Arc<dyn PhysicalExpr>>,
        partition_reducer: StrengthReducedU64,
        hash_buffer: Vec<u64>,
        indices: Vec<Vec<u32>>,
    },
    RoundRobin {
        num_partitions: usize,
        next_idx: usize,
    },
    Range { /* ... */ }
}

Note the shape of the Hash state: there is no per-partition RecordBatch builder, and no per-partition data buffer. There is one hash_buffer: Vec<u64> (one hash per row) and one indices: Vec<Vec<u32>> (one list of row numbers per partition). Everything else is derived.

Round robin: the batch is the unit, not the row

Partitioning::RoundRobinBatch is a bit of a misnomer if you come from Spark: it does not interleave rows. The whole batch goes to one output partition, and the next batch goes to the next partition:

rust 
// datafusion/physical-plan/src/repartition/mod.rs:1240-1247
BatchPartitionerState::RoundRobin {
    num_partitions,
    next_idx,
} => {
    let idx = *next_idx;
    *next_idx = (*next_idx + 1) % *num_partitions;
    Box::new(std::iter::once(Ok((idx, batch))))
}

No hashing, no take, no copy — the batch is handed downstream unchanged. The only subtlety is where you start, and DataFusion solves it neatly so that N inputs feeding M outputs don’t all pile onto partition 0:

rust 
// datafusion/physical-plan/src/repartition/mod.rs:1055-1068
pub fn new_round_robin_partitioner(
    num_partitions: usize,
    timer: metrics::Time,
    input_partition: usize,
    num_input_partitions: usize,
) -> Self {
    Self {
        state: BatchPartitionerState::RoundRobin {
            num_partitions,
            next_idx: (input_partition * num_partitions) / num_input_partitions,
        },
        timer,
    }
}

With 2 inputs and 4 outputs, input 0 starts at partition 0 and input 1 starts at partition 2. Combined with the fact that each input alternates, the batches land evenly.

Hash: route with integers first

rust 
// datafusion/physical-plan/src/repartition/mod.rs:1257-1279
let arrays = evaluate_expressions_to_arrays(exprs.as_slice(), &batch)?;

hash_buffer.clear();
hash_buffer.resize(batch.num_rows(), 0);

create_hashes(
    &arrays,
    REPARTITION_RANDOM_STATE.random_state(),
    hash_buffer,
)?;

for v in indices.iter_mut() {
    v.clear();
}

partition_reducer.partition_indices(hash_buffer, indices);

// Finished building index-arrays for output partitions
timer.done();

let partitioned_batches =
    Self::partition_grouped_take(&batch, indices, &self.timer)?;

Two details worth pausing on.

The hash seed is fixed. REPARTITION_RANDOM_STATE is a SeededRandomState::with_seed(0) (datafusion/physical-plan/src/repartition/mod.rs:675-677). A random seed per query would make the same key land in different partitions on different runs, breaking co-partitioning assumptions between the two sides of a join and making plans non-reproducible. Fixing the seed makes hash partitioning a pure function of the key, which is exactly what the planner’s “both children must be hash-partitioned on the same keys” rule relies on.

The modulo is strength-reduced. Computing hash % num_partitions for every row means an integer division per row, and integer division is expensive. DataFusion precomputes the compiler’s classic constant-divisor optimization once, at partitioner construction time:

rust 
// datafusion/physical-plan/src/repartition/mod.rs:939-986
enum StrengthReducedU64 {
    PowerOfTwo { mask: u64 },
    Reciprocal { divisor: u64, reciprocal: u128 },
}

fn partition_indices(self, hash_buffer: &[u64], indices: &mut [Vec<u32>]) {
    match self {
        Self::PowerOfTwo { mask } => {
            for (index, hash) in hash_buffer.iter().enumerate() {
                indices[(*hash & mask) as usize].push(index as u32);
            }
        }
        Self::Reciprocal { divisor, reciprocal } => {
            for (index, hash) in hash_buffer.iter().enumerate() {
                let quotient = Self::quotient(*hash, reciprocal);
                let partition = *hash - quotient * divisor;
                indices[partition as usize].push(index as u32);
            }
        }
    }
}

Powers of two become a mask; everything else becomes a 128-bit reciprocal multiply (ceil(2^128 / divisor)) with the remainder recovered by subtraction. Either way the inner loop is a multiply plus a subtract — no div instruction at all.

Multi-column keys: hash-then-combine, not hash-of-a-tuple

A multi-column key (PARTITION BY a, b or a join on (a, b)) is hashed one column at a time, with the second and subsequent columns rehashing the running value rather than computing independent hashes and merging them:

rust 
// datafusion/common/src/hash_utils.rs:1264-1279
pub fn create_hashes<'a, I, T>(
    arrays: I,
    random_state: &impl HashState,
    hashes_buffer: &'a mut [u64],
) -> Result<&'a mut [u64]> {
    for (i, array) in arrays.into_iter().enumerate() {
        // combine hashes with `combine_hashes` for all columns besides the first
        let rehash = i >= 1;
        hash_single_array(array.as_dyn_array(), random_state, hashes_buffer, rehash)?;
    }
    Ok(hashes_buffer)
}

For primitive arrays the “rehash” path is especially cheap, because it seeds the hasher with the previous hash rather than hashing the value and then combining:

rust 
// datafusion/common/src/hash_utils.rs:322-333
if array.null_count() == 0 {
    if rehash {
        for (hash, &value) in hashes_buffer.iter_mut().zip(array.values().iter()) {
            let mut hasher = random_state.seeded_state(*hash).build_hasher();
            value.hash_write(&mut hasher);
            *hash = hasher.finish();
        }
    } else { /* ... */ }
}

The alternative — combine_hashes(l, r) = (17*37 + l) * 37 + r (hash_utils.rs:100-103) — is used for the generic (non-primitive) path. Both are order-sensitive, which is what you want: (a=1,b=2) and (a=2,b=1) must not collide.

The payoff: one take, then slices

Here is the whole trick, in one function:

rust 
// datafusion/physical-plan/src/repartition/mod.rs:1328-1391
/// Build repartitioned hash/range output batches using one `take` per input batch.
///
/// The routers first fills one index vector per output partition. This method
/// concatenates those index vectors, performs one grouped `take_arrays`, and
/// then returns each output partition as a slice of the reordered batch.
///
/// For example, given partition indices:
///
/// ```text
/// partition 0: [2, 5]
/// partition 1: []
/// partition 2: [0, 3, 4]
/// ```
///
/// this method takes rows in `[2, 5, 0, 3, 4]` order once, then returns
/// `partition 0 = slice(0, 2)` and `partition 2 = slice(2, 3)`.
fn partition_grouped_take(
    batch: &RecordBatch,
    indices: &mut [Vec<u32>],
    timer: &metrics::Time,
) -> Result<Vec<Result<(usize, RecordBatch)>>> {
    let mut partition_ranges = Vec::with_capacity(indices.len());
    let mut reordered_indices = Vec::with_capacity(batch.num_rows());

    for (partition, p_indices) in indices.iter_mut().enumerate() {
        if p_indices.is_empty() {
            continue;
        }
        let start = reordered_indices.len();
        reordered_indices.extend_from_slice(p_indices);
        partition_ranges.push((partition, start, p_indices.len()));
        p_indices.clear();
    }

    if reordered_indices.is_empty() {
        return Ok(vec![]);
    }

    if partition_ranges.len() == 1 && partition_ranges[0].2 == batch.num_rows() {
        let (partition, _, _) = partition_ranges[0];
        return Ok(vec![Ok((partition, batch.clone()))]);
    }

    let batches = {
        let _timer = timer.timer();
        let indices_array: PrimitiveArray<UInt32Type> = reordered_indices.into();
        let columns = take_arrays(batch.columns(), &indices_array, None)?;

        let mut options = RecordBatchOptions::new();
        options = options.with_row_count(Some(indices_array.len()));
        let reordered_batch =
            RecordBatch::try_new_with_options(batch.schema(), columns, &options)?;

        partition_ranges
            .into_iter()
            .map(|(partition, start, len)| {
                Ok((partition, reordered_batch.slice(start, len)))
            })
            .collect()
    };

    Ok(batches)
}

Read it as three distinct optimizations:

  1. The index vectors are concatenated, not applied one at a time. Partition 0’s rows, then partition 1’s rows, then partition 2’s rows, all in one UInt32Array. That array is exactly the permutation we want.

  2. One take_arrays for the entire batch, not one per partition. take_arrays is a gather: it produces new arrays whose contents are the selected rows in the requested order. Doing it once over 8192 rows × M partitions’ worth of indices costs the same per row as doing it M times over smaller index sets, but it allocates one output array per column instead of M, and it keeps the hot loop tight.

  3. The per-partition batches are slices, not copies. RecordBatch::slice(start, len) produces a new RecordBatch that points at the same Arrow buffers with an offset and a length. No bytes are copied to make partition 2. This is why the function is called “grouped_take” and not “grouped_copy”.

    text 
    original batch  ─┐
                 │  indices per partition: p0=[2,5] p1=[] p2=[0,3,4]
reordered order ─┴─► take() once ─► [2,5,0,3,4]
                                     ├── slice(0,2) ──► partition 0
                                     └── slice(2,3) ──► partition 2
                                         (same buffers, different offset)

There’s also a fast path: if all rows in the batch land in a single partition, the batch is returned as-is with no gather at all (mod.rs:1367-1370). That’s the degenerate case of a completely skewed key.

Why the slices don’t leak memory

A slice keeps the whole reordered batch alive, not just its slice. If partition 0 gets 2 rows and partition 2 gets 3 rows, both slices keep the 5-row buffer alive — which is fine, because together they are the whole buffer, and the buffer is dropped once all its slices are. The pathological case is one big partition plus many tiny ones: the tiny partitions keep a large buffer alive. This is bounded by the batch size, not by the dataset, so it’s a constant-factor concern rather than a leak.

Downstream: coalescing the shreds

A repartitioned batch is tiny — batch_size / M rows on average. Feeding 8192-row shards downstream would wreck every downstream operator’s throughput assumptions. So RepartitionExec coalesces on the producer side, before the channel, using one shared LimitedBatchCoalescer per output partition:

rust 
// datafusion/physical-plan/src/repartition/mod.rs:1461-1480
/// For this, one shared [`LimitedBatchCoalescer`] per output partition is used:
///
/// ```text
///                         ┌───┐                           ┌───┐
///                      ┌─▶│   │────────▶.───────────.     │   │     ┌──────────────────┐
///                      │  └───┘ ┌───┐  ( Coalescer 0 )──▶ ├───┤ ───▶│     Output 0     │
///                      │┌──────▶│   │──▶`───────────'     │   │     └──────────────────┘
///                      ││       └───┘                     └───┘
/// ┌──────────────────┐ ││                                           ┌──────────────────┐
/// │BatchPartitioner 0│─┘│                                           │     Output 1     │
/// └──────────────────┘  │                                           └──────────────────┘
/// ```

The coalescer is shared across all input tasks targeting the same output partition (SharedCoalescer, mod.rs:276-325), with an AtomicUsize counting how many senders are still active so that the last one to finish flushes the residual batch. Coalescing is skipped for unbounded inputs, where withholding a partial batch would stall the stream forever (mod.rs:480).

The rest of the machinery, briefly

The routing above is only the CPU half. RepartitionExec also has to move the resulting batches between threads, and it does that with a deliberately unusual channel design:

text 
// datafusion/physical-plan/src/repartition/distributor_channels.rs:23-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).

A single global “gate” counter (empty_channels: AtomicUsize) implements global backpressure across all M output channels with one atomic increment/decrement per send — rather than a semaphore with M permits. The producer loop then yields to the runtime every num_partitions batches, because an endless input stream would otherwise never give the executor a chance to run (mod.rs:2347-2368).

Under memory pressure the whole path degrades gracefully: OutputChannel::send tries reservation.try_grow(size) and, on failure, spills the batch to a per-(input, output) SpillPool and sends a RepartitionBatch::Spilled marker instead of the data (mod.rs:232-252). The reader side runs a small state machine (StreamState::ReadingMemory / ReadingSpilled, mod.rs:2473-2480) that blocks on the spill stream until the spilled batch arrives, so FIFO order survives the memory/spill boundary.

Trade-offs and failure modes

The take is still a full copy. The slicing is free, but the single take_arrays copies every byte of every column into a new buffer. A shuffle therefore costs one full copy of the data, no matter how many partitions there are. This is the irreducible cost of moving data between threads; the optimization is that you pay it once rather than M times, and that the routing decision costs zero bytes.

Skew is the real enemy. Hash partitioning distributes keys, not rows. If 90% of rows share one key, 90% of the batch lands in one partition. The fast path at mod.rs:1367 only helps when the skew is total. With partial skew you get one huge output partition and M-1 nearly empty ones, and the downstream aggregate on that partition becomes the critical path. The planner’s defense is to hash on the exact key tuple rather than a subset (see the “strictest distribution match” note in datafusion/physical-optimizer/src/ensure_requirements/mod.rs:140-147), because adding more key columns spreads the hot key across more partitions. But no amount of planner cleverness fixes a genuinely skewed key: if country = 'US' is 90% of your data, hashing on (country, city) is the only fix, and only if the query semantics allow it.

Round robin trades locality for balance. RoundRobinBatch is perfectly balanced and needs no key evaluation, but it destroys key locality — every downstream group-by or join has to re-shuffle. That is exactly why the planner only inserts a round-robin when the operator has declared Distribution::UnspecifiedDistribution (i.e. it doesn’t care) but does benefit from more parallelism.

Range partitioning is the third option. Partitioning::Range routes rows by sorted split points rather than by hash (repartition/range.rs), so each partition owns a contiguous key range. It gives perfect balance and preserves sort order, at the cost of needing sampled split points and requiring the input to be sorted by the range key.

Bottom line

The shuffle is the single most performance-sensitive piece of DataFusion’s runtime, and its design is a good illustration of a general principle: decide with integers, move bytes once, and hand out views. The row-index vectors cost 4 bytes per row and are thrown away immediately; the one take_arrays produces a single reordered buffer; the per-partition RecordBatchs are zero-copy windows into it. The result is a shuffle whose cost is O(bytes) and whose routing cost is O(rows), with no hidden O(partitions × bytes) term.

DatafusionExecutionPerformanceInternals