[data] Zarr datasource

.vale/styles/config/vocabularies/Data/accept.txt

@@ -57,5 +57,6 @@ YOLO
 [Ss]harded
 [Pp]arameterization(s)?
 repr
+Zarr
 [Uu]psample(d|s)?
 [Rr]ebatch(ing|ed|es)?

doc/source/data/api/loading_data.rst

@@ -352,6 +352,15 @@ WebDataset
 
    read_webdataset
 
+Zarr
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_zarr
+
 Partitioning API
 ^^^^^^^^^^^^^^^^
 

doc/source/data/loading-data.rst

@@ -150,6 +150,16 @@ To view the full list of supported file formats, see the
             petal.width   float
             sepal.length  float
 
+    .. tab-item:: Zarr
+
+        To read a Zarr v2 store, call :func:`~ray.data.read_zarr`.
+
+        .. code-block:: python
+
+            import ray
+
+            ds = ray.data.read_zarr("s3://anonymous@ray-example-data/mnist-tiny.zarr")
+
 
 Reading files from local disk
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

doc/source/data/user-guide.rst

@@ -24,6 +24,7 @@ shows you how to achieve several tasks.
     working-with-images
     working-with-text
     working-with-tensors
+    working-with-zarr
     working-with-pytorch
     working-with-llms
     how-to-avoid-ooms

doc/source/data/working-with-zarr.md

