iter_batches#

Return an iterable over batches of data.

This method is useful for model training.

Note

This operation will trigger execution of the lazy transformations performed on this dataset.

Examples

>>> import ray
>>> ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
>>> for batch in ds.iter_batches(batch_size=2, batch_format="numpy"):
...     print(batch)
{'image': array([[[[...]]]], dtype=uint8)}
{'image': array([[[[...]]]], dtype=uint8)}

Note

Breaking out of the for-loop above shuts the streaming executor down so it stops producing blocks into the object store. If you keep your own reference to the iterator (it = iter(...)), cleanup is deferred until that reference is dropped — call it.close() to release resources eagerly.

Some libraries (for example PyTorch Lightning’s limit_train_batches) hold an iter() reference internally to cap how many batches are consumed. In those cases prefer ds.limit(n) on the dataset so iteration ends naturally after n rows.

Time complexity: O(1)

Parameters:

prefetch_batches (int) – The number of batches to fetch ahead of the current batch. Ray Data prepares batches in background threads while your code consumes the current one. If greater than 0, a thread pool formats the batches. The pool has prefetch_batches threads, up to a limit of 4. To change the limit, set the RAY_DATA_MAX_FORMAT_THREADPOOL_NUM_WORKERS environment variable. If 0, Ray Data doesn’t prefetch blocks and formats each batch without the thread pool. Defaults to 1.
batch_size (int | None) – The number of rows in each batch, or None to use entire blocks as batches (blocks may contain different numbers of rows). The final batch may include fewer than batch_size rows if drop_last is False. Defaults to 256.
batch_format (str | None) – If "default" or "numpy", batches are Dict[str, numpy.ndarray]. If "pandas", batches are pandas.DataFrame. If "pyarrow", batches are pyarrow.Table. If "cudf" [Experimental], batches are cudf.DataFrame.
drop_last (bool) – Whether to drop the last batch if it’s incomplete.
local_shuffle_buffer_size (int | None) – If not None, Ray Data randomly shuffles the data using a local in-memory shuffle buffer, and this value serves as the minimum number of rows that must be in the buffer to yield a batch. When there are no more rows to add to the buffer, Ray Data drains the remaining rows. A larger buffer makes the shuffle more random but might increase the latency to the first batch. If this value is smaller than batch_size, Ray Data rounds it up to batch_size. Local shuffling raises a ValueError if batch_size is None.
local_shuffle_seed (int | None) – The seed to use for the local random shuffle.
_collate_fn (Callable[[pyarrow.Table | pandas.DataFrame | Dict[str, ndarray] | cudf.DataFrame], CollatedData] | None) – A custom function to collate each batch of data.

Returns:

An iterable over batches of data.

Return type:

Iterable[pyarrow.Table | pandas.DataFrame | Dict[str, ndarray] | cudf.DataFrame]