shuffle_batches#

Shuffle rows within batches using parallel Ray Data tasks.

Each batch receives a uniform random permutation. With batch_size=None, shuffle each input block independently. A positive batch_size uses target-sized batching, as in map_batches(). Batches can span input blocks, but each task’s final batch can be smaller, including in the middle of the dataset. Shuffle batches are independent of consumer batch sizes and output block sizes.

A downstream limit() can reduce the input before shuffling, so its result can differ from a prefix of the unlimited shuffle.

To construct exact windows, call repartition() with target_num_rows_per_block and strict=True, then shuffle each block with batch_size=None. Only the final window can be smaller. Window boundaries follow the upstream execution order.

This operation streams results without materializing the whole dataset. Rows don’t move between shuffle batches, and batches aren’t randomly reordered. Use random_shuffle() for global row shuffling. See Shuffling data for usage and performance considerations.

Examples

>>> import ray
>>> ds = ray.data.range(20, override_num_blocks=4)
>>> shuffled = ds.shuffle_batches(batch_size=8, seed=42)
>>> sum(len(b) for b in shuffled.iter_batches(batch_format="pyarrow"))
20
>>> # Shuffle exact eight-row windows, except the final remainder.
>>> exact = ds.repartition(
...     target_num_rows_per_block=8, strict=True
... ).shuffle_batches(seed=42)
>>> sum(len(b) for b in exact.iter_batches(batch_format="pyarrow"))
20

Parameters:

batch_size (int | None) – Target rows per shuffle batch. If None, shuffle each input block independently. Defaults to None. A value of one returns the input unchanged.
seed (int | RandomSeedConfig | None) – Optional integer seed or RandomSeedConfig. An integer repeats permutations for the same input layout and task grouping. Enable preserve_order on the dataset’s execution options for reproducible output order. None reseeds between executions.
compute (TaskPoolStrategy | None) – Optional task-pool strategy controlling shuffle parallelism.
num_cpus (float | None) – CPU reservation for each shuffle map task.
memory (float | None) – Heap-memory scheduling reservation per shuffle task, in bytes. This is not a hard memory limit.

Returns:

A Dataset with rows shuffled within each batch.

Return type:

Dataset

PublicAPI (alpha): This API is in alpha and may change before becoming stable.