Skip to content

Commit 2bedca4

Browse files
authored
remove gate and add a test (#9958)
1 parent 4e728b8 commit 2bedca4

File tree

3 files changed

+23
-4
lines changed

3 files changed

+23
-4
lines changed

doc/whats-new.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ New Features
5252
~~~~~~~~~~~~
5353
- Relax nanosecond datetime restriction in CF time decoding (:issue:`7493`, :pull:`9618`).
5454
By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_ and `Spencer Clark <https://github.com/spencerkclark>`_.
55+
- Enable the ``compute=False`` option in :py:meth:`DataTree.to_zarr`. (:pull:`9958`).
56+
By `Sam Levang <https://github.com/slevang>`_.
5557
- Improve the error message raised when no key is matching the available variables in a dataset. (:pull:`9943`)
5658
By `Jimmy Westling <https://github.com/illviljan>`_.
5759
- :py:meth:`DatasetGroupBy.first` and :py:meth:`DatasetGroupBy.last` can now use ``flox`` if available. (:issue:`9647`)

xarray/core/datatree_io.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ def _datatree_to_zarr(
8787
consolidated: bool = True,
8888
group: str | None = None,
8989
write_inherited_coords: bool = False,
90-
compute: Literal[True] = True,
90+
compute: bool = True,
9191
**kwargs,
9292
):
9393
"""This function creates an appropriate datastore for writing a datatree
@@ -103,9 +103,6 @@ def _datatree_to_zarr(
103103
"specifying a root group for the tree has not been implemented"
104104
)
105105

106-
if not compute:
107-
raise NotImplementedError("compute=False has not been implemented yet")
108-
109106
if encoding is None:
110107
encoding = {}
111108

@@ -127,6 +124,7 @@ def _datatree_to_zarr(
127124
mode=mode,
128125
encoding=encoding.get(node.path),
129126
consolidated=False,
127+
compute=compute,
130128
**kwargs,
131129
)
132130
if "w" in mode:

xarray/tests/test_backends_datatree.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,25 @@ def test_to_zarr_default_write_mode(self, tmpdir, simple_datatree):
440440
with pytest.raises(zarr.errors.ContainsGroupError):
441441
simple_datatree.to_zarr(tmpdir)
442442

443+
@requires_dask
444+
def test_to_zarr_compute_false(self, tmpdir, simple_datatree):
445+
import dask.array as da
446+
447+
filepath = tmpdir / "test.zarr"
448+
original_dt = simple_datatree.chunk()
449+
original_dt.to_zarr(filepath, compute=False)
450+
451+
for node in original_dt.subtree:
452+
for name, variable in node.dataset.variables.items():
453+
var_dir = filepath / node.path / name
454+
var_files = var_dir.listdir()
455+
assert var_dir / ".zarray" in var_files
456+
assert var_dir / ".zattrs" in var_files
457+
if isinstance(variable.data, da.Array):
458+
assert var_dir / "0" not in var_files
459+
else:
460+
assert var_dir / "0" in var_files
461+
443462
def test_to_zarr_inherited_coords(self, tmpdir):
444463
original_dt = DataTree.from_dict(
445464
{

0 commit comments

Comments
 (0)