Zarr reader

ci/upstream.yml

@@ -26,6 +26,7 @@ dependencies:
   - pytest
   - pooch
   - fsspec
+  - dask
   - pip
   - pip:
       - icechunk>=0.1.0a7 # Installs zarr v3 as dependency

conftest.py

@@ -24,6 +24,22 @@ def pytest_runtest_setup(item):
         )
 
 
+def _xarray_subset():
+    ds = xr.tutorial.open_dataset("air_temperature", chunks={})
+    return ds.isel(time=slice(0, 10), lat=slice(0, 9), lon=slice(0, 18)).chunk(
+        {"time": 5}
+    )
+
+
+@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 empty_netcdf4_file(tmpdir):
     # Set up example xarray dataset

docs/releases.rst

@@ -9,6 +9,9 @@ v1.2.1 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Adds a Zarr reader to ``open_virtual_dataset``, which allows opening Zarr V2 and V3 stores as virtual datasets.
+  (:pull:`#271`) By `Raphael Hagen <https://github.com/norlandrhagen>`_.
+
 - Added a ``.nbytes`` accessor method which displays the bytes needed to hold the virtual references in memory.
   (:issue:`167`, :pull:`227`) By `Tom Nicholas <https://github.com/TomNicholas>`_.
 

virtualizarr/backend.py

@@ -17,15 +17,15 @@
     KerchunkVirtualBackend,
     NetCDF3VirtualBackend,
     TIFFVirtualBackend,
-    ZarrV3VirtualBackend,
+    ZarrVirtualBackend,
 )
 from virtualizarr.readers.common import VirtualBackend
 from virtualizarr.utils import _FsspecFSFromFilepath, check_for_collisions
 
 # TODO add entrypoint to allow external libraries to add to this mapping
 VIRTUAL_BACKENDS = {
     "kerchunk": KerchunkVirtualBackend,
-    "zarr_v3": ZarrV3VirtualBackend,
+    "zarr": ZarrVirtualBackend,
     "dmrpp": DMRPPVirtualBackend,
     # all the below call one of the kerchunk backends internally (https://fsspec.github.io/kerchunk/reference.html#file-format-backends)
     "hdf5": HDF5VirtualBackend,
@@ -72,8 +72,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

@@ -65,9 +65,9 @@ def _get_manifestarray_codecs(
 def _is_zarr_array(array: object) -> bool:
     """Check if the array is an instance of Zarr Array."""
     try:
-        from zarr import Array
+        from zarr import Array, AsyncArray
 
-        return isinstance(array, Array)
+        return isinstance(array, (Array, AsyncArray))
     except ImportError:
         return False
 

virtualizarr/manifests/manifest.py

@@ -48,7 +48,6 @@ def with_validation(
         """
 
         # note: we can't just use `__init__` or a dataclass' `__post_init__` because we need `fs_root` to be an optional kwarg
-
         path = validate_and_normalize_path_to_uri(path, fs_root=fs_root)
         validate_byte_range(offset=offset, length=length)
         return ChunkEntry(path=path, offset=offset, length=length)
@@ -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}"
             )
@@ -357,17 +357,6 @@ def shape_chunk_grid(self) -> tuple[int, ...]:
     def __repr__(self) -> str:
         return f"ChunkManifest<shape={self.shape_chunk_grid}>"
 
-    @property
-    def nbytes(self) -> int:
-        """
-        Size required to hold these references in memory in bytes.
-
-        Note this is not the size of the referenced chunks if they were actually loaded into memory,
-        this is only the size of the pointers to the chunk locations.
-        If you were to load the data into memory it would be ~1e6x larger for 1MB chunks.
-        """
-        return self._paths.nbytes + self._offsets.nbytes + self._lengths.nbytes
-
     def __getitem__(self, key: ChunkKey) -> ChunkEntry:
         indices = split(key)
         path = self._paths[indices]

virtualizarr/readers/__init__.py

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

virtualizarr/readers/common.py

@@ -45,16 +45,24 @@ def maybe_open_loadable_vars_and_indexes(
     # TODO Really we probably want a dedicated backend that iterates over all variables only once
     # TODO See issue #124 for a suggestion of how to avoid calling xarray here.
 
-    fpath = _FsspecFSFromFilepath(
-        filepath=filepath, reader_options=reader_options
-    ).open_file()
+    fpath = _FsspecFSFromFilepath(filepath=filepath, reader_options=reader_options)
+
+    # Updates the Xarray open_dataset kwargs if Zarr
+
+    if fpath.upath.suffix == ".zarr":
+        engine = "zarr"
+        xr_input = fpath.filepath
+
+    else:
+        engine = None
+        xr_input = fpath.open_file()  # type: ignore
 
-    # fpath can be `Any` thanks to fsspec.filesystem(...).open() returning Any.
     ds = open_dataset(
-        fpath,  # type: ignore[arg-type]
+        xr_input,  # type: ignore[arg-type]
         drop_variables=drop_variables,
         group=group,
         decode_times=decode_times,
+        engine=engine,
     )
 
     if indexes is None: