Feature: support rectilinear chunk grid extension #3534

jhamman · 2025-10-20T12:36:21Z

Summary

Adds support for RectilinearChunkGrid extension (Zarr v3), enabling arrays with variable chunk sizes per dimension, plus an enhanced type system for chunk specifications.

Closes: #1595 | Replaces: #1483 | Related: zarr-extensions#25

Key Features

1. RectilinearChunkGrid

arr = zarr.create_array(shape=(60, 100), chunks=[[10, 20, 30], [25, 25, 25, 25]], zarr_format=3)

Zarr v3 only (not compatible with v2, sharding, or from_array())
Supports RLE for chunks in JSON: [[10, 6]] = 6 chunks of size 10
Stored internally in expanded format for fast indexing

2. Enhanced Chunk Types

chunks = arr.chunks  # Returns RegularChunks or RectilinearChunks
isinstance(chunks, RegularChunks)  # Type-safe checking

# perhaps mildly controversial?
chunks.lat           # Named access (if dimension_names set)
dict(chunks.items()) # Dict conversion: {'lat': 10, 'lon': 20}

Design Decisions for `ChunksType`

ABC-based hierarchy: ChunksType is an ABC (not type alias) → enables isinstance() checks and extensibility
Tuple inheritance: Both types inherit from tuple → full backward compatibility, all tuple ops work
Type narrowing: Use isinstance(chunks, RegularChunks) instead of isinstance(chunks[0], int) → cleaner, type-safe
V2 consistency: ArrayV2Metadata.chunks returns RegularChunks (not plain tuple) → consistent API across formats

Review Focus Areas

High Priority

chunk_grids.py: ChunksType ABC, RegularChunks/RectilinearChunks implementations, RLE expansion
metadata/v3.py: chunks property returns correct type based on chunk grid
metadata/v2.py: _chunks_tuple field + cached_property, _replace_with_field_mapping() helper
indexing.py: Rectilinear indexing logic, edge cases

Breaking Changes

None. Fully backward compatible.

TODO:

