# Shuffles

> Shuffles in Dista: TaskMany for a task with several results, TaskWith for a value beside other tasks' results, partitioning again what earlier tasks made, placement by data location, and spilling past the node store.

[magmalake](https://magmalake.dev/) [Dista](https://magmalake.dev/s/dista)

A shuffle moves data between two sets of tasks: map tasks each split their output by key, and each reduce task takes its share from every map task. The partitions stay on the workers that made them, and the reducers fetch them from there.

## Map and reduce

Two kinds of task make a shuffle work. A `TaskMany` returns several results, each an object of its own, stored, fetched and freed separately, like Ray's `num_returns`. A `TaskWith` takes a value beside a list of other tasks' results, gathered in order into a `Seq`:

```
from dista import ObjectRef, Seq, TaskMany, TaskWith

def map_rows(m: Int) -> Seq[Seq[Pair]]:          # one partition per reducer
    ...
def reduce_part(j: Int, parts: Seq[Seq[Pair]]) -> Seq[Pair]:
    ...

comptime MapRows = TaskMany[Int, Seq[Pair], map_rows]
comptime ReducePart = TaskWith[Int, Seq[Pair], Seq[Pair], reduce_part]

var outputs = List[List[ObjectRef[Seq[Pair]]]]()
for m in range(8):
    outputs.append(rt.remote_many[MapRows](m, 4))   # 4 results each
var reduced = List[ObjectRef[Seq[Pair]]]()
for j in range(4):
    var mine = List[ObjectRef[Seq[Pair]]]()
    for m in range(8):
        mine.append(outputs[m][j])
    reduced.append(rt.remote[ReducePart](j, mine))  # partition j of every map
```

A `TaskMany` must return exactly as many values as asked for, or every one of its results fails. The repository's `tests/test_shuffle.mojo` runs a group-by this way, 8 map tasks for 4 reducers, and checks it against one process.

## Partitioning again

A `TaskMany` can wait on other tasks' results too: `rt.remote_many[T](ref, n)` on one result, or `rt.remote_many[T](refs, n)` when it takes a `Seq`. A later stage can then partition again what earlier tasks made, without the driver fetching any of it.

## Placement by data location

The head learns each worker's node from the results it stores: a node is a host name, or `DISTA_NODE` where it's set. When one node holds more than half the bytes of a task's stored inputs, the task prefers it, and a free worker on that node takes it first. A task waits only while a worker on its own node is free to take it, so a busy node's tasks still go elsewhere. A step after a map saves the whole fetch; an all-to-all shuffle gains little, since each reducer's inputs are spread evenly.

## Spilling

A shuffle's intermediate data can be larger than memory. A worker keeps up to `DISTA_STORE_BYTES` of results in its node store (no limit when unset) and writes the rest to `DISTA_SPILL_DIR` on disk, by default `$TMPDIR`, else `/tmp`. A store write that fails spills too. Readers don't change: one on the same host maps the file, one elsewhere fetches it. [Tuning](https://magmalake.dev/s/dista/docs/tuning#store) covers the store's settings.

The shuffle benchmark, on Dista and on Ray, is in [bench/shuffle](https://github.com/magmalake/dista.mojo/tree/main/bench/shuffle), and `JOB=shuffle bench/cloud/bench matrix` runs it on AWS.

[A magmalake project](https://magmalake.dev/)

Apache 2.0 · © 2026 Marius Seritan