Zarr datasource

[alexandrplashchinsky]

Description

This PR introduces ray.data.read_zarrv2() and the backing ZarrV2Datasource for reading Zarr v2 stores with Ray Data.

This adds a dedicated public API and datasource implementation for Zarr v2 so users can read chunk metadata from consolidated Zarr v2 stores through the standard Ray Data read API surface.

Related issues

N/A

Additional information

This PR adds:

• ray.data.read_zarrv2() as a new public Ray Data read API
• ZarrV2Datasource as the datasource implementation used by the API
• unit tests covering the datasource behavior and API wiring

Example usage:

import ray

ds = ray.data.read_zarrv2("/path/to/store")

[gemini-code-assist]

Code Review

This pull request introduces support for reading Zarr v2 stores in Ray Data by adding the ZarrV2Datasource and a public read_zarrv2 API. The implementation includes support for various storage backends (local, S3, Azure) and handles chunk metadata, slice bounds, and padding. Feedback focuses on reducing logic duplication in chunk calculation, simplifying path normalization, converting a utility method to a static method, and adhering to standard Python formatting for keyword arguments.

[ArturNiederfahrenhorst]

I created this dataset and tested with it. So we don't need to skip here.

[ArturNiederfahrenhorst]

This is a larger change from where the PR started:
The key idea here is: In any given zarr dataset, you'll have a bunch of arrays. Some of these arrays will typically be metadata and as such don't represent rows of the dataset. Others do represent rows of the dataset (example: mist dataset would have two arrays, one for images, one for labels, one image and one label belong in one row).
So when reading data from datasets that have metadata etc., you should provide array_paths to read only the arrays together that align on their 0th axis semantically.

[cursor]

Trailing slash not stripped in full-scan array path

Low Severity

In _load_metadata_full_scan, the dirpath from fs.walk is used directly with removeprefix without stripping trailing slashes. The code correctly uses dirpath.rstrip("/") for the root comparison (line 123) and for building zarray_path (line 126), but the array_path key computed via dirpath.removeprefix(store_prefix) could retain a trailing slash if the filesystem's walk yields one. This would produce keys like "sub/" that won't match the normalized (slash-stripped) paths used in filtering and user-facing array_paths.

 

^{Reviewed by Cursor Bugbot for commit 80024bb. Configure here.}

[cursor]

Mixed-rank array chunks in single batch breaks Arrow

Medium Severity

When materialize=True and the store has multiple arrays of different ranks (e.g., 2-D images and 1-D labels), get_read_tasks mixes chunks from all arrays into the same batch. A single batch's chunk column would then contain numpy arrays of different ranks. Ray Data converts DataFrames to Arrow internally, and Arrow's tensor extension requires uniform rank within a column — this will fail at conversion time when parallelism is low enough to group different arrays together. The alternate implementation in the PR correctly batches per-array to avoid this.

 

^{Reviewed by Cursor Bugbot for commit c309fea. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 3 potential issues.

There are 5 total unresolved issues (including 2 from previous reviews).

^{Reviewed by Cursor Bugbot for commit 1f0f1fa. Configure here.}

[cursor]

Test references undefined module-level function _descriptor_row_size_bytes

Medium Severity

The test calls zarrv2_datasource._descriptor_row_size_bytes(array_name, meta) but this function is never defined anywhere in the codebase. The module defines related constants (_PYINT_BYTES, _PYSTR_BASE, _PYTUPLE_BASE, etc.) that are clearly intended for this function, but the implementation was never added. This test will always fail with AttributeError.

Additional Locations (1)

• python/ray/data/_internal/datasource/zarrv2_datasource.py#L41-L52

 

^{Reviewed by Cursor Bugbot for commit 1f0f1fa. Configure here.}

[cursor]

Unused module-level constants from missing function implementation

Low Severity

The constants _PYINT_BYTES, _PYSTR_BASE, _PYTUPLE_BASE, _PYLIST_BASE, _PYPTR_BYTES, _INT_IN_SEQ, and _PAIR_IN_LIST are defined but never referenced anywhere in the module. They appear to be remnants of a _descriptor_row_size_bytes function that was never implemented, leaving dead code that adds confusion for future readers.

 

^{Reviewed by Cursor Bugbot for commit 1f0f1fa. Configure here.}