diff --git a/doc/user-guide/io.md b/doc/user-guide/io.md index 3a9aa5155d6..c6a8a648af5 100644 --- a/doc/user-guide/io.md +++ b/doc/user-guide/io.md @@ -1244,6 +1244,35 @@ 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`. +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 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 864b344c4ab..ab5bc34c57f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -13,6 +13,14 @@ v2026.07.1 (unreleased) New Features ~~~~~~~~~~~~ +- Support reading Zarr V3 arrays with rectilinear (variable-sized) chunk + 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 `_. + - 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 b134921419f..a2db15e6dad 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 @@ -368,6 +382,8 @@ 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 _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. " @@ -383,6 +399,27 @@ 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, ...], ...]: + """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: 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: + 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 @@ -507,6 +544,18 @@ def extract_zarr_variable_encoding( if chunks is None: chunks = "auto" encoding["chunks"] = chunks + + # 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 + ): + 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),) + return encoding @@ -861,16 +910,36 @@ def open_store_variable(self, name): ) attributes = dict(attributes) + try: + chunks = tuple(zarr_array.chunks) + except NotImplementedError: + # 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 + # 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) + ) + 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, } + try: + shards = zarr_array.shards + except NotImplementedError: + # 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( { "compressors": zarr_array.compressors, "filters": zarr_array.filters, - "shards": zarr_array.shards, + "shards": shards, } ) if self.zarr_group.metadata.zarr_format == 3: @@ -1198,6 +1267,21 @@ 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. + # 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, + 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} @@ -1223,19 +1307,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/__init__.py b/xarray/tests/__init__.py index f50e662af64..18b0898f2ed 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -164,6 +164,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 6879b6a9b96..8cd45a3a37f 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -102,6 +102,7 @@ requires_scipy, requires_scipy_or_netCDF4, requires_zarr, + requires_zarr_rectilinear_chunks, requires_zarr_v3, requires_zarr_v3_async_oindex, ) @@ -7279,6 +7280,237 @@ def test_extract_zarr_variable_encoding() -> None: ) +@requires_zarr_rectilinear_chunks +class TestZarrRectilinearChunksRead: + """Reading rectilinear (variable-sized) zarr chunks. + + xarray can't write these yet, so the stores are created with zarr directly. + """ + + @staticmethod + 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) + 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] + shards=shards, # type: ignore[arg-type, unused-ignore] + dtype=dtype, + dimension_names=dimension_names, + ) + + # 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", + [ + pytest.param( + (60,), + ((10, 20, 30),), + ("x",), + "float32", + ((10, 20, 30),), + ((10, 20, 30),), + id="1d-rectilinear", + ), + pytest.param( + (6, 20), + (2, (5, 10, 5)), + ("x", "y"), + "float64", + (2, (5, 10, 5)), + ((2, 2, 2), (5, 10, 5)), + id="mixed-regular-and-rectilinear", + ), + ], + ) + + @cases + def test_read( + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_chunks, + expected_dask_chunks, + ) -> None: + import zarr + + 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 = 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"] == 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) + + @cases + @requires_dask + def test_read_dask( + self, + tmp_path, + shape, + chunks, + dimension_names, + dtype, + expected_chunks, + expected_dask_chunks, + ) -> None: + """Dask arrays get the exact variable-sized chunks.""" + import zarr + + 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 = 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 == 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.""" + 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.""" + 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,) + 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; the error should say so, not just + "must be an int".""" + 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") + + def test_append_does_not_resize_before_erroring(self, tmp_path) -> None: + """A failed append must not leave the existing array resized.""" + 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: + """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,) + 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") + + @pytest.mark.parametrize("shards", [20, (20,), "auto"], ids=repr) + def test_write_regular_shards_still_works(self, tmp_path, shards) -> None: + """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 + 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 + ) + + @pytest.fixture def fsspec_memory_zarr_stores(): """Write two zarr stores to fsspec's global in-memory filesystem."""