feat: unified chunk grid with rectilinear chunk/shard support by maxrjones · Pull Request #3802 · zarr-developers/zarr-python

maxrjones · 2026-03-21T04:31:27Z

Summary

This PR contains an alternative implementation of the rectilinear chunk grid extension, building on the work in #3534 (RLE helpers, validation logic, and test cases were directly adopted). While the core feature of variable-sized chunks is the same, the internal architecture differs in ways that impact extensibility, performance, and release safety.

I appreciate the patience of those who contributed to #3534, and everyone who's been waiting on this feature. I know it's frustrating to see a new PR after #3534 was so close. That PR provided fundamental components, and I hope people will see the value here. I really believe it is worth the churn for the following reasons:

Key differences from #3534

Extensibility. Each dimension is represented by a type implementing the DimensionGrid protocol (FixedDimension, VaryingDimension). Adding a new dimension type (e.g. TiledDimension for periodic patterns like days-per-month) requires implementing that protocol — no changes to indexing, codecs, or the ChunkGrid class. A prototype was built to verify this.
Performance. The indexing pipeline queries each dimension independently with scalar calls rather than constructing N-d coordinate tuples per chunk lookup. This avoids allocation overhead in the inner loop of every indexer. VaryingDimension uses precomputed prefix sums for O(log n) lookups via binary search. See https://github.com/maxrjones/zarr-chunk-grid-tests for a performance comparison.
Feature flag. Rectilinear chunk grids are gated behind zarr.config.set({'array.rectilinear_chunks': True}) (or ZARR_ARRAY__RECTILINEAR_CHUNKS=True), disabled by default. This gives downstream libraries time to adapt before the API is finalized, and us an opportunity to gracefully finalize the API.
Rectilinear sharding. Shard boundaries can be rectilinear while inner chunks remain regular, with validation that each shard edge is divisible by the inner chunk size. This is tested end-to-end and documented in the user guide.

Design document: docs/design/chunk-grid.md covers the full design, rationale, and a suggested PR sequence for splitting this into reviewable increments, if needed.

Downstream POCs (all passing):

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)

This reverts commit 9c0f582.

src/zarr/core/chunk_grids.py

src/zarr/core/metadata/v3.py

tests/test_unified_chunk_grid.py

