Zarr reader

conftest.py

@@ -44,6 +44,31 @@ def pytest_runtest_setup(item):
         pytest.skip("set --run-minio-tests to run tests requiring docker and minio")
 
 
+def _xarray_subset():
+    ds = xr.tutorial.open_dataset("air_temperature")
+    return ds.isel(time=slice(0, 10), lat=slice(0, 90), lon=slice(0, 180))
+
+
+@pytest.fixture(params=[2, 3])
+def zarr_store(tmpdir, request):
+    ds = _xarray_subset()
+    filepath = f"{tmpdir}/air.zarr"
+    ds.to_zarr(filepath, zarr_format=request.param)
+    ds.close()
+    return filepath
+
+
+@pytest.fixture()
+def zarr_store_scalar(tmpdir):
+    import zarr
+
+    # can/should we create a memorystore instead?
+    store = zarr.storage.LocalStore(str(tmpdir + "/tmp.zarr"))
+    zarr_store_scalar = zarr.create_array(store=store, shape=(), dtype="int8")
+    zarr_store_scalar[()] = 42
+    return zarr_store_scalar
+
+
 # Common codec configurations
 DELTA_CODEC = {"name": "numcodecs.delta", "configuration": {"dtype": "<i8"}}
 ARRAYBYTES_CODEC = {"name": "bytes", "configuration": {"endian": "little"}}

docs/faq.md

@@ -77,7 +77,7 @@ vds.virtualize.to_icechunk(icechunkstore)
 
 ### I already have some data in Zarr, do I have to resave it?
 
