From 851c682c1fbdebad4e4ba4234125f093e1fcd993 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Oct 2026 16:53:52 +0200 Subject: [PATCH 1/7] perf(chunk-grids): hold rectilinear chunk edges as run-length encoded runs Opening an array with a rectilinear chunk grid cost time and memory proportional to the repeat counts in its stored metadata: `from_dict` expanded every run into one int per chunk, validation and `VaryingDimension` walked and re-accumulated them, and `to_dict` re-compressed them. A ~300 byte document declaring `[[1, 10**12]]` got the process OOM-killed. Add `RunLengthEdges`, an immutable `Sequence[int]` stored as merged `(size, count)` runs with bisect-over-runs lookups, and use it for the explicit edges of `RectilinearChunkGridMetadata.chunk_shapes` and `VaryingDimension.edges`. Metadata parsing, validation, serialization, `update_shape`, the dimension lookups, resize and the sharding codec's divisibility check now cost O(number of runs). Assisted-by: ClaudeCode:claude-fable-5.1 --- changes/+rle-runs.bugfix.md | 1 + src/zarr/codecs/sharding.py | 5 +- src/zarr/core/chunk_grids.py | 114 +++++++----- src/zarr/core/common.py | 238 ++++++++++++++++++++++--- src/zarr/core/metadata/v3.py | 54 +++--- src/zarr/testing/strategies.py | 4 +- tests/test_metadata/test_repair.py | 2 +- tests/test_unified_chunk_grid.py | 270 +++++++++++++++++++++++++++-- 8 files changed, 580 insertions(+), 108 deletions(-) create mode 100644 changes/+rle-runs.bugfix.md diff --git a/changes/+rle-runs.bugfix.md b/changes/+rle-runs.bugfix.md new file mode 100644 index 0000000000..048631c38f --- /dev/null +++ b/changes/+rle-runs.bugfix.md @@ -0,0 +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 that reads like the tuple of edges and compares equal to it, but is not a `tuple` instance. `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. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index cb972118ce..27d67bf9bd 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -45,6 +45,7 @@ merge_and_encode_chunk, ) from zarr.core.common import ( + RunLengthEdges, ShapeLike, parse_named_configuration, parse_shapelike, @@ -620,8 +621,10 @@ 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 RunLengthEdges.from_edges(s).sizes + for s in chunk_grid.chunk_shapes ) else: raise TypeError( diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 76bf399a65..24ea7a7273 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -1,6 +1,5 @@ from __future__ import annotations -import bisect import itertools import math import operator @@ -24,6 +23,7 @@ import zarr from zarr.core.common import ( + RunLengthEdges, ShapeLike, ceildiv_int, parse_shapelike, @@ -130,55 +130,67 @@ 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)) + 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: + 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)) + object.__setattr__(self, "ngridcells", rle.num_edges) + + @property + def cumulative(self) -> tuple[int, ...]: + """Prefix sums of the edges, one entry per chunk. + + 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)) 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. @@ -186,7 +198,7 @@ def chunk_offset(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.cumulative[chunk_ix - 1] if chunk_ix > 0 else 0 + return self.edges.offset_of(chunk_ix) if chunk_ix > 0 else 0 def chunk_size(self, chunk_ix: int) -> int: """Buffer size for codec processing. @@ -202,25 +214,28 @@ def data_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. """ - 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[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)) @property def _unique_edge_lengths(self) -> Iterable[int]: - """Distinct chunk edge lengths for this dimension (lazily deduplicated). + """Distinct chunk edge lengths for this dimension, in order of first use. 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. + is divisible by the inner chunk size. Reads one size per run of + equal edges, not one per chunk. """ - seen: set[int] = set() - for e in self.edges: - if e not in seen: - seen.add(e) - yield e + return tuple(dict.fromkeys(self.edges.sizes)) def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: - return np.searchsorted(self.cumulative, indices, side="right") + return self.edges.indices_at(indices) def with_extent(self, new_extent: int) -> VaryingDimension: """Re-bind to *new_extent* without modifying edges. @@ -229,7 +244,7 @@ def with_extent(self, new_extent: int) -> VaryingDimension: are already correct. Raises if the existing edges don't cover *new_extent*. """ - edge_sum = self.cumulative[-1] + edge_sum = self.edges.total if edge_sum < new_extent: raise ValueError( f"VaryingDimension edge sum {edge_sum} is less than new extent {new_extent}" @@ -246,15 +261,15 @@ def resize(self, new_extent: int) -> VaryingDimension: """ 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 @@ -391,7 +406,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, ...] @@ -506,7 +521,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 -- diff --git a/src/zarr/core/common.py b/src/zarr/core/common.py index de009d0b50..0a00904c61 100644 --- a/src/zarr/core/common.py +++ b/src/zarr/core/common.py @@ -1,6 +1,8 @@ from __future__ import annotations import asyncio +import bisect +import itertools import math import warnings from collections.abc import Iterable, Mapping, Sequence @@ -25,6 +27,9 @@ if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Iterator + from typing import Self + + import numpy.typing as npt from zarr.core.metadata.v3 import ChunkGridMetadata @@ -320,15 +325,188 @@ 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}") -def expand_rle(data: Sequence[object], axis: int | None = None) -> list[int]: - """Expand a mixed array of bare integers and RLE pairs, the edges of dimension - `axis` (named in errors, when given). +class RunLengthEdges(Sequence[int]): + """An immutable sequence of chunk edge lengths, stored run-length encoded. - Per the rectilinear chunk grid spec, each element can be: - - a bare integer (an explicit edge length) - - a two-element array ``[value, count]`` (run-length encoded) + The sequence is 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 tuple of edges it stands for (`len`, indexing, iteration, and `==` + against a `tuple`), but only iteration and the comparison with a `tuple` visit 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. + + Instances hash by their runs, so an instance does not hash like the `tuple` it + compares equal to. """ - result: list[int] = [] + + __slots__ = ("_index_stops", "_lookup_tables", "_offset_stops", "counts", "sizes") + + sizes: tuple[int, ...] + """The edge length of each run.""" + counts: tuple[int, ...] + """The number of edges in each run.""" + + def __init__(self, runs: Iterable[tuple[int, int]] = ()) -> None: + sizes: list[int] = [] + counts: list[int] = [] + for size, count in runs: + if count < 1: + raise ValueError(f"Run counts must be >= 1, got {count!r}") + if sizes and sizes[-1] == size: + counts[-1] += count + else: + sizes.append(size) + counts.append(count) + self.sizes = tuple(sizes) + self.counts = tuple(counts) + # Per run, the number of edges and the sum of edges up to and including it. + self._index_stops = tuple(itertools.accumulate(counts)) + self._offset_stops = tuple( + itertools.accumulate(size * count for size, count in zip(sizes, counts, strict=True)) + ) + 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): + return edges + return cls((size, sum(1 for _ in group)) for size, group in itertools.groupby(edges)) + + @property + def runs(self) -> tuple[tuple[int, int], ...]: + """The `(size, count)` runs, no two adjacent runs sharing a size.""" + return tuple(zip(self.sizes, self.counts, strict=True)) + + @property + def num_edges(self) -> int: + """The number of edges. Unlike `len`, this is not limited to `sys.maxsize`.""" + return self._index_stops[-1] if self._index_stops else 0 + + @property + def total(self) -> int: + """The sum of all edges.""" + return self._offset_stops[-1] if self._offset_stops else 0 + + 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]: + 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 __eq__(self, other: object) -> bool: + if isinstance(other, RunLengthEdges): + return self.sizes == other.sizes and self.counts == other.counts + if isinstance(other, tuple): + return len(other) == self.num_edges and all( + a == b for a, b in zip(self, other, strict=True) + ) + return NotImplemented + + def __hash__(self) -> int: + return hash((self.sizes, self.counts)) + + 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]`.""" + return [[size, count] if count > 1 else size for size, count in self.runs] + + def with_edge(self, size: int) -> RunLengthEdges: + """A copy with one more edge of length `size` at the end.""" + return RunLengthEdges((*self.runs, (size, 1))) + + def offset_of(self, index: int) -> int: + """The sum of the first `index` edges: where edge `index` starts, or `total` + for `index == num_edges`.""" + if not 0 <= index <= self.num_edges: + raise IndexError(f"Edge index {index} is out of range for {self.num_edges} edges") + if index == self.num_edges: + return self.total + run = bisect.bisect_right(self._index_stops, index) + if run == 0: + return index * self.sizes[0] + return self._offset_stops[run - 1] + (index - self._index_stops[run - 1]) * self.sizes[run] + + def index_at(self, position: int) -> int: + """The index of the edge that covers `position`, a coordinate in `[0, total)` + along the axis the edges tile.""" + if not 0 <= position < self.total: + raise IndexError( + f"Position {position} is out of range for edges summing to {self.total}" + ) + run = bisect.bisect_right(self._offset_stops, position) + if run == 0: + return position // self.sizes[0] + return ( + self._index_stops[run - 1] + (position - self._offset_stops[run - 1]) // self.sizes[run] + ) + + def indices_at(self, positions: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: + """Vectorized `index_at`, without the bounds check: a negative position gives 0, + and a position of `total` or more gives `num_edges`.""" + positions = np.asarray(positions) + if not self.sizes: + return np.zeros(positions.shape, dtype=np.intp) + if self._lookup_tables is None: + # No position reaches a run that starts past the largest intp, and a size + # that large is never exceeded, so the tables always fit intp. + intp_max = int(np.iinfo(np.intp).max) + offsets = [o for o in (0, *self._offset_stops[:-1]) if o <= intp_max] + n = len(offsets) + self._lookup_tables = ( + np.array(offsets, dtype=np.intp), + np.array((0, *self._index_stops[: n - 1]), dtype=np.intp), + np.array([min(size, intp_max) for size in self.sizes[:n]], dtype=np.intp), + np.array(min(self.num_edges, intp_max), dtype=np.intp), + ) + offsets_arr, indices_arr, sizes_arr, last = self._lookup_tables + run = np.searchsorted(offsets_arr[1:], positions, side="right") + found = indices_arr[run] + (positions - offsets_arr[run]) // sizes_arr[run] + return np.clip(found, 0, last).astype(np.intp, copy=False) + + +def _iter_rle_runs(data: Iterable[object], axis: int | None) -> Iterator[tuple[int, int]]: + """The `(size, count)` runs of a mixed array of bare integers and RLE pairs, each + checked as it is read.""" for item in data: if isinstance(item, list): if len(item) != 2: @@ -336,9 +514,34 @@ def expand_rle(data: Sequence[object], axis: int | None = None) -> list[int]: raise ValueError(f"{subject} must be an integer or [size, count], got {item}") size, count = item repeat = _parse_positive_int(count, "RLE repeat count", axis) - result.extend([parse_chunk_edge(size, axis)] * repeat) + yield parse_chunk_edge(size, axis), repeat else: - result.append(parse_chunk_edge(item, axis)) + yield parse_chunk_edge(item, axis), 1 + + +def parse_rle(data: Sequence[object], axis: int | None = None) -> RunLengthEdges: + """Read a mixed array of bare integers and RLE pairs, the edges of dimension `axis` + (named in errors, when given), without expanding it. + + Per the rectilinear chunk grid spec, each element can be: + - a bare integer (an explicit edge length) + - a two-element array ``[value, count]`` (run-length encoded) + + The cost is in the number of elements, whatever the repeat counts. + """ + return RunLengthEdges(_iter_rle_runs(data, axis)) + + +def expand_rle(data: Sequence[object], axis: int | None = None) -> list[int]: + """Expand a mixed array of bare integers and RLE pairs, the edges of dimension + `axis` (named in errors, when given), into one entry per edge. + + The result is as long as the sum of the repeat counts; `parse_rle` reads the same + input without expanding it. + """ + result: list[int] = [] + for size, count in _iter_rle_runs(data, axis): + result.extend([size] * count) return result @@ -350,20 +553,7 @@ def compress_rle(sizes: Sequence[int]) -> list[int | list[int]]: [10, 10, 10, 5] -> [[10, 3], 5] """ - if not sizes: - return [] - result: list[int | list[int]] = [] - current = sizes[0] - count = 1 - for s in sizes[1:]: - if s == current: - count += 1 - else: - result.append([current, count] if count > 1 else current) - current = s - count = 1 - result.append([current, count] if count > 1 else current) - return result + return RunLengthEdges.from_edges(sizes).to_rle() def validate_rectilinear_kind(kind: str | None) -> None: @@ -394,7 +584,7 @@ def validate_rectilinear_edges( for i, (dim_spec, extent) in enumerate(zip(chunk_shapes, array_shape, strict=True)): if isinstance(dim_spec, int): continue - edge_sum = sum(dim_spec) + edge_sum = dim_spec.total if isinstance(dim_spec, RunLengthEdges) else sum(dim_spec) if edge_sum < extent: raise ValueError( f"Rectilinear chunk edges for dimension {i} sum to {edge_sum} " diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2d5d169e6b..9326c2eb75 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -26,11 +26,11 @@ NamedConfig, NamedRequiredConfig, NodeType, - compress_rle, - expand_rle, + RunLengthEdges, parse_chunk_edge, parse_chunk_shape, parse_named_configuration, + parse_rle, parse_shapelike, validate_rectilinear_edges, validate_rectilinear_kind, @@ -246,20 +246,28 @@ class RectilinearChunkGridMetadataConfig(TypedDict): def _validate_chunk_shapes( chunk_shapes: Sequence[int | Sequence[int]], -) -> tuple[int | tuple[int, ...], ...]: +) -> tuple[int | RunLengthEdges, ...]: """Validate per-dimension chunk specifications. Each element is either a bare ``int`` (regular step size, must be >= 1) - or a sequence of explicit edge lengths (all must be >= 1, non-empty). + or a sequence of explicit edge lengths (all must be >= 1, non-empty), + which is returned run-length encoded. Edges that already are run-length + encoded are checked one run at a time, not one edge at a time. """ - result: list[int | tuple[int, ...]] = [] + result: list[int | RunLengthEdges] = [] for dim_idx, dim_spec in enumerate(chunk_shapes): match dim_spec: + case RunLengthEdges(): + for size in dim_spec.sizes: + parse_chunk_edge(size, dim_idx) + if dim_spec.num_edges == 0: + raise ValueError(f"Dimension {dim_idx} has no chunk edges.") + result.append(dim_spec) case list() | tuple(): - edges = tuple(parse_chunk_edge(edge, dim_idx) for edge in dim_spec) + edges = [parse_chunk_edge(edge, dim_idx) for edge in dim_spec] if not edges: raise ValueError(f"Dimension {dim_idx} has no chunk edges.") - result.append(edges) + result.append(RunLengthEdges.from_edges(edges)) case _: result.append(parse_chunk_edge(dim_spec, dim_idx)) return tuple(result) @@ -304,15 +312,18 @@ class RectilinearChunkGridMetadata(Metadata): - A bare ``int`` — a regular step size that repeats to cover the axis (the spec's single-integer shorthand). - - A ``tuple[int, ...]`` — explicit per-chunk edge lengths (already - expanded from any RLE encoding). + - 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 tuple of edges (and compares equal to it) but 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. This distinction matters for faithful round-tripping: a bare int serializes back as a bare int, while a single-element tuple serializes as a list. """ - chunk_shapes: tuple[int | tuple[int, ...], ...] + chunk_shapes: tuple[int | Sequence[int], ...] def __post_init__(self) -> None: object.__setattr__(self, "chunk_shapes", _validate_chunk_shapes(self.chunk_shapes)) @@ -328,12 +339,9 @@ def to_dict(self) -> RectilinearChunkGridMetadataJSON: # type: ignore[override] # Bare int shorthand — serialize as-is serialized_dims.append(dim_spec) else: - rle = compress_rle(dim_spec) - # Use RLE only if it's actually shorter - if len(rle) < len(dim_spec): - serialized_dims.append(rle) - else: - serialized_dims.append(list(dim_spec)) + # 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()) return { "name": "rectilinear", "configuration": { @@ -353,17 +361,17 @@ def update_shape( Otherwise edges are kept as-is (the spec allows trailing edges beyond the array extent). """ - new_chunk_shapes: list[int | tuple[int, ...]] = [] + new_chunk_shapes: list[int | Sequence[int]] = [] 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: - edge_sum = sum(dim_spec) - if new_ext > edge_sum: - new_chunk_shapes.append((*dim_spec, new_ext - edge_sum)) + edges = RunLengthEdges.from_edges(dim_spec) + if new_ext > edges.total: + new_chunk_shapes.append(edges.with_edge(new_ext - edges.total)) else: - new_chunk_shapes.append(dim_spec) + new_chunk_shapes.append(edges) return RectilinearChunkGridMetadata(chunk_shapes=tuple(new_chunk_shapes)) @classmethod @@ -373,7 +381,7 @@ def from_dict(cls, data: RectilinearChunkGridMetadataJSON) -> Self: # type: ign validate_rectilinear_kind(configuration.get("kind")) raw_shapes = configuration["chunk_shapes"] parsed = [ - tuple(expand_rle(dim_spec, axis)) if isinstance(dim_spec, list) else dim_spec + parse_rle(dim_spec, axis) if isinstance(dim_spec, list) else dim_spec for axis, dim_spec in enumerate(raw_shapes) ] return cls(chunk_shapes=tuple(parsed)) @@ -435,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 | tuple[int, ...]] = [] + chunk_shapes: list[int | Sequence[int]] = [] for dim in chunks.dimensions: if isinstance(dim, FixedDimension): chunk_shapes.append(dim.size) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 4c49a2efb0..693fc613f9 100644 --- a/src/zarr/testing/strategies.py +++ b/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 tuples become lists. Bare ints — the spec's step-size + Explicit edge sequences 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 [list(dim) if isinstance(dim, tuple) else dim for dim in meta.chunk_shapes] + return [dim if isinstance(dim, int) else list(dim) for dim in meta.chunk_shapes] # The most chunks a drawn rectilinear grid declares along one axis, and over diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index 9219b5fe31..38ad102c6a 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -422,7 +422,7 @@ def _rectilinear_from_dict(chunk_shapes: list[Any]) -> RectilinearChunkGridMetad def _second_edge(grid: RectilinearChunkGridMetadata) -> int: edges = grid.chunk_shapes[0] - assert isinstance(edges, tuple) + assert not isinstance(edges, int) return edges[1] diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 1262af7484..59d34054a4 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -8,6 +8,9 @@ from __future__ import annotations +import itertools +import json +import pickle import re from functools import partial from typing import TYPE_CHECKING, Any @@ -25,7 +28,7 @@ VaryingDimension, guess_chunks, ) -from zarr.core.common import compress_rle, expand_rle +from zarr.core.common import RunLengthEdges, compress_rle, expand_rle, parse_rle from zarr.core.dtype import UInt8 from zarr.core.metadata.v3 import ( ArrayV3Metadata, @@ -575,10 +578,13 @@ def test_rle_roundtrip() -> None: "rle-entry-of-three", ], ) -def test_rle_expand_rejects_invalid(rle_input: list[Any], match: str) -> None: - """expand_rle raises ValueError for zero/negative edge lengths or repeat counts""" +@pytest.mark.parametrize("read", [expand_rle, parse_rle]) +def test_rle_expand_rejects_invalid( + read: Callable[..., object], rle_input: list[Any], match: str +) -> None: + """expand_rle and parse_rle raise ValueError for zero/negative edge lengths or repeat counts""" with pytest.raises(ValueError, match=match): - expand_rle(rle_input) + read(rle_input) @pytest.mark.parametrize( @@ -604,10 +610,13 @@ def test_rle_expand_rejects_invalid(rle_input: list[Any], match: str) -> None: "numpy-int-count", ], ) -def test_rle_expand_rejects_non_int(rle_input: list[Any], match: str) -> None: - """expand_rle reads `int`s (and `bool`s) only, not floats or NumPy integers.""" +@pytest.mark.parametrize("read", [expand_rle, parse_rle]) +def test_rle_expand_rejects_non_int( + read: Callable[..., object], rle_input: list[Any], match: str +) -> None: + """expand_rle and parse_rle read `int`s (and `bool`s) only, not floats or NumPy integers.""" with pytest.raises(TypeError, match=re.escape(match)): - expand_rle(rle_input) + read(rle_input) @pytest.mark.parametrize( @@ -620,10 +629,251 @@ def test_rle_expand_rejects_non_int(rle_input: list[Any], match: str) -> None: ], ids=["zero-edge", "fractional-edge", "zero-rle-count", "rle-entry-of-three"], ) -def test_rle_expand_names_dimension(rle_input: list[Any], match: str) -> None: - """Given the dimension `axis` the edges belong to, every error of expand_rle names it.""" +@pytest.mark.parametrize("read", [expand_rle, parse_rle]) +def test_rle_expand_names_dimension( + read: Callable[..., object], rle_input: list[Any], match: str +) -> None: + """Given the dimension `axis` the edges belong to, every error of expand_rle and + parse_rle names it.""" with pytest.raises((TypeError, ValueError), match=f"^Dimension 2: {match}"): - expand_rle(rle_input, axis=2) + read(rle_input, axis=2) + + +@pytest.mark.parametrize( + "rle_input", + [[[10, 3]], [[10, 2], [20, 1]], [[True, 2], [3, 1]], [5, 5, [5, 2], 7], [1, 2, 3]], +) +def test_parse_rle(rle_input: list[Any]) -> None: + """parse_rle reads the edges expand_rle expands, as merged runs""" + parsed = parse_rle(rle_input) + assert parsed == tuple(expand_rle(rle_input)) + assert parsed.to_rle() == compress_rle(expand_rle(rle_input)) + + +# --------------------------------------------------------------------------- +# RunLengthEdges tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "runs", + [ + [(10, 1)], + [(10, 3)], + [(10, 2), (20, 1), (5, 4)], + [(1, 1), (2, 1), (3, 1)], + [(4, 2), (4, 3), (1, 1), (4, 1)], + ], + ids=["one-edge", "one-run", "mixed-runs", "no-repeats", "adjacent-equal-runs"], +) +def test_run_length_edges(runs: list[tuple[int, int]]) -> None: + """RunLengthEdges reads like the tuple of edges it encodes: length, indexing, slicing, + iteration, equality, prefix sums and position lookups all agree with the expanded + edges, and adjacent runs of one size are merged""" + expanded = tuple(size for size, count in runs for _ in range(count)) + offsets = (0, *itertools.accumulate(expanded)) + edges = RunLengthEdges(runs) + + assert edges == expanded + assert edges == RunLengthEdges.from_edges(expanded) + assert hash(edges) == hash(RunLengthEdges.from_edges(expanded)) + assert edges != (*expanded, 1) + assert pickle.loads(pickle.dumps(edges)) == edges + assert all(a != b for (a, _), (b, _) in itertools.pairwise(edges.runs)) + assert len(edges) == edges.num_edges == len(expanded) + assert edges.total == sum(expanded) + assert tuple(edges) == expanded + assert [edges[i] for i in range(-len(expanded), len(expanded))] == [*expanded, *expanded] + assert edges[1:] == expanded[1:] + assert edges[:-1] == expanded[:-1] + assert edges[::2] == expanded[::2] + assert all(edges.count(size) == expanded.count(size) for size in {*expanded, 99}) + assert all((size in edges) == (size in expanded) for size in {*expanded, 99}) + assert edges.to_rle() == compress_rle(expanded) + assert edges.with_edge(7) == (*expanded, 7) + assert [edges.offset_of(i) for i in range(len(expanded) + 1)] == list(offsets) + positions = np.arange(-2, edges.total + 3, dtype=np.intp) + expected_indices = np.searchsorted(offsets[1:], positions, side="right") + assert [edges.index_at(p) for p in range(edges.total)] == expected_indices[2:-3].tolist() + np.testing.assert_array_equal(edges.indices_at(positions), expected_indices) + + +def test_run_length_edges_rejects_count_below_one() -> None: + """RunLengthEdges raises ValueError for a run of fewer than one edge""" + with pytest.raises(ValueError, match="Run counts must be >= 1, got 0"): + RunLengthEdges([(5, 0)]) + + +@pytest.mark.parametrize("index", [3, -4]) +def test_run_length_edges_getitem_out_of_range(index: int) -> None: + """Indexing RunLengthEdges past either end raises IndexError""" + with pytest.raises(IndexError, match=f"Edge index {index} is out of range for 3 edges"): + RunLengthEdges([(5, 3)])[index] + + +@pytest.mark.parametrize("index", [4, -1]) +def test_run_length_edges_offset_of_out_of_range(index: int) -> None: + """offset_of raises IndexError for an index that is neither an edge nor the end""" + with pytest.raises(IndexError, match=f"Edge index {index} is out of range for 3 edges"): + RunLengthEdges([(5, 3)]).offset_of(index) + + +@pytest.mark.parametrize("position", [15, -1]) +def test_run_length_edges_index_at_out_of_range(position: int) -> None: + """index_at raises IndexError for a position no edge covers""" + with pytest.raises(IndexError, match=f"Position {position} is out of range"): + RunLengthEdges([(5, 3)]).index_at(position) + + +# --------------------------------------------------------------------------- +# Grids with very many chunks per run +# --------------------------------------------------------------------------- + +HUGE = 2**40 +"""A repeat count no test could afford to expand: one entry per chunk would take +terabytes of memory.""" + + +def test_varying_dimension_huge_runs() -> None: + """A VaryingDimension of runs of 2**40 chunks is built, queried, resized and shown + from its runs alone, with the values the expanded edges would give""" + # chunks [0, HUGE) have size 3, chunk HUGE has size 7, chunks (HUGE, 2 * HUGE] size 2 + edge_sum = 3 * HUGE + 7 + 2 * HUGE + dim = VaryingDimension(RunLengthEdges([(3, HUGE), (7, 1), (2, HUGE)]), extent=edge_sum - 3) + + assert dim.ngridcells == 2 * HUGE + 1 + assert dim.nchunks == 2 * HUGE + assert tuple(dim._unique_edge_lengths) == (3, 7, 2) + assert dim._size_repr == f"[[3, {HUGE}], 7, [2, {HUGE}]]" + lookups = { + 0: 0, + 3 * HUGE - 1: HUGE - 1, + 3 * HUGE: HUGE, + 3 * HUGE + 6: HUGE, + 3 * HUGE + 7: HUGE + 1, + edge_sum - 4: 2 * HUGE - 1, + } + for idx, chunk in lookups.items(): + assert dim.index_to_chunk(idx) == chunk + np.testing.assert_array_equal( + dim.indices_to_chunks(np.array(list(lookups), dtype=np.intp)), list(lookups.values()) + ) + assert [dim.chunk_offset(c) for c in (0, HUGE - 1, HUGE, HUGE + 1, 2 * HUGE)] == [ + 0, + 3 * HUGE - 3, + 3 * HUGE, + 3 * HUGE + 7, + edge_sum - 2, + ] + assert [dim.chunk_size(c) for c in (0, HUGE, HUGE + 1, 2 * HUGE)] == [3, 7, 2, 2] + # the extent ends one element into the last chunk that holds data + assert [dim.data_size(c) for c in (0, 2 * HUGE - 1, 2 * HUGE)] == [3, 1, 0] + + grown = dim.resize(edge_sum + 5) + assert grown.edges.to_rle() == [[3, HUGE], 7, [2, HUGE], 5] + assert dim.resize(10).edges == dim.edges + assert dim.with_extent(edge_sum).nchunks == 2 * HUGE + 1 + + +def _store_rectilinear_array( + shape: list[int], chunk_shapes: list[Any], codecs: list[dict[str, Any]] +) -> MemoryStore: + """A store whose array `a` has the stored rectilinear chunk grid `chunk_shapes`.""" + document = { + "zarr_format": 3, + "node_type": "array", + "shape": shape, + "data_type": "int8", + "fill_value": 0, + "chunk_grid": { + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": chunk_shapes}, + }, + "chunk_key_encoding": {"name": "default"}, + "codecs": codecs, + "attributes": {}, + } + store = MemoryStore() + buffer = default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + zarr.core.sync.sync(store.set("a/zarr.json", buffer)) + return store + + +@pytest.mark.parametrize( + ("edge", "codecs"), + [ + (1, [{"name": "bytes"}]), + ( + 4, + [ + { + "name": "sharding_indexed", + "configuration": { + "chunk_shape": [2], + "codecs": [{"name": "bytes"}], + "index_codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "crc32c"}, + ], + }, + } + ], + ), + ], + ids=["unsharded", "sharded"], +) +def test_open_rectilinear_huge_repeat_count(edge: int, codecs: list[dict[str, Any]]) -> None: + """An array whose stored rectilinear grid repeats one edge 2**40 times opens, reads, + writes, grows and re-serializes at the cost of its run-length encoded metadata: + expanding the edges would exhaust memory long before this test could pass""" + extent = edge * HUGE + store = _store_rectilinear_array([extent], [[[edge, HUGE]]], codecs) + + arr = zarr.open_array(store, path="a", mode="r+") + + assert arr.shape == (extent,) + assert arr._chunk_grid.grid_shape == (HUGE,) + assert arr.metadata.to_dict()["chunk_grid"] == { + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": ([[edge, HUGE]],)}, + } + assert f"[[{edge}, {HUGE}]]" in repr(arr._chunk_grid) + np.testing.assert_array_equal(arr[extent - 3 :], [0, 0, 0]) + np.testing.assert_array_equal(arr[[0, extent // 2, extent - 1]], [0, 0, 0]) + arr[extent // 2 : extent // 2 + 2] = [1, 2] + np.testing.assert_array_equal(arr[extent // 2 - 1 : extent // 2 + 3], [0, 1, 2, 0]) + + arr.resize((extent + 8,)) + assert arr._chunk_grid.grid_shape == (HUGE + 1,) + stored_buffer = zarr.core.sync.sync(store.get("a/zarr.json", default_buffer_prototype())) + assert stored_buffer is not None + stored = json.loads(stored_buffer.to_bytes()) + assert stored["chunk_grid"]["configuration"]["chunk_shapes"] == [[[edge, HUGE], 8]] + # Shrinking the stored grid keeps the trailing edges. (`Array.resize` to a smaller + # shape visits every chunk it drops, whatever the chunk grid, so it is not run here.) + grid = arr.metadata.chunk_grid + assert isinstance(grid, RectilinearChunkGridMetadata) + assert grid.update_shape((extent + 8,), (8,)) == grid + + +def test_rectilinear_metadata_rejects_invalid_run_length_edge() -> None: + """RectilinearChunkGridMetadata rejects run-length encoded edges whose size is below + 1, naming the dimension""" + with pytest.raises(ValueError, match="^Dimension 1: chunk edge length must be >= 1, got 0"): + RectilinearChunkGridMetadata(chunk_shapes=(4, RunLengthEdges([(0, HUGE)]))) + + +def test_rectilinear_metadata_rejects_empty_run_length_edges() -> None: + """RectilinearChunkGridMetadata rejects run-length encoded edges with no runs""" + with pytest.raises(ValueError, match="Dimension 0 has no chunk edges"): + RectilinearChunkGridMetadata(chunk_shapes=(RunLengthEdges(),)) + + +def test_varying_dimension_rejects_invalid_huge_run() -> None: + """VaryingDimension reports a non-positive edge in a long run without writing out + every edge""" + with pytest.raises(ValueError, match=re.escape(f"got [[0, {HUGE}]]")): + VaryingDimension(RunLengthEdges([(0, HUGE)]), extent=0) # --------------------------------------------------------------------------- From 7ddb469803d28371dec73d4a9ee4f71f272d64f9 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Oct 2026 16:54:22 +0200 Subject: [PATCH 2/7] docs: name the changelog fragment after the pull request Assisted-by: ClaudeCode:claude-fable-5.1 --- changes/{+rle-runs.bugfix.md => 4479.bugfix.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{+rle-runs.bugfix.md => 4479.bugfix.md} (100%) diff --git a/changes/+rle-runs.bugfix.md b/changes/4479.bugfix.md similarity index 100% rename from changes/+rle-runs.bugfix.md rename to changes/4479.bugfix.md From 26a62878cdd3fc52995976fb5b62719252aaa41c Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Oct 2026 17:11:50 +0200 Subject: [PATCH 3/7] refactor(chunk-grids): compare RunLengthEdges only with RunLengthEdges A `RunLengthEdges` is not a `tuple`, so `==` against one is now false instead of an edge-by-edge comparison. This also removes the mismatch between equality and hashing. Tests compare `tuple(edges)` or whole metadata objects instead. Assisted-by: ClaudeCode:claude-fable-5.1 --- changes/4479.bugfix.md | 2 +- src/zarr/core/common.py | 13 ++++-------- src/zarr/core/metadata/v3.py | 4 ++-- tests/test_chunk_grids.py | 2 +- tests/test_properties.py | 4 ++-- tests/test_unified_chunk_grid.py | 35 ++++++++++++++++---------------- 6 files changed, 28 insertions(+), 32 deletions(-) diff --git a/changes/4479.bugfix.md b/changes/4479.bugfix.md index 048631c38f..8a9b1955c1 100644 --- a/changes/4479.bugfix.md +++ 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 that reads like the tuple of edges and compares equal to it, but is not a `tuple` instance. `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 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. diff --git a/src/zarr/core/common.py b/src/zarr/core/common.py index 0a00904c61..324f624523 100644 --- a/src/zarr/core/common.py +++ b/src/zarr/core/common.py @@ -331,15 +331,14 @@ class RunLengthEdges(Sequence[int]): The sequence is 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 tuple of edges it stands for (`len`, indexing, iteration, and `==` - against a `tuple`), but only iteration and the comparison with a `tuple` visit every - edge. + reads like the sequence of edges it stands for (`len`, indexing, iteration), but + only iteration visits 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. - Instances hash by their runs, so an instance does not hash like the `tuple` it - compares equal to. + It is not a `tuple` and equals only another `RunLengthEdges` with the same edges; + compare `tuple(edges)` against a tuple. """ __slots__ = ("_index_stops", "_lookup_tables", "_offset_stops", "counts", "sizes") @@ -430,10 +429,6 @@ def __getitem__(self, index: int | slice) -> int | RunLengthEdges: def __eq__(self, other: object) -> bool: if isinstance(other, RunLengthEdges): return self.sizes == other.sizes and self.counts == other.counts - if isinstance(other, tuple): - return len(other) == self.num_edges and all( - a == b for a, b in zip(self, other, strict=True) - ) return NotImplemented def __hash__(self) -> int: diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 9326c2eb75..098a2db15c 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -314,8 +314,8 @@ class RectilinearChunkGridMetadata(Metadata): (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 tuple of edges (and compares equal to it) but stores - one entry per run of equal edges, so a grid costs time and memory in + 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. This distinction matters for faithful round-tripping: a bare int diff --git a/tests/test_chunk_grids.py b/tests/test_chunk_grids.py index 5d96e5fe30..61153dff24 100644 --- a/tests/test_chunk_grids.py +++ b/tests/test_chunk_grids.py @@ -44,7 +44,7 @@ def _assert_chunks_equal( assert isinstance(a, VaryingDimension), ( f"axis {axis}: expected VaryingDimension, got {a!r}" ) - assert a.edges == tuple(e), f"axis {axis}: edges {a.edges} != {tuple(e)}" + assert tuple(a.edges) == tuple(e), f"axis {axis}: edges {a.edges} != {tuple(e)}" assert a.extent == span, f"axis {axis}: extent {a.extent} != {span}" diff --git a/tests/test_properties.py b/tests/test_properties.py index b448a524aa..ea8575d535 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -618,7 +618,7 @@ def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: src = zarr.create_array(MemoryStore(), shape=(3, 3), chunks=([1, 2], 1), dtype="uint8") grid = src.metadata.chunk_grid # type: ignore[union-attr] assert isinstance(grid, RectilinearChunkGridMetadata) - assert grid.chunk_shapes == ((1, 2), 1) + assert grid == RectilinearChunkGridMetadata(chunk_shapes=((1, 2), 1)) chunks = chunks_param_from_rectilinear(grid) # a list of lists, not tuples (`[1, 2] != (1, 2)`) assert chunks == [[1, 2], 1] @@ -646,7 +646,7 @@ def test_rectilinear_chunk_grid_declarations(data: st.DataObject) -> None: "configuration": {"kind": "inline", "chunk_shapes": declaration}, } meta = RectilinearChunkGridMetadata.from_dict(stored) - assert meta.chunk_shapes == chunk_shapes + assert meta == RectilinearChunkGridMetadata(chunk_shapes=chunk_shapes) serialized = json.loads(json.dumps(meta.to_dict())) assert serialized["name"] == "rectilinear" diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 59d34054a4..6b4b063eeb 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -250,7 +250,7 @@ def test_fixed_dimension_rejects_invalid(size: int, extent: int, match: str) -> def test_varying_dimension_construction() -> None: """VaryingDimension stores edges, cumulative sums, nchunks, and extent correctly""" d = VaryingDimension([10, 20, 30], extent=60) - assert d.edges == (10, 20, 30) + assert tuple(d.edges) == (10, 20, 30) assert d.cumulative == (10, 30, 60) assert d.nchunks == 3 assert d.extent == 60 @@ -646,7 +646,7 @@ def test_rle_expand_names_dimension( def test_parse_rle(rle_input: list[Any]) -> None: """parse_rle reads the edges expand_rle expands, as merged runs""" parsed = parse_rle(rle_input) - assert parsed == tuple(expand_rle(rle_input)) + assert list(parsed) == expand_rle(rle_input) assert parsed.to_rle() == compress_rle(expand_rle(rle_input)) @@ -667,30 +667,31 @@ def test_parse_rle(rle_input: list[Any]) -> None: ids=["one-edge", "one-run", "mixed-runs", "no-repeats", "adjacent-equal-runs"], ) def test_run_length_edges(runs: list[tuple[int, int]]) -> None: - """RunLengthEdges reads like the tuple of edges it encodes: length, indexing, slicing, - iteration, equality, prefix sums and position lookups all agree with the expanded - edges, and adjacent runs of one size are merged""" + """RunLengthEdges reads like the sequence of edges it encodes: length, indexing, + slicing, iteration, prefix sums and position lookups all agree with the expanded + edges, adjacent runs of one size are merged, and it equals only a RunLengthEdges of + the same edges, never the tuple of them""" expanded = tuple(size for size, count in runs for _ in range(count)) offsets = (0, *itertools.accumulate(expanded)) edges = RunLengthEdges(runs) - assert edges == expanded + assert edges != expanded assert edges == RunLengthEdges.from_edges(expanded) assert hash(edges) == hash(RunLengthEdges.from_edges(expanded)) - assert edges != (*expanded, 1) + assert edges != RunLengthEdges.from_edges((*expanded, 1)) assert pickle.loads(pickle.dumps(edges)) == edges assert all(a != b for (a, _), (b, _) in itertools.pairwise(edges.runs)) assert len(edges) == edges.num_edges == len(expanded) assert edges.total == sum(expanded) assert tuple(edges) == expanded assert [edges[i] for i in range(-len(expanded), len(expanded))] == [*expanded, *expanded] - assert edges[1:] == expanded[1:] - assert edges[:-1] == expanded[:-1] - assert edges[::2] == expanded[::2] + assert tuple(edges[1:]) == expanded[1:] + assert tuple(edges[:-1]) == expanded[:-1] + assert tuple(edges[::2]) == expanded[::2] assert all(edges.count(size) == expanded.count(size) for size in {*expanded, 99}) assert all((size in edges) == (size in expanded) for size in {*expanded, 99}) assert edges.to_rle() == compress_rle(expanded) - assert edges.with_edge(7) == (*expanded, 7) + assert tuple(edges.with_edge(7)) == (*expanded, 7) assert [edges.offset_of(i) for i in range(len(expanded) + 1)] == list(offsets) positions = np.arange(-2, edges.total + 3, dtype=np.intp) expected_indices = np.searchsorted(offsets[1:], positions, side="right") @@ -1352,7 +1353,7 @@ def test_mixed_scalar_and_list_dims_keep_shorthand(tmp_path: Path) -> None: assert isinstance(arr.metadata, ArrayV3Metadata) grid = arr.metadata.chunk_grid assert isinstance(grid, RectilinearChunkGridMetadata) - assert grid.chunk_shapes == (5, (10, 20, 70)) + assert grid == RectilinearChunkGridMetadata(chunk_shapes=(5, (10, 20, 70))) # The chunk/shard spec matrix for `from_array` round-trip tests. Every kind of @@ -1599,7 +1600,7 @@ def test_from_array_keep_is_o1_in_chunk_count() -> None: dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=False) grid = dst.metadata.chunk_grid assert isinstance(grid, RectilinearChunkGridMetadata) - assert grid.chunk_shapes == (3, (10, 20)) + assert grid == RectilinearChunkGridMetadata(chunk_shapes=(3, (10, 20))) @pytest.mark.parametrize( @@ -1624,7 +1625,7 @@ def test_resize_uniform_rectilinear_appends_edge( assert isinstance(rect.metadata, ArrayV3Metadata) rect_grid = rect.metadata.chunk_grid assert isinstance(rect_grid, RectilinearChunkGridMetadata) - assert rect_grid.chunk_shapes == ((10, 10, 10, 15),) + assert rect_grid == RectilinearChunkGridMetadata(chunk_shapes=((10, 10, 10, 15),)) assert rect.nchunks == 4 # the appended region is exactly the new edge chunk rect[30:45] = 1 @@ -2770,7 +2771,7 @@ def test_update_shape_shrink_creates_boundary() -> None: new_grid = grid.update_shape((45,)) dim = new_grid._dimensions[0] assert isinstance(dim, VaryingDimension) - assert dim.edges == (10, 20, 30) + assert tuple(dim.edges) == (10, 20, 30) assert dim.extent == 45 assert dim.chunk_size(2) == 30 assert dim.data_size(2) == 15 @@ -2782,7 +2783,7 @@ def test_update_shape_shrink_to_exact_boundary() -> None: new_grid = grid.update_shape((30,)) dim = new_grid._dimensions[0] assert isinstance(dim, VaryingDimension) - assert dim.edges == (10, 20, 30) + assert tuple(dim.edges) == (10, 20, 30) assert dim.nchunks == 2 assert dim.ngridcells == 3 assert dim.extent == 30 @@ -3314,7 +3315,7 @@ def test_rectilinear_from_dict( ) -> None: """RectilinearChunkGridMetadata.from_dict correctly parses all spec forms.""" grid = RectilinearChunkGridMetadata.from_dict(json_input) - assert grid.chunk_shapes == expected_chunk_shapes + assert grid == RectilinearChunkGridMetadata(chunk_shapes=expected_chunk_shapes) @pytest.mark.parametrize( From 6808f75292c558eedd6b797ab372d8e32f6ec198 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Oct 2026 17:47:12 +0200 Subject: [PATCH 4/7] 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 --- changes/4479.bugfix.md | 2 +- src/zarr/codecs/sharding.py | 4 +- src/zarr/core/chunk_grids.py | 16 +++---- src/zarr/core/common.py | 70 ++++++++++-------------------- src/zarr/core/metadata/v3.py | 36 +++++++-------- src/zarr/testing/strategies.py | 4 +- tests/test_chunk_grids.py | 4 +- tests/test_metadata/test_repair.py | 2 +- tests/test_properties.py | 10 +++-- tests/test_unified_chunk_grid.py | 51 ++++++++++++---------- 10 files changed, 92 insertions(+), 107 deletions(-) diff --git a/changes/4479.bugfix.md b/changes/4479.bugfix.md index 8a9b1955c1..9fab9fc88a 100644 --- a/changes/4479.bugfix.md +++ 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. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 27d67bf9bd..53a1605699 100644 --- a/src/zarr/codecs/sharding.py +++ b/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( diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 24ea7a7273..d693e2856e 100644 --- a/src/zarr/core/chunk_grids.py +++ b/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,7 +206,7 @@ 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. @@ -214,7 +214,7 @@ def data_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 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, ...]: @@ -222,7 +222,7 @@ def _data_sizes(self) -> tuple[int, ...]: 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`. """ diff --git a/src/zarr/core/common.py b/src/zarr/core/common.py index 324f624523..547134006a 100644 --- a/src/zarr/core/common.py +++ b/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. diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 098a2db15c..b6ddbfe8ef 100644 --- a/src/zarr/core/metadata/v3.py +++ b/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,7 +304,7 @@ 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. @@ -312,21 +312,22 @@ class RectilinearChunkGridMetadata(Metadata): - 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) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 693fc613f9..6bcaf78cb2 100644 --- a/src/zarr/testing/strategies.py +++ b/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 diff --git a/tests/test_chunk_grids.py b/tests/test_chunk_grids.py index 61153dff24..d4227dea5b 100644 --- a/tests/test_chunk_grids.py +++ b/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}" diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index 38ad102c6a..b5c28c0618 100644 --- a/tests/test_metadata/test_repair.py +++ b/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]] = { diff --git a/tests/test_properties.py b/tests/test_properties.py index ea8575d535..f0d5092b6a 100644 --- a/tests/test_properties.py +++ b/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) ) diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 6b4b063eeb..c59ee31dc4 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -60,7 +60,7 @@ def _edges(grid: ChunkGrid, dim: int) -> tuple[int, ...]: if isinstance(d, FixedDimension): return tuple(d.size for _ in range(d.nchunks)) if isinstance(d, VaryingDimension): - return tuple(d.edges) + return tuple(d.edges.expand()) raise TypeError(f"Unexpected dimension type: {type(d)}") @@ -250,7 +250,7 @@ def test_fixed_dimension_rejects_invalid(size: int, extent: int, match: str) -> def test_varying_dimension_construction() -> None: """VaryingDimension stores edges, cumulative sums, nchunks, and extent correctly""" d = VaryingDimension([10, 20, 30], extent=60) - assert tuple(d.edges) == (10, 20, 30) + assert tuple(d.edges.expand()) == (10, 20, 30) assert d.cumulative == (10, 30, 60) assert d.nchunks == 3 assert d.extent == 60 @@ -646,7 +646,7 @@ def test_rle_expand_names_dimension( def test_parse_rle(rle_input: list[Any]) -> None: """parse_rle reads the edges expand_rle expands, as merged runs""" parsed = parse_rle(rle_input) - assert list(parsed) == expand_rle(rle_input) + assert list(parsed.expand()) == expand_rle(rle_input) assert parsed.to_rle() == compress_rle(expand_rle(rle_input)) @@ -667,10 +667,10 @@ def test_parse_rle(rle_input: list[Any]) -> None: ids=["one-edge", "one-run", "mixed-runs", "no-repeats", "adjacent-equal-runs"], ) def test_run_length_edges(runs: list[tuple[int, int]]) -> None: - """RunLengthEdges reads like the sequence of edges it encodes: length, indexing, - slicing, iteration, prefix sums and position lookups all agree with the expanded - edges, adjacent runs of one size are merged, and it equals only a RunLengthEdges of - the same edges, never the tuple of them""" + """RunLengthEdges answers for the edges it encodes: the edge count, each edge's size, + the expansion, prefix sums and position lookups all agree with the expanded edges, + adjacent runs of one size are merged, and it equals only a RunLengthEdges of the + same edges, never the tuple of them""" expanded = tuple(size for size, count in runs for _ in range(count)) offsets = (0, *itertools.accumulate(expanded)) edges = RunLengthEdges(runs) @@ -681,17 +681,12 @@ def test_run_length_edges(runs: list[tuple[int, int]]) -> None: assert edges != RunLengthEdges.from_edges((*expanded, 1)) assert pickle.loads(pickle.dumps(edges)) == edges assert all(a != b for (a, _), (b, _) in itertools.pairwise(edges.runs)) - assert len(edges) == edges.num_edges == len(expanded) + assert edges.num_edges == len(expanded) assert edges.total == sum(expanded) - assert tuple(edges) == expanded - assert [edges[i] for i in range(-len(expanded), len(expanded))] == [*expanded, *expanded] - assert tuple(edges[1:]) == expanded[1:] - assert tuple(edges[:-1]) == expanded[:-1] - assert tuple(edges[::2]) == expanded[::2] - assert all(edges.count(size) == expanded.count(size) for size in {*expanded, 99}) - assert all((size in edges) == (size in expanded) for size in {*expanded, 99}) + assert tuple(edges.expand()) == expanded + assert [edges.size_of(i) for i in range(len(expanded))] == list(expanded) assert edges.to_rle() == compress_rle(expanded) - assert tuple(edges.with_edge(7)) == (*expanded, 7) + assert tuple(edges.with_edge(7).expand()) == (*expanded, 7) assert [edges.offset_of(i) for i in range(len(expanded) + 1)] == list(offsets) positions = np.arange(-2, edges.total + 3, dtype=np.intp) expected_indices = np.searchsorted(offsets[1:], positions, side="right") @@ -705,11 +700,23 @@ def test_run_length_edges_rejects_count_below_one() -> None: RunLengthEdges([(5, 0)]) -@pytest.mark.parametrize("index", [3, -4]) -def test_run_length_edges_getitem_out_of_range(index: int) -> None: - """Indexing RunLengthEdges past either end raises IndexError""" +@pytest.mark.parametrize("index", [3, -1]) +def test_run_length_edges_size_of_out_of_range(index: int) -> None: + """size_of raises IndexError for an index that is not an edge""" with pytest.raises(IndexError, match=f"Edge index {index} is out of range for 3 edges"): - RunLengthEdges([(5, 3)])[index] + RunLengthEdges([(5, 3)]).size_of(index) + + +@pytest.mark.parametrize( + "use", + [len, iter, lambda edges: edges[0], lambda edges: 5 in edges], + ids=["len", "iter", "getitem", "contains"], +) +def test_run_length_edges_is_not_a_sequence(use: Callable[[Any], object]) -> None: + """RunLengthEdges has no length, iteration, indexing or membership test, so no caller + visits every edge without asking for `expand`""" + with pytest.raises(TypeError): + use(RunLengthEdges([(5, 3)])) @pytest.mark.parametrize("index", [4, -1]) @@ -2771,7 +2778,7 @@ def test_update_shape_shrink_creates_boundary() -> None: new_grid = grid.update_shape((45,)) dim = new_grid._dimensions[0] assert isinstance(dim, VaryingDimension) - assert tuple(dim.edges) == (10, 20, 30) + assert tuple(dim.edges.expand()) == (10, 20, 30) assert dim.extent == 45 assert dim.chunk_size(2) == 30 assert dim.data_size(2) == 15 @@ -2783,7 +2790,7 @@ def test_update_shape_shrink_to_exact_boundary() -> None: new_grid = grid.update_shape((30,)) dim = new_grid._dimensions[0] assert isinstance(dim, VaryingDimension) - assert tuple(dim.edges) == (10, 20, 30) + assert tuple(dim.edges.expand()) == (10, 20, 30) assert dim.nchunks == 2 assert dim.ngridcells == 3 assert dim.extent == 30 From 034444270101f675a7d04c65981ba2d992fa9a91 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Oct 2026 17:53:48 +0200 Subject: [PATCH 5/7] refactor(chunk-grids): remove VaryingDimension.cumulative The per-chunk prefix sums were no longer used by any lookup and cost time and memory in the number of chunks on every access. `chunk_offset` gives the start of a chunk from the runs. Assisted-by: ClaudeCode:claude-fable-5.1 --- changes/4479.bugfix.md | 2 +- src/zarr/core/chunk_grids.py | 9 --------- tests/test_unified_chunk_grid.py | 3 +-- 3 files changed, 2 insertions(+), 12 deletions(-) diff --git a/changes/4479.bugfix.md b/changes/4479.bugfix.md index 9fab9fc88a..649e3c1b3f 100644 --- a/changes/4479.bugfix.md +++ 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 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. +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 removed; `VaryingDimension.chunk_offset` gives the start of a chunk. The repr of a chunk grid shows a dimension of more than 100 chunks in its run-length encoded form. diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index d693e2856e..bc5e2df816 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -178,15 +178,6 @@ def __init__(self, edges: Sequence[int] | RunLengthEdges, extent: int) -> None: object.__setattr__(self, "nchunks", n) object.__setattr__(self, "ngridcells", rle.num_edges) - @property - def cumulative(self) -> tuple[int, ...]: - """Prefix sums of the edges, one entry per chunk. - - 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.expand())) - 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}") diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index c59ee31dc4..901595632b 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -248,10 +248,9 @@ def test_fixed_dimension_rejects_invalid(size: int, extent: int, match: str) -> def test_varying_dimension_construction() -> None: - """VaryingDimension stores edges, cumulative sums, nchunks, and extent correctly""" + """VaryingDimension stores edges, nchunks, and extent correctly""" d = VaryingDimension([10, 20, 30], extent=60) assert tuple(d.edges.expand()) == (10, 20, 30) - assert d.cumulative == (10, 30, 60) assert d.nchunks == 3 assert d.extent == 60 From e529ad748fa416191e1e76530ba658dd028aed8f Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Oct 2026 18:05:52 +0200 Subject: [PATCH 6/7] refactor(chunk-grids): remove unused members of the dimension grids Nothing in the library called `with_extent`, `ngridcells` or `_unique_edge_lengths` on `FixedDimension` / `VaryingDimension`, did an `isinstance` check against `DimensionGrid`, or used `is_regular_1d` / `is_regular_nd`. Remove them along with the tests that only exercised them. `VaryingDimension.chunk_offset` no longer maps a negative chunk index to 0, a guard left over from indexing a tuple of prefix sums. Both dimension types raise the same `IndexError` message from `index_to_chunk`, and docstrings that described the old tuple-backed behavior are corrected. Assisted-by: ClaudeCode:claude-fable-5.1 --- changes/4479.bugfix.md | 2 +- src/zarr/core/chunk_grids.py | 116 +++++-------------------------- tests/test_metadata/test_v3.py | 43 +----------- tests/test_unified_chunk_grid.py | 67 ++++-------------- 4 files changed, 34 insertions(+), 194 deletions(-) diff --git a/changes/4479.bugfix.md b/changes/4479.bugfix.md index 649e3c1b3f..46c595ed2e 100644 --- a/changes/4479.bugfix.md +++ 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 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 removed; `VaryingDimension.chunk_offset` gives the start of a chunk. 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 removed; `VaryingDimension.chunk_offset` gives the start of a chunk. The repr of a chunk grid shows a dimension of more than 100 chunks in its run-length encoded form. The unused `with_extent` and `ngridcells` members of `FixedDimension` and `VaryingDimension`, and the unused helpers `zarr.core.chunk_grids.is_regular_1d` and `is_regular_nd`, are removed; both dimension types now raise the same `IndexError` message from `index_to_chunk`. diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index bc5e2df816..94a8a273f3 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -14,7 +14,6 @@ Protocol, SupportsIndex, cast, - runtime_checkable, ) import numpy as np @@ -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,34 +89,12 @@ 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 @@ -159,7 +131,6 @@ class VaryingDimension: 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 def __init__(self, edges: Sequence[int] | RunLengthEdges, extent: int) -> None: rle = RunLengthEdges.from_edges(edges) @@ -176,7 +147,6 @@ def __init__(self, edges: Sequence[int] | RunLengthEdges, extent: int) -> None: # Cache nchunks: number of chunks that overlap [0, extent) n = 0 if extent == 0 else rle.index_at(extent - 1) + 1 object.__setattr__(self, "nchunks", n) - object.__setattr__(self, "ngridcells", rle.num_edges) def index_to_chunk(self, idx: int) -> int: if idx < 0 or idx >= self.extent: @@ -184,26 +154,26 @@ def index_to_chunk(self, idx: int) -> int: 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.edges.offset_of(chunk_ix) 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.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)``. """ return max(0, min(self.edges.size_of(chunk_ix), self.extent - self.chunk_offset(chunk_ix))) @@ -215,38 +185,14 @@ def _data_sizes(self) -> tuple[int, ...]: last = self.nchunks - 1 return (*itertools.islice(self.edges.expand(), last), self.data_size(last)) - @property - def _unique_edge_lengths(self) -> Iterable[int]: - """Distinct chunk edge lengths for this dimension, in order of first use. - - Used by shard validation to check that every unique edge length - is divisible by the inner chunk size. Reads one size per run of - equal edges, not one per chunk. - """ - return tuple(dict.fromkeys(self.edges.sizes)) - def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: return self.edges.indices_at(indices) - 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.edges.total - 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) - 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). """ @@ -263,24 +209,18 @@ def _size_repr(self) -> str: 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: ... @@ -365,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: """ diff --git a/tests/test_metadata/test_v3.py b/tests/test_metadata/test_v3.py index 94240c36b9..680197717f 100644 --- a/tests/test_metadata/test_v3.py +++ b/tests/test_metadata/test_v3.py @@ -11,7 +11,7 @@ from tests.conftest import Expect, ExpectFail from tests.test_metadata.conftest import minimal_metadata_dict_v3 from zarr.core.buffer import default_buffer_prototype -from zarr.core.chunk_grids import ChunkGrid, is_regular_1d, is_regular_nd +from zarr.core.chunk_grids import ChunkGrid from zarr.core.config import config from zarr.core.dtype import Float64, UInt8 from zarr.core.group import GroupMetadata, parse_node_type @@ -108,47 +108,6 @@ def test_parse_codecs_unknown_raises(monkeypatch: pytest.MonkeyPatch) -> None: parse_codecs([{"name": "unknown"}]) -# --------------------------------------------------------------------------- -# Chunk-grid regularity helpers -# --------------------------------------------------------------------------- - -_REGULAR_1D_CASES: list[Expect[list[int], bool]] = [ - Expect(input=[], output=True, id="empty"), - Expect(input=[10], output=True, id="single-chunk"), - Expect(input=[10, 10, 10], output=True, id="all-equal"), - Expect(input=[10, 10, 10, 7], output=True, id="smaller-boundary"), - Expect(input=[10, 10, 10, 10], output=True, id="exact-multiple"), - Expect(input=[10, 5, 10], output=False, id="middle-mismatch"), - Expect(input=[10, 10, 10, 12], output=False, id="last-larger"), - # The first chunk anchors the size; later mismatches in the middle fail - # before the boundary check. - Expect(input=[5, 10, 5], output=False, id="middle-larger"), -] - - -@pytest.mark.parametrize("case", _REGULAR_1D_CASES, ids=lambda c: c.id) -def test_is_regular_1d_sequence(case: Expect[list[int], bool]) -> None: - """`is_regular_1d` accepts plain Python sequences.""" - assert is_regular_1d(case.input) is case.output - assert is_regular_1d(tuple(case.input)) is case.output - - -@pytest.mark.parametrize( - "case", - [ - Expect(input=[[10, 10, 10], [5, 5]], output=True, id="all-regular"), - Expect(input=[[10, 10, 7], [5, 5, 5, 3]], output=True, id="all-regular-with-boundary"), - Expect(input=[[10, 10, 10], [5, 8, 5]], output=False, id="second-dim-irregular"), - Expect(input=[[10, 5, 10], [5, 5]], output=False, id="first-dim-irregular"), - Expect(input=[], output=True, id="zero-dims"), - ], - ids=lambda c: c.id, -) -def test_is_regular_nd_sequence(case: Expect[list[list[int]], bool]) -> None: - """`is_regular_nd` returns True iff every per-dim spec is regular.""" - assert is_regular_nd(case.input) is case.output - - def test_create_chunk_grid_metadata_unknown_dimension_type() -> None: """`create_chunk_grid_metadata` rejects dimension grids it does not recognize.""" grid = ChunkGrid(dimensions=(object(),)) # type: ignore[arg-type] diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 901595632b..9c219fdd6b 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -75,13 +75,15 @@ def _edges(grid: ChunkGrid, dim: int) -> tuple[int, ...]: (VaryingDimension([10, 20, 30], extent=60), 60, "out of bounds"), (VaryingDimension([10, 20, 30], extent=60), 100, "out of bounds"), (FixedDimension(size=10, extent=95), 95, "out of bounds"), - (FixedDimension(size=10, extent=95), -1, "Negative"), + (FixedDimension(size=10, extent=95), -1, "out of bounds"), + (VaryingDimension([10, 20, 30], extent=60), -1, "out of bounds"), ], ids=[ "varying-at-extent", "varying-past-extent", "fixed-at-extent", "fixed-negative", + "varying-negative", ], ) def test_dimension_index_to_chunk_bounds( @@ -748,9 +750,9 @@ def test_varying_dimension_huge_runs() -> None: edge_sum = 3 * HUGE + 7 + 2 * HUGE dim = VaryingDimension(RunLengthEdges([(3, HUGE), (7, 1), (2, HUGE)]), extent=edge_sum - 3) - assert dim.ngridcells == 2 * HUGE + 1 + assert dim.edges.num_edges == 2 * HUGE + 1 assert dim.nchunks == 2 * HUGE - assert tuple(dim._unique_edge_lengths) == (3, 7, 2) + assert dim.edges.sizes == (3, 7, 2) assert dim._size_repr == f"[[3, {HUGE}], 7, [2, {HUGE}]]" lookups = { 0: 0, @@ -779,7 +781,7 @@ def test_varying_dimension_huge_runs() -> None: grown = dim.resize(edge_sum + 5) assert grown.edges.to_rle() == [[3, HUGE], 7, [2, HUGE], 5] assert dim.resize(10).edges == dim.edges - assert dim.with_extent(edge_sum).nchunks == 2 * HUGE + 1 + assert VaryingDimension(dim.edges, extent=edge_sum).nchunks == 2 * HUGE + 1 def _store_rectilinear_array( @@ -1017,28 +1019,6 @@ def test_spec_example() -> None: # --------------------------------------------------------------------------- -def test_parse_chunk_grid_varying_extent_mismatch_raises() -> None: - """Reconstructing a ChunkGrid with mismatched extents raises ValueError""" - g = ChunkGrid.from_sizes((60, 100), [[10, 20, 30], [50, 50]]) - with pytest.raises(ValueError, match="extent"): - ChunkGrid( - dimensions=tuple( - dim.with_extent(ext) for dim, ext in zip(g._dimensions, (100, 100), strict=True) - ) - ) - - -def test_parse_chunk_grid_varying_extent_match_ok() -> None: - """Reconstructing a ChunkGrid with matching extents succeeds""" - g = ChunkGrid.from_sizes((60, 100), [[10, 20, 30], [50, 50]]) - g2 = ChunkGrid( - dimensions=tuple( - dim.with_extent(ext) for dim, ext in zip(g._dimensions, (60, 100), strict=True) - ) - ) - assert g2._dimensions[0].extent == 60 - - @pytest.mark.parametrize( ("chunk_shapes", "array_shape", "match"), [ @@ -1099,17 +1079,6 @@ def test_parse_chunk_grid_rectilinear_rle_extent_validated() -> None: ChunkGrid.from_sizes((100, 50), meta.chunk_shapes) -def test_parse_chunk_grid_varying_dimension_extent_mismatch_on_chunkgrid_input() -> None: - """ChunkGrid constructor rejects VaryingDimension with extent exceeding sum of edges""" - g = ChunkGrid.from_sizes((60, 50), [[10, 20, 30], [25, 25]]) - with pytest.raises(ValueError, match="less than"): - ChunkGrid( - dimensions=tuple( - dim.with_extent(ext) for dim, ext in zip(g._dimensions, (100, 50), strict=True) - ) - ) - - # --------------------------------------------------------------------------- # Rectilinear indexing tests # --------------------------------------------------------------------------- @@ -1770,10 +1739,9 @@ def test_fixed_dimension_zero_extent(size: int) -> None: """A zero-length axis has zero chunks and behaves like an empty grid.""" d = FixedDimension(size=size, extent=0) assert d.nchunks == 0 - assert d.ngridcells == 0 assert d.data_size(0) == 0 - assert d.with_extent(0) == d - assert d.with_extent(3) == FixedDimension(size=size, extent=3) + assert d.resize(0) == d + assert d.resize(3) == FixedDimension(size=size, extent=3) empty = np.array([], dtype=np.intp) np.testing.assert_array_equal(d.indices_to_chunks(empty), empty) with pytest.raises(IndexError): @@ -1826,9 +1794,7 @@ def test_parse_chunk_grid_preserves_varying_extent() -> None: assert g._dimensions[0].extent == 60 g2 = ChunkGrid( - dimensions=tuple( - dim.with_extent(ext) for dim, ext in zip(g._dimensions, (60, 100), strict=True) - ) + dimensions=tuple(dim.resize(ext) for dim, ext in zip(g._dimensions, (60, 100), strict=True)) ) assert isinstance(g2._dimensions[0], VaryingDimension) assert g2._dimensions[0].extent == 60 @@ -1840,9 +1806,7 @@ def test_parse_chunk_grid_rebinds_fixed_extent() -> None: assert g._dimensions[0].extent == 100 g2 = ChunkGrid( - dimensions=tuple( - dim.with_extent(ext) for dim, ext in zip(g._dimensions, (50, 100), strict=True) - ) + dimensions=tuple(dim.resize(ext) for dim, ext in zip(g._dimensions, (50, 100), strict=True)) ) assert isinstance(g2._dimensions[0], FixedDimension) assert g2._dimensions[0].extent == 50 @@ -2447,7 +2411,7 @@ def test_varying_dimension_zero_extent() -> None: """VaryingDimension with extent=0 has zero active chunks but retains all grid cells.""" d = VaryingDimension([10, 20], extent=0) assert d.nchunks == 0 - assert d.ngridcells == 2 + assert d.edges.num_edges == 2 # No chunks overlap [0, 0), so the grid is structurally non-empty but logically empty g = ChunkGrid(dimensions=(d,)) assert g.grid_shape == (0,) @@ -2483,7 +2447,8 @@ def test_overflow_multiple_chunks_past_extent() -> None: """Edges past extent are structural; nchunks counts active only.""" g = ChunkGrid.from_sizes((50,), [[10, 20, 30, 40]]) d = g._dimensions[0] - assert d.ngridcells == 4 + assert isinstance(d, VaryingDimension) + assert d.edges.num_edges == 4 assert d.nchunks == 3 assert d.data_size(0) == 10 assert d.data_size(1) == 20 @@ -2791,7 +2756,7 @@ def test_update_shape_shrink_to_exact_boundary() -> None: assert isinstance(dim, VaryingDimension) assert tuple(dim.edges.expand()) == (10, 20, 30) assert dim.nchunks == 2 - assert dim.ngridcells == 3 + assert dim.edges.num_edges == 3 assert dim.extent == 30 assert dim.data_size(1) == 20 @@ -2811,9 +2776,7 @@ def test_update_shape_parse_chunk_grid_rebinds_extent() -> None: """parse_chunk_grid re-binds VaryingDimension extent to array shape.""" g = ChunkGrid.from_sizes((60,), [[10, 20, 30]]) g2 = ChunkGrid( - dimensions=tuple( - dim.with_extent(ext) for dim, ext in zip(g._dimensions, (50,), strict=True) - ) + dimensions=tuple(dim.resize(ext) for dim, ext in zip(g._dimensions, (50,), strict=True)) ) dim = g2._dimensions[0] assert isinstance(dim, VaryingDimension) From 1c1d0f3794039b5ec537415f014b96efc3da22de Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Oct 2026 21:18:25 +0200 Subject: [PATCH 7/7] docs: rewrite the changelog entry for users and file it as misc The entry named private internals and read as a bug fix. Describe only the user-visible effect, and classify it as `misc` since nothing was incorrect before. Assisted-by: ClaudeCode:claude-fable-5.1 --- changes/4479.bugfix.md | 1 - changes/4479.misc.md | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) delete mode 100644 changes/4479.bugfix.md create mode 100644 changes/4479.misc.md diff --git a/changes/4479.bugfix.md b/changes/4479.bugfix.md deleted file mode 100644 index 46c595ed2e..0000000000 --- a/changes/4479.bugfix.md +++ /dev/null @@ -1 +0,0 @@ -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 removed; `VaryingDimension.chunk_offset` gives the start of a chunk. The repr of a chunk grid shows a dimension of more than 100 chunks in its run-length encoded form. The unused `with_extent` and `ngridcells` members of `FixedDimension` and `VaryingDimension`, and the unused helpers `zarr.core.chunk_grids.is_regular_1d` and `is_regular_nd`, are removed; both dimension types now raise the same `IndexError` message from `index_to_chunk`. diff --git a/changes/4479.misc.md b/changes/4479.misc.md new file mode 100644 index 0000000000..b7c51cc3e6 --- /dev/null +++ b/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.