feat: rectilinear chunks in Zarr backend

.github/workflows/upstream-dev-ci.yaml

@@ -158,12 +158,9 @@ jobs:
   mypy-upstream-dev:
     name: mypy-upstream-dev
     runs-on: ubuntu-latest
-    needs: [detect-ci-trigger, cache-pixi-lock]
+    needs: cache-pixi-lock
     if: |
-      always()
-      && (
-          contains( github.event.pull_request.labels.*.name, 'run-upstream')
-      )
+      always() && needs.cache-pixi-lock.result == 'success'
     defaults:
       run:
         shell: bash -l {0}

doc/user-guide/io.rst

@@ -1241,6 +1241,33 @@ shape of each coordinate array in the ``encoding`` argument:
 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. Requires ``zarr-python >= 3.2``, and the experimental feature must be
+enabled for both reading and writing:
+
+.. code-block:: python
+
+    import zarr
+
+    with zarr.config.set({"array.rectilinear_chunks": True}):
+        ds = xr.Dataset({"var": ("x", np.arange(60))}).chunk({"x": (10, 20, 30)})
+        ds.to_zarr("rectilinear.zarr", zarr_format=3, mode="w")
+
+        roundtrip = xr.open_zarr("rectilinear.zarr", zarr_format=3)
+        roundtrip.chunks["x"]  # (10, 20, 30)
+
+Rectilinear chunks can also be specified through the ``encoding`` argument
+(one sequence per dimension), e.g. ``encoding={"var": {"chunks": ((10, 20, 30),)}}``.
+Writing non-uniform chunks to a zarr v2 store raises ``ValueError`` because the
+feature is Zarr Format 3-only.
+
 Groups
 ~~~~~~
 

doc/whats-new.rst

@@ -14,6 +14,11 @@ v2026.05.0 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Support reading and writing Zarr V3 arrays with rectilinear (variable-sized)
+  chunk grids. Requires zarr-python >= 3.2 with
+  ``zarr.config.set({"array.rectilinear_chunks": True})``, which must be set
+  for both reading and writing rectilinear-chunked stores. (:pull:`11279`).
+  By `Max Jones <https://github.com/maxrjones>`_.
 
 Breaking Changes
 ~~~~~~~~~~~~~~~~

xarray/backends/chunks.py

@@ -1,3 +1,5 @@
+import itertools
+
 import numpy as np
 
 from xarray.core.datatree import Variable
