Check for aligned chunks when writing to existing variables

doc/whats-new.rst

@@ -64,6 +64,10 @@ Deprecations
 Bug fixes
 ~~~~~~~~~
 
+- Writing to unaligned zarr chunks with ``region`` will now raise an error
+  (unless ``safe_chunks=False``); previously an error would only be raised on
+  new variables. (:pull:`8459`, :issue:`8371`)
+  By `Maximilian Roos <https://github.com/max-sixty>`_.
 - Port `bug fix from pandas <https://github.com/pandas-dev/pandas/pull/55283>`_
   to eliminate the adjustment of resample bin edges in the case that the
   resampling frequency has units of days and is greater than one day

xarray/backends/zarr.py

@@ -679,6 +679,17 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
             if v.encoding == {"_FillValue": None} and fill_value is None:
                 v.encoding = {}
 
+            # We need to do this for both new and existing variables to ensure we're not
+            # writing to a partial chunk, even though we don't use the `encoding` value
+            # when writing to an existing variable. See
+            # https://github.com/pydata/xarray/issues/8371 for details.
+            encoding = extract_zarr_variable_encoding(
+                v,
+                raise_on_invalid=check,
+                name=vn,
+                safe_chunks=self._safe_chunks,
+            )
+
             if name in self.zarr_group:
                 # existing variable
                 # TODO: if mode="a", consider overriding the existing variable
@@ -709,9 +720,6 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
                     zarr_array = self.zarr_group[name]
             else:
                 # new variable
-                encoding = extract_zarr_variable_encoding(
-                    v, raise_on_invalid=check, name=vn, safe_chunks=self._safe_chunks
-                )
                 encoded_attrs = {}
                 # the magic for storing the hidden dimension data
                 encoded_attrs[DIMENSION_KEY] = dims

xarray/tests/test_backends.py

@@ -5434,3 +5434,22 @@
     ds_region.to_zarr(
         tmp_path / "test.zarr", region={"x": slice(0, 1), "y": slice(0, 1)}
     )
+
+
+@requires_zarr
+def test_zarr_region_chunk_align(tmp_path):
+    """
+    Check that writing to unaligned chunks with `region` fails, assuming `safe_chunks=False`.
+    """
+    ds = (
+        xr.DataArray(np.arange(120).reshape(4, 3, -1), dims=list("abc"))
+        .rename("var1")
+        .to_dataset()
+    )
+
+    ds.chunk(5).to_zarr(tmp_path / "foo.zarr", compute=False, mode="w")
+    with pytest.raises(ValueError):
+        for r in range(ds.sizes["a"]):
+            ds.chunk(3).isel(a=[r]).to_zarr(
+                tmp_path / "foo.zarr", region=dict(a=slice(r, r + 1))
+            )