Fix: Chunksize when opening sharded zarr

doc/whats-new.rst

@@ -312,6 +312,11 @@ Bug Fixes
   This was due to a bug in the reduction method generation introduced in :pull:`11556`
   (:pull:`11686`).
   By `Andrew Scherer <https://github.com/andrew-s28>`_.
+- Fix the ``preferred_chunks`` encoding of sharded Zarr variables read with
+  :py:func:`open_zarr`: it now holds the shard (outer/storage chunk) sizes
+  instead of the inner chunk sizes, so the default dask chunks align with the
+  stored shards rather than splitting them (:issue:`11460`, :pull:`11464`).
+  By `Taus Møller <https://github.com/taus>`_.
 - Fix indexing with a 0-dimensional non-NumPy array, e.g. a CuPy array, which
   raised ``ValueError`` instead of dropping the indexed dimension (:issue:`8820`).
   By `Trinh Quoc Anh <https://github.com/tqa236>`_.

xarray/backends/zarr.py

@@ -953,14 +953,6 @@ def open_store_variable(self, name):
         chunks = _compact_chunk_sizes(
             tuple(x if isinstance(x, tuple) else (x,) for x in chunks)
         )
-        preferred_chunks = dict(zip(dimensions, chunks, strict=True))
-
-        encoding = {
-            "chunks": chunks,
-            "preferred_chunks": preferred_chunks,
-        }
-        if array_wrapper.dtype.kind == "T":
-            encoding["dtype"] = array_wrapper.dtype
 
         try:
             shards = zarr_array.shards
@@ -969,13 +961,20 @@ def open_store_variable(self, name):
             # (i.e. outer/storage chunk) sizes instead, e.g. ((1, 2),).
             shards = _compact_chunk_sizes(zarr_array.write_chunk_sizes)
 
-        encoding.update(
-            {
-                "compressors": zarr_array.compressors,
-                "filters": zarr_array.filters,
-                "shards": shards,
-            }
-        )
+        outer_chunks = shards if shards is not None else chunks
+        preferred_chunks = dict(zip(dimensions, outer_chunks, strict=True))
+
+        encoding = {
+            "chunks": chunks,
+            "preferred_chunks": preferred_chunks,
+            "compressors": zarr_array.compressors,
+            "filters": zarr_array.filters,
+            "shards": shards,
+        }
+
+        if array_wrapper.dtype.kind == "T":
+            encoding["dtype"] = array_wrapper.dtype
+
         if self.zarr_group.metadata.zarr_format == 3:
             encoding.update({"serializer": zarr_array.serializer})
 

xarray/tests/backends/test_zarr.py

@@ -3235,3 +3235,75 @@ def test_dataset_to_zarr_align_chunks_true(self, tmp_store) -> None:
                 region="auto",
             )
             assert_identical(ds, xr.open_zarr(store))
+
+
+@requires_zarr
+@requires_dask
+@pytest.mark.parametrize(
+    "shape, enc_chunks, enc_shards",
+    [
+        (
+            (16, 16),
+            (2, 2),
+            (8, 4),
+        ),
+        (
+            (10, 10),
+            (2, 2),
+            (4, 4),
+        ),
+        (
+            (100, 10),
+            (2, 2),
+            (4, 4),
+        ),
+        (
+            (100, 10, 1),
+            (2, 2, 1),
+            (4, 4, 1),
+        ),
+    ],
+)
+def test_zarr_shard_roundtrip(
+    shape: tuple[int, ...],
+    enc_chunks: tuple[int, ...],
+    enc_shards: tuple[int, ...],
+    tmp_path: Path,
+) -> None:
+
+    dimensions = tuple(f"dim_{i}" for i in range(len(shape)))
+    sizes = dict(zip(dimensions, shape, strict=True))
+
+    data_values = np.arange(np.prod(shape)).reshape(shape)
+    coords = {dim: np.arange(sizes[dim]) for dim in dimensions}
+
+    chunks = dict(zip(dimensions, enc_chunks, strict=True))
+    shards = dict(zip(dimensions, enc_shards, strict=True))
+
+    ds = xr.Dataset(
+        {
+            "data": (dimensions, data_values),
+        },
+        coords=coords,
+    )
+
+    # Make the xarray chunks match the zarr shards
+    ds_sharded = ds.chunk(shards)
+    encoding: dict[str, dict[str, tuple[int, ...]]] = {
+        dim: {"chunks": (chunks[dim],), "shards": (shards[dim],)} for dim in dimensions
+    }
+    encoding["data"] = {"chunks": enc_chunks, "shards": enc_shards}
+
+    store = tmp_path / "test.zarr"
+    ds_sharded.to_zarr(
+        store,
+        mode="w",
+        encoding=encoding,
+        zarr_format=3,
+        consolidated=False,
+    )
+
+    # reopen dataset and store it back
+    reopened_ds = xr.open_zarr(store, consolidated=False)
+    reopened_ds.to_zarr(store, mode="a")
+    assert reopened_ds.chunks == ds_sharded.chunks