@@ -133,9 +135,12 @@ def align_nd_chunks(
 
 def build_grid_chunks(
     size: int,
-    chunk_size: int,
+    chunk_size: int | tuple[int, ...],
     region: slice | None = None,
 ) -> tuple[int, ...]:
+    if isinstance(chunk_size, tuple):
+        return _build_rectilinear_grid_chunks(chunk_size, region)
+
     if region is None:
         region = slice(0, size)
 
@@ -153,9 +158,34 @@ def build_grid_chunks(
     return tuple(chunks_on_region)
 
 
+def _build_rectilinear_grid_chunks(
+    chunk_sizes: tuple[int, ...],
+    region: slice | None = None,
+) -> tuple[int, ...]:
+    """Build grid chunks for a rectilinear dimension within a region."""
+    if region is None or region == slice(None):
+        return tuple(chunk_sizes)
+
+    region_start = region.start or 0
+    region_stop = region.stop or sum(chunk_sizes)
+
+    boundaries = [0, *itertools.accumulate(chunk_sizes, initial=0)]
+
+    result = []
+    for chunk_start, chunk_end in itertools.pairwise(boundaries):
+        if chunk_end <= region_start or chunk_start >= region_stop:
+            continue
+
+        effective_start = max(chunk_start, region_start)
+        effective_end = min(chunk_end, region_stop)
+        result.append(effective_end - effective_start)
+
+    return tuple(result)
+
+
 def grid_rechunk(
     v: Variable,
-    enc_chunks: tuple[int, ...],
+    enc_chunks: tuple[int | tuple[int, ...], ...],
     region: tuple[slice, ...],
 ) -> Variable:
     nd_v_chunks = v.chunks
@@ -181,9 +211,36 @@ def grid_rechunk(
     return v
 
 
+def _validate_rectilinear_chunk_alignment(
+    dask_chunks: tuple[int, ...],
+    enc_chunks: tuple[int, ...],
+    axis: int,
+    name: str,
+    region: slice = slice(None),
+) -> None:
+    """Validate dask chunks align with rectilinear encoding chunk boundaries."""
+    enc_stops = set(itertools.accumulate(enc_chunks))
+    region_start = region.start or 0
+    dask_stops = {region_start + s for s in itertools.accumulate(dask_chunks)}
+    # The final stop (total size) always matches. Exclude it.
+    total = sum(enc_chunks)
+    enc_stops.discard(total)
+    dask_stops.discard(total)
+    bad = dask_stops - enc_stops
+    if bad:
+        raise ValueError(
+            f"Specified rectilinear encoding chunks {enc_chunks!r} for variable "
+            f"named {name!r} would overlap multiple Dask chunks on axis {axis}. "
+            f"Dask chunk boundaries at positions {sorted(bad)} do not align with "
+            f"encoding chunk boundaries at {sorted(enc_stops)}. "
+            "Writing this array in parallel with Dask could lead to corrupted data. "
+            "Consider rechunking using `chunk()` or setting `safe_chunks=False`."
+        )
+
+
 def validate_grid_chunks_alignment(
     nd_v_chunks: tuple[tuple[int, ...], ...] | None,
-    enc_chunks: tuple[int, ...],
+    enc_chunks: tuple[int | tuple[int, ...], ...],
     backend_shape: tuple[int, ...],
     region: tuple[slice, ...],
     allow_partial_chunks: bool,
@@ -213,6 +270,18 @@ def validate_grid_chunks_alignment(
         backend_shape,
         strict=True,
     ):
+        if isinstance(chunk_size, tuple):
+            # Rectilinear dimension: use boundary-based validation
+            _validate_rectilinear_chunk_alignment(
+                dask_chunks=v_chunks,
+                enc_chunks=chunk_size,
+                axis=axis,
+                name=name,
+                region=interval,
+            )
+            continue
+
+        # Regular dimension: use existing validation logic
         for i, chunk in enumerate(v_chunks[1:-1]):
             if chunk % chunk_size:
                 raise ValueError(

xarray/backends/zarr.py

@@ -46,6 +46,10 @@
     from xarray.core.types import ZarrArray, ZarrGroup
 
 
+def _has_rectilinear_chunks() -> bool:
+    return module_available("zarr", minversion="3.2")
+
+
 def _get_mappers(*, storage_options, store, chunk_store):
     # expand str and path-like arguments
     store = _normalize_path(store)
@@ -333,7 +337,7 @@ async def async_getitem(self, key):
         )
 
 
-def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name):
+def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name, zarr_format):
     """
     Given encoding chunks (possibly None or []) and variable chunks
     (possibly None or []).
@@ -355,18 +359,34 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name):
     # while dask chunks can be variable sized
     # https://dask.pydata.org/en/latest/array-design.html#chunks
     if var_chunks and not enc_chunks:
+        if zarr_format == 3 and _has_rectilinear_chunks():
+            # Check if dask chunks are regular (uniform except for last chunk)
+            has_varying_interior = any(
+                len(set(chunks[:-1])) > 1 for chunks in var_chunks
+            )
+            has_larger_final = any(chunks[0] < chunks[-1] for chunks in var_chunks)
+            if has_varying_interior or has_larger_final:
+                # Truly rectilinear: return dask-style tuples of per-chunk sizes.
+                # Requires zarr config: array.rectilinear_chunks = True
+                return tuple(var_chunks)
+            # Regular chunks: return the first chunk size per dimension
+            return tuple(chunk[0] for chunk in var_chunks)
+
         if any(len(set(chunks[:-1])) > 1 for chunks in var_chunks):
             raise ValueError(
-                "Zarr requires uniform chunk sizes except for final chunk. "
+                "Zarr v2 requires uniform chunk sizes except for the final chunk. "
                 f"Variable named {name!r} has incompatible dask chunks: {var_chunks!r}. "
-                "Consider rechunking using `chunk()`."
+                "Consider rechunking using `chunk()`, or switching to the "
+                "zarr v3 format with zarr-python>=3.2."
             )
         if any((chunks[0] < chunks[-1]) for chunks in var_chunks):
             raise ValueError(
-                "Final chunk of Zarr array must be the same size or smaller "
-                f"than the first. Variable named {name!r} has incompatible Dask chunks {var_chunks!r}."
-                "Consider either rechunking using `chunk()` or instead deleting "
-                "or modifying `encoding['chunks']`."
+                "The final chunk of a Zarr v2 array or a Zarr v3 array without the "
+                "rectilinear chunks extension must be the same size or smaller "
+                f"than the first. Variable named {name!r} has incompatible Dask "
+                f"chunks {var_chunks!r}. "
+                "Consider switching to Zarr v3 with the rectilinear chunks extension, "
+                "rechunking using `chunk()` or deleting or modifying `encoding['chunks']`."
             )
         # return the first chunk for each dimension
         return tuple(chunk[0] for chunk in var_chunks)
@@ -380,7 +400,12 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name):
     if isinstance(enc_chunks, integer_types):
         enc_chunks_tuple = ndim * (enc_chunks,)
     else:
-        enc_chunks_tuple = tuple(enc_chunks)
+        # Normalize per-dimension rectilinear specs (possibly lists, e.g.
+        # encoding={"chunks": [[10, 20, 30], 25]}) to tuples, so downstream
+        # code only deals with elements of type int | tuple[int, ...]
+        enc_chunks_tuple = tuple(
+            tuple(x) if isinstance(x, (list, tuple)) else x for x in enc_chunks
+        )
 
     if len(enc_chunks_tuple) != ndim:
         # throw away encoding chunks, start over
@@ -389,8 +414,17 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name):
             var_chunks,
             ndim,
             name,
+            zarr_format,
         )
 
+    # Rectilinear chunks: each element is a sequence of per-chunk edge lengths
+    if (
+        zarr_format == 3
+        and _has_rectilinear_chunks()
+        and any(not isinstance(x, int) for x in enc_chunks_tuple)
+    ):
+        return enc_chunks_tuple
+
     for x in enc_chunks_tuple:
         if not isinstance(x, int):
             raise TypeError(
@@ -532,6 +566,7 @@ def extract_zarr_variable_encoding(
         var_chunks=variable.chunks,
         ndim=variable.ndim,
         name=name,
+        zarr_format=zarr_format,
     )
     if _zarr_v3() and chunks is None:
         chunks = "auto"
@@ -910,9 +945,20 @@ 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 >= 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,
         }
 
         if _zarr_v3():

xarray/tests/__init__.py

@@ -134,23 +134,9 @@ def _importorskip(
 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")
-if has_zarr:
-    import zarr
-
-    # manual update by checking attrs for now
-    # TODO: use version specifier
-    # installing from git main is giving me a lower version than the
-    # most recently released zarr
-    has_zarr_v3_dtypes = hasattr(zarr.core, "dtype")
-    has_zarr_v3_async_oindex = hasattr(zarr.AsyncArray, "oindex")
-
-    requires_zarr_v3_dtypes = pytest.mark.skipif(
-        not has_zarr_v3_dtypes, reason="requires zarr>3.1.0"
-    )
-    requires_zarr_v3_async_oindex = pytest.mark.skipif(
-        not has_zarr_v3_async_oindex, reason="requires zarr>3.1.1"
-    )
-
+has_zarr_rectilinear_chunks, requires_zarr_rectilinear_chunks = _importorskip(
+    "zarr", "3.2.0"
+)
 
 has_fsspec, requires_fsspec = _importorskip("fsspec")
 has_iris, requires_iris = _importorskip("iris")