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

changes/4479.misc.md

@@ -0,0 +1 @@
+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.

src/zarr/codecs/sharding.py

@@ -627,8 +627,9 @@ def validate(
         if isinstance(chunk_grid, RegularChunkGridMetadata):
             edges_per_dim: tuple[tuple[int, ...], ...] = tuple((s,) for s in chunk_grid.chunk_shape)
         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 s 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

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import bisect
 import itertools
 import math
 import operator
@@ -15,7 +14,6 @@
     Protocol,
     SupportsIndex,
     cast,
-    runtime_checkable,
 )
 
 import numpy as np
@@ -24,6 +22,7 @@
 
 import zarr
 from zarr.core.common import (
+    RunLengthEdges,
     ShapeLike,
     ceildiv_int,
     parse_shapelike,
@@ -53,26 +52,21 @@ class FixedDimension:
     size: int  # chunk edge length (>= 1)
     extent: int  # array dimension length (>= 0)
     nchunks: int = field(init=False, repr=False)
-    ngridcells: int = field(init=False, repr=False)
 
     def __post_init__(self) -> None:
         if self.size < 1:
             raise ValueError(f"FixedDimension size must be >= 1, got {self.size}")
         if self.extent < 0:
             raise ValueError(f"FixedDimension extent must be >= 0, got {self.extent}")
-        n = ceildiv_int(self.extent, self.size)
-        object.__setattr__(self, "nchunks", n)
-        object.__setattr__(self, "ngridcells", n)
+        object.__setattr__(self, "nchunks", ceildiv_int(self.extent, self.size))
 
     def index_to_chunk(self, idx: int) -> int:
-        if idx < 0:
-            raise IndexError(f"Negative index {idx} is not allowed")
-        if idx >= self.extent:
-            raise IndexError(f"Index {idx} is out of bounds for extent {self.extent}")
+        if idx < 0 or idx >= self.extent:
+            raise IndexError(f"Index {idx} out of bounds for dimension with extent {self.extent}")
         return idx // self.size
 
     def chunk_offset(self, chunk_ix: int) -> int:
-        """Byte-aligned start position of chunk *chunk_ix* in array coordinates.
+        """Start position of chunk *chunk_ix* in array coordinates.
 
         Does not validate *chunk_ix* — callers must ensure it is in
         ``[0, nchunks)``. Use ``ChunkGrid.__getitem__`` for safe access.
@@ -95,186 +89,138 @@ def data_size(self, chunk_ix: int) -> int:
         """
         return max(0, min(self.size, self.extent - chunk_ix * self.size))
 
-    @property
-    def _unique_edge_lengths(self) -> Iterable[int]:
-        """Distinct chunk edge lengths for this dimension.
-
-        Used by shard validation to check that every unique edge length
-        is divisible by the inner chunk size. O(1) for fixed dimensions
-        since there is only one edge length.
-        """
-        return (self.size,)
-
     def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]:
         return indices // self.size
 
-    def with_extent(self, new_extent: int) -> FixedDimension:
-        """Re-bind to *new_extent* without modifying edges.
-
-        Used when constructing a grid from existing metadata where edges
-        are already correct. Raises on
-        ``VaryingDimension`` if edges don't cover the new extent.
-        """
-        return FixedDimension(size=self.size, extent=new_extent)
-
     def resize(self, new_extent: int) -> FixedDimension:
-        """Adapt for a user-initiated array resize, growing edges if needed.
-
-        For ``FixedDimension`` this is identical to ``with_extent`` since
-        regular grids don't store explicit edges.
-        """
+        """Adapt for a user-initiated array resize: the chunk size is kept and
+        re-bound to *new_extent*."""
         return FixedDimension(size=self.size, extent=new_extent)
 
     @property
     def _size_repr(self) -> str:
         return str(self.size)
 
 
-@dataclass(frozen=True)
+_MAX_EDGES_IN_REPR = 100
+"""The most chunk edges of one dimension written out one by one in a repr or an error
+message; a longer dimension is shown run-length encoded."""
+
+
+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.expand()))
+    return repr(edges.to_rle())
+
+
+@dataclass(frozen=True, init=False)
 class VaryingDimension:
     """Explicit per-chunk sizes. The last chunk may extend past the array
     extent (``extent < sum(edges)``), in which case ``data_size`` clips to
     the valid region while ``chunk_size`` returns the full edge length for
     codec processing. This underflow is allowed to match how regular grids
     handle boundary chunks, and to support shrinking an array without
-    rewriting chunk edges (the spec allows trailing edges beyond the extent)."""
+    rewriting chunk edges (the spec allows trailing edges beyond the extent).
+
+    The edges are held run-length encoded (see `RunLengthEdges`), so building
+    the dimension, resizing it and every lookup cost time and memory in the
+    number of runs of equal edges, not in the number of chunks: a dimension of
+    ``2**40`` equal chunks is one run."""
 
-    edges: tuple[int, ...]  # per-chunk edge lengths (all > 0)
-    cumulative: tuple[int, ...]  # prefix sums for O(log n) lookup
+    edges: RunLengthEdges  # per-chunk edge lengths (all > 0), run-length encoded
     extent: int  # array dimension length (may be < sum(edges) after resize)
     nchunks: int = field(init=False, repr=False)  # cached at construction
-    ngridcells: int = field(init=False, repr=False)  # cached at construction
-
-    # TODO(perf): for long dimensions (O(million chunks)):
-    # - with_extent/resize recompute cumulative sums and nchunks from scratch;
-    #   add a fast path that reuses the existing cumulative tuple.
-    # - Consider storing cumulative as ndarray so bisect calls can use
-    #   np.searchsorted. Scalar lookups (chunk_offset, index_to_chunk)
-    #   would need benchmarking to confirm no regression.
-    def __init__(self, edges: Sequence[int], extent: int) -> None:
-        edges_tuple = tuple(edges)
-        if not edges_tuple:
+
+    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")
-        if any(e <= 0 for e in edges_tuple):
-            raise ValueError(f"All edge lengths must be > 0, got {edges_tuple}")
-        cumulative = tuple(itertools.accumulate(edges_tuple))
+        if any(size <= 0 for size in rle.sizes):
+            raise ValueError(f"All edge lengths must be > 0, got {_edges_repr(rle)}")
         if extent < 0:
             raise ValueError(f"VaryingDimension extent must be >= 0, got {extent}")
-        if extent > cumulative[-1]:
-            raise ValueError(
-                f"VaryingDimension extent {extent} exceeds sum of edges {cumulative[-1]}"
-            )
-        object.__setattr__(self, "edges", edges_tuple)
-        object.__setattr__(self, "cumulative", cumulative)
+        if extent > rle.total:
+            raise ValueError(f"VaryingDimension extent {extent} exceeds sum of edges {rle.total}")
+        object.__setattr__(self, "edges", rle)
         object.__setattr__(self, "extent", extent)
         # Cache nchunks: number of chunks that overlap [0, extent)
-        if extent == 0:
-            n = 0
-        else:
-            n = bisect.bisect_left(cumulative, extent) + 1
+        n = 0 if extent == 0 else rle.index_at(extent - 1) + 1
         object.__setattr__(self, "nchunks", n)
-        object.__setattr__(self, "ngridcells", len(edges_tuple))
 
     def index_to_chunk(self, idx: int) -> int:
         if idx < 0 or idx >= self.extent:
             raise IndexError(f"Index {idx} out of bounds for dimension with extent {self.extent}")
-        return bisect.bisect_right(self.cumulative, idx)
+        return self.edges.index_at(idx)
 
     def chunk_offset(self, chunk_ix: int) -> int:
-        """Start position of chunk *chunk_ix* in array coordinates.
+        """Start position of chunk *chunk_ix* in array coordinates; for
+        ``chunk_ix == edges.num_edges``, the sum of all edges.
 
-        Does not validate *chunk_ix* — callers must ensure it is in
-        ``[0, ngridcells)``. Use ``ChunkGrid.__getitem__`` for safe access.
+        Raises ``IndexError`` for any other *chunk_ix* outside
+        ``[0, edges.num_edges)``.
         """
-        return self.cumulative[chunk_ix - 1] if chunk_ix > 0 else 0
+        return self.edges.offset_of(chunk_ix)
 
     def chunk_size(self, chunk_ix: int) -> int:
         """Buffer size for codec processing.
 
-        Does not validate *chunk_ix* — callers must ensure it is in
-        ``[0, ngridcells)``. Use ``ChunkGrid.__getitem__`` for safe access.
+        Raises ``IndexError`` for *chunk_ix* outside ``[0, edges.num_edges)``.
         """
-        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.
+        """Valid data region within the buffer — clipped at extent; 0 for an
+        edge that lies entirely past the extent.
 
-        Does not validate *chunk_ix* — callers must ensure it is in
-        ``[0, ngridcells)``. Use ``ChunkGrid.__getitem__`` for safe access.
+        Raises ``IndexError`` for *chunk_ix* outside ``[0, edges.num_edges)``.
         """
-        offset = self.cumulative[chunk_ix - 1] if chunk_ix > 0 else 0
-        return max(0, min(self.edges[chunk_ix], self.extent - offset))
+        return max(0, min(self.edges.size_of(chunk_ix), self.extent - self.chunk_offset(chunk_ix)))
 
     @property
-    def _unique_edge_lengths(self) -> Iterable[int]:
-        """Distinct chunk edge lengths for this dimension (lazily deduplicated).
-
-        Used by shard validation to check that every unique edge length
-        is divisible by the inner chunk size. Lazy deduplication avoids
-        materializing all edges for dimensions with many repeated sizes.
-        """
-        seen: set[int] = set()
-        for e in self.edges:
-            if e not in seen:
-                seen.add(e)
-                yield e
+    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.expand(), last), self.data_size(last))
 
     def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]:
-        return np.searchsorted(self.cumulative, indices, side="right")
-
-    def with_extent(self, new_extent: int) -> VaryingDimension:
-        """Re-bind to *new_extent* without modifying edges.
-
-        Used when constructing a grid from existing metadata where edges
-        are already correct. Raises if the
-        existing edges don't cover *new_extent*.
-        """
-        edge_sum = self.cumulative[-1]
-        if edge_sum < new_extent:
-            raise ValueError(
-                f"VaryingDimension edge sum {edge_sum} is less than new extent {new_extent}"
-            )
-        return VaryingDimension(self.edges, extent=new_extent)
+        return self.edges.indices_at(indices)
 
     def resize(self, new_extent: int) -> VaryingDimension:
         """Adapt for a user-initiated array resize, growing edges if needed.
 
-        Unlike ``with_extent``, this never fails — if *new_extent* exceeds
-        the current edge sum, a new chunk is appended to cover the gap.
+        If *new_extent* exceeds the current edge sum, a new chunk is
+        appended to cover the gap.
         Shrinking preserves all edges (the spec allows trailing edges
         beyond the array extent).
         """
         if new_extent == self.extent:
             return self
-        elif new_extent > self.cumulative[-1]:
-            expanded_edges = list(self.edges) + [new_extent - self.cumulative[-1]]
-            return VaryingDimension(expanded_edges, extent=new_extent)
+        elif new_extent > self.edges.total:
+            grown = self.edges.with_edge(new_extent - self.edges.total)
+            return VaryingDimension(grown, extent=new_extent)
         else:
             return VaryingDimension(self.edges, extent=new_extent)
 
     @property
     def _size_repr(self) -> str:
-        return repr(tuple(self.edges))
+        return _edges_repr(self.edges)
 
 
-@runtime_checkable
 class DimensionGrid(Protocol):
     """Structural interface shared by FixedDimension and VaryingDimension."""
 
     @property
     def nchunks(self) -> int: ...
     @property
-    def ngridcells(self) -> int: ...
-    @property
     def extent(self) -> int: ...
     def index_to_chunk(self, idx: int) -> int: ...
     def chunk_offset(self, chunk_ix: int) -> int: ...
     def chunk_size(self, chunk_ix: int) -> int: ...
     def data_size(self, chunk_ix: int) -> int: ...
     def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: ...
-    @property
-    def _unique_edge_lengths(self) -> Iterable[int]: ...
-    def with_extent(self, new_extent: int) -> DimensionGrid: ...
     def resize(self, new_extent: int) -> DimensionGrid: ...
     @property
     def _size_repr(self) -> str: ...
@@ -359,28 +305,6 @@ def _chunk_list(chunks: object) -> list[Any]:
         ) from None
 
 
-def is_regular_1d(dim_chunks: Sequence[int]) -> bool:
-    """Check if a single dimension's chunk sizes represent a regular grid.
-
-    A regular dimension has either all chunks the same size, or all
-    but the last chunk the same size with the last chunk smaller
-    (boundary chunk).
-    """
-    if len(dim_chunks) <= 1:
-        return True
-    first = dim_chunks[0]
-    for c in dim_chunks[1:-1]:
-        if c != first:
-            return False
-    # Last chunk must be the same size or a smaller boundary chunk
-    return dim_chunks[-1] <= first
-
-
-def is_regular_nd(chunks: Iterable[Sequence[int]]) -> bool:
-    """Check if an N-dimensional chunk specification represents a regular grid."""
-    return all(is_regular_1d(d) for d in chunks)
-
-
 @dataclass(frozen=True)
 class ChunkGrid:
     """
@@ -391,7 +315,7 @@ class ChunkGrid:
     ``grid[coords]`` to return a ``ChunkSpec`` without external parameters.
 
     Internally represents each dimension as either FixedDimension (uniform chunks)
-    or VaryingDimension (per-chunk edge lengths with prefix sums).
+    or VaryingDimension (per-chunk edge lengths, run-length encoded).
     """
 
     _dimensions: tuple[DimensionGrid, ...]
@@ -432,7 +356,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.
 
@@ -444,7 +368,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`.
         """
@@ -506,7 +430,12 @@ def chunk_sizes(self) -> tuple[tuple[int, ...], ...]:
             One inner tuple per dimension, each containing the data size
             of every chunk along that dimension.
         """
-        return tuple(tuple(d.data_size(i) for i in range(d.nchunks)) for d in self._dimensions)
+        return tuple(
+            d._data_sizes
+            if isinstance(d, VaryingDimension)
+            else tuple(d.data_size(i) for i in range(d.nchunks))
+            for d in self._dimensions
+        )
 
     # -- Collection interface --