From 2b2413aa319ddac05fd683ba9161e74365a8bb32 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 16 Sep 2026 14:28:04 -0400 Subject: [PATCH 01/19] Support reading Zarr V3 rectilinear (variable-sized) chunks Extracts the read-only half of #11279 (by Max Jones) into its own change. The read path just needs to catch the NotImplementedError that zarr_array.chunks raises for a rectilinear grid and fall back to write_chunk_sizes for the per-dimension chunk tuples. Write support (#11279) still has open correctness issues flagged in review -- align_chunks/safe_chunks validation reuses uniform-grid logic that doesn't handle partial chunks or region-start alignment -- so it's left for a separate follow-up PR rather than gating this on resolving those. Co-Authored-By: Claude Sonnet 5 --- doc/user-guide/io.md | 25 +++++++++++++++ doc/whats-new.rst | 7 +++++ xarray/backends/zarr.py | 16 ++++++++-- xarray/tests/__init__.py | 3 ++ xarray/tests/test_backends.py | 57 +++++++++++++++++++++++++++++++++++ 5 files changed, 106 insertions(+), 2 deletions(-) diff --git a/doc/user-guide/io.md b/doc/user-guide/io.md index 839a76935f1..239d1c90e9d 100644 --- a/doc/user-guide/io.md +++ b/doc/user-guide/io.md @@ -1231,6 +1231,31 @@ ds.to_zarr( The number of chunks on Tair matches our dask chunks, while there is now only a single chunk in the directory stores of each coordinate. +(io.zarr.rectilinear-chunks)= + +### Variable-sized (rectilinear) chunks + +Zarr v3 supports _rectilinear_ chunk grids, where chunk sizes vary along one +or more dimensions. This is useful when natural data boundaries (yearly +chunks of a daily time series, per-tile spatial extents) don't align to a +regular grid. Reading such arrays requires `zarr-python >= 3.2`, with the +experimental feature enabled: + +```python +import zarr + +with zarr.config.set({"array.rectilinear_chunks": True}): + roundtrip = xr.open_zarr("rectilinear.zarr", zarr_format=3) + roundtrip.chunks["x"] # e.g. (10, 20, 30) +``` + +```{note} +xarray can currently only *read* rectilinear-chunked Zarr V3 arrays, not +write them. Rectilinear arrays must be created with `zarr-python` directly +(or another tool, such as Icechunk) before being opened with `xr.open_zarr`. +Write support is tracked in [GH11279](https://github.com/pydata/xarray/pull/11279). +``` + ### Groups Nested groups in zarr stores can be represented by loading the store as a diff --git a/doc/whats-new.rst b/doc/whats-new.rst index da7259da069..f3009149fa4 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -13,6 +13,13 @@ v2026.07.1 (unreleased) New Features ~~~~~~~~~~~~ +- Support reading Zarr V3 arrays with rectilinear (variable-sized) chunk + grids, requiring ``zarr-python >= 3.2`` with + ``zarr.config.set({"array.rectilinear_chunks": True})``. Writing + rectilinear chunks from xarray is not yet supported (:pull:`11279`). + By `Max Jones `_ and + `Tom Nicholas `_. + - Added `PyArrowCapsule interface `_ to :py:class:`DataArray` (``__arrow_c_schema__`` and ``__arrow_c_stream__``), enabling near zero-copy export to pyarrow, polars or duckdb. diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index f0d0d081a5d..58245a1a0d2 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -860,9 +860,21 @@ def open_store_variable(self, name): ) attributes = dict(attributes) + try: + # Regular chunk grid: a single uniform chunk size per dimension, + # e.g. (10,) for a 60-element axis chunked into 10s. + chunks = tuple(zarr_array.chunks) + except NotImplementedError: + # Rectilinear chunk grid (zarr-python >= 3.2): chunk sizes vary + # along an axis, so there is no single chunk size. `.chunks` + # raises and we instead read the explicit per-chunk listing, + # e.g. ((10, 20, 30),). + chunks = zarr_array.write_chunk_sizes + preferred_chunks = dict(zip(dimensions, chunks, strict=True)) + encoding = { - "chunks": zarr_array.chunks, - "preferred_chunks": dict(zip(dimensions, zarr_array.chunks, strict=True)), + "chunks": chunks, + "preferred_chunks": preferred_chunks, } encoding.update( diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 15326b1d7e7..049039d7342 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -159,6 +159,9 @@ def get_dask_chunkmanager(): requires_zarr_v3 = requires_zarr has_zarr_v3_dtypes, requires_zarr_v3_dtypes = _importorskip("zarr", "3.1.0") has_zarr_v3_async_oindex, requires_zarr_v3_async_oindex = _importorskip("zarr", "3.1.2") +has_zarr_rectilinear_chunks, requires_zarr_rectilinear_chunks = _importorskip( + "zarr", "3.2.0" +) if has_zarr: import zarr diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index a0b93abf671..0d3c4b73989 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -101,6 +101,7 @@ requires_scipy, requires_scipy_or_netCDF4, requires_zarr, + requires_zarr_rectilinear_chunks, requires_zarr_v3, ) from xarray.tests.test_coding_times import ( @@ -7181,6 +7182,62 @@ def test_extract_zarr_variable_encoding() -> None: ) +@requires_zarr_rectilinear_chunks +def test_rectilinear_chunks_read(tmp_path: Path) -> None: + """Read a rectilinear (variable-sized chunk) array created directly by zarr. + + xarray does not yet support *writing* rectilinear chunks, so the store + here is created with zarr-python directly. + """ + import zarr + + store_path = tmp_path / "zarr_native.zarr" + data = np.arange(60, dtype="float32") + + with zarr.config.set({"array.rectilinear_chunks": True}): + root = zarr.open_group(store_path, mode="w", zarr_format=3) + arr = root.create( + "var", + shape=(60,), + # older zarr stubs (<3.2) don't include rectilinear chunk types + chunks=((10, 20, 30),), # type: ignore[arg-type, unused-ignore] + dtype="float32", + dimension_names=("x",), + ) + arr[:] = data + + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + assert roundtrip.chunks["x"] == (10, 20, 30) + assert roundtrip["var"].encoding["chunks"] == ((10, 20, 30),) + assert roundtrip["var"].encoding["preferred_chunks"] == {"x": (10, 20, 30)} + np.testing.assert_array_equal(roundtrip["var"].values, data) + + +@requires_zarr_rectilinear_chunks +def test_rectilinear_chunks_read_multidim(tmp_path: Path) -> None: + """Read a multi-dimensional array with a mix of regular and rectilinear dims.""" + import zarr + + store_path = tmp_path / "zarr_native_2d.zarr" + data = np.arange(120, dtype="float64").reshape(6, 20) + + with zarr.config.set({"array.rectilinear_chunks": True}): + root = zarr.open_group(store_path, mode="w", zarr_format=3) + arr = root.create( + "var", + shape=(6, 20), + chunks=(2, (5, 10, 5)), # type: ignore[arg-type, unused-ignore] + dtype="float64", + dimension_names=("x", "y"), + ) + arr[:] = data + + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + assert roundtrip.chunks["x"] == (2, 2, 2) + assert roundtrip.chunks["y"] == (5, 10, 5) + np.testing.assert_array_equal(roundtrip["var"].values, data) + + @requires_zarr @requires_fsspec @pytest.mark.filterwarnings("ignore:deallocating CachingFileManager") From c06afede5f30c0c9c58a9ee587491284cb0a34e3 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 16 Sep 2026 14:38:54 -0400 Subject: [PATCH 02/19] Use read_chunk_sizes instead of write_chunk_sizes for rectilinear reads Addresses review feedback on #11279 (same line of code, carried over in the split): write_chunk_sizes gives outer/shard sizes, but .chunks (used in the regular-grid branch just above) gives the inner chunk shape when sharding is enabled. read_chunk_sizes keeps the two branches consistent for sharded rectilinear arrays. https://github.com/pydata/xarray/pull/11279#discussion_r3349959314 https://github.com/pydata/xarray/pull/11279#discussion_r3349941077 Co-Authored-By: Claude Sonnet 5 --- xarray/backends/zarr.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 58245a1a0d2..509176051e2 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -868,8 +868,10 @@ def open_store_variable(self, name): # Rectilinear chunk grid (zarr-python >= 3.2): chunk sizes vary # along an axis, so there is no single chunk size. `.chunks` # raises and we instead read the explicit per-chunk listing, - # e.g. ((10, 20, 30),). - chunks = zarr_array.write_chunk_sizes + # e.g. ((10, 20, 30),). Use read_chunk_sizes (not + # write_chunk_sizes) so this matches `.chunks` above in + # returning the inner chunk shape when sharding is used. + chunks = zarr_array.read_chunk_sizes preferred_chunks = dict(zip(dimensions, chunks, strict=True)) encoding = { From 3fb191b45bfb266805f3e317a3c9bd14a60a3896 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 16 Sep 2026 14:48:49 -0400 Subject: [PATCH 03/19] Add tests for mixed regular/rectilinear dims and dask-backed reads Mirrors coverage from #11279 that was missing here: a dimension mix of regular and variable-length chunks, and reading into dask arrays (checking the dask array's own .chunks, not just Dataset.chunks). Co-Authored-By: Claude Sonnet 5 --- xarray/tests/test_backends.py | 115 ++++++++++++++++++++++++++++------ 1 file changed, 95 insertions(+), 20 deletions(-) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 0d3c4b73989..b97f89603e3 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7182,59 +7182,134 @@ def test_extract_zarr_variable_encoding() -> None: ) -@requires_zarr_rectilinear_chunks -def test_rectilinear_chunks_read(tmp_path: Path) -> None: - """Read a rectilinear (variable-sized chunk) array created directly by zarr. +def _create_rectilinear_zarr_array(store_path, shape, chunks, dimension_names, dtype): + """Create an array with the given (possibly rectilinear) chunk grid. - xarray does not yet support *writing* rectilinear chunks, so the store - here is created with zarr-python directly. + xarray does not yet support *writing* rectilinear chunks, so tests create + the store with zarr-python directly and only exercise xarray's read path. """ import zarr + root = zarr.open_group(store_path, mode="w", zarr_format=3) + arr = root.create( + "var", + shape=shape, + # older zarr stubs (<3.2) don't include rectilinear chunk types + chunks=chunks, # type: ignore[arg-type, unused-ignore] + dtype=dtype, + dimension_names=dimension_names, + ) + return arr + + +@requires_zarr_rectilinear_chunks +def test_rectilinear_chunks_read(tmp_path: Path) -> None: + """Read a rectilinear (variable-sized chunk) array, without dask.""" + import zarr + store_path = tmp_path / "zarr_native.zarr" data = np.arange(60, dtype="float32") with zarr.config.set({"array.rectilinear_chunks": True}): - root = zarr.open_group(store_path, mode="w", zarr_format=3) - arr = root.create( - "var", + arr = _create_rectilinear_zarr_array( + store_path, shape=(60,), - # older zarr stubs (<3.2) don't include rectilinear chunk types - chunks=((10, 20, 30),), # type: ignore[arg-type, unused-ignore] + chunks=((10, 20, 30),), dtype="float32", dimension_names=("x",), ) arr[:] = data - roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) - assert roundtrip.chunks["x"] == (10, 20, 30) + roundtrip = xr.open_zarr( + store_path, zarr_format=3, consolidated=False, chunks=None + ) assert roundtrip["var"].encoding["chunks"] == ((10, 20, 30),) assert roundtrip["var"].encoding["preferred_chunks"] == {"x": (10, 20, 30)} + assert isinstance(roundtrip["var"].data, np.ndarray) np.testing.assert_array_equal(roundtrip["var"].values, data) @requires_zarr_rectilinear_chunks -def test_rectilinear_chunks_read_multidim(tmp_path: Path) -> None: - """Read a multi-dimensional array with a mix of regular and rectilinear dims.""" +def test_rectilinear_chunks_read_mixed_dims(tmp_path: Path) -> None: + """Read a multi-dimensional array with a mix of regular and rectilinear dims, + without dask.""" import zarr store_path = tmp_path / "zarr_native_2d.zarr" data = np.arange(120, dtype="float64").reshape(6, 20) with zarr.config.set({"array.rectilinear_chunks": True}): - root = zarr.open_group(store_path, mode="w", zarr_format=3) - arr = root.create( - "var", + arr = _create_rectilinear_zarr_array( + store_path, + shape=(6, 20), + chunks=(2, (5, 10, 5)), + dtype="float64", + dimension_names=("x", "y"), + ) + arr[:] = data + + roundtrip = xr.open_zarr( + store_path, zarr_format=3, consolidated=False, chunks=None + ) + assert roundtrip["var"].encoding["chunks"] == (2, (5, 10, 5)) + assert roundtrip["var"].encoding["preferred_chunks"] == { + "x": 2, + "y": (5, 10, 5), + } + assert isinstance(roundtrip["var"].data, np.ndarray) + np.testing.assert_array_equal(roundtrip["var"].values, data) + + +@requires_zarr_rectilinear_chunks +@requires_dask +def test_rectilinear_chunks_read_dask(tmp_path: Path) -> None: + """Reading a rectilinear array into dask arrays picks up the exact, + variable-sized chunks as the dask chunks (not just xarray's `chunks` + property, but the underlying dask array's own chunking).""" + import zarr + + store_path = tmp_path / "zarr_native_dask.zarr" + data = np.arange(60, dtype="float32") + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = _create_rectilinear_zarr_array( + store_path, + shape=(60,), + chunks=((10, 20, 30),), + dtype="float32", + dimension_names=("x",), + ) + arr[:] = data + + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + assert isinstance(roundtrip["var"].data, dask_array_type) + assert roundtrip["var"].data.chunks == ((10, 20, 30),) + np.testing.assert_array_equal(roundtrip["var"].values, data) + + +@requires_zarr_rectilinear_chunks +@requires_dask +def test_rectilinear_chunks_read_dask_mixed_dims(tmp_path: Path) -> None: + """Same as above, but for a 2D array with one regular and one rectilinear + dimension.""" + import zarr + + store_path = tmp_path / "zarr_native_2d_dask.zarr" + data = np.arange(120, dtype="float64").reshape(6, 20) + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = _create_rectilinear_zarr_array( + store_path, shape=(6, 20), - chunks=(2, (5, 10, 5)), # type: ignore[arg-type, unused-ignore] + chunks=(2, (5, 10, 5)), dtype="float64", dimension_names=("x", "y"), ) arr[:] = data roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) - assert roundtrip.chunks["x"] == (2, 2, 2) - assert roundtrip.chunks["y"] == (5, 10, 5) + assert isinstance(roundtrip["var"].data, dask_array_type) + assert roundtrip["var"].data.chunks == ((2, 2, 2), (5, 10, 5)) np.testing.assert_array_equal(roundtrip["var"].values, data) From 76a4042f01c341eab23ffcdab6f41b3e45f4a752 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 16 Sep 2026 14:53:31 -0400 Subject: [PATCH 04/19] Organize rectilinear chunk read tests into a parametrized class Groups the four tests into TestZarrRectilinearChunksRead, sharing the zarr-array-creation helper and parametrizing the two shape/chunk cases (pure rectilinear, mixed regular+rectilinear) across a plain read test and a dask-backed read test. Co-Authored-By: Claude Sonnet 5 --- xarray/tests/test_backends.py | 212 +++++++++++++++------------------- 1 file changed, 96 insertions(+), 116 deletions(-) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index b97f89603e3..c57bd3d0de9 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7182,135 +7182,115 @@ def test_extract_zarr_variable_encoding() -> None: ) -def _create_rectilinear_zarr_array(store_path, shape, chunks, dimension_names, dtype): - """Create an array with the given (possibly rectilinear) chunk grid. - - xarray does not yet support *writing* rectilinear chunks, so tests create - the store with zarr-python directly and only exercise xarray's read path. - """ - import zarr - - root = zarr.open_group(store_path, mode="w", zarr_format=3) - arr = root.create( - "var", - shape=shape, - # older zarr stubs (<3.2) don't include rectilinear chunk types - chunks=chunks, # type: ignore[arg-type, unused-ignore] - dtype=dtype, - dimension_names=dimension_names, - ) - return arr - - -@requires_zarr_rectilinear_chunks -def test_rectilinear_chunks_read(tmp_path: Path) -> None: - """Read a rectilinear (variable-sized chunk) array, without dask.""" - import zarr - - store_path = tmp_path / "zarr_native.zarr" - data = np.arange(60, dtype="float32") - - with zarr.config.set({"array.rectilinear_chunks": True}): - arr = _create_rectilinear_zarr_array( - store_path, - shape=(60,), - chunks=((10, 20, 30),), - dtype="float32", - dimension_names=("x",), - ) - arr[:] = data - - roundtrip = xr.open_zarr( - store_path, zarr_format=3, consolidated=False, chunks=None - ) - assert roundtrip["var"].encoding["chunks"] == ((10, 20, 30),) - assert roundtrip["var"].encoding["preferred_chunks"] == {"x": (10, 20, 30)} - assert isinstance(roundtrip["var"].data, np.ndarray) - np.testing.assert_array_equal(roundtrip["var"].values, data) - - @requires_zarr_rectilinear_chunks -def test_rectilinear_chunks_read_mixed_dims(tmp_path: Path) -> None: - """Read a multi-dimensional array with a mix of regular and rectilinear dims, - without dask.""" - import zarr - - store_path = tmp_path / "zarr_native_2d.zarr" - data = np.arange(120, dtype="float64").reshape(6, 20) +class TestZarrRectilinearChunksRead: + """Tests for reading (not writing) rectilinear (variable-sized) zarr chunks. - with zarr.config.set({"array.rectilinear_chunks": True}): - arr = _create_rectilinear_zarr_array( - store_path, - shape=(6, 20), - chunks=(2, (5, 10, 5)), - dtype="float64", - dimension_names=("x", "y"), - ) - arr[:] = data + xarray does not yet support *writing* rectilinear chunks, so these stores + are created directly with zarr-python rather than through `Dataset.to_zarr`. + """ - roundtrip = xr.open_zarr( - store_path, zarr_format=3, consolidated=False, chunks=None + @staticmethod + def create_zarr_array(store_path, shape, chunks, dimension_names, dtype): + import zarr + + root = zarr.open_group(store_path, mode="w", zarr_format=3) + return root.create( + "var", + shape=shape, + # older zarr stubs (<3.2) don't include rectilinear chunk types + chunks=chunks, # type: ignore[arg-type, unused-ignore] + dtype=dtype, + dimension_names=dimension_names, ) - assert roundtrip["var"].encoding["chunks"] == (2, (5, 10, 5)) - assert roundtrip["var"].encoding["preferred_chunks"] == { - "x": 2, - "y": (5, 10, 5), - } - assert isinstance(roundtrip["var"].data, np.ndarray) - np.testing.assert_array_equal(roundtrip["var"].values, data) - -@requires_zarr_rectilinear_chunks -@requires_dask -def test_rectilinear_chunks_read_dask(tmp_path: Path) -> None: - """Reading a rectilinear array into dask arrays picks up the exact, - variable-sized chunks as the dask chunks (not just xarray's `chunks` - property, but the underlying dask array's own chunking).""" - import zarr + cases = pytest.mark.parametrize( + "shape,chunks,dimension_names,dtype,expected_preferred_chunks,expected_dask_chunks", + [ + pytest.param( + (60,), + ((10, 20, 30),), + ("x",), + "float32", + {"x": (10, 20, 30)}, + ((10, 20, 30),), + id="1d-rectilinear", + ), + pytest.param( + (6, 20), + (2, (5, 10, 5)), + ("x", "y"), + "float64", + {"x": 2, "y": (5, 10, 5)}, + ((2, 2, 2), (5, 10, 5)), + id="mixed-regular-and-rectilinear", + ), + ], + ) - store_path = tmp_path / "zarr_native_dask.zarr" - data = np.arange(60, dtype="float32") + @cases + def test_read( + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_preferred_chunks, + expected_dask_chunks, + ) -> None: + import zarr - with zarr.config.set({"array.rectilinear_chunks": True}): - arr = _create_rectilinear_zarr_array( - store_path, - shape=(60,), - chunks=((10, 20, 30),), - dtype="float32", - dimension_names=("x",), - ) - arr[:] = data + data = np.arange(np.prod(shape), dtype=dtype).reshape(shape) + store_path = tmp_path / "source.zarr" - roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) - assert isinstance(roundtrip["var"].data, dask_array_type) - assert roundtrip["var"].data.chunks == ((10, 20, 30),) - np.testing.assert_array_equal(roundtrip["var"].values, data) + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, shape, chunks, dimension_names, dtype + ) + arr[:] = data + roundtrip = xr.open_zarr( + store_path, zarr_format=3, consolidated=False, chunks=None + ) + assert roundtrip["var"].encoding["chunks"] == chunks + assert ( + roundtrip["var"].encoding["preferred_chunks"] + == expected_preferred_chunks + ) + assert isinstance(roundtrip["var"].data, np.ndarray) + np.testing.assert_array_equal(roundtrip["var"].values, data) -@requires_zarr_rectilinear_chunks -@requires_dask -def test_rectilinear_chunks_read_dask_mixed_dims(tmp_path: Path) -> None: - """Same as above, but for a 2D array with one regular and one rectilinear - dimension.""" - import zarr + @cases + @requires_dask + def test_read_dask( + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_preferred_chunks, + expected_dask_chunks, + ) -> None: + """Reading into dask arrays picks up the exact, variable-sized chunks + as the dask chunks (not just `Dataset.chunks`, but the underlying + dask array's own chunking).""" + import zarr - store_path = tmp_path / "zarr_native_2d_dask.zarr" - data = np.arange(120, dtype="float64").reshape(6, 20) + data = np.arange(np.prod(shape), dtype=dtype).reshape(shape) + store_path = tmp_path / "source.zarr" - with zarr.config.set({"array.rectilinear_chunks": True}): - arr = _create_rectilinear_zarr_array( - store_path, - shape=(6, 20), - chunks=(2, (5, 10, 5)), - dtype="float64", - dimension_names=("x", "y"), - ) - arr[:] = data + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, shape, chunks, dimension_names, dtype + ) + arr[:] = data - roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) - assert isinstance(roundtrip["var"].data, dask_array_type) - assert roundtrip["var"].data.chunks == ((2, 2, 2), (5, 10, 5)) - np.testing.assert_array_equal(roundtrip["var"].values, data) + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + assert isinstance(roundtrip["var"].data, dask_array_type) + assert roundtrip["var"].data.chunks == expected_dask_chunks + np.testing.assert_array_equal(roundtrip["var"].values, data) @requires_zarr From 54d6fc5fca8ec3eb51613ae6876ab0499684d0b8 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 16 Sep 2026 15:04:52 -0400 Subject: [PATCH 05/19] Fix rectilinear test expectations for zarr's actual regular-grid check CI failed on test-py314-no-dask: my local zarr (3.2.1) let .chunks return a compact per-dim value (2, (5, 10, 5)) for a mixed regular/rectilinear array without raising, but zarr 3.3.0 (what CI installs) correctly raises NotImplementedError there per the docstring, so xarray falls back to read_chunk_sizes and reports the fully-expanded form (2, 2, 2), (5, 10, 5)) instead. Verified against zarr 3.3.0 locally and updated the expected values to match the documented/correct behavior rather than the older version's quirk. Also collapsed the parametrize down to a single expected_chunks field, since it's the same fully-expanded tuple in both the encoding check and the dask-chunks check. Co-Authored-By: Claude Sonnet 5 --- xarray/tests/test_backends.py | 36 ++++++++++++----------------------- 1 file changed, 12 insertions(+), 24 deletions(-) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index c57bd3d0de9..b5271be4500 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7204,15 +7204,19 @@ def create_zarr_array(store_path, shape, chunks, dimension_names, dtype): dimension_names=dimension_names, ) + # `expected_chunks` is the fully-expanded, dask-style tuple-of-tuples that + # zarr-python reports (via `read_chunk_sizes`) for any array whose chunk + # grid isn't purely regular, even along dimensions that happen to be + # uniformly chunked -- e.g. dimension "x" below is uniformly chunked but + # still reported as (2, 2, 2), not the compact form 2. cases = pytest.mark.parametrize( - "shape,chunks,dimension_names,dtype,expected_preferred_chunks,expected_dask_chunks", + "shape,chunks,dimension_names,dtype,expected_chunks", [ pytest.param( (60,), ((10, 20, 30),), ("x",), "float32", - {"x": (10, 20, 30)}, ((10, 20, 30),), id="1d-rectilinear", ), @@ -7221,7 +7225,6 @@ def create_zarr_array(store_path, shape, chunks, dimension_names, dtype): (2, (5, 10, 5)), ("x", "y"), "float64", - {"x": 2, "y": (5, 10, 5)}, ((2, 2, 2), (5, 10, 5)), id="mixed-regular-and-rectilinear", ), @@ -7230,14 +7233,7 @@ def create_zarr_array(store_path, shape, chunks, dimension_names, dtype): @cases def test_read( - self, - tmp_path, - shape, - chunks, - dimension_names, - dtype, - expected_preferred_chunks, - expected_dask_chunks, + self, tmp_path, shape, chunks, dimension_names, dtype, expected_chunks ) -> None: import zarr @@ -7253,10 +7249,9 @@ def test_read( roundtrip = xr.open_zarr( store_path, zarr_format=3, consolidated=False, chunks=None ) - assert roundtrip["var"].encoding["chunks"] == chunks - assert ( - roundtrip["var"].encoding["preferred_chunks"] - == expected_preferred_chunks + assert roundtrip["var"].encoding["chunks"] == expected_chunks + assert roundtrip["var"].encoding["preferred_chunks"] == dict( + zip(dimension_names, expected_chunks, strict=True) ) assert isinstance(roundtrip["var"].data, np.ndarray) np.testing.assert_array_equal(roundtrip["var"].values, data) @@ -7264,14 +7259,7 @@ def test_read( @cases @requires_dask def test_read_dask( - self, - tmp_path, - shape, - chunks, - dimension_names, - dtype, - expected_preferred_chunks, - expected_dask_chunks, + self, tmp_path, shape, chunks, dimension_names, dtype, expected_chunks ) -> None: """Reading into dask arrays picks up the exact, variable-sized chunks as the dask chunks (not just `Dataset.chunks`, but the underlying @@ -7289,7 +7277,7 @@ def test_read_dask( roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) assert isinstance(roundtrip["var"].data, dask_array_type) - assert roundtrip["var"].data.chunks == expected_dask_chunks + assert roundtrip["var"].data.chunks == expected_chunks np.testing.assert_array_equal(roundtrip["var"].values, data) From 8e83b52af219b74f25726c3fe1c3e5d468d25938 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 16 Sep 2026 15:38:12 -0400 Subject: [PATCH 06/19] Handle rectilinear shard grids when reading (not just rectilinear chunks) zarr_array.shards raises NotImplementedError for a rectilinear shard grid the same way zarr_array.chunks does for a rectilinear chunk grid, but only the latter was handled. Regular inner chunks with rectilinear shards (e.g. chunks=(1,), shards=((1, 2),)) crashed xr.open_zarr with an uncaught NotImplementedError. Falls back to write_chunk_sizes (the outer/storage-level sizes, matching what "shards" means) the same way the chunks fallback uses read_chunk_sizes (the inner/read-level sizes). Reported at https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 Co-Authored-By: Claude Sonnet 5 --- xarray/backends/zarr.py | 14 +++++++++++++- xarray/tests/test_backends.py | 35 ++++++++++++++++++++++++++++++++++- 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 509176051e2..468f6b1f86f 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -879,11 +879,23 @@ def open_store_variable(self, name): "preferred_chunks": preferred_chunks, } + try: + # Regular shard grid (or no sharding at all, in which case this is + # None): a single uniform shard size per dimension. + shards = zarr_array.shards + except NotImplementedError: + # Rectilinear shard grid: shard sizes vary along an axis, so + # there is no single shard size. `.shards` raises and we instead + # read the explicit per-shard listing, e.g. ((1, 2),). + # write_chunk_sizes (not read_chunk_sizes) gives the outer, + # storage-level (shard) sizes here. + shards = zarr_array.write_chunk_sizes + encoding.update( { "compressors": zarr_array.compressors, "filters": zarr_array.filters, - "shards": zarr_array.shards, + "shards": shards, } ) if self.zarr_group.metadata.zarr_format == 3: diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index b5271be4500..f7b4e9c1fa6 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7191,7 +7191,9 @@ class TestZarrRectilinearChunksRead: """ @staticmethod - def create_zarr_array(store_path, shape, chunks, dimension_names, dtype): + def create_zarr_array( + store_path, shape, chunks, dimension_names, dtype, shards=None + ): import zarr root = zarr.open_group(store_path, mode="w", zarr_format=3) @@ -7200,6 +7202,7 @@ def create_zarr_array(store_path, shape, chunks, dimension_names, dtype): shape=shape, # older zarr stubs (<3.2) don't include rectilinear chunk types chunks=chunks, # type: ignore[arg-type, unused-ignore] + shards=shards, # type: ignore[arg-type, unused-ignore] dtype=dtype, dimension_names=dimension_names, ) @@ -7280,6 +7283,36 @@ def test_read_dask( assert roundtrip["var"].data.chunks == expected_chunks np.testing.assert_array_equal(roundtrip["var"].values, data) + def test_read_rectilinear_shards(self, tmp_path) -> None: + """Regular (uniform) inner chunks, but rectilinear (variable-sized) + shards. `.shards` raises NotImplementedError here, same as `.chunks` + does for rectilinear chunk grids, so it needs the same fallback. + + https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 + """ + import zarr + + data = np.array([1.0, 2.0, 3.0], dtype="float32") + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, + shape=(3,), + chunks=(1,), + shards=((1, 2),), + dimension_names=("x",), + dtype="float32", + ) + arr[:] = data + + roundtrip = xr.open_zarr( + store_path, zarr_format=3, consolidated=False, chunks=None + ) + assert roundtrip["var"].encoding["chunks"] == ((1, 1, 1),) + assert roundtrip["var"].encoding["shards"] == ((1, 2),) + np.testing.assert_array_equal(roundtrip["var"].values, data) + @requires_zarr @requires_fsspec From eac23405d1d5286e6beea71ef8092ed1cd7adce8 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 16 Sep 2026 15:44:09 -0400 Subject: [PATCH 07/19] Address review feedback: PR reference, error message, region-write docs 1. whats-new.rst cited #11279 (the write-side PR) instead of this PR (#11592); fixed to reference both. 2/3. Trying to write a Dataset read from a rectilinear-chunked store (via round-trip to_zarr, or a region write) hit a generic "must be an int or a tuple of ints" TypeError that gave no indication this was about rectilinear chunks or how to work around it. Both cases go through the same _determine_zarr_chunks code path, so a single fix covers both: the error now names rectilinear chunks and states the two-step workaround (clear the chunk encoding *and* rechunk to a uniform size -- verified both are actually needed). Documented the region-write limitation in the rectilinear chunks docs note. 4. CI coverage: confirmed this is an upstream constraint, not something to fix here -- zarr>=3.2 requires Python>=3.12 (pypi.org/pypi/zarr/3.2.0/json), so xarray's Python-3.11-pinned environments can never resolve a zarr new enough to exercise this feature and correctly skip via requires_zarr_rectilinear_chunks. https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 (Ian Hunt-Isaak) Co-Authored-By: Claude Sonnet 5 --- doc/user-guide/io.md | 6 +++++- doc/whats-new.rst | 3 ++- xarray/backends/zarr.py | 14 ++++++++++++++ xarray/tests/test_backends.py | 27 +++++++++++++++++++++++++++ 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/doc/user-guide/io.md b/doc/user-guide/io.md index 239d1c90e9d..55c14d9ce77 100644 --- a/doc/user-guide/io.md +++ b/doc/user-guide/io.md @@ -1253,7 +1253,11 @@ with zarr.config.set({"array.rectilinear_chunks": True}): xarray can currently only *read* rectilinear-chunked Zarr V3 arrays, not write them. Rectilinear arrays must be created with `zarr-python` directly (or another tool, such as Icechunk) before being opened with `xr.open_zarr`. -Write support is tracked in [GH11279](https://github.com/pydata/xarray/pull/11279). +This also means a Dataset opened from a rectilinear-chunked store cannot be +written back with {py:meth}`Dataset.to_zarr`, including writing to a +`region` of the store -- both raise a `TypeError` naming the affected +variable. Write support is tracked in +[GH11279](https://github.com/pydata/xarray/pull/11279). ``` ### Groups diff --git a/doc/whats-new.rst b/doc/whats-new.rst index f3009149fa4..e8380f5b794 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -16,7 +16,8 @@ New Features - Support reading Zarr V3 arrays with rectilinear (variable-sized) chunk grids, requiring ``zarr-python >= 3.2`` with ``zarr.config.set({"array.rectilinear_chunks": True})``. Writing - rectilinear chunks from xarray is not yet supported (:pull:`11279`). + rectilinear chunks from xarray is not yet supported (:pull:`11592`, + extracted from :pull:`11279`). By `Max Jones `_ and `Tom Nicholas `_. diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 468f6b1f86f..5e962402d24 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -367,6 +367,20 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): for x in enc_chunks_tuple: if not isinstance(x, int): + if isinstance(x, list | tuple): + raise TypeError( + f"encoding['chunks']={enc_chunks_tuple!r} for variable " + f"named {name!r} looks like a rectilinear (variable-sized) " + "chunk grid, e.g. from a store opened with " + "zarr.config.set({'array.rectilinear_chunks': True}). " + "Writing rectilinear chunks is not yet supported by xarray " + "(only reading them is). To write this variable with a " + "regular chunk grid instead, clear its chunk encoding " + "(`ds.drop_encoding()`, or " + f"`del ds[{name!r}].encoding['chunks']`) *and* make sure " + "its own chunks are uniform too, e.g. with " + "`ds.chunk({dim: size})`." + ) raise TypeError( "zarr chunk sizes specified in `encoding['chunks']` " "must be an int or a tuple of ints. " diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index f7b4e9c1fa6..d4621aef3a5 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7313,6 +7313,33 @@ def test_read_rectilinear_shards(self, tmp_path) -> None: assert roundtrip["var"].encoding["shards"] == ((1, 2),) np.testing.assert_array_equal(roundtrip["var"].values, data) + def test_write_after_read_gives_helpful_error(self, tmp_path) -> None: + """Writing isn't supported yet, so round-tripping (and region writes, + which go through the same encoding-extraction code path) should fail + with an error that names rectilinear chunks and a workaround, not a + generic "must be an int" message. + + https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 + """ + import zarr + + data = np.arange(60, dtype="float32") + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, + shape=(60,), + chunks=((10, 20, 30),), + dimension_names=("x",), + dtype="float32", + ) + arr[:] = data + + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + with pytest.raises(TypeError, match=r"rectilinear"): + roundtrip.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + @requires_zarr @requires_fsspec From ef8385bf28fc250cfd7ece20c730ed2c667387d6 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 10:45:42 -0400 Subject: [PATCH 08/19] Validate encoding['shards'] on write; it bypassed the chunks guard Investigated the claim that reading a rectilinear-shard store and writing it back "succeeds cleanly" with valid rectilinear output. Traced zarr-python's ArrayV3Metadata.chunks/.shards properties: both gate on the identical `isinstance(self.chunk_grid, RegularChunkGridMetadata)` check, so for a store actually read through our open_store_variable fallback, encoding["chunks"] and encoding["shards"] always agree -- if the shard grid is genuinely rectilinear, both raise NotImplementedError together, and our existing encoding["chunks"] guard already catches it (verified: a round-trip from a real rectilinear-sharded store still raises TypeError, not a silent write). The part that *is* real: encoding["shards"] is passed straight through to zarr_group.create() in _create_new_array with no validation at all, unlike encoding["chunks"] (_determine_zarr_chunks). Confirmed this is pre-existing on main, unrelated to anything in this PR: hand-setting ds['var'].encoding = {"chunks": (10,), "shards": ((20, 10, 30),)} and calling to_zarr() writes a genuine rectilinear shard grid on main today, with no rectilinear reading involved at all. Added the same style of validation for encoding["shards"] that encoding["chunks"] already has, plus a regression test reproducing the hand-set-encoding case. https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 Co-Authored-By: Claude Sonnet 5 --- xarray/backends/zarr.py | 21 +++++++++++++++++++++ xarray/tests/test_backends.py | 21 +++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 5e962402d24..9f9ff79f623 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -520,6 +520,27 @@ def extract_zarr_variable_encoding( if chunks is None: chunks = "auto" encoding["chunks"] = chunks + + # `encoding["shards"]` isn't run through `_determine_zarr_chunks` above, + # so a rectilinear (variable-sized) shard spec -- e.g. read from a + # rectilinear-sharded store via `open_store_variable`, or set by hand -- + # would otherwise pass straight through to zarr's array creation + # unchecked, silently writing a rectilinear shard grid that xarray + # doesn't otherwise support writing. + shards = encoding.get("shards") + if shards is not None and any(not isinstance(x, int) for x in shards): + raise TypeError( + f"encoding['shards']={shards!r} for variable named {name!r} " + "looks like a rectilinear (variable-sized) shard grid, e.g. " + "from a store opened with " + "zarr.config.set({'array.rectilinear_chunks': True}). Writing " + "rectilinear shards is not yet supported by xarray (only " + "reading them is). To write this variable with a regular shard " + "grid instead, clear its shard encoding (`ds.drop_encoding()`, " + f"or `del ds[{name!r}].encoding['shards']`) *and* make sure its " + "own chunks are uniform too, e.g. with `ds.chunk({dim: size})`." + ) + return encoding diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index d4621aef3a5..360259437d5 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7340,6 +7340,27 @@ def test_write_after_read_gives_helpful_error(self, tmp_path) -> None: with pytest.raises(TypeError, match=r"rectilinear"): roundtrip.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: + """encoding["shards"] must be validated the same way encoding["chunks"] + is: it is passed straight through to zarr's array creation otherwise, + so a rectilinear (variable-sized) shard spec -- whether read from a + store or set by hand -- would silently write a rectilinear-sharded + array, even though xarray doesn't support writing rectilinear + anything. + + https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 + """ + data = np.arange(60, dtype="float32") + ds = xr.Dataset({"var": ("x", data)}) + ds["var"].encoding["chunks"] = (10,) + ds["var"].encoding["shards"] = ((20, 10, 30),) + + import zarr + + with zarr.config.set({"array.rectilinear_chunks": True}): + with pytest.raises(TypeError, match=r"rectilinear"): + ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + @requires_zarr @requires_fsspec From d1891747156302362817c4a3d3cc78130347605c Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 15:58:14 -0400 Subject: [PATCH 09/19] Fix regression: int encoding['shards'] crashed the new write guard any(not isinstance(x, int) for x in shards) iterated shards without checking it was iterable first. A bare int shards value (applies to every dimension, same convention as encoding['chunks']) crashed with TypeError: 'int' object is not iterable instead of writing normally. Guard it the same way _determine_zarr_chunks already guards chunks: skip validation when shards is an int. Co-Authored-By: Claude Sonnet 5 --- xarray/backends/zarr.py | 6 +++++- xarray/tests/test_backends.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 9f9ff79f623..2211d0457b3 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -528,7 +528,11 @@ def extract_zarr_variable_encoding( # unchecked, silently writing a rectilinear shard grid that xarray # doesn't otherwise support writing. shards = encoding.get("shards") - if shards is not None and any(not isinstance(x, int) for x in shards): + if ( + shards is not None + and not isinstance(shards, integer_types) + and any(not isinstance(x, int) for x in shards) + ): raise TypeError( f"encoding['shards']={shards!r} for variable named {name!r} " "looks like a rectilinear (variable-sized) shard grid, e.g. " diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 360259437d5..3607acfc7e4 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7361,6 +7361,20 @@ def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: with pytest.raises(TypeError, match=r"rectilinear"): ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + def test_write_regular_shards_still_works(self, tmp_path) -> None: + """The rectilinear-shards guard must not choke on a plain regular + shard spec, including a bare int (applies to every dimension) rather + than a tuple -- iterating that directly raises `TypeError: 'int' + object is not iterable`. + """ + data = np.arange(60, dtype="float32") + + for shards in (20, (20,)): + ds = xr.Dataset({"var": ("x", data)}) + ds["var"].encoding["chunks"] = 10 + ds["var"].encoding["shards"] = shards + ds.to_zarr(tmp_path / f"dest-{shards}.zarr", zarr_format=3, mode="w") + @requires_zarr @requires_fsspec From 9fa0678a3e2f6d655b39cc9be6f28ce78675cb31 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 16:49:31 -0400 Subject: [PATCH 10/19] Validate encoding before resizing an existing array on append extract_zarr_variable_encoding ran after zarr_array.resize() in set_variables, so if it raised (as it now does for rectilinear chunk encodings, which xarray can't write yet) the store was left with a grown array whose new region was never written. Move the encoding extraction ahead of the resize; it only depends on the variable, not on the region or the existing array. Co-authored-by: Claude --- xarray/backends/zarr.py | 30 +++++++++++++++++------------- xarray/tests/test_backends.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 2211d0457b3..bb5b2c1a8cd 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -1262,6 +1262,23 @@ def set_variables( else: del v.encoding["_FillValue"] + # 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. + # Note: Ideally there should be two functions, one for validating the chunks and + # another one for extracting the encoding. + # This must happen *before* any resize of an existing array below: if + # the encoding is rejected (e.g. rectilinear chunks, which xarray can't + # write yet) after resizing, the store would be left with a grown array + # whose new region was never written. + encoding = extract_zarr_variable_encoding( + v, + raise_on_invalid=vn in check_encoding_set, + name=vn, + zarr_format=3 if is_zarr_v3_format else 2, + ) + zarr_shape = None write_region = self._write_region if self._write_region is not None else {} write_region = {dim: write_region.get(dim, slice(None)) for dim in dims} @@ -1287,19 +1304,6 @@ def set_variables( zarr_shape = zarr_array.shape region = tuple(write_region[dim] for dim in dims) - # 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. - # Note: Ideally there should be two functions, one for validating the chunks and - # another one for extracting the encoding. - encoding = extract_zarr_variable_encoding( - v, - raise_on_invalid=vn in check_encoding_set, - name=vn, - zarr_format=3 if is_zarr_v3_format else 2, - ) - # 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 diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 3607acfc7e4..a8024878ec3 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7340,6 +7340,34 @@ def test_write_after_read_gives_helpful_error(self, tmp_path) -> None: with pytest.raises(TypeError, match=r"rectilinear"): roundtrip.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + def test_append_does_not_resize_before_erroring(self, tmp_path) -> None: + """Appending along a dimension of a rectilinear variable must fail + *before* the existing zarr array is resized. Otherwise the store is + left with a grown array whose new region was never written. + """ + import zarr + + data = np.arange(60, dtype="float32") + store_path = tmp_path / "source.zarr" + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = self.create_zarr_array( + store_path, + shape=(60,), + chunks=((10, 20, 30),), + dimension_names=("x",), + dtype="float32", + ) + arr[:] = data + + ds = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + with pytest.raises(TypeError, match=r"rectilinear"): + ds.to_zarr( + store_path, append_dim="x", zarr_format=3, consolidated=False + ) + + assert zarr.open_array(store_path / "var").shape == (60,) + def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: """encoding["shards"] must be validated the same way encoding["chunks"] is: it is passed straight through to zarr's array creation otherwise, From 84c0ef33893b7e55cf586465e35cbff1c2f5d3e3 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 16:50:31 -0400 Subject: [PATCH 11/19] Only flag encoding['shards'] as rectilinear when it nests sequences The new shards guard iterated whatever it was given and rejected any non-int element. zarr also accepts shards="auto" (which worked before this branch) and a config dict, both of which iterate to non-int values and so hit a misleading "looks like a rectilinear shard grid" error. Restrict the check to a list/tuple containing lists/tuples, mirroring the chunks guard, and cover "auto" in the regression test. Co-authored-by: Claude --- xarray/backends/zarr.py | 9 +++++---- xarray/tests/test_backends.py | 25 ++++++++++++++----------- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index bb5b2c1a8cd..c205db00c9e 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -527,11 +527,12 @@ def extract_zarr_variable_encoding( # would otherwise pass straight through to zarr's array creation # unchecked, silently writing a rectilinear shard grid that xarray # doesn't otherwise support writing. + # Only a sequence *of sequences* is rectilinear; zarr also accepts an int, + # a tuple of ints, the string "auto" and a config dict here, none of which + # should trip this check. shards = encoding.get("shards") - if ( - shards is not None - and not isinstance(shards, integer_types) - and any(not isinstance(x, int) for x in shards) + if isinstance(shards, list | tuple) and any( + isinstance(x, list | tuple) for x in shards ): raise TypeError( f"encoding['shards']={shards!r} for variable named {name!r} " diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index a8024878ec3..df63eac6222 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7389,19 +7389,22 @@ def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: with pytest.raises(TypeError, match=r"rectilinear"): ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") - def test_write_regular_shards_still_works(self, tmp_path) -> None: - """The rectilinear-shards guard must not choke on a plain regular - shard spec, including a bare int (applies to every dimension) rather - than a tuple -- iterating that directly raises `TypeError: 'int' - object is not iterable`. + @pytest.mark.parametrize("shards", [20, (20,), "auto"], ids=repr) + def test_write_regular_shards_still_works(self, tmp_path, shards) -> None: + """The rectilinear-shards guard must only fire on a sequence of + sequences. zarr also accepts a bare int, a tuple of ints and the + string "auto" for shards; iterating an int raises TypeError and + iterating "auto" yields non-int characters, so neither may be fed + to a naive element check. """ data = np.arange(60, dtype="float32") - - for shards in (20, (20,)): - ds = xr.Dataset({"var": ("x", data)}) - ds["var"].encoding["chunks"] = 10 - ds["var"].encoding["shards"] = shards - ds.to_zarr(tmp_path / f"dest-{shards}.zarr", zarr_format=3, mode="w") + ds = xr.Dataset({"var": ("x", data)}) + ds["var"].encoding["chunks"] = 10 + ds["var"].encoding["shards"] = shards + ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") + assert ( + xr.open_zarr(tmp_path / "dest.zarr")["var"].encoding["shards"] is not None + ) @requires_zarr From f706ca5a3d8fa9c54591f72dbebcf5e965e4398e Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 16:54:08 -0400 Subject: [PATCH 12/19] Normalise rectilinear chunk encodings across zarr-python versions What zarr-python reports for an array whose chunk grid is regular along only some dimensions, or for regular inner chunks inside rectilinear shards, differs by version: 3.2 returns a mixed tuple like (2, (5, 10, 5)) from Array.chunks, 3.3 raises and read_chunk_sizes expands the regular dimension to (2, 2, 2), and 3.4 returns the compact inner chunk shape for sharded arrays where 3.3 raises. The tests were pinned to 3.3 and failed on 3.2.x (the locally installed version) and 3.4.0 (what unpinned CI environments resolve to). Compact zarr's per-chunk size listings back to a single int along every dimension that describes a regular grid, so encoding['chunks'] and encoding['shards'] come out the same regardless of zarr version: an int per regular dimension, a tuple of sizes per rectilinear one. Test expectations updated to that representation, with the fully expanded form kept for the dask-chunks assertion. zarr-python 3.2.x also crashes on a bare-int shards spec inside its own metadata parsing, so that test case is skipped below 3.3. Co-authored-by: Claude --- xarray/backends/zarr.py | 37 +++++++++++++++++++++- xarray/tests/__init__.py | 1 + xarray/tests/test_backends.py | 58 ++++++++++++++++++++++++++++------- 3 files changed, 84 insertions(+), 12 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index c205db00c9e..b65389cc6c8 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -396,6 +396,34 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): return enc_chunks_tuple +def _compact_chunk_sizes( + chunk_sizes: tuple[tuple[int, ...], ...], +) -> tuple[int | tuple[int, ...], ...]: + """Compact zarr's per-chunk size listing back to a chunk *shape* where possible. + + zarr-python's ``read_chunk_sizes`` / ``write_chunk_sizes`` list every + chunk's size along each dimension, dask-style, e.g. ``((10, 10, 5),)``. + Along any dimension where that listing describes a regular grid (all + chunks equal except a possibly smaller final one) it is replaced by the + single chunk size, e.g. ``10``; genuinely variable-sized dimensions keep + their listing. This gives the same representation regardless of which + zarr-python version is installed: whether ``Array.chunks`` raises for a + given array, or returns a compact or partially compact tuple, has changed + between 3.2, 3.3 and 3.4. + """ + compacted: list[int | tuple[int, ...]] = [] + for sizes in chunk_sizes: + sizes = tuple(sizes) + if not sizes: + # zero-length dimension: nothing to compact + compacted.append(sizes) + elif len(set(sizes[:-1])) <= 1 and sizes[-1] <= sizes[0]: + compacted.append(sizes[0]) + else: + compacted.append(sizes) + return tuple(compacted) + + def _get_zarr_dims_and_attrs(zarr_obj, dimension_key, try_nczarr): # Check for attributes and dimension name metadata as discussed in the Zarr encoding # specification https://docs.xarray.dev/en/stable/internals/zarr-encoding-spec.html @@ -912,6 +940,13 @@ def open_store_variable(self, name): # write_chunk_sizes) so this matches `.chunks` above in # returning the inner chunk shape when sharding is used. chunks = zarr_array.read_chunk_sizes + # Which arrays make `.chunks` raise, and what it returns for a grid + # that is regular along only some dimensions, differs between + # zarr-python versions, so normalise to one representation: an int per + # regular dimension, a tuple of sizes per rectilinear one. + 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 = { @@ -929,7 +964,7 @@ def open_store_variable(self, name): # read the explicit per-shard listing, e.g. ((1, 2),). # write_chunk_sizes (not read_chunk_sizes) gives the outer, # storage-level (shard) sizes here. - shards = zarr_array.write_chunk_sizes + shards = _compact_chunk_sizes(zarr_array.write_chunk_sizes) encoding.update( { diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 049039d7342..3028841ad9d 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -162,6 +162,7 @@ def get_dask_chunkmanager(): has_zarr_rectilinear_chunks, requires_zarr_rectilinear_chunks = _importorskip( "zarr", "3.2.0" ) +has_zarr_v3_3, requires_zarr_v3_3 = _importorskip("zarr", "3.3.0") if has_zarr: import zarr diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index df63eac6222..bf8c4f0507d 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -80,6 +80,7 @@ has_numpy_2, has_scipy, has_zarr, + has_zarr_v3_3, has_zarr_v3_async_oindex, has_zarr_v3_dtypes, mock, @@ -7207,13 +7208,15 @@ def create_zarr_array( dimension_names=dimension_names, ) - # `expected_chunks` is the fully-expanded, dask-style tuple-of-tuples that - # zarr-python reports (via `read_chunk_sizes`) for any array whose chunk - # grid isn't purely regular, even along dimensions that happen to be - # uniformly chunked -- e.g. dimension "x" below is uniformly chunked but - # still reported as (2, 2, 2), not the compact form 2. + # `expected_chunks` is what xarray reports in `encoding["chunks"]`: an int + # for every regularly-chunked dimension and the explicit per-chunk sizes + # for every rectilinear one. zarr-python itself reports these arrays + # differently from version to version (3.2 returns the mixed form from + # `.chunks`, 3.3 raises and `read_chunk_sizes` expands "x" to (2, 2, 2)), + # so xarray normalises. `expected_dask_chunks` is the fully expanded form + # dask uses. cases = pytest.mark.parametrize( - "shape,chunks,dimension_names,dtype,expected_chunks", + "shape,chunks,dimension_names,dtype,expected_chunks,expected_dask_chunks", [ pytest.param( (60,), @@ -7221,6 +7224,7 @@ def create_zarr_array( ("x",), "float32", ((10, 20, 30),), + ((10, 20, 30),), id="1d-rectilinear", ), pytest.param( @@ -7228,6 +7232,7 @@ def create_zarr_array( (2, (5, 10, 5)), ("x", "y"), "float64", + (2, (5, 10, 5)), ((2, 2, 2), (5, 10, 5)), id="mixed-regular-and-rectilinear", ), @@ -7236,7 +7241,14 @@ def create_zarr_array( @cases def test_read( - self, tmp_path, shape, chunks, dimension_names, dtype, expected_chunks + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_chunks, + expected_dask_chunks, ) -> None: import zarr @@ -7262,7 +7274,14 @@ def test_read( @cases @requires_dask def test_read_dask( - self, tmp_path, shape, chunks, dimension_names, dtype, expected_chunks + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_chunks, + expected_dask_chunks, ) -> None: """Reading into dask arrays picks up the exact, variable-sized chunks as the dask chunks (not just `Dataset.chunks`, but the underlying @@ -7280,13 +7299,16 @@ def test_read_dask( roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) assert isinstance(roundtrip["var"].data, dask_array_type) - assert roundtrip["var"].data.chunks == expected_chunks + assert roundtrip["var"].data.chunks == expected_dask_chunks np.testing.assert_array_equal(roundtrip["var"].values, data) def test_read_rectilinear_shards(self, tmp_path) -> None: """Regular (uniform) inner chunks, but rectilinear (variable-sized) shards. `.shards` raises NotImplementedError here, same as `.chunks` does for rectilinear chunk grids, so it needs the same fallback. + Whether `.chunks` raises for this array too depends on the zarr + version (it does on 3.2/3.3, not on 3.4), so the regular inner chunk + must be reported the same way, as a single int, either way. https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 """ @@ -7309,7 +7331,7 @@ def test_read_rectilinear_shards(self, tmp_path) -> None: roundtrip = xr.open_zarr( store_path, zarr_format=3, consolidated=False, chunks=None ) - assert roundtrip["var"].encoding["chunks"] == ((1, 1, 1),) + assert roundtrip["var"].encoding["chunks"] == (1,) assert roundtrip["var"].encoding["shards"] == ((1, 2),) np.testing.assert_array_equal(roundtrip["var"].values, data) @@ -7389,7 +7411,21 @@ def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: with pytest.raises(TypeError, match=r"rectilinear"): ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") - @pytest.mark.parametrize("shards", [20, (20,), "auto"], ids=repr) + @pytest.mark.parametrize( + "shards", + [ + pytest.param( + 20, + marks=pytest.mark.skipif( + not has_zarr_v3_3, + reason="zarr-python < 3.3 crashes on a bare int shards spec", + ), + ), + (20,), + "auto", + ], + ids=repr, + ) def test_write_regular_shards_still_works(self, tmp_path, shards) -> None: """The rectilinear-shards guard must only fire on a sequence of sequences. zarr also accepts a bare int, a tuple of ints and the From ae7568714f4b7bca105ca25e069c5d877615f224 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 16:54:41 -0400 Subject: [PATCH 13/19] Clarify rectilinear write errors for region and append writes The messages told users to clear the variable's chunk/shard encoding, but for writes into an existing array (region= or append_dim=) the store's own encoding is re-applied to the variable, so that advice cannot work. Say explicitly that such writes are unsupported and that the workaround only applies when writing to a new array. Co-authored-by: Claude --- xarray/backends/zarr.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index b65389cc6c8..9ac4f2c1676 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -374,7 +374,9 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): "chunk grid, e.g. from a store opened with " "zarr.config.set({'array.rectilinear_chunks': True}). " "Writing rectilinear chunks is not yet supported by xarray " - "(only reading them is). To write this variable with a " + "(only reading them is), so this variable cannot be written " + "into an existing rectilinear-chunked array (with `region` " + "or `append_dim`). To write it to a *new* array with a " "regular chunk grid instead, clear its chunk encoding " "(`ds.drop_encoding()`, or " f"`del ds[{name!r}].encoding['chunks']`) *and* make sure " @@ -568,8 +570,10 @@ def extract_zarr_variable_encoding( "from a store opened with " "zarr.config.set({'array.rectilinear_chunks': True}). Writing " "rectilinear shards is not yet supported by xarray (only " - "reading them is). To write this variable with a regular shard " - "grid instead, clear its shard encoding (`ds.drop_encoding()`, " + "reading them is), so this variable cannot be written into an " + "existing rectilinear-sharded array (with `region` or " + "`append_dim`). To write it to a *new* array with a regular " + "shard grid instead, clear its shard encoding (`ds.drop_encoding()`, " f"or `del ds[{name!r}].encoding['shards']`) *and* make sure its " "own chunks are uniform too, e.g. with `ds.chunk({dim: size})`." ) From 2ebea815e72351cdabb26bc6ceb8c9bdc5f3f950 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 17:37:28 -0400 Subject: [PATCH 14/19] Expand a bare int encoding['shards'] to a per-dimension tuple zarr-python 3.2.x crashes on an int shard spec ("'int' object is not iterable" from its metadata parsing), so expand it to one size per dimension in extract_zarr_variable_encoding, as zarr itself does for chunks. This keeps the documented zarr >= 3.2 minimum working without a version skip in the tests. As a side effect the dask/shard alignment checks in set_variables, which only run for a tuple spec, now also cover int shard specs. Co-authored-by: Claude --- xarray/backends/zarr.py | 8 ++++++++ xarray/tests/__init__.py | 1 - xarray/tests/test_backends.py | 20 ++++---------------- 3 files changed, 12 insertions(+), 17 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 9ac4f2c1676..2647f74a411 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -577,6 +577,14 @@ def extract_zarr_variable_encoding( f"or `del ds[{name!r}].encoding['shards']`) *and* make sure its " "own chunks are uniform too, e.g. with `ds.chunk({dim: size})`." ) + if isinstance(shards, integer_types): + # Expand a bare int to one size per dimension, as zarr does for + # chunks. zarr-python 3.2.x crashes on an int shard spec + # ("'int' object is not iterable" from its metadata parsing), and a + # tuple also lets the dask/shard alignment checks in + # ZarrStore.set_variables see the shard grid, which they skip for + # anything that isn't a tuple. + encoding["shards"] = variable.ndim * (int(shards),) return encoding diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 3028841ad9d..049039d7342 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -162,7 +162,6 @@ def get_dask_chunkmanager(): has_zarr_rectilinear_chunks, requires_zarr_rectilinear_chunks = _importorskip( "zarr", "3.2.0" ) -has_zarr_v3_3, requires_zarr_v3_3 = _importorskip("zarr", "3.3.0") if has_zarr: import zarr diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index bf8c4f0507d..31f603276ec 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -80,7 +80,6 @@ has_numpy_2, has_scipy, has_zarr, - has_zarr_v3_3, has_zarr_v3_async_oindex, has_zarr_v3_dtypes, mock, @@ -7411,27 +7410,16 @@ def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: with pytest.raises(TypeError, match=r"rectilinear"): ds.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") - @pytest.mark.parametrize( - "shards", - [ - pytest.param( - 20, - marks=pytest.mark.skipif( - not has_zarr_v3_3, - reason="zarr-python < 3.3 crashes on a bare int shards spec", - ), - ), - (20,), - "auto", - ], - ids=repr, - ) + @pytest.mark.parametrize("shards", [20, (20,), "auto"], ids=repr) def test_write_regular_shards_still_works(self, tmp_path, shards) -> None: """The rectilinear-shards guard must only fire on a sequence of sequences. zarr also accepts a bare int, a tuple of ints and the string "auto" for shards; iterating an int raises TypeError and iterating "auto" yields non-int characters, so neither may be fed to a naive element check. + + A bare int is expanded to a tuple by xarray before reaching zarr, + which also sidesteps zarr-python 3.2.x crashing on an int shard spec. """ data = np.arange(60, dtype="float32") ds = xr.Dataset({"var": ("x", data)}) From ea0da51fe647f332fe0e05441ae611eec2686505 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Thu, 17 Sep 2026 17:42:04 -0400 Subject: [PATCH 15/19] Shorten comments and test docstrings Co-authored-by: Claude --- xarray/backends/zarr.py | 67 +++++++++-------------------------- xarray/tests/test_backends.py | 64 +++++++-------------------------- 2 files changed, 29 insertions(+), 102 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 2647f74a411..f4a31e2bd6a 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -401,23 +401,15 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): def _compact_chunk_sizes( chunk_sizes: tuple[tuple[int, ...], ...], ) -> tuple[int | tuple[int, ...], ...]: - """Compact zarr's per-chunk size listing back to a chunk *shape* where possible. - - zarr-python's ``read_chunk_sizes`` / ``write_chunk_sizes`` list every - chunk's size along each dimension, dask-style, e.g. ``((10, 10, 5),)``. - Along any dimension where that listing describes a regular grid (all - chunks equal except a possibly smaller final one) it is replaced by the - single chunk size, e.g. ``10``; genuinely variable-sized dimensions keep - their listing. This gives the same representation regardless of which - zarr-python version is installed: whether ``Array.chunks`` raises for a - given array, or returns a compact or partially compact tuple, has changed - between 3.2, 3.3 and 3.4. + """Replace a dask-style listing of chunk sizes with a single int along any + dimension where it describes a regular grid, e.g. ((10, 10, 5),) -> (10,). + + Gives a consistent representation across zarr-python versions. """ compacted: list[int | tuple[int, ...]] = [] for sizes in chunk_sizes: sizes = tuple(sizes) - if not sizes: - # zero-length dimension: nothing to compact + if not sizes: # zero-length dimension compacted.append(sizes) elif len(set(sizes[:-1])) <= 1 and sizes[-1] <= sizes[0]: compacted.append(sizes[0]) @@ -551,15 +543,8 @@ def extract_zarr_variable_encoding( chunks = "auto" encoding["chunks"] = chunks - # `encoding["shards"]` isn't run through `_determine_zarr_chunks` above, - # so a rectilinear (variable-sized) shard spec -- e.g. read from a - # rectilinear-sharded store via `open_store_variable`, or set by hand -- - # would otherwise pass straight through to zarr's array creation - # unchecked, silently writing a rectilinear shard grid that xarray - # doesn't otherwise support writing. - # Only a sequence *of sequences* is rectilinear; zarr also accepts an int, - # a tuple of ints, the string "auto" and a config dict here, none of which - # should trip this check. + # Reject rectilinear shards (a sequence of sequences) like we do for chunks. + # Zarr also accepts an int, tuple of ints, "auto" or a dict here. shards = encoding.get("shards") if isinstance(shards, list | tuple) and any( isinstance(x, list | tuple) for x in shards @@ -578,12 +563,7 @@ def extract_zarr_variable_encoding( "own chunks are uniform too, e.g. with `ds.chunk({dim: size})`." ) if isinstance(shards, integer_types): - # Expand a bare int to one size per dimension, as zarr does for - # chunks. zarr-python 3.2.x crashes on an int shard spec - # ("'int' object is not iterable" from its metadata parsing), and a - # tuple also lets the dask/shard alignment checks in - # ZarrStore.set_variables see the shard grid, which they skip for - # anything that isn't a tuple. + # Expand to a tuple: zarr-python 3.2.x crashes on an int shard spec. encoding["shards"] = variable.ndim * (int(shards),) return encoding @@ -941,21 +921,13 @@ def open_store_variable(self, name): attributes = dict(attributes) try: - # Regular chunk grid: a single uniform chunk size per dimension, - # e.g. (10,) for a 60-element axis chunked into 10s. chunks = tuple(zarr_array.chunks) except NotImplementedError: - # Rectilinear chunk grid (zarr-python >= 3.2): chunk sizes vary - # along an axis, so there is no single chunk size. `.chunks` - # raises and we instead read the explicit per-chunk listing, - # e.g. ((10, 20, 30),). Use read_chunk_sizes (not - # write_chunk_sizes) so this matches `.chunks` above in - # returning the inner chunk shape when sharding is used. + # Rectilinear chunk grid (zarr-python >= 3.2): `.chunks` raises, so + # read the per-chunk sizes instead, e.g. ((10, 20, 30),). chunks = zarr_array.read_chunk_sizes - # Which arrays make `.chunks` raise, and what it returns for a grid - # that is regular along only some dimensions, differs between - # zarr-python versions, so normalise to one representation: an int per - # regular dimension, a tuple of sizes per rectilinear one. + # Normalise to an int per regular dim and a tuple per rectilinear dim, + # since what zarr-python returns above varies between versions. chunks = _compact_chunk_sizes( tuple(x if isinstance(x, tuple) else (x,) for x in chunks) ) @@ -967,15 +939,10 @@ def open_store_variable(self, name): } try: - # Regular shard grid (or no sharding at all, in which case this is - # None): a single uniform shard size per dimension. shards = zarr_array.shards except NotImplementedError: - # Rectilinear shard grid: shard sizes vary along an axis, so - # there is no single shard size. `.shards` raises and we instead - # read the explicit per-shard listing, e.g. ((1, 2),). - # write_chunk_sizes (not read_chunk_sizes) gives the outer, - # storage-level (shard) sizes here. + # Rectilinear shard grid: `.shards` raises, so read the per-shard + # (i.e. outer/storage chunk) sizes instead, e.g. ((1, 2),). shards = _compact_chunk_sizes(zarr_array.write_chunk_sizes) encoding.update( @@ -1316,10 +1283,8 @@ def set_variables( # https://github.com/pydata/xarray/issues/8371 for details. # Note: Ideally there should be two functions, one for validating the chunks and # another one for extracting the encoding. - # This must happen *before* any resize of an existing array below: if - # the encoding is rejected (e.g. rectilinear chunks, which xarray can't - # write yet) after resizing, the store would be left with a grown array - # whose new region was never written. + # Must run before any resize below, so a rejected encoding can't + # leave behind a resized array that was never written to. encoding = extract_zarr_variable_encoding( v, raise_on_invalid=vn in check_encoding_set, diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 31f603276ec..87d11497eaf 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7184,10 +7184,9 @@ def test_extract_zarr_variable_encoding() -> None: @requires_zarr_rectilinear_chunks class TestZarrRectilinearChunksRead: - """Tests for reading (not writing) rectilinear (variable-sized) zarr chunks. + """Reading rectilinear (variable-sized) zarr chunks. - xarray does not yet support *writing* rectilinear chunks, so these stores - are created directly with zarr-python rather than through `Dataset.to_zarr`. + xarray can't write these yet, so the stores are created with zarr directly. """ @staticmethod @@ -7207,13 +7206,9 @@ def create_zarr_array( dimension_names=dimension_names, ) - # `expected_chunks` is what xarray reports in `encoding["chunks"]`: an int - # for every regularly-chunked dimension and the explicit per-chunk sizes - # for every rectilinear one. zarr-python itself reports these arrays - # differently from version to version (3.2 returns the mixed form from - # `.chunks`, 3.3 raises and `read_chunk_sizes` expands "x" to (2, 2, 2)), - # so xarray normalises. `expected_dask_chunks` is the fully expanded form - # dask uses. + # expected_chunks is what xarray puts in encoding["chunks"]: an int per + # regular dim, a tuple of sizes per rectilinear dim (zarr-python itself + # differs between versions here). expected_dask_chunks is the expanded form. cases = pytest.mark.parametrize( "shape,chunks,dimension_names,dtype,expected_chunks,expected_dask_chunks", [ @@ -7282,9 +7277,7 @@ def test_read_dask( expected_chunks, expected_dask_chunks, ) -> None: - """Reading into dask arrays picks up the exact, variable-sized chunks - as the dask chunks (not just `Dataset.chunks`, but the underlying - dask array's own chunking).""" + """Dask arrays get the exact variable-sized chunks.""" import zarr data = np.arange(np.prod(shape), dtype=dtype).reshape(shape) @@ -7302,15 +7295,8 @@ def test_read_dask( np.testing.assert_array_equal(roundtrip["var"].values, data) def test_read_rectilinear_shards(self, tmp_path) -> None: - """Regular (uniform) inner chunks, but rectilinear (variable-sized) - shards. `.shards` raises NotImplementedError here, same as `.chunks` - does for rectilinear chunk grids, so it needs the same fallback. - Whether `.chunks` raises for this array too depends on the zarr - version (it does on 3.2/3.3, not on 3.4), so the regular inner chunk - must be reported the same way, as a single int, either way. - - https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 - """ + """Regular inner chunks with rectilinear shards. The regular chunks + must be reported as an int regardless of zarr-python version.""" import zarr data = np.array([1.0, 2.0, 3.0], dtype="float32") @@ -7335,13 +7321,8 @@ def test_read_rectilinear_shards(self, tmp_path) -> None: np.testing.assert_array_equal(roundtrip["var"].values, data) def test_write_after_read_gives_helpful_error(self, tmp_path) -> None: - """Writing isn't supported yet, so round-tripping (and region writes, - which go through the same encoding-extraction code path) should fail - with an error that names rectilinear chunks and a workaround, not a - generic "must be an int" message. - - https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 - """ + """Writing isn't supported yet; the error should say so, not just + "must be an int".""" import zarr data = np.arange(60, dtype="float32") @@ -7362,10 +7343,7 @@ def test_write_after_read_gives_helpful_error(self, tmp_path) -> None: roundtrip.to_zarr(tmp_path / "dest.zarr", zarr_format=3, mode="w") def test_append_does_not_resize_before_erroring(self, tmp_path) -> None: - """Appending along a dimension of a rectilinear variable must fail - *before* the existing zarr array is resized. Otherwise the store is - left with a grown array whose new region was never written. - """ + """A failed append must not leave the existing array resized.""" import zarr data = np.arange(60, dtype="float32") @@ -7390,15 +7368,7 @@ def test_append_does_not_resize_before_erroring(self, tmp_path) -> None: assert zarr.open_array(store_path / "var").shape == (60,) def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: - """encoding["shards"] must be validated the same way encoding["chunks"] - is: it is passed straight through to zarr's array creation otherwise, - so a rectilinear (variable-sized) shard spec -- whether read from a - store or set by hand -- would silently write a rectilinear-sharded - array, even though xarray doesn't support writing rectilinear - anything. - - https://github.com/pydata/xarray/pull/11592#issuecomment-5703342460 - """ + """Rectilinear shards must be rejected on write, like rectilinear chunks.""" data = np.arange(60, dtype="float32") ds = xr.Dataset({"var": ("x", data)}) ds["var"].encoding["chunks"] = (10,) @@ -7412,15 +7382,7 @@ def test_write_rectilinear_shards_blocked(self, tmp_path) -> None: @pytest.mark.parametrize("shards", [20, (20,), "auto"], ids=repr) def test_write_regular_shards_still_works(self, tmp_path, shards) -> None: - """The rectilinear-shards guard must only fire on a sequence of - sequences. zarr also accepts a bare int, a tuple of ints and the - string "auto" for shards; iterating an int raises TypeError and - iterating "auto" yields non-int characters, so neither may be fed - to a naive element check. - - A bare int is expanded to a tuple by xarray before reaching zarr, - which also sidesteps zarr-python 3.2.x crashing on an int shard spec. - """ + """Every non-rectilinear shard spec zarr accepts must still write.""" data = np.arange(60, dtype="float32") ds = xr.Dataset({"var": ("x", data)}) ds["var"].encoding["chunks"] = 10 From d4ab93bb48d5c8bf5bcc64cb48b7622e658abbf5 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 29 Sep 2026 11:24:02 -0400 Subject: [PATCH 16/19] Share one helper for the rectilinear chunks/shards write errors Co-authored-by: Claude --- xarray/backends/zarr.py | 44 +++++++++++++++-------------------------- 1 file changed, 16 insertions(+), 28 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 3d6cc98d1b6..8022a2e0b9c 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -308,6 +308,20 @@ async def async_getitem(self, key): ) +def _rectilinear_encoding_error(key: str, value, name) -> TypeError: + """Error for a rectilinear (variable-sized) chunk or shard spec, which xarray + can read but not yet write.""" + return TypeError( + f"encoding[{key!r}]={value!r} for variable {name!r} is a rectilinear " + "(variable-sized) grid, e.g. read from a store opened with " + "zarr.config.set({'array.rectilinear_chunks': True}). xarray can read " + "rectilinear grids but not yet write them, including with `region` or " + "`append_dim`. To write a regular grid to a new array instead, clear " + f"the encoding (`del ds[{name!r}].encoding[{key!r}]`) and make sure the " + "variable's own chunks are uniform (e.g. `ds.chunk({dim: size})`)." + ) + + def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): """ Given encoding chunks (possibly None or []) and variable chunks @@ -369,21 +383,7 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name): for x in enc_chunks_tuple: if not isinstance(x, int): if isinstance(x, list | tuple): - raise TypeError( - f"encoding['chunks']={enc_chunks_tuple!r} for variable " - f"named {name!r} looks like a rectilinear (variable-sized) " - "chunk grid, e.g. from a store opened with " - "zarr.config.set({'array.rectilinear_chunks': True}). " - "Writing rectilinear chunks is not yet supported by xarray " - "(only reading them is), so this variable cannot be written " - "into an existing rectilinear-chunked array (with `region` " - "or `append_dim`). To write it to a *new* array with a " - "regular chunk grid instead, clear its chunk encoding " - "(`ds.drop_encoding()`, or " - f"`del ds[{name!r}].encoding['chunks']`) *and* make sure " - "its own chunks are uniform too, e.g. with " - "`ds.chunk({dim: size})`." - ) + raise _rectilinear_encoding_error("chunks", enc_chunks_tuple, name) raise TypeError( "zarr chunk sizes specified in `encoding['chunks']` " "must be an int or a tuple of ints. " @@ -550,19 +550,7 @@ def extract_zarr_variable_encoding( if isinstance(shards, list | tuple) and any( isinstance(x, list | tuple) for x in shards ): - raise TypeError( - f"encoding['shards']={shards!r} for variable named {name!r} " - "looks like a rectilinear (variable-sized) shard grid, e.g. " - "from a store opened with " - "zarr.config.set({'array.rectilinear_chunks': True}). Writing " - "rectilinear shards is not yet supported by xarray (only " - "reading them is), so this variable cannot be written into an " - "existing rectilinear-sharded array (with `region` or " - "`append_dim`). To write it to a *new* array with a regular " - "shard grid instead, clear its shard encoding (`ds.drop_encoding()`, " - f"or `del ds[{name!r}].encoding['shards']`) *and* make sure its " - "own chunks are uniform too, e.g. with `ds.chunk({dim: size})`." - ) + raise _rectilinear_encoding_error("shards", shards, name) if isinstance(shards, integer_types): # Expand to a tuple: zarr-python 3.2.x crashes on an int shard spec. encoding["shards"] = variable.ndim * (int(shards),) From 4f9b2eb66d8f0128350224d0d0130e25232e390b Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 29 Sep 2026 12:51:44 -0400 Subject: [PATCH 17/19] Fix opening a rectilinear array with a zero-length dimension into dask zarr's read_chunk_sizes lists no chunks at all along a zero-length dimension, and passing that empty tuple through as the preferred chunk made dask raise "Empty tuples are not allowed in chunks" on zarr >= 3.3. Report it as (0,) instead, which is what dask itself uses. Co-authored-by: Claude --- xarray/backends/zarr.py | 5 +++-- xarray/tests/test_backends.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 8022a2e0b9c..a2db15e6dad 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -410,8 +410,9 @@ def _compact_chunk_sizes( compacted: list[int | tuple[int, ...]] = [] for sizes in chunk_sizes: sizes = tuple(sizes) - if not sizes: # zero-length dimension - compacted.append(sizes) + if not sizes: + # zero-length dimension: dask rejects an empty tuple here + compacted.append((0,)) elif len(set(sizes[:-1])) <= 1 and sizes[-1] <= sizes[0]: compacted.append(sizes[0]) else: diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index e6476ded983..1c06b235b5f 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7361,6 +7361,25 @@ def test_read_dask( assert roundtrip["var"].data.chunks == expected_dask_chunks np.testing.assert_array_equal(roundtrip["var"].values, data) + def test_read_zero_length_dim(self, tmp_path) -> None: + """zarr reports no chunk sizes at all along a zero-length dimension, + which must not be passed to dask as an empty tuple.""" + import zarr + + store_path = tmp_path / "source.zarr" + with zarr.config.set({"array.rectilinear_chunks": True}): + self.create_zarr_array( + store_path, + shape=(0, 20), + chunks=(2, (5, 10, 5)), + dimension_names=("x", "y"), + dtype="float32", + ) + roundtrip = xr.open_zarr(store_path, zarr_format=3, consolidated=False) + assert roundtrip["var"].shape == (0, 20) + assert roundtrip["var"].encoding["chunks"][1] == (5, 10, 5) + assert roundtrip["var"].data.chunks == ((0,), (5, 10, 5)) + def test_read_rectilinear_shards(self, tmp_path) -> None: """Regular inner chunks with rectilinear shards. The regular chunks must be reported as an int regardless of zarr-python version.""" From f7954f949e18410563df6aca840354ac46373860 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 29 Sep 2026 13:26:17 -0400 Subject: [PATCH 18/19] Clarify that rectilinear reads don't raise the minimum zarr version Co-authored-by: Claude --- doc/whats-new.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 29fa34d9e47..ab5bc34c57f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -14,10 +14,10 @@ v2026.07.1 (unreleased) New Features ~~~~~~~~~~~~ - Support reading Zarr V3 arrays with rectilinear (variable-sized) chunk - grids, requiring ``zarr-python >= 3.2`` with - ``zarr.config.set({"array.rectilinear_chunks": True})``. Writing - rectilinear chunks from xarray is not yet supported (:pull:`11592`, - extracted from :pull:`11279`). + grids. Using this feature needs ``zarr-python >= 3.2`` with + ``zarr.config.set({"array.rectilinear_chunks": True})``; xarray's minimum + supported zarr version is unchanged. Writing rectilinear chunks from + xarray is not yet supported (:pull:`11592`, extracted from :pull:`11279`). By `Max Jones `_ and `Tom Nicholas `_. From 67572289cc39d04e6f9f37ec758536c3316ec6a4 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 29 Sep 2026 14:47:32 -0400 Subject: [PATCH 19/19] Mark zero-length rectilinear test as requiring dask Co-authored-by: Claude --- xarray/tests/test_backends.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index aa93b75a40f..8cd45a3a37f 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7392,6 +7392,7 @@ def test_read_dask( assert roundtrip["var"].data.chunks == expected_dask_chunks np.testing.assert_array_equal(roundtrip["var"].values, data) + @requires_dask def test_read_zero_length_dim(self, tmp_path) -> None: """zarr reports no chunk sizes at all along a zero-length dimension, which must not be passed to dask as an empty tuple."""