refactor(chunk-grids): make RunLengthEdges a plain value, not a sequence

Per-chunk access belongs to `VaryingDimension`, which binds edges to an
extent; the metadata model only needs runs, their total and the RLE
form. A `Sequence[int]` interface on `RunLengthEdges` invited
`for e in edges` or `list(edges)` on a grid with 2**40 chunks.

Drop `len`, indexing, slicing, iteration, membership and `count`. Add
`size_of(index)` for one edge and `expand()` as the single, explicit way
to visit every edge. `RectilinearChunkGridMetadata.chunk_shapes` is now
annotated as what it holds, `tuple[int | RunLengthEdges, ...]`, while its
constructor still accepts tuples and lists of edges.

Assisted-by: ClaudeCode:claude-fable-5.1

Author: d-v-b

changes/4479.bugfix.md

@@ -1 +1 @@
-Opening an array with a rectilinear chunk grid, and indexing it, now costs time and memory proportional to the size of the grid's run-length encoded metadata rather than to its number of chunks: a stored grid such as `[[1, 1000000000000]]`, a few hundred bytes of JSON, previously took seconds per ten million chunks to open and could exhaust memory. The explicit edges in `RectilinearChunkGridMetadata.chunk_shapes` and `VaryingDimension.edges` are now held as `zarr.core.common.RunLengthEdges`, an immutable sequence of the edges that is not a `tuple` and does not compare equal to one: use `tuple(edges)` to compare against a tuple. `VaryingDimension.cumulative` is now computed on access. The repr of a chunk grid shows a dimension of more than 100 chunks in its run-length encoded form.
+Opening an array with a rectilinear chunk grid, and indexing it, now costs time and memory proportional to the size of the grid's run-length encoded metadata rather than to its number of chunks: a stored grid such as `[[1, 1000000000000]]`, a few hundred bytes of JSON, previously took seconds per ten million chunks to open and could exhaust memory. The explicit edges in `RectilinearChunkGridMetadata.chunk_shapes` and `VaryingDimension.edges` are now held as `zarr.core.common.RunLengthEdges`, an immutable run-length encoded value that is not a sequence: it has no `len`, indexing or iteration and does not compare equal to a `tuple`. Its `expand()` method yields one entry per chunk. `VaryingDimension.cumulative` is now computed on access. The repr of a chunk grid shows a dimension of more than 100 chunks in its run-length encoded form.

src/zarr/codecs/sharding.py