-No! VirtualiZarr can (well, [soon will be able to](https://github.com/zarr-developers/VirtualiZarr/issues/262)) create virtual references pointing to existing Zarr stores in the same way as for other file formats.
+No! VirtualiZarr can create virtual references pointing to existing Zarr stores in the same way as for other file formats. Note: Currently only reading Zarr V3 is supported.
 
 ### Can I add a new reader for my custom file format?
 
@@ -119,7 +119,7 @@ Users of Kerchunk may find the following comparison table useful, which shows wh
 | From a netCDF3 file                                                      | `kerchunk.netCDF3.NetCDF3ToZarr`                                                                                                    | `open_virtual_dataset(..., filetype='netcdf3')`, via `kerchunk.netCDF3.NetCDF3ToZarr`                                                                                     |
 | From a COG / tiff file                                                   | `kerchunk.tiff.tiff_to_zarr`                                                                                                        | `open_virtual_dataset(..., filetype='tiff')`, via `kerchunk.tiff.tiff_to_zarr` or potentially `tifffile` (❌ Not yet implemented - see [issue #291](https://github.com/zarr-developers/VirtualiZarr/issues/291))                                                               |
 | From a Zarr v2 store                                                     | `kerchunk.zarr.ZarrToZarr`                                                                                                          | `open_virtual_dataset(..., filetype='zarr')` (❌ Not yet implemented - see [issue #262](https://github.com/zarr-developers/VirtualiZarr/issues/262))                                                                                       |
-| From a Zarr v3 store                                                     | ❌                                                                                                          | `open_virtual_dataset(..., filetype='zarr')` (❌ Not yet implemented - see [issue #262](https://github.com/zarr-developers/VirtualiZarr/issues/262))                                                                                       |
+| From a Zarr v3 store                                                     |                                                                                                          | `open_virtual_dataset(..., filetype='zarr')`                                                                                        |
 | From a GRIB2 file                                                        | `kerchunk.grib2.scan_grib`                                                                                                          | `open_virtual_datatree(..., filetype='grib')` (❌ Not yet implemented - see [issue #11](https://github.com/zarr-developers/VirtualiZarr/issues/11))                                                                                |
 | From a FITS file                                                         | `kerchunk.fits.process_file`                                                                                                        | `open_virtual_dataset(..., filetype='fits')`, via `kerchunk.fits.process_file`                                                                                      |
 | From a HDF4 file                                                         | `kerchunk.hdf4.HDF4ToZarr`                                                                                                        | `open_virtual_dataset(..., filetype='hdf4')`, via `kerchunk.hdf4.HDF4ToZarr` (❌ Not yet implemented - see [issue #216](https://github.com/zarr-developers/VirtualiZarr/issues/216))                                                        |

docs/releases.rst

@@ -9,6 +9,8 @@ v1.3.3 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Adds a Zarr reader to ``open_virtual_dataset``, which allows opening Zarr V3 stores as virtual datasets.
+  (:pull:`#271`) By `Raphael Hagen <https://github.com/norlandrhagen>`_.
 - Added experimental ManifestStore (:pull:`490`).
 - Added :py:meth:`ManifestStore.to_virtual_dataset()` method (:pull:`522`).
   By `Tom Nicholas <https://github.com/TomNicholas>`_.

virtualizarr/backend.py

@@ -28,6 +28,7 @@
     KerchunkVirtualBackend,
     NetCDF3VirtualBackend,
     TIFFVirtualBackend,
+    ZarrVirtualBackend,
 )
 from virtualizarr.readers.api import VirtualBackend
 from virtualizarr.utils import _FsspecFSFromFilepath
@@ -43,6 +44,7 @@
 # TODO add entrypoint to allow external libraries to add to this mapping
 VIRTUAL_BACKENDS = {
     "kerchunk": KerchunkVirtualBackend,
+    "zarr": ZarrVirtualBackend,
     "dmrpp": DMRPPVirtualBackend,
     "hdf5": HDFVirtualBackend,
     "netcdf4": HDFVirtualBackend,  # note this is the same as for hdf5
@@ -70,6 +72,7 @@ class FileType(AutoName):
     fits = auto()
     dmrpp = auto()
     kerchunk = auto()
+    zarr = auto()
 
 
 def automatically_determine_filetype(
@@ -87,8 +90,7 @@ def automatically_determine_filetype(
 
     # TODO how do we handle kerchunk json / parquet here?
     if Path(filepath).suffix == ".zarr":
-        # TODO we could imagine opening an existing zarr store, concatenating it, and writing a new virtual one...
-        raise NotImplementedError()
+        return FileType.zarr
 
     # Read magic bytes from local or remote file
     fpath = _FsspecFSFromFilepath(

virtualizarr/codecs.py

@@ -94,10 +94,20 @@ def get_codec_config(codec: ZarrCodec) -> dict[str, Any]:
     """
     Extract configuration from a codec, handling both zarr-python and numcodecs codecs.
     """
+
     if hasattr(codec, "codec_config"):
         return codec.codec_config
     elif hasattr(codec, "get_config"):
         return codec.get_config()
+    elif hasattr(codec, "_zstd_codec"):
+        # related issue: https://github.com/zarr-developers/VirtualiZarr/issues/514
+        # very silly workaround. codec.to_dict for zstd gives:
+        # {'name': 'zstd', 'configuration': {'level': 0, 'checksum': False}}
+        # which when passed through ArrayV2Metadata -> numcodecs.get_codec gives the error:
+        # *** numcodecs.errors.UnknownCodecError: codec not available: 'None'
+        # if codec._zstd_codec.get_config() : {'id': 'zstd', 'level': 0, 'checksum': False}
+        # is passed to numcodecs.get_codec. It works fine.
+        return codec._zstd_codec.get_config()
     elif hasattr(codec, "to_dict"):
         return codec.to_dict()
     else:

virtualizarr/manifests/manifest.py

@@ -1,4 +1,3 @@
-import json
 import re
 from collections.abc import ItemsView, Iterable, Iterator, KeysView, ValuesView
 from pathlib import PosixPath
@@ -84,7 +83,8 @@ def validate_and_normalize_path_to_uri(path: str, fs_root: str | None = None) ->
         return urlunparse(components)
 
     elif any(path.startswith(prefix) for prefix in VALID_URI_PREFIXES):
-        if not PosixPath(path).suffix:
+        # Question: This feels fragile, is there a better way to ID a Zarr
+        if not PosixPath(path).suffix and "zarr" not in path:
             raise ValueError(
                 f"entries in the manifest must be paths to files, but this path has no file suffix: {path}"
             )
@@ -96,7 +96,7 @@ def validate_and_normalize_path_to_uri(path: str, fs_root: str | None = None) ->
         # using PosixPath here ensures a clear error would be thrown on windows (whose paths and platform are not officially supported)
         _path = PosixPath(path)
 
-        if not _path.suffix:
+        if not _path.suffix and "zarr" not in path:
             raise ValueError(
                 f"entries in the manifest must be paths to files, but this path has no file suffix: {path}"
             )
@@ -436,21 +436,6 @@ def __eq__(self, other: Any) -> bool:
         lengths_equal = (self._lengths == other._lengths).all()
         return paths_equal and offsets_equal and lengths_equal
 
-    @classmethod
-    def from_zarr_json(cls, filepath: str) -> "ChunkManifest":
-        """Create a ChunkManifest from a Zarr manifest.json file."""
-
-        with open(filepath, "r") as manifest_file:
-            entries = json.load(manifest_file)
-
-        return cls(entries=entries)
-
-    def to_zarr_json(self, filepath: str) -> None:
-        """Write the manifest to a Zarr manifest.json file."""
-        entries = self.dict()
-        with open(filepath, "w") as json_file:
-            json.dump(entries, json_file, indent=4, separators=(", ", ": "))
-
     def rename_paths(
         self,
         new: str | Callable[[str], str],

virtualizarr/manifests/store.py

@@ -18,7 +18,7 @@
 
 from virtualizarr.manifests.array import ManifestArray
 from virtualizarr.manifests.group import ManifestGroup
-from virtualizarr.vendor.zarr.metadata import dict_to_buffer
+from virtualizarr.vendor.zarr.core.metadata import dict_to_buffer
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Iterable, Mapping

virtualizarr/readers/__init__.py

@@ -5,6 +5,9 @@
 from virtualizarr.readers.kerchunk import KerchunkVirtualBackend
 from virtualizarr.readers.netcdf3 import NetCDF3VirtualBackend
 from virtualizarr.readers.tiff import TIFFVirtualBackend
+from virtualizarr.readers.zarr import (
+    ZarrVirtualBackend,
+)
 
 __all__ = [
     "DMRPPVirtualBackend",
@@ -14,4 +17,5 @@
     "KerchunkVirtualBackend",
     "NetCDF3VirtualBackend",
     "TIFFVirtualBackend",
+    "ZarrVirtualBackend",
 ]

virtualizarr/readers/zarr.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path  # noqa
+from typing import (
+    Any,
+    Hashable,
+    Iterable,
+    Mapping,
+    Optional,
+)
+
+import numpy as np
+from xarray import Dataset, Index
+from zarr.api.asynchronous import open_group as open_group_async
+from zarr.core.metadata import ArrayV3Metadata
+
+from virtualizarr.manifests import (
+    ChunkManifest,
+    ManifestArray,
+    ManifestGroup,
+    ManifestStore,
+)
+from virtualizarr.manifests.manifest import validate_and_normalize_path_to_uri  # noqa
+from virtualizarr.readers.api import VirtualBackend
+from virtualizarr.vendor.zarr.core.common import _concurrent_map
+
+FillValueT = bool | str | float | int | list | None
+
+ZARR_DEFAULT_FILL_VALUE: dict[str, FillValueT] = {
+    # numpy dtypes's hierarchy lets us avoid checking for all the widths
+    # https://numpy.org/doc/stable/reference/arrays.scalars.html
+    np.dtype("bool").kind: False,
+    np.dtype("int").kind: 0,
+    np.dtype("float").kind: 0.0,
+    np.dtype("complex").kind: [0.0, 0.0],
+    np.dtype("datetime64").kind: 0,
+}
+
+
+import zarr
+
+
+async def get_chunk_mapping_prefix(zarr_array: zarr.AsyncArray, filepath: str) -> dict:
+    """Create a dictionary to pass into ChunkManifest __init__"""
+
+    # TODO: For when we want to support reading V2 we should parse the /c/ and "/" between chunks
+
+    prefix = zarr_array.name.lstrip("/") + "/c/"
+    prefix_keys = [(x,) async for x in zarr_array.store.list_prefix(prefix)]
+    _lengths = await _concurrent_map(prefix_keys, zarr_array.store.getsize)
+
+    chunk_keys = [x[0].split(prefix)[1] for x in prefix_keys]
+    _dict_keys = [key.replace("/", ".") for key in chunk_keys]
+    _paths = [filepath + "/" + prefix + key for key in chunk_keys]
+
+    _offsets = [0] * len(_lengths)
+    return {
+        key: {"path": path, "offset": offset, "length": length}
+        for key, path, offset, length in zip(
+            _dict_keys,
+            _paths,
+            _offsets,
+            _lengths,
+        )
+    }
+
+
+async def build_chunk_manifest(
+    zarr_array: zarr.AsyncArray, filepath: str
+) -> ChunkManifest:
+    """Build a ChunkManifest from a dictionary"""
+    chunk_map = await get_chunk_mapping_prefix(zarr_array=zarr_array, filepath=filepath)
+    return ChunkManifest(chunk_map)
+
+
+def get_metadata(zarr_array: zarr.AsyncArray[Any]) -> ArrayV3Metadata:
+    fill_value = zarr_array.metadata.fill_value
+    if fill_value is not None:
+        fill_value = ZARR_DEFAULT_FILL_VALUE[zarr_array.metadata.fill_value.dtype.kind]
+
+    zarr_format = zarr_array.metadata.zarr_format
+
+    if zarr_format == 2:
+        # TODO: Once we want to support V2, we will have to deconstruct the
+        # zarr_array codecs etc. and reconstruct them with create_v3_array_metadata
+        raise NotImplementedError("Reading Zarr V2 currently not supported.")
+
+    elif zarr_format == 3:
+        return zarr_array.metadata
+
+    else:
+        raise NotImplementedError("Zarr format is not recognized as v2 or v3.")
+
+
+async def _construct_manifest_array(zarr_array: zarr.AsyncArray[Any], filepath: str):
+    array_metadata = get_metadata(zarr_array=zarr_array)
+
+    chunk_manifest = await build_chunk_manifest(zarr_array, filepath=filepath)
+    return ManifestArray(metadata=array_metadata, chunkmanifest=chunk_manifest)
+
+
+async def _construct_manifest_group(
+    filepath: str,
+    *,
+    reader_options: Optional[dict] = None,
+    drop_variables: str | Iterable[str] | None = None,
+    group: str | None = None,
+):
+    reader_options = reader_options or {}
+    zarr_group = await open_group_async(
+        filepath,
+        storage_options=reader_options.get("storage_options"),
+        path=group,
+        mode="r",
+    )
+
+    zarr_array_keys = [key async for key in zarr_group.array_keys()]
+
+    _drop_vars: list[Hashable] = [] if drop_variables is None else list(drop_variables)
+
+    zarr_arrays = await asyncio.gather(
+        *[zarr_group.getitem(var) for var in zarr_array_keys if var not in _drop_vars]
+    )
+
+    manifest_arrays = await asyncio.gather(
+        *[
+            _construct_manifest_array(zarr_array=array, filepath=filepath)  # type: ignore[arg-type]
+            for array in zarr_arrays
+        ]
+    )
+
+    manifest_dict = {
+        array.basename: result for array, result in zip(zarr_arrays, manifest_arrays)
+    }
+    return ManifestGroup(manifest_dict, attributes=zarr_group.attrs)
+
+
+def _construct_manifest_store(
+    filepath: str,
+    *,
+    reader_options: Optional[dict] = None,
+    drop_variables: str | Iterable[str] | None = None,
+    group: str | None = None,
+) -> ManifestStore:
+    import asyncio
+
+    manifest_group = asyncio.run(
+        _construct_manifest_group(
+            filepath=filepath,
+            group=group,
+            drop_variables=drop_variables,
+            reader_options=reader_options,
+        )
+    )
+    return ManifestStore(manifest_group)
+
+
+class ZarrVirtualBackend(VirtualBackend):
+    @staticmethod
+    def open_virtual_dataset(
+        filepath: str,
+        group: str | None = None,
+        drop_variables: str | Iterable[str] | None = None,
+        loadable_variables: Iterable[str] | None = None,
+        decode_times: bool | None = None,
+        indexes: Mapping[str, Index] | None = None,
+        virtual_backend_kwargs: Optional[dict] = None,
+        reader_options: Optional[dict] = None,
+    ) -> Dataset:
+        filepath = validate_and_normalize_path_to_uri(
+            filepath, fs_root=Path.cwd().as_uri()
+        )
+
+        manifest_store = _construct_manifest_store(
+            filepath=filepath,
+            group=group,
+            drop_variables=drop_variables,
+            reader_options=reader_options,
+        )
+
+        ds = manifest_store.to_virtual_dataset(
+            loadable_variables=loadable_variables,
+            decode_times=decode_times,
+            indexes=indexes,
+        )
+        return ds