variable-sized chunks with zarr v3 by keewis · Pull Request #10880 · pydata/xarray

keewis · 2025-10-27T16:32:41Z

Closes #xxxx
Tests added
User visible changes (including notable bug fixes) are documented in whats-new.rst

Building on top of zarr-developers/zarr-python#3534, this is a draft PR that allows writing variable-sized chunks to zarr.

To see this in action, try:

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "xarray @ git+https://github.com/keewis/xarray.git@variable-chunking",
#   "zarr @ git+https://github.com/jhamman/zarr-python.git@feature/rectilinear-chunk-grid",
# ]
# ///

import numpy as np
import xarray as xr

rng = np.random.default_rng(seed=0)
values = rng.normal(size=(365, 20))

ds = xr.Dataset(
    {"a": (["time", "x"], values)},
    coords={"time": xr.date_range("2025-01-01", freq="d", periods=365)}
)
chunked = ds.chunk({"time": xr.groupers.TimeResampler(freq="ME"), "x": 10})

chunked.to_zarr(
    "variable_chunks.zarr",
    mode="w",
    safe_chunks=False,
    zarr_format=3,
    consolidated=False,
)

ds = xr.open_dataset(store, engine="zarr", chunks={})
print(ds.chunksizes)
# Frozen({'time': (31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31), 'x': (10, 10)})

At the moment, this requires safe_chunks=False because I didn't change the chunk alignment machinery, yet.

cc @d-v-b, @jhamman, @dcherian

keewis · 2025-10-27T16:34:05Z

xarray/backends/zarr.py

    # 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:


this check is probably not sufficient

jhamman · 2025-10-27T16:54:53Z

We need zarr-python>=3, which doesn't work with @jhamman's fork because it doesn't have tags for versions above 3.0.0b2

I just pushed tags to my fork!

keewis · 2025-10-27T16:57:29Z

thanks, I've changed the example back to using your fork

TomNicholas · 2025-11-05T16:53:10Z

xarray/backends/zarr.py

        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 final chunk. "
                f"Variable named {name!r} has incompatible dask chunks: {var_chunks!r}. "
                "Consider rechunking using `chunk()`."
            )
        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 "
+                "Final chunk of a Zarr v2 array must be the same size or smaller "


This is not correct - it's unfortunately not as simple as "Zarr V3 supports variable-length chunking but Zarr V2 doesn't".

jhamman · 2026-01-08T18:23:28Z

@keewis - zarr-developers/zarr-python#3534 is approaching a mergable state. Curious if you want to take another pass through this PR before we merge it and provide any feedback.

keewis · 2026-01-08T21:43:58Z

sure.

Do you know if there's a specific tell on whether rectilinear chunks are available? So far I've been using zarr_format == 3, but that's clearly not sufficient (and might even become wrong, should this ever be backported to the zarr 2 format). I can extend that by also checking if the RectilinearChunks and RegularChunks classes are exposed (see the most recent commit), but not sure if that's a sufficient check, either. Basically, what I'd like to check is whether the format extension is available.

I've posted a comment to the zarr PR (which doesn't seem to really affect the code here).

Finally, I still need to figure out how to change validate_grid_chunks_alignment, and we also need tests, both for reading and writing.

keewis · 2026-01-08T21:45:57Z

xarray/backends/zarr.py

                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."


I still struggle with accurately expressing the prerequisites for rectilinear chunk support. Maybe this is fine, but we could also ask for "rectilinear chunk support"?

Suggested change

"zarr v3 format with zarr-python>=3.2."

"zarr v3 format with enabled rectilinear chunk support."

keewis · 2026-01-08T21:49:15Z

pixi.toml

 dask = { git = "https://github.com/dask/dask" }
 distributed = { git = "https://github.com/dask/distributed" }
-zarr = { git = "https://github.com/zarr-developers/zarr-python" }
+zarr = { git = "https://github.com/jhamman/zarr-python", branch = "feature/rectilinear-chunk-grid" }


revert before merging:

Suggested change

zarr = { git = "https://github.com/jhamman/zarr-python", branch = "feature/rectilinear-chunk-grid" }

zarr = { git = "https://github.com/zarr-developers/zarr-python" }

maxrjones · 2026-02-11T16:35:23Z