Add unit tests and/or doctests in docstrings
Add docstrings and API docs for any new/modified user-facing classes and functions
New/modified features documented in docs/user-guide/*.md
Changes documented as a new file in changes/
GitHub Actions have all passed
Test coverage is 100% (Codecov passes)

…ature/rectilinear-chunk-grid

codecov · 2025-10-20T12:47:43Z

Codecov Report

❌ Patch coverage is 69.20943% with 222 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.13%. Comparing base (1c05c1a) to head (9024817).

Files with missing lines	Patch %	Lines
src/zarr/core/chunk_grids.py	57.75%	169 Missing ⚠️
src/zarr/core/indexing.py	85.07%	20 Missing ⚠️
src/zarr/core/array.py	79.36%	13 Missing ⚠️
src/zarr/testing/strategies.py	89.28%	9 Missing ⚠️
src/zarr/core/metadata/v2.py	69.56%	7 Missing ⚠️
src/zarr/core/metadata/v3.py	85.71%	2 Missing ⚠️
src/zarr/__init__.py	0.00%	1 Missing ⚠️
src/zarr/core/_info.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3534      +/-   ##
==========================================
+ Coverage   60.89%   61.13%   +0.23%     
==========================================
  Files          86       86              
  Lines       10182    10767     +585     
==========================================
+ Hits         6200     6582     +382     
- Misses       3982     4185     +203

Files with missing lines	Coverage Δ
src/zarr/api/asynchronous.py	`72.20% <100.00%> (ø)`
src/zarr/api/synchronous.py	`36.61% <ø> (ø)`
src/zarr/core/group.py	`70.27% <ø> (ø)`
src/zarr/__init__.py	`63.04% <0.00%> (-1.41%)`	⬇️
src/zarr/core/_info.py	`51.80% <0.00%> (ø)`
src/zarr/core/metadata/v3.py	`60.08% <85.71%> (+2.07%)`	⬆️
src/zarr/core/metadata/v2.py	`60.31% <69.56%> (+2.17%)`	⬆️
src/zarr/testing/strategies.py	`94.18% <89.28%> (-3.66%)`	⬇️
src/zarr/core/array.py	`68.63% <79.36%> (+0.15%)`	⬆️
src/zarr/core/indexing.py	`70.19% <85.07%> (+0.73%)`	⬆️
... and 1 more

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…ature/rectilinear-chunk-grid

docs/user-guide/arrays.md

d-v-b · 2025-10-20T17:55:16Z

docs/user-guide/arrays.md

+With variable chunking, the standard `.chunks` property is not available since chunks
+have different sizes. Instead, access chunk information through the chunk grid:


I think it would be better if .chunks just had a different type (tuple of tuples of ints)

d-v-b · 2025-10-20T21:28:48Z

src/zarr/core/chunk_grids.py

+
+
+@dataclass(frozen=True)
+class RectilinearChunkGrid(ChunkGrid):


thoughts on just calling this class Rectilinear, and renaming the RegularChunkGrid to Regular? We could keep around a RegularChunkGrid class for compatibility. But I feel like people know these are chunk grids when they import them

50/50. I think the more descriptive class is useful when looking at a tracebacks. Plus, this is currently in .core so its not meant to be used directly by users.

src/zarr/core/chunk_grids.py

d-v-b · 2025-10-20T21:33:33Z

src/zarr/core/chunk_grids.py

+                )
+
+    @cached_property
+    def _cumulative_sizes(self) -> tuple[tuple[int, ...], ...]:


nice, this is cached

src/zarr/core/indexing.py

dcherian · 2025-10-27T21:31:42Z

src/zarr/testing/strategies.py

+        chunk_shapes_rle = [
+            [[c, r] for c, r in zip(draw(dim_chunks), draw(repeats), strict=True)]
+            for _ in range(ndim)
+        ]
+        return RectilinearChunkGrid(chunk_shapes=chunk_shapes_rle)


This needs fixing

dcherian · 2025-10-27T21:32:15Z

tests/test_properties.py

-@given(data=st.data())
-async def test_basic_indexing(data: st.DataObject) -> None:
-    zarray = data.draw(simple_arrays())
+@given(data=st.data(), zarray=st.one_of([simple_arrays(), complex_chunked_arrays()]))


Because the search space for the standard arrays strategy is so large, i made a different one complex_chunked_arrays that purely checks different chunk grids
with simple_arrays() we are only spending 10% of our time trying RectilinearChunkGrid so using this approach. We should boost number of examples too.

shoyer · 2025-10-28T03:11:15Z

docs/user-guide/arrays.md

+2. **Not compatible with sharding**: You cannot use variable chunking together with
+   the sharding feature. Arrays must use either variable chunking or sharding, but not both.


I hope this is a temporary limitation! There's a natural extension of rectilinear chunk grids to rectilinear shard grids.

xref: zarr-developers/zarr-extensions#34

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

…arr-python into feature/rectilinear-chunk-grid

…ature/rectilinear-chunk-grid

jhamman · 2026-01-08T18:22:30Z

@d-v-b, @maxrjones, et al - this is ready for a real review. I'd be happy if we merged this in using an "experimental" warning of some kind and let it bake from 3.2->3.3. Suggestions on the best way to communicate the experimental nature of this are welcome.

keewis · 2026-01-08T21:13:23Z

is it intentional that the chunk grid is RectilinearChunkGrid when explicitly enumerating regular chunk sizes?

import zarr

store = zarr.storage.MemoryStore()
arr = zarr.create_array(
    store,
    shape=(60, 100),
    chunks=[[20, 20, 20], [25, 25, 25, 25]],
    zarr_format=3,
    dtype="uint8"
)
arr.metadata.chunk_grid # RectilinearChunkGrid

Or is RegularChunkGrid only chosen if I specify a tuple of values, in this case (20, 25)?

Not sure how much this matters in practice (xarray's code does not appear to be affected by the difference), but that's what dask normalizes the chunks to.

jhamman · 2026-01-14T16:58:17Z

is it intentional that the chunk grid is RectilinearChunkGrid when explicitly enumerating regular chunk sizes?

@keewis - yes, this is intentional. We decided this was preferable because it will allow the user to add variable lengths later (through an extend/append workflow).

codspeed-hq · 2026-01-16T18:43:13Z

CodSpeed Performance Report

Merging this PR will not alter performance

_{Comparing jhamman:feature/rectilinear-chunk-grid (665847a) with main (e94504e)}

Summary

✅ 30 untouched benchmarks

d-v-b · 2026-01-26T18:46:52Z

src/zarr/core/chunk_grids.py

+
+    @property
+    @abstractmethod
+    def dimension_names(self) -> tuple[str | None, ...] | None:


where is this used? I feel like integer indexing should be fine for the chunk grid.

d-v-b · 2026-01-26T18:52:41Z

src/zarr/core/chunk_grids.py

+    chunk_shapes: Sequence[Sequence[ChunkEdgeLength]]
+
+
+class ChunksType(tuple[Any, ...], ABC):  # noqa: SLOT001


I'm not sold on subclassing tuple here. It's going to lead to a lot of surprises unless we carefully implement a lot of stuff that we could just use tuple for:

>>> RegularChunks((10,10,1))[:2] (10, 10) # tuple, not RegularChunks >>> RegularChunks((10,10,1)) + (1,) (10, 10, 1, 1) # tuple, not RegularChunks

I would much rather make a breaking change than introduce something super confusing

d-v-b · 2026-01-26T18:53:13Z

src/zarr/core/chunk_grids.py

+    @abstractmethod
+    def keys(self) -> Iterator[str]:
+        """Return dimension names, enabling dict(chunks) conversion."""
+        ...
+
+    @abstractmethod
+    def values(self) -> Iterator[Any]:
+        """Return chunk values, enabling dict(chunks) conversion."""
+        ...
+
+    @abstractmethod
+    def items(self) -> Iterator[tuple[str, Any]]:
+        """Return (name, value) pairs, enabling dict(chunks) conversion."""
+        ...


Why do we need key: value access?

d-v-b · 2026-01-26T19:25:47Z

src/zarr/core/chunk_grids.py

+
+
+@dataclass(frozen=True)
+class ResolvedChunkSpec:


does this need to be a class? It has no methods. Seems like either a TypedDict or just a tuple would work

d-v-b · 2026-01-26T19:26:57Z

src/zarr/core/chunk_grids.py

+    shards: tuple[int, ...] | None
+
+
+def _validate_zarr_format_compatibility(


why do we need this here? Can't we check when we construct the array metadata document?

d-v-b · 2026-01-26T19:29:41Z

src/zarr/core/chunk_grids.py

+            )
+
+
+def _validate_sharding_compatibility(


shouldn't this be defined over in the array v3 metadata module, and used there to check that the array metadata document is valid?

d-v-b · 2026-01-26T19:31:26Z

src/zarr/core/chunk_grids.py

+            )
+
+
+def _validate_data_compatibility(


does this need to be stand-alone function? are we going to use this logic anywhere other than it's current usage site (resolve_chunk_spec)

d-v-b · 2026-01-26T19:40:32Z

src/zarr/core/metadata/v3.py


    @property
-    def chunks(self) -> tuple[int, ...]:
+    def chunks(self) -> ChunksType:


I'd be willing to drop this attribute entirely if it's too complicated to support in a clean way.

d-v-b · 2026-01-26T19:51:50Z

src/zarr/core/metadata/v3.py

            out_dict["data_type"] = dtype_meta.to_json(zarr_format=3)  # type: ignore[unreachable]
        return out_dict

    def update_shape(self, shape: tuple[int, ...]) -> Self:


IMO this needs to take a parameter that can define the new chunks (which can default to None or some other sentinel flagging "default behavior" semantics)

jhamman added 5 commits October 16, 2025 13:32

feature(chunkgrids): add rectillinear chunk grid metadata support

0fd633f

implementation (wip)

31c831d

fixup

d40dd30

fixup types

5f84198

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

b96a650

…ature/rectilinear-chunk-grid

github-actions bot added the needs release notes Automatically applied to PRs which haven't added release notes label Oct 20, 2025

jhamman added 3 commits October 20, 2025 09:42

more tests

6deb98e

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

fc07af1

…ature/rectilinear-chunk-grid

docs

527f7cf

github-actions bot removed the needs release notes Automatically applied to PRs which haven't added release notes label Oct 20, 2025

fixups

cea5fc2

d-v-b reviewed Oct 20, 2025

View reviewed changes

docs/user-guide/arrays.md Show resolved Hide resolved

d-v-b reviewed Oct 20, 2025

View reviewed changes

src/zarr/core/chunk_grids.py Show resolved Hide resolved

d-v-b reviewed Oct 20, 2025

View reviewed changes

src/zarr/core/chunk_grids.py

)

@cached_property

def _cumulative_sizes(self) -> tuple[tuple[int, ...], ...]:

Copy link

Contributor

d-v-b Oct 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice, this is cached

d-v-b reviewed Oct 20, 2025

View reviewed changes

src/zarr/core/indexing.py Show resolved Hide resolved

keewis mentioned this pull request Oct 27, 2025

variable-sized chunks with zarr v3 pydata/xarray#10880

Draft

3 tasks

dcherian added 2 commits October 27, 2025 14:34

Updated strategies

cef873f

better approach

79ecee9

dcherian reviewed Oct 27, 2025

View reviewed changes

shoyer reviewed Oct 28, 2025

View reviewed changes

jhamman and others added 6 commits October 28, 2025 09:57

handle run length encoded chunk grid spec in the top level api

e116f63

fixups

6eab392

fixups + support partial edge chunks / resize

1935d19

Update docs/user-guide/arrays.md

4f36ffb

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

Merge branch 'main' into feature/rectilinear-chunk-grid

4c33eac

fixups + support partial edge chunks / resize

99fd0a8

jhamman added 5 commits December 20, 2025 16:14

Merge branch 'feature/rectilinear-chunk-grid' of github.com:jhamman/z…

db3253d

…arr-python into feature/rectilinear-chunk-grid

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

3187229

…ature/rectilinear-chunk-grid

fixups

77fb91a

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

d3a8997

…ature/rectilinear-chunk-grid

fixup mypy

b55cb8e

jhamman marked this pull request as ready for review December 30, 2025 21:41

jhamman and others added 5 commits January 4, 2026 12:21

special type for array.chunks

372704e

Merge branch 'main' into feature/rectilinear-chunk-grid

41d4a91

special type for array.chunks

d4e4d2b

fixup mypy

38c8b4f

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

376d97b

…ature/rectilinear-chunk-grid

jhamman requested review from d-v-b and dcherian January 8, 2026 18:52

Merge branch 'main' into feature/rectilinear-chunk-grid

9024817

doc notes about this being experimental until 3.3

84f47d3

Merge branch 'main' into feature/rectilinear-chunk-grid

665847a

d-v-b reviewed Jan 26, 2026

View reviewed changes

		With variable chunking, the standard `.chunks` property is not available since chunks
		have different sizes. Instead, access chunk information through the chunk grid:



		@dataclass(frozen=True)
		class RectilinearChunkGrid(ChunkGrid):

		2. Not compatible with sharding: You cannot use variable chunking together with
		the sharding feature. Arrays must use either variable chunking or sharding, but not both.

		chunk_shapes: Sequence[Sequence[ChunkEdgeLength]]


		class ChunksType(tuple[Any, ...], ABC): # noqa: SLOT001

		shards: tuple[int, ...] \| None


		def _validate_zarr_format_compatibility(

Uh oh!

Feature: support rectilinear chunk grid extension #3534

Are you sure you want to change the base?

Feature: support rectilinear chunk grid extension #3534

Conversation

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

Uh oh!

Summary

Key Features

1. RectilinearChunkGrid

2. Enhanced Chunk Types

Design Decisions for ChunksType

Review Focus Areas

Breaking Changes

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcherian 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!

Choose a reason for hiding this comment

Uh oh!

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!

jhamman commented Jan 14, 2026

Uh oh!

codspeed-hq bot commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will not alter performance

Summary

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

jhamman commented Oct 20, 2025 •

edited

Loading

Design Decisions for `ChunksType`

codecov bot commented Oct 20, 2025 •

edited

Loading

dcherian Oct 27, 2025 •

edited

Loading

keewis commented Jan 8, 2026 •

edited

Loading

codspeed-hq bot commented Jan 16, 2026 •

edited

Loading