Allow chunk_store argument when opening Zarr datasets (#3804)

* Allow chunk store for Zarr datasets

* Add test for open_zarr() chunk_store argument

* Add "chunk_store" argument to to_zarr()

* Simplify chunk_store argument handling

* blacken

* Add minimum zarr version requirement in docstring

* Update xarray/tests/test_backends.py

Co-authored-by: Ryan Abernathey <[email protected]>
Co-authored-by: dcherian <[email protected]>
Co-authored-by: Deepak Cherian <[email protected]>

xarray/backends/api.py

@@ -1307,6 +1307,7 @@ def _validate_append_dim_and_encoding(
 def to_zarr(
     dataset,
     store=None,
+    chunk_store=None,
     mode=None,
     synchronizer=None,
     group=None,
@@ -1322,6 +1323,8 @@ def to_zarr(
     """
     if isinstance(store, Path):
         store = str(store)
+    if isinstance(chunk_store, Path):
+        chunk_store = str(store)
     if encoding is None:
         encoding = {}
 
@@ -1346,6 +1349,7 @@ def to_zarr(
         synchronizer=synchronizer,
         group=group,
         consolidate_on_close=consolidated,
+        chunk_store=chunk_store,
     )
     zstore.append_dim = append_dim
     writer = ArrayWriter()

xarray/backends/zarr.py

@@ -278,10 +278,14 @@ def open_group(
         group=None,
         consolidated=False,
         consolidate_on_close=False,
+        chunk_store=None,
     ):
         import zarr
 
         open_kwargs = dict(mode=mode, synchronizer=synchronizer, path=group)
+        if chunk_store:
+            open_kwargs["chunk_store"] = chunk_store
+
         if consolidated:
             # TODO: an option to pass the metadata_key keyword
             zarr_group = zarr.open_consolidated(store, **open_kwargs)
@@ -505,6 +509,7 @@ def open_zarr(
     drop_variables=None,
     consolidated=False,
     overwrite_encoded_chunks=False,
+    chunk_store=None,
     decode_timedelta=None,
     **kwargs,
 ):
@@ -565,6 +570,8 @@ def open_zarr(
     consolidated : bool, optional
         Whether to open the store using zarr's consolidated metadata
         capability. Only works for stores that have already been consolidated.
+    chunk_store : MutableMapping, optional
+        A separate Zarr store only for chunk data.
     decode_timedelta : bool, optional
         If True, decode variables and coordinates with time units in
         {'days', 'hours', 'minutes', 'seconds', 'milliseconds', 'microseconds'}
@@ -646,6 +653,7 @@ def maybe_decode_store(store, lock=False):
         synchronizer=synchronizer,
         group=group,
         consolidated=consolidated,
+        chunk_store=chunk_store,
     )
     ds = maybe_decode_store(zarr_store)
 

xarray/core/dataset.py

@@ -1571,6 +1571,7 @@ def to_netcdf(
     def to_zarr(
         self,
         store: Union[MutableMapping, str, Path] = None,
+        chunk_store: Union[MutableMapping, str, Path] = None,
         mode: str = None,
         synchronizer=None,
         group: str = None,
@@ -1589,6 +1590,9 @@ def to_zarr(
         ----------
         store : MutableMapping, str or Path, optional
             Store or path to directory in file system.
+        chunk_store : MutableMapping, str or Path, optional
+            Store or path to directory in file system only for Zarr array chunks.
+            Requires zarr-python v2.4.0 or later.
         mode : {"w", "w-", "a", None}, optional
             Persistence mode: "w" means create (overwrite if exists);
             "w-" means create (fail if exists);
@@ -1649,6 +1653,7 @@ def to_zarr(
         return to_zarr(
             self,
             store=store,
+            chunk_store=chunk_store,
             mode=mode,
             synchronizer=synchronizer,
             group=group,

xarray/tests/test_backends.py

@@ -1564,6 +1564,15 @@ def test_roundtrip_consolidated(self):
             self.check_dtypes_roundtripped(expected, actual)
             assert_identical(expected, actual)
 
+    def test_with_chunkstore(self):
+        expected = create_test_data()
+        with self.create_zarr_target() as store_target, self.create_zarr_target() as chunk_store:
+            save_kwargs = {"chunk_store": chunk_store}
+            self.save(expected, store_target, **save_kwargs)
+            open_kwargs = {"chunk_store": chunk_store}
+            with self.open(store_target, **open_kwargs) as ds:
+                assert_equal(ds, expected)
+
     @requires_dask
     def test_auto_chunk(self):
         original = create_test_data().chunk()