Iterating Over Data

Consume Ray Datasets in training and inference loops with iter_batches, iter_torch_batches, and iter_rows.

Ray Data’s iterators stream blocks from a dataset, decoding them into the format your code consumes.

iter_batches

The general-purpose iterator. Returns a dict of column-arrays in the requested batch format.

for batch in ds.iter_batches(batch_size=64, batch_format="numpy"):
    process(batch["features"], batch["label"])

batch_format	Type returned
"numpy" (default)	dict[str, np.ndarray]
"pandas"	pd.DataFrame
"pyarrow"	pa.Table

iter_torch_batches

Yields PyTorch tensors directly, with optional dtype and device casts.

for batch in ds.iter_torch_batches(batch_size=64, dtypes={"x": torch.float32}):
    out = model(batch["x"].cuda())

iter_tf_batches

Equivalent for TensorFlow.

for batch in ds.iter_tf_batches(batch_size=64):
    ...

iter_rows

For row-by-row iteration. Slower but useful for debugging or for non-vectorized libraries.

for row in ds.iter_rows():
    print(row)

Prefetching

Ray Data prefetches blocks ahead of the consumer to overlap IO with compute. Tune with prefetch_batches:

ds.iter_torch_batches(batch_size=64, prefetch_batches=4)

Sharding for distributed training

Use streaming_split to produce one independent iterator per training worker.

shards = ds.streaming_split(n=4, equal=True)

@ray.remote(num_gpus=1)
def train(shard):
    for batch in shard.iter_torch_batches(batch_size=64):
        ...

ray.get([train.remote(s) for s in shards])

When you pass a Dataset to TorchTrainer or TFTrainer, sharding happens automatically.

Local shuffle

For randomization during training, enable a local shuffle buffer per iterator:

ds.iter_torch_batches(batch_size=64, local_shuffle_buffer_size=10_000)

Next steps

Shuffling

Choose between local and global shuffles.

Performance tips

Tune iterator throughput.