@@ -45,7 +45,6 @@
     merge_and_encode_chunk,
 )
 from zarr.core.common import (
-    RunLengthEdges,
     ShapeLike,
     parse_named_configuration,
     parse_shapelike,
@@ -623,8 +622,7 @@ def validate(
         elif isinstance(chunk_grid, RectilinearChunkGridMetadata):
             # One size per run of equal edges, not one per chunk.
             edges_per_dim = tuple(
-                (s,) if isinstance(s, int) else RunLengthEdges.from_edges(s).sizes
-                for s in chunk_grid.chunk_shapes
+                (s,) if isinstance(s, int) else s.sizes for s in chunk_grid.chunk_shapes
             )
         else:
             raise TypeError(

src/zarr/core/chunk_grids.py

@@ -138,7 +138,7 @@ def _size_repr(self) -> str:
 def _edges_repr(edges: RunLengthEdges) -> str:
     """The edges as a tuple, or run-length encoded if there are too many to write out."""
     if edges.num_edges <= _MAX_EDGES_IN_REPR:
-        return repr(tuple(edges))
+        return repr(tuple(edges.expand()))
     return repr(edges.to_rle())
 
 
@@ -161,7 +161,7 @@ class VaryingDimension:
     nchunks: int = field(init=False, repr=False)  # cached at construction
     ngridcells: int = field(init=False, repr=False)  # cached at construction
 
-    def __init__(self, edges: Sequence[int], extent: int) -> None:
+    def __init__(self, edges: Sequence[int] | RunLengthEdges, extent: int) -> None:
         rle = RunLengthEdges.from_edges(edges)
         if rle.num_edges == 0:
             raise ValueError("VaryingDimension edges must not be empty")
@@ -185,7 +185,7 @@ def cumulative(self) -> tuple[int, ...]:
         Lookups do not use this: it is computed on each access, in time and
         memory proportional to the number of chunks.
         """
-        return tuple(itertools.accumulate(self.edges))
+        return tuple(itertools.accumulate(self.edges.expand()))
 
     def index_to_chunk(self, idx: int) -> int:
         if idx < 0 or idx >= self.extent:
@@ -206,23 +206,23 @@ def chunk_size(self, chunk_ix: int) -> int:
         Does not validate *chunk_ix* — callers must ensure it is in
         ``[0, ngridcells)``. Use ``ChunkGrid.__getitem__`` for safe access.
         """
-        return self.edges[chunk_ix]
+        return self.edges.size_of(chunk_ix)
 
     def data_size(self, chunk_ix: int) -> int:
         """Valid data region within the buffer — clipped at extent.
 
         Does not validate *chunk_ix* — callers must ensure it is in
         ``[0, ngridcells)``. Use ``ChunkGrid.__getitem__`` for safe access.
         """
-        return max(0, min(self.edges[chunk_ix], self.extent - self.chunk_offset(chunk_ix)))
+        return max(0, min(self.edges.size_of(chunk_ix), self.extent - self.chunk_offset(chunk_ix)))
 
     @property
     def _data_sizes(self) -> tuple[int, ...]:
         """`data_size` of every chunk that holds data, in order."""
         if self.nchunks == 0:
             return ()
         last = self.nchunks - 1
-        return (*itertools.islice(self.edges, last), self.data_size(last))
+        return (*itertools.islice(self.edges.expand(), last), self.data_size(last))
 
     @property
     def _unique_edge_lengths(self) -> Iterable[int]:
@@ -447,7 +447,7 @@ def from_metadata(cls, metadata: ArrayMetadata) -> ChunkGrid:
     def from_sizes(
         cls,
         array_shape: ShapeLike,
-        chunk_sizes: Sequence[int | Sequence[int]],
+        chunk_sizes: Sequence[int | Sequence[int] | RunLengthEdges],
     ) -> ChunkGrid:
         """Create a ChunkGrid from per-dimension chunk size specifications.
 
@@ -459,7 +459,7 @@ def from_sizes(
             Per-dimension chunk sizes. Each element is either:
 
             - An ``int`` — regular (fixed) chunk size for that dimension.
-            - A ``Sequence[int]`` — explicit per-chunk edge lengths, kept as a
+            - A ``Sequence[int]`` or a `RunLengthEdges` — explicit per-chunk edge lengths, kept as a
               ``VaryingDimension`` even when the edges are uniform: the spec
               form declares the grid kind, as in `normalize_chunks_1d`.
         """

src/zarr/core/common.py

@@ -27,7 +27,6 @@
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Iterator
-    from typing import Self
 
     import numpy.typing as npt
 
@@ -325,20 +324,22 @@ def parse_chunk_shape(data: object) -> tuple[int, ...]:
     raise TypeError(f"A chunk shape must be an iterable of chunk edge lengths, got {data!r}")
 
 
-class RunLengthEdges(Sequence[int]):
-    """An immutable sequence of chunk edge lengths, stored run-length encoded.
+class RunLengthEdges:
+    """The chunk edge lengths of one dimension, stored run-length encoded.
 
-    The sequence is held as `(size, count)` runs, with adjacent runs of the same size
+    The edges are held as `(size, count)` runs, with adjacent runs of the same size
     merged, so construction, lookups and prefix sums cost time and memory in the number
-    of runs, not in the number of edges: `RunLengthEdges([(1, 2**40)])` is one run. It
-    reads like the sequence of edges it stands for (`len`, indexing, iteration), but
-    only iteration visits every edge.
+    of runs, not in the number of edges: `RunLengthEdges([(1, 2**40)])` is one run.
+
+    This is deliberately not a sequence: there is no `len`, indexing or iteration, so
+    nothing visits every edge by accident. `size_of`, `offset_of` and `index_at` answer
+    per-edge questions from the runs, and `expand` is the one method that yields every
+    edge.
 
     Sizes are not checked here: what a valid edge length is, and how to report an invalid
     one, is up to the caller (see `parse_chunk_edge`). Each count must be at least 1.
 
-    It is not a `tuple` and equals only another `RunLengthEdges` with the same edges;
-    compare `tuple(edges)` against a tuple.
+    Two instances are equal if they encode the same edges.
     """
 
     __slots__ = ("_index_stops", "_lookup_tables", "_offset_stops", "counts", "sizes")
@@ -369,10 +370,10 @@ def __init__(self, runs: Iterable[tuple[int, int]] = ()) -> None:
         self._lookup_tables: tuple[npt.NDArray[np.intp], ...] | None = None
 
     @classmethod
-    def from_edges(cls, edges: Iterable[int]) -> Self:
-        """Run-length encode `edges`, one entry per edge. An instance of this class is
-        returned as it is."""
-        if isinstance(edges, cls):
+    def from_edges(cls, edges: Iterable[int] | RunLengthEdges) -> RunLengthEdges:
+        """Run-length encode `edges`, one entry per edge. A `RunLengthEdges` is returned
+        as it is."""
+        if isinstance(edges, RunLengthEdges):
             return edges
         return cls((size, sum(1 for _ in group)) for size, group in itertools.groupby(edges))
 
@@ -383,7 +384,7 @@ def runs(self) -> tuple[tuple[int, int], ...]:
 
     @property
     def num_edges(self) -> int:
-        """The number of edges. Unlike `len`, this is not limited to `sys.maxsize`."""
+        """The number of edges."""
         return self._index_stops[-1] if self._index_stops else 0
 
     @property
@@ -394,37 +395,15 @@ def total(self) -> int:
     def __reduce__(self) -> tuple[type[RunLengthEdges], tuple[tuple[tuple[int, int], ...]]]:
         return type(self), (self.runs,)
 
-    def __len__(self) -> int:
-        return self.num_edges
-
-    def __iter__(self) -> Iterator[int]:
+    def expand(self) -> Iterator[int]:
+        """Every edge, in order: as many items as there are edges, however few the runs."""
         return itertools.chain.from_iterable(map(itertools.repeat, self.sizes, self.counts))
 
-    def __contains__(self, value: object) -> bool:
-        return value in self.sizes
-
-    @overload
-    def __getitem__(self, index: int) -> int: ...
-    @overload
-    def __getitem__(self, index: slice) -> RunLengthEdges: ...
-    def __getitem__(self, index: int | slice) -> int | RunLengthEdges:
-        num_edges = self.num_edges
-        if isinstance(index, slice):
-            start, stop, step = index.indices(num_edges)
-            if step != 1:
-                return RunLengthEdges.from_edges(self[i] for i in range(start, stop, step))
-            runs: list[tuple[int, int]] = []
-            run_start = 0
-            for size, run_stop in zip(self.sizes, self._index_stops, strict=True):
-                count = min(stop, run_stop) - max(start, run_start)
-                if count > 0:
-                    runs.append((size, count))
-                run_start = run_stop
-            return RunLengthEdges(runs)
-        position = index + num_edges if index < 0 else index
-        if not 0 <= position < num_edges:
-            raise IndexError(f"Edge index {index} is out of range for {num_edges} edges")
-        return self.sizes[bisect.bisect_right(self._index_stops, position)]
+    def size_of(self, index: int) -> int:
+        """The length of edge `index`."""
+        if not 0 <= index < self.num_edges:
+            raise IndexError(f"Edge index {index} is out of range for {self.num_edges} edges")
+        return self.sizes[bisect.bisect_right(self._index_stops, index)]
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, RunLengthEdges):
@@ -437,9 +416,6 @@ def __hash__(self) -> int:
     def __repr__(self) -> str:
         return f"{type(self).__name__}({self.to_rle()})"
 
-    def count(self, value: object) -> int:
-        return sum(count for size, count in self.runs if size == value)
-
     def to_rle(self) -> list[int | list[int]]:
         """The mixed run-length encoding of the rectilinear chunk grid spec: a run of
         one edge is its bare size, a longer run is `[size, count]`."""
@@ -569,7 +545,7 @@ def validate_rectilinear_kind(kind: str | None) -> None:
 
 
 def validate_rectilinear_edges(
-    chunk_shapes: Sequence[int | Sequence[int]], array_shape: Sequence[int]
+    chunk_shapes: Sequence[int | Sequence[int] | RunLengthEdges], array_shape: Sequence[int]
 ) -> None:
     """Validate that rectilinear chunk edges cover the array extent per dimension.
 

src/zarr/core/metadata/v3.py

@@ -245,7 +245,7 @@ class RectilinearChunkGridMetadataConfig(TypedDict):
 
 
 def _validate_chunk_shapes(
-    chunk_shapes: Sequence[int | Sequence[int]],
+    chunk_shapes: Sequence[int | Sequence[int] | RunLengthEdges],
 ) -> tuple[int | RunLengthEdges, ...]:
     """Validate per-dimension chunk specifications.
 
@@ -304,29 +304,30 @@ def from_dict(cls, data: RegularChunkGridMetadataJSON) -> Self:  # type: ignore[
         return cls(chunk_shape=parse_chunk_shape(configuration["chunk_shape"]))
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(frozen=True, kw_only=True, init=False)
 class RectilinearChunkGridMetadata(Metadata):
     """Metadata-only description of a rectilinear chunk grid.
 
     Each element of ``chunk_shapes`` is either:
 
     - A bare ``int`` — a regular step size that repeats to cover the axis
       (the spec's single-integer shorthand).
-    - A sequence of ``int`` — explicit per-chunk edge lengths. It is given
-      as a ``tuple`` or a ``list`` and held as a `RunLengthEdges`, which
-      reads like the sequence of edges (but is not a ``tuple``, and does
-      not compare equal to one) and stores one entry per run of equal edges, so a grid costs time and memory in
-      the size of its run-length encoded form, not in its number of chunks.
+    - A `RunLengthEdges` — explicit per-chunk edge lengths. They may be given
+      as a ``tuple`` or a ``list`` of ``int``, and are held run-length
+      encoded, one entry per run of equal edges, so a grid costs time and
+      memory in the size of its run-length encoded form, not in its number
+      of chunks. A `RunLengthEdges` is not a sequence: use its ``expand``
+      method for one entry per chunk.
 
     This distinction matters for faithful round-tripping: a bare int
-    serializes back as a bare int, while a single-element tuple serializes
+    serializes back as a bare int, while a single explicit edge serializes
     as a list.
     """
 
-    chunk_shapes: tuple[int | Sequence[int], ...]
+    chunk_shapes: tuple[int | RunLengthEdges, ...]
 
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "chunk_shapes", _validate_chunk_shapes(self.chunk_shapes))
+    def __init__(self, *, chunk_shapes: Sequence[int | Sequence[int] | RunLengthEdges]) -> None:
+        object.__setattr__(self, "chunk_shapes", _validate_chunk_shapes(chunk_shapes))
 
     @property
     def ndim(self) -> int:
@@ -341,7 +342,7 @@ def to_dict(self) -> RectilinearChunkGridMetadataJSON:  # type: ignore[override]
             else:
                 # A run of one edge is written as its bare size, so edges without
                 # repeats come out as a plain list.
-                serialized_dims.append(RunLengthEdges.from_edges(dim_spec).to_rle())
+                serialized_dims.append(dim_spec.to_rle())
         return {
             "name": "rectilinear",
             "configuration": {
@@ -361,17 +362,16 @@ def update_shape(
           Otherwise edges are kept as-is (the spec allows trailing edges
           beyond the array extent).
         """
-        new_chunk_shapes: list[int | Sequence[int]] = []
+        new_chunk_shapes: list[int | RunLengthEdges] = []
         for dim_spec, new_ext in zip(self.chunk_shapes, new_shape, strict=True):
             if isinstance(dim_spec, int):
                 # Bare int covers any extent — no change needed
                 new_chunk_shapes.append(dim_spec)
             else:
-                edges = RunLengthEdges.from_edges(dim_spec)
-                if new_ext > edges.total:
-                    new_chunk_shapes.append(edges.with_edge(new_ext - edges.total))
+                if new_ext > dim_spec.total:
+                    new_chunk_shapes.append(dim_spec.with_edge(new_ext - dim_spec.total))
                 else:
-                    new_chunk_shapes.append(edges)
+                    new_chunk_shapes.append(dim_spec)
         return RectilinearChunkGridMetadata(chunk_shapes=tuple(new_chunk_shapes))
 
     @classmethod
@@ -443,7 +443,7 @@ def create_chunk_grid_metadata(
         return RegularChunkGridMetadata(chunk_shape=chunks.chunk_shape)
     # Uniform dimensions stay bare ints — the rectilinear grid spec treats
     # a bare int as a step size repeating to cover the axis.
-    chunk_shapes: list[int | Sequence[int]] = []
+    chunk_shapes: list[int | RunLengthEdges] = []
     for dim in chunks.dimensions:
         if isinstance(dim, FixedDimension):
             chunk_shapes.append(dim.size)

src/zarr/testing/strategies.py

@@ -553,12 +553,12 @@ def chunks_param_from_rectilinear(
 ) -> list[int | list[int]]:
     """Convert rectilinear chunk grid metadata into a `chunks=` argument.
 
-    Explicit edge sequences become lists. Bare ints — the spec's step-size
+    Explicit edges become lists. Bare ints — the spec's step-size
     shorthand meaning "repeat to cover the axis" — pass through unchanged;
     wrapping one in a single-element list would instead declare exactly one
     chunk, which fails normalization whenever the axis needs more than one.
     """
-    return [dim if isinstance(dim, int) else list(dim) for dim in meta.chunk_shapes]
+    return [dim if isinstance(dim, int) else list(dim.expand()) for dim in meta.chunk_shapes]
 
 
 # The most chunks a drawn rectilinear grid declares along one axis, and over

tests/test_chunk_grids.py

@@ -44,7 +44,9 @@ def _assert_chunks_equal(
             assert isinstance(a, VaryingDimension), (
                 f"axis {axis}: expected VaryingDimension, got {a!r}"
             )
-            assert tuple(a.edges) == tuple(e), f"axis {axis}: edges {a.edges} != {tuple(e)}"
+            assert tuple(a.edges.expand()) == tuple(e), (
+                f"axis {axis}: edges {a.edges} != {tuple(e)}"
+            )
         assert a.extent == span, f"axis {axis}: extent {a.extent} != {span}"
 
 

tests/test_metadata/test_repair.py

@@ -423,7 +423,7 @@ def _rectilinear_from_dict(chunk_shapes: list[Any]) -> RectilinearChunkGridMetad
 def _second_edge(grid: RectilinearChunkGridMetadata) -> int:
     edges = grid.chunk_shapes[0]
     assert not isinstance(edges, int)
-    return edges[1]
+    return edges.size_of(1)
 
 
 CHUNK_EDGE_SITES: dict[str, Callable[[Any], object]] = {

tests/test_properties.py

@@ -2,7 +2,7 @@
 import json
 import numbers
 import warnings
-from collections.abc import Generator
+from collections.abc import Generator, Sequence
 from typing import Any
 
 import numpy as np
@@ -306,9 +306,11 @@ def test_block_indexing(data: st.DataObject) -> None:
     # The block grid is worked out from the stored declaration, not by zarr's grid code.
     assert isinstance(zarray.metadata, ArrayV3Metadata)
     grid = zarray.metadata.chunk_grid
-    declared = (
-        grid.chunk_shapes if isinstance(grid, RectilinearChunkGridMetadata) else grid.chunk_shape
-    )
+    declared: Sequence[int | tuple[int, ...]]
+    if isinstance(grid, RectilinearChunkGridMetadata):
+        declared = [d if isinstance(d, int) else tuple(d.expand()) for d in grid.chunk_shapes]
+    else:
+        declared = grid.chunk_shape
     chunk_sizes = tuple(
         declared_chunk_data_sizes(d, n) for d, n in zip(declared, zarray.shape, strict=True)
     )