Test support for regions in to_zarr

Author: tomwhite

cubed_xarray/cubedmanager.py

@@ -210,23 +210,12 @@ def store(
         if lock:
             raise NotImplementedError("Locking is not supported.")
 
-        regions = kwargs.pop("regions", None)
-        if regions:
-            # regions is either a tuple of slices or a collection of tuples of slices
-            if isinstance(regions, tuple):
-                regions = [regions]
-            for t in regions:
-                if not all(r == slice(None) for r in t):
-                    raise NotImplementedError(
-                        "Only whole slices are supported for regions."
-                    )
-
         kwargs.pop("flush", None)  # not used
 
         return store(
             sources,
             targets,
-            **kwargs,
+            **kwargs,  # regions passed through here
         )
 
 

cubed_xarray/tests/test_wrapping.py

@@ -1,6 +1,7 @@
 import sys
 
 import cubed
+import numpy as np
 import pytest
 import xarray as xr
 from cubed.runtime.create import create_executor
@@ -75,6 +76,29 @@ def test_to_zarr(tmpdir, executor):
         assert_allclose(original, computed)
 
 
+# based on test_write_region
+def test_write_region(tmpdir):
+    zeros = xr.Dataset({"u": (("x",), np.zeros(10))}).chunk(
+        2, chunked_array_type="cubed"
+    )
+    nonzeros = xr.Dataset({"u": (("x",), np.arange(1, 11))}).chunk(
+        2, chunked_array_type="cubed"
+    )
+
+    store = tmpdir / "out.zarr"
+    zeros.to_zarr(
+        store,
+        encoding={"u": dict(chunks=2)},
+    )
+    with xr.open_zarr(store) as actual:
+        assert_identical(actual, zeros)
+    for i in range(0, 10, 2):
+        region = {"x": slice(i, i + 2)}
+        nonzeros.isel(region).to_zarr(store, region=region)
+    with xr.open_zarr(store) as actual:
+        assert_identical(actual, nonzeros)
+
+
 def test_dataset_accessor_visualize(tmp_path):
     spec = cubed.Spec(allowed_mem="200MB")