diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 5c1126fbc4d..54f9fec28d9 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -52,6 +52,8 @@ New Features ~~~~~~~~~~~~ - Relax nanosecond datetime restriction in CF time decoding (:issue:`7493`, :pull:`9618`). By `Kai Mühlbauer `_ and `Spencer Clark `_. +- Enable the ``compute=False`` option in :py:meth:`DataTree.to_zarr`. (:pull:`9958`). + By `Sam Levang `_. - Improve the error message raised when no key is matching the available variables in a dataset. (:pull:`9943`) By `Jimmy Westling `_. - :py:meth:`DatasetGroupBy.first` and :py:meth:`DatasetGroupBy.last` can now use ``flox`` if available. (:issue:`9647`) diff --git a/xarray/core/datatree_io.py b/xarray/core/datatree_io.py index 3a57ebc7d2d..31f5a93dd5e 100644 --- a/xarray/core/datatree_io.py +++ b/xarray/core/datatree_io.py @@ -87,7 +87,7 @@ def _datatree_to_zarr( consolidated: bool = True, group: str | None = None, write_inherited_coords: bool = False, - compute: Literal[True] = True, + compute: bool = True, **kwargs, ): """This function creates an appropriate datastore for writing a datatree @@ -103,9 +103,6 @@ def _datatree_to_zarr( "specifying a root group for the tree has not been implemented" ) - if not compute: - raise NotImplementedError("compute=False has not been implemented yet") - if encoding is None: encoding = {} @@ -127,6 +124,7 @@ def _datatree_to_zarr( mode=mode, encoding=encoding.get(node.path), consolidated=False, + compute=compute, **kwargs, ) if "w" in mode: diff --git a/xarray/tests/test_backends_datatree.py b/xarray/tests/test_backends_datatree.py index 9cbe3f2247a..bca528ca042 100644 --- a/xarray/tests/test_backends_datatree.py +++ b/xarray/tests/test_backends_datatree.py @@ -440,6 +440,25 @@ def test_to_zarr_default_write_mode(self, tmpdir, simple_datatree): with pytest.raises(zarr.errors.ContainsGroupError): simple_datatree.to_zarr(tmpdir) + @requires_dask + def test_to_zarr_compute_false(self, tmpdir, simple_datatree): + import dask.array as da + + filepath = tmpdir / "test.zarr" + original_dt = simple_datatree.chunk() + original_dt.to_zarr(filepath, compute=False) + + for node in original_dt.subtree: + for name, variable in node.dataset.variables.items(): + var_dir = filepath / node.path / name + var_files = var_dir.listdir() + assert var_dir / ".zarray" in var_files + assert var_dir / ".zattrs" in var_files + if isinstance(variable.data, da.Array): + assert var_dir / "0" not in var_files + else: + assert var_dir / "0" in var_files + def test_to_zarr_inherited_coords(self, tmpdir): original_dt = DataTree.from_dict( {