perf(chunk-grids): hold rectilinear chunk edges as run-length encoded runs

[d-v-b]

This PR fixes a scaling bug in our model of rectilinear chunk edges by avoiding premature materialization of each chunk declared by a run-length-encoded sequence. Fixing this amounts to making the VaryingDimension class dumber, which means we can actually remove some methods from that class, which is nice.

🤖 AI text below 🤖

Opening and indexing an array with a rectilinear chunk grid is now much faster and uses much less memory when the grid has long runs of equal-sized chunks. Previously the time and memory needed grew with the number of chunks, so an array with many millions of chunks could take seconds to open, or run out of memory.

Problem

With array.rectilinear_chunks enabled, a ~300 byte zarr.json declaring "chunk_shapes": [[[1, N]]] opened in 0.16 s for N=106, 1.6 s for N=107 and 4.8 s for N=3*107, at roughly 50 bytes per chunk; N=1012 got the process OOM-killed. RectilinearChunkGridMetadata.from_dict expanded every run into one int per chunk, validation and VaryingDimension walked and re-accumulated them, and to_dict re-compressed them.

Change

• RunLengthEdges (zarr.core.common): an immutable value stored as merged (size, count) runs, with per-run cumulative chunk counts and offsets, bisect-over-runs scalar lookups (size_of, offset_of, index_at) and a vectorized numpy lookup. It is deliberately not a sequence; expand() is the one method that yields every edge. parse_rle reads stored JSON into it without expanding; expand_rle and compress_rle remain, with the same error messages.
• VaryingDimension holds its edges as RunLengthEdges; construction, index_to_chunk, chunk_offset, chunk_size, data_size, indices_to_chunks, and resize cost O(runs).
• RectilinearChunkGridMetadata parses, validates, serializes and resizes per run. Serialized output is unchanged.
• The sharding codec's validate checks one edge size per run.
• zarr.core.metadata.repair is untouched: it works on the JSON document and never expanded edges.

After the change the document above opens, slices and fancy-indexes in about 5 ms with ~0.05 MB peak memory for N = 106, 3*107, 240 and 1012.

Behavior changes

• chunk_shapes[i] and VaryingDimension.edges are RunLengthEdges, not tuples: len, indexing, iteration and in raise TypeError, and == against a tuple is false. A RunLengthEdges equals only another RunLengthEdges with the same edges. tuple(edges.expand()) gives the per-chunk tuple. Existing tests that compared against nested tuples now compare the expansion or whole metadata objects.
• RectilinearChunkGridMetadata.chunk_shapes is annotated tuple[int | RunLengthEdges, ...]; the constructor still accepts tuples and lists of edges.
• VaryingDimension.cumulative is removed. No lookup used it any more, and chunk_offset(i) gives the start of chunk i from the runs.
• ChunkGrid.__repr__ and the "All edge lengths must be > 0" error show a dimension of more than 100 chunks in RLE form; shorter dimensions are shown as before.
• VaryingDimension.chunk_offset past the last grid cell raises IndexError with a different message, and a negative chunk index raises instead of returning 0.
• Removed because nothing in the library used them: with_extent, ngridcells and _unique_edge_lengths on both dimension types and the DimensionGrid protocol, @runtime_checkable on that protocol, and is_regular_1d / is_regular_nd. Tests that only exercised them are removed or rewritten against resize and edges.
• FixedDimension.index_to_chunk and VaryingDimension.index_to_chunk raise the same IndexError message.

Not addressed

• ChunkGrid.chunk_sizes (Array.read_chunk_sizes / write_chunk_sizes) returns one entry per chunk by contract and stays O(chunks).
• Array.resize to a smaller shape enumerates every dropped chunk, for regular grids as well.
• packages/zarr-indexing has its own VaryingDimension with per-chunk tuples and is unchanged.

Tests

• RunLengthEdges: one parametrized test against the expanded tuple as oracle, one test per error case, and one test that it is not a sequence.
• VaryingDimension with runs of 2**40 chunks, checked against closed-form values.
• test_open_rectilinear_huge_repeat_count (unsharded and sharded) opens, reads, writes, grows and re-serializes an array whose grid repeats one edge 2**40 times. It has no timing assertion: it can only pass if nothing expands the edges.
• The existing expand_rle error tests also run against parse_rle.

The agent ran prek run --all-files (passing, including mypy) and the full test suite locally with tests/test_store/test_fsspec.py deselected (12423 passed).

Author attestation

• I am a human, these are my changes, and I have reviewed and understood every change and can explain why each is correct.

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)

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 98.67550% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.75%. Comparing base (a53d51f) to head (1c1d0f3).

Files with missing lines	Patch %	Lines
src/zarr/core/chunk_grids.py	97.14%	1 Missing ⚠️
src/zarr/core/common.py	98.96%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4479      +/-   ##
==========================================
+ Coverage   94.69%   94.75%   +0.06%     
==========================================
  Files          94       94              
  Lines       13598    13649      +51     
==========================================
+ Hits        12876    12933      +57     
+ Misses        722      716       -6     

Files with missing lines	Coverage Δ
src/zarr/codecs/sharding.py	95.73% <ø> (ø)
src/zarr/core/metadata/v3.py	96.92% <100.00%> (+0.01%)	⬆️
src/zarr/testing/strategies.py	97.03% <100.00%> (ø)
src/zarr/core/chunk_grids.py	97.85% <97.14%> (+1.50%)	⬆️
src/zarr/core/common.py	94.65% <98.96%> (+2.30%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[d-v-b]

cc @maxrjones for visibility

[d-v-b]

this looks better @maxrjones

[maxrjones]

wow, reads perfect! thanks for fixing!