Fix silent data corruption when writing dask arrays to sharded zarr stores (#11117)

When writing dask-backed arrays to zarr with sharding enabled, the chunk
alignment validation was checking against zarr's internal chunk size instead
of the shard size. This allowed configurations where dask chunk boundaries
didn't align with shard boundaries, causing data corruption during parallel
writes.

This fix uses shard boundaries (when shards are specified) for both
grid_rechunk() and validate_grid_chunks_alignment(), ensuring that parallel
writes don't cross shard boundaries.

Fixes #10831

Co-authored-by: Claude <noreply@anthropic.com>

Author: kkollsga

doc/whats-new.rst

@@ -26,6 +26,10 @@ Deprecations
 Bug Fixes
 ~~~~~~~~~
 
+- Fix silent data corruption when writing dask arrays to sharded Zarr stores.
+  Dask chunk boundaries must now align with shard boundaries, not just internal
+  Zarr chunk boundaries (:issue:`10831`).
+
 
 Documentation
 ~~~~~~~~~~~~~

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,