Allow to_zarr to be a lazy operation by passing compute=False (#884)

Author: tomwhite

cubed/core/ops.py

@@ -277,7 +277,9 @@ def __iter__(self):
         return out
 
 
-def to_zarr(x: "Array", store, path=None, region=None, executor=None, **kwargs):
+def to_zarr(
+    x: "Array", store, path=None, region=None, compute=True, *, executor=None, **kwargs
+):
     """Save an array to Zarr storage.
 
     Note that this operation is eager, and will run the computation
@@ -293,12 +295,17 @@ def to_zarr(x: "Array", store, path=None, region=None, executor=None, **kwargs):
         Group path
     region : tuple of slices, optional
         The region of data that should be written to in target.
+    compute : boolean, optional
+        If True compute immediately, return array otherwise.
     executor : cubed.runtime.types.Executor, optional
         The executor to use to run the computation.
         Defaults to using the in-process Python executor.
     """
     out = _store_array(x, store, path=path, region=region)
-    out.compute(executor=executor, _return_in_memory_array=False, **kwargs)
+    if compute:
+        out.compute(executor=executor, _return_in_memory_array=False, **kwargs)
+    else:
+        return out
 
 
 def blockwise(

cubed/tests/test_core.py

@@ -224,6 +224,25 @@ def test_to_zarr_array(tmp_path, spec, executor):
     assert_array_equal(res[:], np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
 
 
+def test_to_zarr_lazy_compute(tmp_path, spec):
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    store = tmp_path / "output.zarr"
+    z = create_zarr(
+        np.zeros(a.shape, dtype=a.dtype),
+        chunks=(2, 2),
+        store=store,
+    )
+    b = cubed.to_zarr(a, z, compute=False)
+    # target has not been computed yet
+    res = open_storage_array(store, mode="r")
+    with pytest.raises(AssertionError):
+        assert_array_equal(res[:], np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+
+    b.compute()
+    res = open_storage_array(store, mode="r")
+    assert_array_equal(res[:], np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+
+
 def test_to_zarr_region(tmp_path, spec, executor):
     a = xp.asarray(
         [