Human gold patch for pydata__xarray-11117

PR: pydata#11117
Issue: pydata#10831
Base commit: 7a63b5a
Changed lines: 53

Author: SWE-bench Complex

xarray/backends/zarr.py

@@ -1240,16 +1240,22 @@ def set_variables(
                 zarr_format=3 if is_zarr_v3_format else 2,
             )
 
-            if self._align_chunks and isinstance(encoding["chunks"], tuple):
+            # When shards are specified, dask chunks must align with shard boundaries
+            # (not just zarr chunk boundaries) to avoid data corruption during
+            # parallel writes. See https://github.com/pydata/xarray/issues/10831
+            effective_write_chunks = encoding.get("shards") or encoding["chunks"]
+
+            if self._align_chunks and isinstance(effective_write_chunks, tuple):
                 v = grid_rechunk(
                     v=v,
-                    enc_chunks=encoding["chunks"],
+                    enc_chunks=effective_write_chunks,
                     region=region,
                 )
 
-            if self._safe_chunks and isinstance(encoding["chunks"], tuple):
+            if self._safe_chunks and isinstance(effective_write_chunks, tuple):
                 # the hard case
                 # DESIGN CHOICE: do not allow multiple dask chunks on a single zarr chunk
+                # (or shard, when sharding is enabled)
                 # this avoids the need to get involved in zarr synchronization / locking
                 # From zarr docs:
                 #  "If each worker in a parallel computation is writing to a
@@ -1260,7 +1266,7 @@ def set_variables(
                 shape = zarr_shape or v.shape
                 validate_grid_chunks_alignment(
                     nd_v_chunks=v.chunks,
-                    enc_chunks=encoding["chunks"],
+                    enc_chunks=effective_write_chunks,
                     region=region,
                     allow_partial_chunks=self._mode != "r+",
                     name=name,

xarray/tests/test_backends.py

@@ -2869,6 +2869,41 @@ def test_shard_encoding(self) -> None:
                 with self.roundtrip(data) as actual:
                     pass
 
+    @requires_dask
+    def test_shard_encoding_with_dask(self) -> None:
+        # Test that dask chunks must align with shard boundaries.
+        # See https://github.com/pydata/xarray/issues/10831
+        if not (has_zarr_v3 and zarr.config.config["default_zarr_format"] == 3):
+            pytest.skip("sharding requires zarr v3 format")
+
+        ds = xr.DataArray(np.arange(12), dims="x", name="var1").to_dataset()
+
+        # Case 1: Dask chunks equal to shards should work
+        # (zarr chunk=3, shard=6, dask chunk=6)
+        ds1 = ds.chunk({"x": 6})
+        ds1["var1"].encoding = {"chunks": (3,), "shards": (6,)}
+        with self.roundtrip(ds1) as actual:
+            assert_identical(ds, actual)
+
+        # Case 2: Dask chunks that are multiples of shards should work
+        # (zarr chunk=1, shard=3, dask chunk=6)
+        ds2 = ds.chunk({"x": 6})
+        ds2["var1"].encoding = {"chunks": (1,), "shards": (3,)}
+        with self.roundtrip(ds2) as actual:
+            assert_identical(ds, actual)
+
+        # Case 3: Dask chunks smaller than shards should fail
+        # (zarr chunk=2, shard=4, dask chunk=3) - dask chunk doesn't align with shard
+        ds3 = ds.chunk({"x": 3})
+        ds3["var1"].encoding = {"chunks": (2,), "shards": (4,)}
+        with pytest.raises(ValueError, match=r"would overlap"):
+            with self.roundtrip(ds3) as actual:
+                pass
+
+        # Case 4: Can bypass with safe_chunks=False (but data may be corrupted)
+        with self.roundtrip(ds3, save_kwargs={"safe_chunks": False}) as actual:
+            pass
+
     @requires_dask
     @pytest.mark.skipif(
         ON_WINDOWS,