@@ -0,0 +1,155 @@
+(working_with_zarr)=
+
+# Working with Zarr
+
+Ray Data reads [Zarr v2](https://zarr.readthedocs.io/) stores — chunked, compressed,
+N-dimensional arrays on local disk or cloud object storage — with
+{func}`ray.data.read_zarr` (`zarr-python` 2.x / Zarr v2 stores).
+
+This guide covers:
+
+- [The two output schemas](#output-schemas) — long-form (default) and aligned wide-form
+- [Selecting arrays and metadata discovery](#selecting-arrays-and-metadata-discovery)
+- [Controlling chunk size](#controlling-chunk-size)
+- [Reading row-aligned arrays](#reading-row-aligned-arrays)
+- [Custom codecs](#custom-codecs)
+- [Zarr's .zattrs](#zarr-zattrs)
+
+For the full parameter reference, see {func}`ray.data.read_zarr`.
+
+## Output schemas
+
+`read_zarr` produces one of two schemas, selected by `align_axis_0`.
+
+### Long-form (default)
+
+By default each output row is **one chunk of one array**, with columns:
+
+- `array` — the array's path in the store (for example `"data/camera0_rgb"`, or `""` for a root-level array).
+- `chunk_index` — the N-D index of the chunk in its array's chunk grid.
+- `chunk_slices` — per-axis `(start, stop)` of the chunk in the array's coordinate space.
+- `chunk` — the chunk's data at its natural shape (trailing-edge chunks may be shorter; no padding).
+
+Arrays read in the same call need not share any dimension — different ranks, shapes,
+dtypes, and native chunk sizes coexist as separate rows.
+
+```python
+import ray
+
+ds = ray.data.read_zarr("s3://anonymous@ray-example-data/mnist-tiny.zarr")
+```
+
+```{note}
+The `chunk` column is a tensor, and tensors of different rank or dtype can't be
+combined into one batch. Consume long-form **per array** (filter on the `array`
+column first), or — when arrays are row-aligned (share `shape[0]`) — use
+`align_axis_0=True` so each array becomes its own column, which is batch-safe.
+```
+
+### Aligned wide-form (`align_axis_0=True`)
+
+With `align_axis_0=True` each row is **one axis-0 chunk shared across the selected
+arrays**, with columns:
+
+- `t_start`, `t_stop` — the global axis-0 range of the row.
+- one column per selected array, holding that array's `[t_start:t_stop, ...]` slice.
+
+All selected arrays must share `shape[0]` and resolve to the same axis-0 chunk size
+(after any `chunk_shapes` override). Use `array_paths` to choose which arrays participate —
+`align_axis_0` itself doesn't filter.
+
+```python
+ds = ray.data.read_zarr(
+    "s3://anonymous@ray-example-data/mnist-tiny.zarr",
+    align_axis_0=True,
+    chunk_shapes=[50],
+)
+```
+
+## Selecting arrays and metadata discovery
+
+By default `read_zarr` reads every array it discovers. Pass `array_paths` to read a
+subset:
+
+```python
+ds = ray.data.read_zarr(store_uri, array_paths=["images", "labels"])
+```
+
+Discovery follows these rules:
+
+- If the store has consolidated `.zmetadata`, it's the canonical array list (filtered by
+  `array_paths` if given). This is the fast path.
+- Otherwise, if `array_paths` is given, each requested array's metadata is read directly
+  — no `.zmetadata` required.
+- Otherwise, if `allow_full_metadata_scan=True`, the store is recursively scanned for
+  arrays. This can be slow or costly on large remote stores, so it's off by default;
+  prefer consolidating metadata with `zarr.consolidate_metadata` ahead of time.
+
+## Controlling chunk size
+
+Zarr stores are often chunked finely (for example one image per chunk).
+You can use `chunk_shapes` to chunk the leading axes **at read
+time** to coarsen (or refine) the granularity at which reading happens.
+Note that this doesn't affect downstream batch sizes and is internal to the reading operation.
+Finely chunked reading can hurt performance.
+
+- A **sequence** applies as a shared prefix across all selected arrays, overriding the
+  leading axes and keeping trailing axes native. `chunk_shapes=[16]` turns native chunks
+  `(1, 224, 224, 3)` into `(16, 224, 224, 3)` and `(50,)` into `(16,)`.
+- A **dict** overrides per array; arrays absent from it keep native chunks.
+
+```python
+# Coarsen every array's axis 0 to 16-element chunks.
+ds = ray.data.read_zarr(store_uri, chunk_shapes=[16])
+
+# Different overrides per array.
+ds = ray.data.read_zarr(store_uri, chunk_shapes={"images": [16], "labels": [64]})
+```
+
+## Reading row-aligned arrays
+
+When arrays share an axis-0 (for example a timestep axis), `align_axis_0=True`
+co-iterates them as the [wide-form schema](#output-schemas) — one row per axis-0
+chunk, one column per array.
+
+For sliding-window pipelines, `overlap` extends each row's per-array data forward by `N`
+timesteps from the next row's range (clipped at the end of the store). With
+`overlap=K-1`, any window of length `K` that starts in a row's owned `[t_start, t_stop)`
+fits entirely within that row's slice.
+
+```python
+ds = ray.data.read_zarr(
+    store_uri,
+    align_axis_0=True,
+    chunk_shapes=[50],
+    overlap=9,  # length-10 windows fit within a row
+)
+```
+
+## Custom codecs
+
+Stores compressed with non-stdlib codecs (for example `imagecodecs` `JPEG-XL`) need the
+codec package imported and registered **in every Ray worker**, not just the driver.
+Register it with a `worker_process_setup_hook` — pass an importable callable or its
+dotted path (a string of code isn't accepted; a string is interpreted as an import
+path):
+
+```python
+import ray
+
+ray.init(runtime_env={
+    "worker_process_setup_hook": "imagecodecs.numcodecs.register_codecs",
+})
+```
+
+This is a particularity of the underlying Zarr library.
+
+
+(zarr-zattrs)=
+## Zarr's .zattrs
+
+`read_zarr` doesn't surface each array's `.zattrs` (Zarr user attributes) in the row
+schema — they're invariant per array, so repeating them on every row would just bloat
+the output. Read them separately (for example with the `zarr` package) if your job
+needs them.
+

pyrefly.toml

@@ -65,4 +65,15 @@ ignore-missing-imports = [
     "rapidsmpf.*",
     "rmm.*",
     "confluent_kafka.*",
+    "zarr.*",
+]
+
+# numcodecs' ``abc.Codec`` ABC declares its abstract ``encode``/``decode`` with
+# empty bodies, so pyrefly infers their return type as ``None`` and flags any
+# real codec subclass (like the one in test_zarrv2) as an inconsistent override
+# -- a numcodecs typing defect we can't satisfy without lying about the return
+# type. So treat numcodecs as Any. (``ignore-missing-imports`` only covers
+# modules that can't be found, not installed ones, so it belongs here.)
+replace-imports-with-any = [
+    "numcodecs.*",
 ]

python/deplocks/ci/data-base-ci_depset_py3.10.lock

@@ -313,6 +313,11 @@ arro3-core==0.8.0 \
     # via
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   deltalake
+asciitree==0.3.3 ; python_full_version < '3.11' \
+    --hash=sha256:4aa4b9b649f85e3fcb343363d97564aa1fb62e249677f2e18a96765145cc0f6e
+    # via
+    #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
+    #   zarr
 asn1crypto==1.5.1 \
     --hash=sha256:13ae38502be632115abf8a24cbe5f4da52e3b5231990aff31123c805306ccb9c \
     --hash=sha256:db4e40728b728508912cbb3d44f19ce188f218e9eba635821bb4b68564f8fd67
@@ -1339,6 +1344,12 @@ fastavro==1.12.1 \
     # via
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   -r python/requirements/ml/py313/data-test-requirements.txt
+fasteners==0.20 ; python_full_version < '3.11' and sys_platform != 'emscripten' \
+    --hash=sha256:55dce8792a41b56f727ba6e123fcaee77fd87e638a6863cec00007bfea84c8d8 \
+    --hash=sha256:9422c40d1e350e4259f509fb2e608d6bc43c0136f79a00db1b49046029d0b3b7
+    # via
+    #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
+    #   zarr
 fastrlock==0.8.3 ; sys_platform != 'darwin' \
     --hash=sha256:001fd86bcac78c79658bac496e8a17472d64d558cd2227fdc768aa77f877fe40 \
     --hash=sha256:04bb5eef8f460d13b8c0084ea5a9d3aab2c0573991c880c0a34a56bb14951d30 \
@@ -2745,6 +2756,27 @@ networkx==3.2.1 \
     --hash=sha256:9f1bb5cf3409bf324e0a722c20bdb4c20ee39bf1c30ce8ae499c8502b0b5e0c6 \
     --hash=sha256:f18c69adc97877c42332c170849c96cefa91881c99a7cb3e95b7c659ebdc1ec2
     # via torch
+numcodecs==0.13.1 ; python_full_version < '3.11' \
+    --hash=sha256:233bc7f26abce24d57e44ea8ebeb5cd17084690b4e7409dd470fdb75528d615f \
+    --hash=sha256:237b7171609e868a20fd313748494444458ccd696062f67e198f7f8f52000c15 \
+    --hash=sha256:2a86f5367af9168e30f99727ff03b27d849c31ad4522060dde0bce2923b3a8bc \
+    --hash=sha256:2eda97dd2f90add98df6d295f2c6ae846043396e3d51a739ca5db6c03b5eb666 \
+    --hash=sha256:3501a848adaddce98a71a262fee15cd3618312692aa419da77acd18af4a6a3f6 \
+    --hash=sha256:3f593c7506b0ab248961a3b13cb148cc6e8355662ff124ac591822310bc55ecf \
+    --hash=sha256:5195bea384a6428f8afcece793860b1ab0ae28143c853f0b2b20d55a8947c917 \
+    --hash=sha256:796b3e6740107e4fa624cc636248a1580138b3f1c579160f260f76ff13a4261b \
+    --hash=sha256:7a60d75179fd6692e301ddfb3b266d51eb598606dcae7b9fc57f986e8d65cb43 \
+    --hash=sha256:80d3071465f03522e776a31045ddf2cfee7f52df468b977ed3afdd7fe5869701 \
+    --hash=sha256:90d3065ae74c9342048ae0046006f99dcb1388b7288da5a19b3bddf9c30c3176 \
+    --hash=sha256:96add4f783c5ce57cc7e650b6cac79dd101daf887c479a00a29bc1487ced180b \
+    --hash=sha256:96e42f73c31b8c24259c5fac6adba0c3ebf95536e37749dc6c62ade2989dca28 \
+    --hash=sha256:a3cf37881df0898f3a9c0d4477df88133fe85185bffe57ba31bcc2fa207709bc \
+    --hash=sha256:da2230484e6102e5fa3cc1a5dd37ca1f92dfbd183d91662074d6f7574e3e8f53 \
+    --hash=sha256:e5db4824ebd5389ea30e54bc8aeccb82d514d28b6b68da6c536b8fa4596f4bca \
+    --hash=sha256:eda7d7823c9282e65234731fd6bd3986b1f9e035755f7fed248d7d366bb291ab
+    # via
+    #   -r python/requirements/ml/py313/data-test-requirements.txt
+    #   zarr
 numexpr==2.14.1 ; python_full_version < '3.12' \
     --hash=sha256:03130afa04edf83a7b590d207444f05a00363c9b9ea5d81c0f53b1ea13fad55a \
     --hash=sha256:05f9366d23a2e991fd5a8b5e61a17558f028ba86158a4552f8f239b005cdf83c \
@@ -2880,6 +2912,7 @@ numpy==2.2.6 \
     #   keras
     #   ml-dtypes
     #   modin
+    #   numcodecs
     #   numexpr
     #   pandas
     #   pylance
@@ -2899,6 +2932,7 @@ numpy==2.2.6 \
     #   torchtext
     #   torchvision
     #   webdataset
+    #   zarr
 nvidia-nccl-cu12==2.27.5 ; platform_machine != 'aarch64' and sys_platform == 'linux' \
     --hash=sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457
     # via
@@ -5771,6 +5805,10 @@ yarl==1.23.0 \
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   aiohttp
     #   delta-sharing
+zarr==2.18.3 ; python_full_version < '3.11' \
+    --hash=sha256:2580d8cb6dd84621771a10d31c4d777dca8a27706a1a89b29f42d2d37e2df5ce \
+    --hash=sha256:b1f7dfd2496f436745cdd4c7bcf8d3b4bc1dceef5fdd0d589c87130d842496dd
+    # via -r python/requirements/ml/py313/data-test-requirements.txt
 zict==3.0.0 \
     --hash=sha256:5796e36bd0e0cc8cf0fbc1ace6a68912611c1dbd74750a3f3026b9b9d6a327ae \
     --hash=sha256:e321e263b6a97aafc0790c3cfb3c04656b7066e6738c37fffcca95d803c9fba5

python/deplocks/ci/data-base-ci_depset_py3.11.lock

@@ -313,6 +313,11 @@ arro3-core==0.8.0 \
     # via
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   deltalake
+asciitree==0.3.3 \
+    --hash=sha256:4aa4b9b649f85e3fcb343363d97564aa1fb62e249677f2e18a96765145cc0f6e
+    # via
+    #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
+    #   zarr
 asn1crypto==1.5.1 \
     --hash=sha256:13ae38502be632115abf8a24cbe5f4da52e3b5231990aff31123c805306ccb9c \
     --hash=sha256:db4e40728b728508912cbb3d44f19ce188f218e9eba635821bb4b68564f8fd67
@@ -1214,6 +1219,7 @@ deprecated==1.3.1 \
     # via
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   deltalake
+    #   numcodecs
 dill==0.4.1 \
     --hash=sha256:1e1ce33e978ae97fcfcff5638477032b801c46c7c65cf717f95fbc2248f79a9d \
     --hash=sha256:423092df4182177d4d8ba8290c8a5b640c66ab35ec7da59ccfa00f6fa3eea5fa
@@ -1343,6 +1349,12 @@ fastavro==1.12.1 \
     # via
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   -r python/requirements/ml/py313/data-test-requirements.txt
+fasteners==0.20 ; sys_platform != 'emscripten' \
+    --hash=sha256:55dce8792a41b56f727ba6e123fcaee77fd87e638a6863cec00007bfea84c8d8 \
+    --hash=sha256:9422c40d1e350e4259f509fb2e608d6bc43c0136f79a00db1b49046029d0b3b7
+    # via
+    #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
+    #   zarr
 fastrlock==0.8.3 ; sys_platform != 'darwin' \
     --hash=sha256:001fd86bcac78c79658bac496e8a17472d64d558cd2227fdc768aa77f877fe40 \
     --hash=sha256:04bb5eef8f460d13b8c0084ea5a9d3aab2c0573991c880c0a34a56bb14951d30 \
@@ -2754,6 +2766,24 @@ networkx==3.6.1 \
     # via
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   torch
+numcodecs==0.15.1 \
+    --hash=sha256:1d471a1829ce52d3f365053a2bd1379e32e369517557c4027ddf5ac0d99c591e \
+    --hash=sha256:1dfdea4a67108205edfce99c1cb6cd621343bc7abb7e16a041c966776920e7de \
+    --hash=sha256:698f1d59511488b8fe215fadc1e679a4c70d894de2cca6d8bf2ab770eed34dfd \
+    --hash=sha256:a34f0fe5e5f3b837bbedbeb98794a6d4a12eeeef8d4697b523905837900b5e1c \
+    --hash=sha256:a4f7bdb26f1b34423cb56d48e75821223be38040907c9b5954eeb7463e7eb03c \
+    --hash=sha256:b0a9d9cd29a0088220682dda4a9898321f7813ff7802be2bbb545f6e3d2f10ff \
+    --hash=sha256:bef8c8e64fab76677324a07672b10c31861775d03fc63ed5012ca384144e4bb9 \
+    --hash=sha256:c3a09e22140f2c691f7df26303ff8fa2dadcf26d7d0828398c0bc09b69e5efa3 \
+    --hash=sha256:cdfaef9f5f2ed8f65858db801f1953f1007c9613ee490a1c56233cd78b505ed5 \
+    --hash=sha256:daed6066ffcf40082da847d318b5ab6123d69ceb433ba603cb87c323a541a8bc \
+    --hash=sha256:e2547fa3a7ffc9399cfd2936aecb620a3db285f2630c86c8a678e477741a4b3c \
+    --hash=sha256:e3d82b70500cf61e8d115faa0d0a76be6ecdc24a16477ee3279d711699ad85f3 \
+    --hash=sha256:eeed77e4d6636641a2cc605fbc6078c7a8f2cc40f3dfa2b3f61e52e6091b04ff
+    # via
+    #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
+    #   -r python/requirements/ml/py313/data-test-requirements.txt
+    #   zarr
 numexpr==2.14.1 ; python_full_version < '3.12' \
     --hash=sha256:03130afa04edf83a7b590d207444f05a00363c9b9ea5d81c0f53b1ea13fad55a \
     --hash=sha256:05f9366d23a2e991fd5a8b5e61a17558f028ba86158a4552f8f239b005cdf83c \
@@ -2889,6 +2919,7 @@ numpy==2.2.6 \
     #   keras
     #   ml-dtypes
     #   modin
+    #   numcodecs
     #   numexpr
     #   pandas
     #   pylance
@@ -2908,6 +2939,7 @@ numpy==2.2.6 \
     #   torchtext
     #   torchvision
     #   webdataset
+    #   zarr
 nvidia-nccl-cu12==2.27.5 ; platform_machine != 'aarch64' and sys_platform == 'linux' \
     --hash=sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457
     # via
@@ -5799,6 +5831,12 @@ yarl==1.23.0 \
     #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
     #   aiohttp
     #   delta-sharing
+zarr==2.18.7 \
+    --hash=sha256:ac3dc4033e9ae4e9d7b5e27c97ea3eaf1003cc0a07f010bd83d5134bf8c4b223 \
+    --hash=sha256:b2b8f66f14dac4af66b180d2338819981b981f70e196c9a66e6bfaa9e59572f5
+    # via
+    #   -c /tmp/ray-deps/requirements_compiled_py3.13.txt
+    #   -r python/requirements/ml/py313/data-test-requirements.txt
 zict==3.0.0 \
     --hash=sha256:5796e36bd0e0cc8cf0fbc1ace6a68912611c1dbd74750a3f3026b9b9d6a327ae \
     --hash=sha256:e321e263b6a97aafc0790c3cfb3c04656b7066e6738c37fffcca95d803c9fba5