From ee95fa443a7f9e82a50dd2af140bf8ed22ed41b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Taus=20M=C3=B8ller?= Date: Mon, 20 Jul 2026 12:29:58 +0200 Subject: [PATCH 1/3] Add unit test for sharded zarr roundtrip --- xarray/tests/backends/test_zarr.py | 71 ++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/xarray/tests/backends/test_zarr.py b/xarray/tests/backends/test_zarr.py index 46973d69598..33184e67013 100644 --- a/xarray/tests/backends/test_zarr.py +++ b/xarray/tests/backends/test_zarr.py @@ -3219,3 +3219,74 @@ def test_dataset_to_zarr_align_chunks_true(self, tmp_store) -> None: region="auto", ) assert_identical(ds, xr.open_zarr(store)) + + +@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 = {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 From 22cfdc6fddb9b2642e27d9e15171975b302a80f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Taus=20M=C3=B8ller?= Date: Thu, 15 Oct 2026 11:48:58 +0200 Subject: [PATCH 2/3] set preferred_chunks to shards if available --- doc/whats-new.rst | 6 ++++++ xarray/backends/zarr.py | 29 ++++++++++++++--------------- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index fc3da497a15..0f27558c060 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -252,6 +252,12 @@ Bug Fixes This was due to a bug in the reduction method generation introduced in :pull:`11556` (:pull:`11686`). By `Andrew Scherer `_. +- 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 chunking + is aligned with the storage layout instead of splitting the stored shards. + (:issue:`11460`, :pull:`11464`). + By `Taus Møller `_. Documentation diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index c10b28a5ee8..1463f8b696c 100644 --- a/xarray/backends/zarr.py +++ b/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}) From ca1effe9ff4f7d6fce7b3d9874b18b0b54ac9722 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Fri, 16 Oct 2026 13:42:46 +0200 Subject: [PATCH 3/3] Fix whats-new formatting and test dimension order Co-Authored-By: Claude Opus 5.5 --- doc/whats-new.rst | 5 ++--- xarray/tests/backends/test_zarr.py | 3 ++- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index a63587dd928..47f474bda80 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -270,9 +270,8 @@ Bug Fixes By `Andrew Scherer `_. - 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 chunking - is aligned with the storage layout instead of splitting the stored shards. - (:issue:`11460`, :pull:`11464`). + 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 `_. - Fix DatasetWeighted dropping coordinates/dimensions which are not affected by the reduction nor in the weights array (:issue:`11560`, :pull:`11562`). diff --git a/xarray/tests/backends/test_zarr.py b/xarray/tests/backends/test_zarr.py index 33184e67013..e7911a1d846 100644 --- a/xarray/tests/backends/test_zarr.py +++ b/xarray/tests/backends/test_zarr.py @@ -3221,6 +3221,7 @@ def test_dataset_to_zarr_align_chunks_true(self, tmp_store) -> None: assert_identical(ds, xr.open_zarr(store)) +@requires_zarr @requires_dask @pytest.mark.parametrize( "shape, enc_chunks, enc_shards", @@ -3254,7 +3255,7 @@ def test_zarr_shard_roundtrip( tmp_path: Path, ) -> None: - dimensions = {f"dim_{i}" for i in range(len(shape))} + 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)