@keewis would you be willing to push xarray tags to your fork (https://github.com/keewis/xarray/tags)? This would enable trying out your PR's integration in VirtualiZarr

keewis · 2026-02-11T17:29:51Z

would you be willing to push xarray tags to your fork

done

maxrjones

Some rough suggestions from trying this out in virtual-tiff/VirtualiZarr regarding using chunk_grid.chunk_shape rather than .chunks

maxrjones · 2026-02-17T01:46:08Z

xarray/backends/zarr.py

+        chunks = tuple(zarr_array.chunks)
+        preferred_chunks = dict(zip(dimensions, chunks, strict=True))


Suggested change

chunks = tuple(zarr_array.chunks)

preferred_chunks = dict(zip(dimensions, chunks, strict=True))

chunk_grid = zarr_array.metadata.chunk_grid

if has_chunk_grid_support and isinstance(chunk_grid, RegularChunkGrid):

chunks = chunk_grid.chunk_shape

preferred_chunks = dict(zip(dimensions, chunks, strict=True))

elif has_chunk_grid_support:

# RectilinearChunkGrid or other non-regular grids — store the

# full chunk_grid and skip preferred_chunks since there's no

# single chunk size per dimension

chunks = chunk_grid

preferred_chunks = {}

else:

# Fallback for older zarr-python without chunk_grid support

chunks = tuple(zarr_array.chunks)

preferred_chunks = dict(zip(dimensions, chunks, strict=True))

This suggestion adds support for RectilinearChunkGrids, which do not implement .chunks. It relies on private API, so probably worth advocating for Regular/RectilinearChunkGrid to be made public after the main PR lands

maxrjones · 2026-02-17T01:46:57Z

xarray/backends/zarr.py

+try:
+    from zarr import RectilinearChunks, RegularChunks  # noqa: F401
+
+    has_variable_chunk_support = True
+except ImportError:
+    has_variable_chunk_support = False


Suggested change

try:

from zarr import RectilinearChunks, RegularChunks # noqa: F401

has_variable_chunk_support = True

except ImportError:

has_variable_chunk_support = False

try:

from zarr.core.chunk_grids import RegularChunkGrid

has_chunk_grid_support = True

except ImportError:

has_chunk_grid_support = False

Used for the variable chunk grid support later on, see note there about making it public API

maxrjones · 2026-02-17T01:47:21Z

xarray/backends/zarr.py

    # 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_variable_chunk_support:


Suggested change

if zarr_format == 3 and has_variable_chunk_support:

if zarr_format == 3 and has_chunk_grid_support:

keewis added 3 commits October 27, 2025 17:07

allow variable chunk sizes for zarr v3

9e54553

change the error message to mention zarr v2

ebfe0be

infer the chunk shape(s) from the zarr v3 metadata

e2fd4c3

github-actions bot added topic-backends topic-zarr Related to zarr storage library io labels Oct 27, 2025

keewis commented Oct 27, 2025

View reviewed changes

keewis marked this pull request as draft October 27, 2025 16:34

TomNicholas reviewed Nov 5, 2025

View reviewed changes

keewis added run-upstream Run upstream CI skip-ci labels Jan 8, 2026

keewis added 5 commits January 8, 2026 21:05

Merge branch 'main' into variable-chunking

c39476d

temporarily install from the zarr branch

c99881c

try to detect if variable chunk sizes are supported

52ea952

add a bit more detail to the error messages

960a18a

simplify the construction of the chunks / preferred_chunks

f4a01b4

keewis commented Jan 8, 2026

View reviewed changes

maxrjones reviewed Feb 17, 2026

View reviewed changes

maxrjones mentioned this pull request Feb 17, 2026

feat: Support variable chunking zarr-developers/VirtualiZarr#877

Draft

8 tasks

	"zarr v3 format with zarr-python>=3.2."
	"zarr v3 format with enabled rectilinear chunk support."

	zarr = { git = "https://github.com/jhamman/zarr-python", branch = "feature/rectilinear-chunk-grid" }
	zarr = { git = "https://github.com/zarr-developers/zarr-python" }

		chunks = tuple(zarr_array.chunks)
		preferred_chunks = dict(zip(dimensions, chunks, strict=True))

-        chunks = tuple(zarr_array.chunks)
-        preferred_chunks = dict(zip(dimensions, chunks, strict=True))
+        chunk_grid = zarr_array.metadata.chunk_grid
+        if has_chunk_grid_support and isinstance(chunk_grid, RegularChunkGrid):
+            chunks = chunk_grid.chunk_shape
+            preferred_chunks = dict(zip(dimensions, chunks, strict=True))
+        elif has_chunk_grid_support:
+            # RectilinearChunkGrid or other non-regular grids — store the
+            # full chunk_grid and skip preferred_chunks since there's no
+            # single chunk size per dimension
+            chunks = chunk_grid
+            preferred_chunks = {}
+        else:
+            # Fallback for older zarr-python without chunk_grid support
+            chunks = tuple(zarr_array.chunks)
+            preferred_chunks = dict(zip(dimensions, chunks, strict=True))

	if zarr_format == 3 and has_variable_chunk_support:
	if zarr_format == 3 and has_chunk_grid_support:

Uh oh!

Conversation

keewis commented Oct 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

keewis Oct 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jhamman commented Oct 27, 2025

Uh oh!

keewis commented Oct 27, 2025

Uh oh!

TomNicholas Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

jhamman commented Jan 8, 2026

Uh oh!

keewis commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

keewis Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

keewis Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

maxrjones commented Feb 11, 2026

Uh oh!

keewis commented Feb 11, 2026

Uh oh!

maxrjones left a comment

Choose a reason for hiding this comment

Uh oh!

maxrjones Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

maxrjones Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

maxrjones Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

keewis commented Oct 27, 2025 •

edited

Loading

keewis Oct 27, 2025 •

edited

Loading

keewis commented Jan 8, 2026 •

edited

Loading

keewis Jan 8, 2026 •

edited

Loading