Enable DataTree.to_zarr(compute=False)

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 <https://github.com/kmuehlbauer>`_ and `Spencer Clark <https://github.com/spencerkclark>`_.
+- Enable the ``compute=False`` option in :py:meth:`DataTree.to_zarr`. (:pull:`9958`).
+  By `Sam Levang <https://github.com/slevang>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/core/datatree_io.py

@@ -84,7 +84,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
@@ -100,9 +100,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 = {}
 
@@ -124,6 +121,7 @@ def _datatree_to_zarr(
             mode=mode,
             encoding=encoding.get(node.path),
             consolidated=False,
+            compute=compute,
             **kwargs,
         )
         if "w" in mode:

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(
             {