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.

maxrjones · 2026-03-30T02:25:29Z

I added a real-world example based on the data that @tinaok shared in #3534 at https://zarr--3802.org.readthedocs.build/en/3802/user-guide/examples/rectilinear_chunks/. I could adjust the source branches for the juv-compliant notebook before we merge.

@d-v-b would you have any time for a review on this PR?

src/zarr/core/array.py

src/zarr/core/chunk_grids.py

d-v-b · 2026-03-30T13:15:14Z

src/zarr/core/chunk_grids.py

    return _shards_out, _chunks_out
+
+
+class _RegularChunkGridMeta(type):


what's this for? I'm guessing it would not be necessary if we had a clear abstraction boundary between the in-memory ChunkGrid and the metadata chunk grids.

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

abarciauskas-bgse · 2026-03-30T14:43:48Z

docs/design/chunk-grid.md

 | Fully variable | Arbitrary per-chunk sizes | `[5, 12, 3, 20]` |

-A registry-based plugin system adds complexity without clear benefit.
+Prior iterations on the chunk grid design were based on the Zarr V3 spec's definition of chunk grids as an extension point alongside codecs, dtypes, etc. Therefore, we started designing the chunk grid implementation following a similar registry based approach. However, in practice chunk grids are fundamentally different than codecs. Codecs are independent; supporting `zstd` tells you nothing about `gzip`. Chunk grids are not: every regular grid is a valid rectilinear grid. A registry-based plugin system makes sense for codecs but adds complexity without clear benefit for chunk grids. Here we start from some basic goals and propose a more fitting design for supporting different chunk grids in zarr-python.


👍🏽

Nit: registry based --> registry-based

…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>

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

Improve design doc

f2ed1c4

github-actions bot removed the needs release notes Automatically applied to PRs which haven't added release notes label Mar 29, 2026

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

a857bac

github-actions bot added the needs release notes Automatically applied to PRs which haven't added release notes label Mar 29, 2026

maxrjones added 4 commits March 29, 2026 21:05

Add demo notebook

7318255

Add release notes

662ceef

Fix indexing empty slices

994e329

Add open support

f1c5182

github-actions bot removed the needs release notes Automatically applied to PRs which haven't added release notes label Mar 30, 2026

Improve release note

c4f7cf4