…metadata (#7) * chore: simplify sharding codec validation against varying chunk grid metadata * test: restore test strength

docs/design/chunk-grid.md

…chunk grid (#8) * refactor: allow regular-style chunk grid declaration for rectilinear chunk grid The rectilinear chunk grid spec allows bare integers per dimension (meaning "regular step size"), distinct from explicit single-element edge lists. This commit widens `RectilinearChunkGrid.chunk_shapes` to `tuple[int | tuple[int, ...], ...]` so bare ints are preserved for faithful JSON round-tripping. Additionally: - unifies `_validate_chunk_shapes` to handle both regular and rectilinear validation; `_parse_chunk_shape` now delegates to it - adds `from_sizes` method to `ChunkGrid`, accepting `int | Sequence[int]` per dimension - removes `from_regular` and `from_rectilinear` methods from `ChunkGrid` - removes `parse_chunk_grid` from `chunk_grids.py` (JSON → ChunkGrid shortcut that bypassed the metadata layer) - removes `serialize_chunk_grid`, `_infer_chunk_grid_name`, and serialization helpers from `chunk_grids.py` (ChunkGrid never needs to be serialized; metadata DTOs handle it) - renames `parse_chunk_grid` in `v3.py` to `parse_chunk_grid_metadata` to disambiguate - moves the rectilinear feature flag to `RectilinearChunkGrid.__post_init__` - simplifies sharding codec validation into a single divisibility check for both regular and rectilinear grids - updates `validate_rectilinear_edges` to skip bare-int dimensions - refactors chunk grid tests to functional style with parametrization - adds docstrings to all test functions * chore: remove .claude * refactor: rename chunk_grid parsing function --------- Co-authored-by: Max Jones <14077947+maxrjones@users.noreply.github.com>

design/chunk-grid.md

docs/user-guide/config.md

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

src/zarr/core/chunk_grids.py

d-v-b · 2026-03-30T16:45:30Z

src/zarr/core/chunk_grids.py

+            if spec is not None:
+                yield spec
+
+    def all_chunk_coords(


imo the name all_chunk_coords doesn't convey that this is an iterator, but iter_chunk_coords would. if we need all_chunk_coords for backwards compat, can we deprecate that method and make it call iter_chunk_coords instead?

d-v-b · 2026-03-30T16:47:23Z

src/zarr/core/chunk_grids.py

+            if spec is not None:
+                yield spec.slices
+
+    def get_nchunks(self) -> int:


nit: i imagine this name is here for backwards compatibility, but nchunks (like ndim) feels like a tighter name.

src/zarr/core/chunk_grids.py

d-v-b · 2026-03-30T16:57:46Z

src/zarr/core/common.py


 BytesLike = bytes | bytearray | memoryview
 ShapeLike = Iterable[int | np.integer[Any]] | int | np.integer[Any]
+ChunksLike = ShapeLike | Sequence[Sequence[int]] | None


what would we lose if we used Iterable[Iterable[int]]? we'd have to call tuple before we could get a length, but we could also accept more inputs.

src/zarr/core/indexing.py

maxrjones · 2026-03-30T23:57:56Z

Thanks for the thorough reviews, @d-v-b!

I'll make DimensionGrid an internal implementation detail by making _dimensions private.

I'd prefer to address the naming/API comments on all_chunk_coords, get_nchunks, and the Iterable[Iterable[int]] type in common.py in follow-ups to keep the scope here focused because those all predate this PR.

maxrjones added 30 commits March 9, 2026 11:15

Add prospectus

4f63867

Initial prospectus POC

fb96207

V2 prospectus

b3e72ec

V3 prospectus

d6d551a

Fastforward POC to V3

8b8af74

Remove prospectus

784f4e7

Fix sharding

f1a1bc3

Fix bugs

30fa867

Support sequence in array functions

1282c7b

Add end-to-end tests

ea89f33

Collapse indexing paths

2eba460

Add DimensionGrid protocol

fa07396

Remove the try/except escape hatch from ChunkGrid.chunk_shape

a0acd95

Cache is_regular

ce0527d

Produce RLE directly

f433668

Fix bugs

02fd7c5

Separate chunk grid serialization

42ef639

Retain comments

55e720b

Update block and coordinate indexing

1b7871d

POC: TiledDimension

9c0f582

Support rectilinear shards

1e2fa97

Revert "POC: TiledDimension"

1f424b0

This reverts commit 9c0f582.

Fix __getitem__ for 1d chunk grids

0967e53

Implement resize

67a684d

Fix spec compliance

61c48a4

Fix .info

7d5ebb8

Fix typing

e74586a

Adopt joe's property testing strategy

8ab3ca8

Remove RegularChunkGrid

80d8280

Use none rather than sentinel value

ffc7805

d-v-b reviewed Mar 30, 2026

View reviewed changes

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

maxrjones added 3 commits March 30, 2026 09:18

Make chunk_grid property private

26b4760

Improve docstrings

5a7280b

Update notebook

60ad5cb

d-v-b reviewed Mar 30, 2026

View reviewed changes

src/zarr/core/metadata/v3.py Outdated Show resolved Hide resolved

d-v-b reviewed Mar 30, 2026

View reviewed changes

tests/test_unified_chunk_grid.py Outdated Show resolved Hide resolved

d-v-b reviewed Mar 30, 2026

View reviewed changes

tests/test_unified_chunk_grid.py Outdated Show resolved Hide resolved

d-v-b reviewed Mar 30, 2026

View reviewed changes

tests/test_unified_chunk_grid.py Outdated Show resolved Hide resolved

maxrjones and others added 2 commits March 30, 2026 09:47

Remove dead code

f2ec718

chore: simplify sharding codec validation against varying chunk grid …

2c06fb2

…metadata (#7) * chore: simplify sharding codec validation against varying chunk grid metadata * test: restore test strength

abarciauskas-bgse reviewed Mar 30, 2026

View reviewed changes

docs/design/chunk-grid.md Outdated Show resolved Hide resolved

d-v-b and others added 2 commits March 30, 2026 10:45

Fix typo

b3b5933

abarciauskas-bgse reviewed Mar 30, 2026

View reviewed changes

design/chunk-grid.md Show resolved Hide resolved

maxrjones added 4 commits March 30, 2026 10:59

Normalize

f80f798

Remove shim

3327cc1

Consistent typing

5642e03

Move design doc outside public docs

6a4c01a

d-v-b reviewed Mar 30, 2026

View reviewed changes

docs/user-guide/config.md Outdated Show resolved Hide resolved

Update config.md

e3ba71f

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

d-v-b reviewed Mar 30, 2026

View reviewed changes

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

d-v-b reviewed Mar 30, 2026

View reviewed changes

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

d-v-b reviewed Mar 30, 2026

View reviewed changes

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

maxrjones added 3 commits March 30, 2026 20:01

Remove get_chunk_shape

6667fda

Make dimensions private

0f0576e

Merge branch 'main' into poc/unified-chunk-grid

e1d57d6

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: unified chunk grid with rectilinear chunk/shard support#3802

feat: unified chunk grid with rectilinear chunk/shard support#3802
maxrjones wants to merge 123 commits intozarr-developers:mainfrom
maxrjones:poc/unified-chunk-grid

maxrjones commented Mar 21, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

d-v-b Mar 30, 2026

Uh oh!

d-v-b Mar 30, 2026

Uh oh!

Uh oh!

d-v-b Mar 30, 2026

Uh oh!

Uh oh!

maxrjones commented Mar 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

maxrjones commented Mar 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Key differences from #3534

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

d-v-b Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

d-v-b Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

d-v-b Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

maxrjones commented Mar 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

maxrjones commented Mar 21, 2026 •

edited

Loading