test(chunk-grids): sample the full rectilinear chunk grid declaration space

changes/4377.misc.md

@@ -0,0 +1 @@
+The hypothesis strategies in `zarr.testing.strategies` now sample the full space of rectilinear chunk grid declarations: `arrays()` passes a rectilinear grid to `create_array` either as metadata or as a `chunks=` specification mixing bare-int steps and explicit edge lists in any arrangement, and the new `rectilinear_chunk_shape_declarations` and `rectilinear_chunk_grids` cover stored metadata, including run-length encoded edge lists in canonical and arbitrary groupings and edges overhanging the array extent. `rectilinear_chunks` and `chunk_grids` also draw zero-length dimensions.

src/zarr/testing/strategies.py

@@ -34,7 +34,11 @@
 from zarr.core.dtype.npy.structured import Struct
 from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
-from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata
+from zarr.core.metadata.v3 import (
+    RectilinearChunkGridMetadata,
+    RectilinearDimSpecJSON,
+    RegularChunkGridMetadata,
+)
 from zarr.core.sync import sync
 from zarr.storage import MemoryStore, StoreLike
 from zarr.storage._utils import _join_paths, normalize_path
@@ -444,18 +448,27 @@ def arrays(
     array_path = _join_paths([path, name])
     root = zarr.open_group(store, mode=open_mode, zarr_format=zarr_format)
 
-    # Convert chunk grid metadata to a form create_array accepts:
-    # - RegularChunkGridMetadata -> flat tuple of ints
-    # - RectilinearChunkGridMetadata -> nested list of ints (triggers rectilinear path)
-    # - v2 -> flat tuple of ints
-    chunks_param: tuple[int, ...] | list[int | list[int]]
+    chunks_param: tuple[int, ...] | list[int | list[int]] | RectilinearChunkGridMetadata
     shard_shape = None
     dim_names = None
     if zarr_format == 3:
         chunk_grid_meta = draw(st.none() | chunk_grids(shape=nparray.shape), label="chunk grid")
         dim_names = draw(dimension_names(ndim=nparray.ndim), label="dimension names")
         if isinstance(chunk_grid_meta, RectilinearChunkGridMetadata):
-            chunks_param = chunks_param_from_rectilinear(chunk_grid_meta)
+            # A rectilinear grid is passed either as the metadata object or in
+            # the list form of `chunks=`, drawn from its own strategy. A 0-d
+            # array has no dimension to hold an edge list.
+            if nparray.ndim > 0 and draw(st.booleans(), label="chunks as lists"):
+                event("rectilinear chunks= as lists")
+                chunks_param = draw(_rectilinear_chunks(shape=nparray.shape), label="chunks")
+                chunk_grid_meta = RectilinearChunkGridMetadata(
+                    chunk_shapes=tuple(
+                        dim if isinstance(dim, int) else tuple(dim) for dim in chunks_param
+                    )
+                )
+            else:
+                event("rectilinear chunks= as metadata")
+                chunks_param = chunk_grid_meta
         elif isinstance(chunk_grid_meta, RegularChunkGridMetadata):
             chunks_param = chunk_grid_meta.chunk_shape
         else:
@@ -505,7 +518,9 @@ def arrays(
             )
             assert shard_shape == a.shards
         else:
-            assert isinstance(a.metadata.chunk_grid, RectilinearChunkGridMetadata)
+            # The stored grid is exactly the declared one: bare ints stay bare
+            # ints, edge lists keep their edges.
+            assert a.metadata.chunk_grid == chunk_grid_meta
             assert shard_shape is None
 
     assert a.basename == name, (a.basename, name)
@@ -546,59 +561,190 @@ def chunks_param_from_rectilinear(
     return [list(dim) if isinstance(dim, tuple) else dim for dim in meta.chunk_shapes]
 
 
+# The most chunks a drawn rectilinear grid declares along one axis, and over
+# the array, all axes together. Indexing tests visit every chunk, so the
+# product is what costs time.
+_RECTILINEAR_MAX_CHUNKS_PER_DIM = 20
+_RECTILINEAR_CHUNK_BUDGET = 400
+
+
+def _max_chunks_per_dim(ndim: int) -> int:
+    """At most `_RECTILINEAR_MAX_CHUNKS_PER_DIM` chunks per dimension, and at
+    most `_RECTILINEAR_CHUNK_BUDGET` in total over `ndim` dimensions."""
+    return max(
+        k
+        for k in range(1, _RECTILINEAR_MAX_CHUNKS_PER_DIM + 1)
+        if k**ndim <= _RECTILINEAR_CHUNK_BUDGET
+    )
+
+
+@st.composite
+def rectilinear_dim_edges(
+    draw: st.DrawFn, *, extent: int, max_chunks: int = _RECTILINEAR_MAX_CHUNKS_PER_DIM
+) -> list[int]:
+    """Explicit chunk edge lengths summing exactly to `extent`.
+
+    A zero `extent` has no chunks for edges to cover, and no non-empty list
+    of positive edges sums to 0; its edges are the chunks the axis grows into
+    on `append` or `resize`, so any non-empty list of positive edges is drawn.
+
+    Two modes: "uneven" cuts the extent at random dividers; "uniform" repeats
+    one size with an optional remainder, optionally shuffled so equal edges
+    are not all adjacent. At most `max_chunks` chunks keeps property tests
+    fast.
+    """
+    assert extent >= 0
+    if extent == 0:
+        event("rectilinear edges: zero extent")
+        return draw(st.lists(st.integers(min_value=1, max_value=10), min_size=1, max_size=5))
+    if draw(st.booleans(), label="uneven edges"):
+        nchunks = draw(st.integers(min_value=1, max_value=min(extent, max_chunks)))
+        # Draw distinct dividers by index into the unused positions: no
+        # rejection, unlike `st.lists(..., unique=True)`.
+        positions = list(range(1, extent))
+        dividers = sorted(
+            positions.pop(draw(st.integers(min_value=0, max_value=len(positions) - 1)))
+            for _ in range(nchunks - 1)
+        )
+        return [b - a for a, b in zip([0, *dividers], [*dividers, extent], strict=True)]
+    size = draw(st.integers(min_value=math.ceil(extent / max_chunks), max_value=extent))
+    edges = [size] * (extent // size)
+    if extent % size:
+        edges.append(extent % size)
+    if draw(st.booleans(), label="shuffle uniform edges"):
+        return list(draw(st.permutations(edges)))
+    return edges
+
+
+def _rectilinear_step(draw: st.DrawFn, *, extent: int, max_chunks: int) -> int:
+    """A bare-int chunk size for one dimension: a step that repeats to cover
+    the extent, with the last chunk possibly smaller. A step larger than the
+    extent (one overhanging chunk) is allowed, as for a regular grid."""
+    step = draw(st.integers(min_value=max(1, math.ceil(extent / max_chunks)), max_value=extent + 3))
+    if step > extent:
+        event("rectilinear step: larger than extent")
+    return step
+
+
 @st.composite
 def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[int]]:
-    """Generate valid rectilinear chunk shapes for a given array shape.
+    """A `chunks=` specification declaring a rectilinear grid over `shape`, as an
+    explicit edge list per dimension summing to the extent (any edges, for a zero
+    extent). A 0-d `shape` gives `[]`.
+
+    To also draw bare-int steps, see `_rectilinear_chunks`.
+    """
+    max_chunks = _max_chunks_per_dim(len(shape))
+    return [draw(rectilinear_dim_edges(extent=e, max_chunks=max_chunks)) for e in shape]
 
-    Uses two modes per dimension:
-    - "expanded": random divider points create arbitrary chunk sizes
-    - "rle": uniform chunks with optional remainder, optionally shuffled
 
-    Keeps max chunks per dimension <= 20 to avoid performance issues
-    in property tests. With higher dimensions, the total chunk count
-    grows multiplicatively.
+@st.composite
+def _rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[int | list[int]]:
+    """A `chunks=` specification declaring a rectilinear grid over `shape`.
+
+    Each dimension is either a bare int (a step size; the last chunk may be
+    smaller) or an explicit edge list summing to the extent (any edges, for a
+    zero extent). At least one dimension is an edge list, since bare ints
+    alone declare a regular grid.
+    Run-length encoding is not part of the `chunks=` syntax; it belongs to
+    stored metadata, see `rectilinear_chunk_shape_declarations`.
+
+    `shape` must have at least one dimension: a 0-d array has no dimension
+    to hold an edge list, so it cannot have a rectilinear grid.
     """
-    chunk_shapes: list[list[int]] = []
-    for size in shape:
-        assert size > 0
-        if size > 1:
-            mode = draw(st.sampled_from(["expanded", "rle"]))
-            if mode == "expanded":
-                event("rectilinear expanded")
-                max_chunks = min(size - 1, 20)
-                nchunks = draw(st.integers(min_value=1, max_value=max_chunks))
-                dividers = sorted(
-                    draw(
-                        st.lists(
-                            st.integers(min_value=1, max_value=size - 1),
-                            min_size=nchunks - 1,
-                            max_size=nchunks - 1,
-                            unique=True,
-                        )
-                    )
-                )
-                chunk_shapes.append(
-                    [a - b for a, b in zip(dividers + [size], [0] + dividers, strict=False)]
-                )
-            else:
-                # RLE mode: uniform chunks with optional remainder
-                max_chunk_size = min(size, 20)
-                chunk_size = draw(st.integers(min_value=1, max_value=max_chunk_size))
-                n_full = size // chunk_size
-                remainder = size % chunk_size
-                chunks_list = [chunk_size] * n_full
-                if remainder > 0:
-                    chunks_list.append(remainder)
-                # Optionally shuffle to create non-contiguous duplicate patterns
-                if draw(st.booleans()):
-                    event("rectilinear rle shuffled")
-                    chunks_list = draw(st.permutations(chunks_list))
-                else:
-                    event("rectilinear rle")
-                chunk_shapes.append(list(chunks_list))
+    assert shape, "a rectilinear grid needs at least one dimension"
+    max_chunks = _max_chunks_per_dim(len(shape))
+    forced_list = draw(st.integers(min_value=0, max_value=len(shape) - 1))
+    chunks: list[int | list[int]] = []
+    for i, extent in enumerate(shape):
+        if i != forced_list and draw(st.booleans(), label="bare int"):
+            chunks.append(_rectilinear_step(draw, extent=extent, max_chunks=max_chunks))
         else:
-            chunk_shapes.append([1])
-    return chunk_shapes
+            chunks.append(draw(rectilinear_dim_edges(extent=extent, max_chunks=max_chunks)))
+    return chunks
+
+
+def _rle_encode(draw: st.DrawFn, edges: list[int]) -> list[int | list[int]]:
+    """Run-length encode `edges` as the spec allows: a mix of bare ints and
+    `[size, count]` pairs. Either the canonical form (each run as one pair,
+    a run of one as a bare int) or an arbitrary grouping, which may split a
+    run across pairs and use `count == 1`."""
+    canonical = draw(st.booleans(), label="canonical rle")
+    if not canonical:
+        event("rectilinear rle: arbitrary grouping")
+    encoded: list[int | list[int]] = []
+    i = 0
+    while i < len(edges):
+        run = 1
+        while i + run < len(edges) and edges[i + run] == edges[i]:
+            run += 1
+        if canonical:
+            count = run
+            bare = run == 1
+        else:
+            count = draw(st.integers(min_value=1, max_value=run))
+            bare = count == 1 and draw(st.booleans(), label="bare edge")
+        encoded.append(edges[i] if bare else [edges[i], count])
+        i += count
+    return encoded
+
+
+def _rectilinear_chunk_shapes(
+    draw: st.DrawFn, shape: tuple[int, ...]
+) -> tuple[int | tuple[int, ...], ...]:
+    """The `chunk_shapes` of a stored rectilinear grid over `shape`: per
+    dimension a bare-int step, or explicit edges. Edges may sum beyond the
+    extent, which the spec allows and a shrinking resize produces."""
+    max_chunks = _max_chunks_per_dim(len(shape))
+    chunk_shapes: list[int | tuple[int, ...]] = []
+    for extent in shape:
+        if draw(st.booleans(), label="bare int"):
+            chunk_shapes.append(_rectilinear_step(draw, extent=extent, max_chunks=max_chunks))
+            continue
+        edges = draw(rectilinear_dim_edges(extent=extent, max_chunks=max_chunks))
+        if extent > 0 and draw(st.booleans(), label="overhang"):
+            event("rectilinear edges: overhang")
+            if draw(st.booleans(), label="trailing edge"):
+                edges = [*edges, draw(st.integers(min_value=1, max_value=5))]
+            else:
+                edges[-1] += draw(st.integers(min_value=1, max_value=5))
+        chunk_shapes.append(tuple(edges))
+    return tuple(chunk_shapes)
+
+
+@st.composite
+def rectilinear_chunk_shape_declarations(
+    draw: st.DrawFn, *, shape: tuple[int, ...]
+) -> tuple[list[RectilinearDimSpecJSON], tuple[int | tuple[int, ...], ...]]:
+    """The `chunk_shapes` of a stored rectilinear chunk grid, with its meaning.
+
+    Samples the whole declaration space of the spec. Per dimension: a bare
+    int step, or an edge list written in full or run-length encoded
+    (canonically, or with arbitrary grouping). Edge lists may sum beyond the
+    extent.
+
+    Returns `(declaration, chunk_shapes)`: the JSON value to store, and the
+    `chunk_shapes` that parsing it must produce.
+    """
+    chunk_shapes = _rectilinear_chunk_shapes(draw, shape)
+    declaration: list[RectilinearDimSpecJSON] = [
+        dim
+        if isinstance(dim, int)
+        else list(dim)
+        if draw(st.booleans(), label="write edges in full")
+        else _rle_encode(draw, list(dim))
+        for dim in chunk_shapes
+    ]
+    return declaration, chunk_shapes
+
+
+@st.composite
+def rectilinear_chunk_grids(
+    draw: st.DrawFn, *, shape: tuple[int, ...]
+) -> RectilinearChunkGridMetadata:
+    """A `RectilinearChunkGridMetadata` over `shape`, per dimension a bare-int
+    step or explicit edges, which may sum beyond the extent."""
+    return RectilinearChunkGridMetadata(chunk_shapes=_rectilinear_chunk_shapes(draw, shape))
 
 
 @st.composite
@@ -613,16 +759,9 @@ def chunk_grids(
 
     This allows property tests to exercise both chunk grid types.
     """
-    # RectilinearChunkGridMetadata doesn't support zero-sized dimensions,
-    # so use RegularChunkGridMetadata if any dimension is 0
-    if any(s == 0 for s in shape):
-        event("using RegularChunkGridMetadata (zero-sized dimensions)")
-        return RegularChunkGridMetadata(chunk_shape=draw(chunk_shapes(shape=shape)))
-
     if zarr.config.get("array.rectilinear_chunks") and draw(st.booleans()):
-        chunks = draw(rectilinear_chunks(shape=shape))
         event("using RectilinearChunkGridMetadata")
-        return RectilinearChunkGridMetadata(chunk_shapes=tuple(tuple(dim) for dim in chunks))
+        return draw(rectilinear_chunk_grids(shape=shape))
     else:
         event("using RegularChunkGridMetadata")
         return RegularChunkGridMetadata(chunk_shape=draw(chunk_shapes(shape=shape)))
@@ -640,7 +779,7 @@ def rectilinear_arrays(
 ) -> Any:
     """Generate a zarr v3 array with rectilinear (variable) chunk grid."""
     shape = draw(shapes)
-    chunk_shapes = draw(rectilinear_chunks(shape=shape))
+    chunk_shapes = draw(_rectilinear_chunks(shape=shape))
 
     np_dtype = draw(dtypes())
     nparray = draw(numpy_arrays(shapes=st.just(shape), dtype=np_dtype))
@@ -943,7 +1082,7 @@ def block_test_arrays(
     ``zarray.write_chunk_sizes`` — the array's *outer* (block / shard) grid, which
     is exactly the grid ``Array.blocks`` addresses; the caller reads it directly.
     """
-    chunks: tuple[int, ...] | list[list[int]]
+    chunks: tuple[int, ...] | list[int | list[int]]
     if draw(st.booleans()):
         # regular arm, optionally sharded
         nparray, chunks = draw(
@@ -960,7 +1099,7 @@ def block_test_arrays(
         # rectilinear arm, always unsharded
         event("block rectilinear")
         shape = draw(_rectilinear_shapes)
-        chunks = draw(rectilinear_chunks(shape=shape))
+        chunks = draw(_rectilinear_chunks(shape=shape))
         nparray = draw(numpy_arrays(shapes=st.just(shape), dtype=draw(dtypes())))
         shards, rectilinear = None, True
 

tests/conftest.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import itertools
 import math
 import os
 import pathlib
@@ -574,6 +575,21 @@ def deep_nan_equal(a: object, b: object) -> bool:
     return nan_equal(a, b)
 
 
+def declared_chunk_data_sizes(declared: int | Sequence[int], extent: int) -> tuple[int, ...]:
+    """The data sizes of the chunks one declared chunk grid dimension places
+    over `extent`, worked out from the declaration alone: a bare int repeats
+    to cover the extent, explicit edges are clipped to it. An oracle that does
+    not go through zarr's chunk grid code."""
+    sizes: list[int] = []
+    offset = 0
+    for edge in itertools.repeat(declared) if isinstance(declared, int) else declared:
+        if offset >= extent:
+            break
+        sizes.append(min(edge, extent - offset))
+        offset += edge
+    return tuple(sizes)
+
+
 def gzip_streams_equal_except_mtime(a: bytes, b: bytes) -> bool:
     """Compare two gzip streams, ignoring the MTIME field of the header.