diff --git a/changes/4377.misc.md b/changes/4377.misc.md new file mode 100644 index 0000000000..6c3f3ce5ea --- /dev/null +++ b/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. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 00932c96ce..4c49a2efb0 100644 --- a/src/zarr/testing/strategies.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index e1016f0156..3e87a724e0 100644 --- a/tests/conftest.py +++ b/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. diff --git a/tests/test_array_stateful.py b/tests/test_array_stateful.py index a25b056c7d..3361d967e3 100644 --- a/tests/test_array_stateful.py +++ b/tests/test_array_stateful.py @@ -40,6 +40,7 @@ from zarr.core.sync import sync from zarr.errors import ZarrUserWarning from zarr.storage import MemoryStore +from zarr.testing.strategies import _rectilinear_chunks pytestmark = [ pytest.mark.slow_hypothesis, @@ -51,24 +52,6 @@ MAX_SIDE = 6 -def _rectilinear_dim(extent: int) -> st.SearchStrategy[int | list[int]]: - """A bare step, or an edge list covering `extent` (any edges for extent 0). - - A small local copy of what `zarr.testing.strategies` draws for rectilinear - declarations, so this test does not depend on that module's experimental API. - """ - steps = st.integers(min_value=1, max_value=MAX_SIDE) - if extent == 0: - return steps | st.lists(steps, min_size=1, max_size=3) - if extent == 1: - return steps | st.just([1]) - cuts = st.lists(st.integers(min_value=1, max_value=extent - 1), unique=True, max_size=3) - edges = cuts.map( - lambda c: [b - a for a, b in zip([0, *sorted(c)], [*sorted(c), extent], strict=True)] - ) - return steps | edges - - async def _list(store: MemoryStore, prefix: str) -> list[str]: return [key async for key in store.list_prefix(prefix)] @@ -110,9 +93,7 @@ def create(self, data: st.DataObject) -> None: elif spelling == "auto": chunks = "auto" elif spelling == "rectilinear": - chunks = [data.draw(_rectilinear_dim(s)) for s in shape] - if not any(isinstance(c, list) for c in chunks): - chunks[0] = [chunks[0]] if shape[0] == 0 else [shape[0]] + chunks = data.draw(_rectilinear_chunks(shape=shape), label="rectilinear chunks") else: chunks = tuple(data.draw(st.integers(1, 3)) for _ in shape) if spelling == "sharded": diff --git a/tests/test_properties.py b/tests/test_properties.py index fcb1076441..b448a524aa 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -19,6 +19,7 @@ import hypothesis.strategies as st from hypothesis import assume, event, given, settings +from tests.conftest import declared_chunk_data_sizes from zarr.abc.store import Store from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON from zarr.core.dtype import get_data_type_from_json, get_data_type_from_native_dtype @@ -26,9 +27,12 @@ from zarr.core.dtype.npy.structured import Struct from zarr.core.dtype.wrapper import ZDType from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RectilinearChunkGridMetadataJSON from zarr.core.sync import sync from zarr.errors import ZarrUserWarning +from zarr.storage import MemoryStore from zarr.testing.strategies import ( + _rectilinear_chunks, array_metadata, arrays, basic_indices, @@ -38,6 +42,8 @@ numpy_arrays, orthogonal_indices, rectilinear_arrays, + rectilinear_chunk_shape_declarations, + rectilinear_chunks, sharded_arrays, simple_arrays, stores, @@ -297,7 +303,18 @@ def test_block_indexing(data: st.DataObject) -> None: # across that matrix (rectilinear + sharded is unsupported and not drawn). zarray, nparray = data.draw(block_test_arrays()) - block_indexer, array_indexer = data.draw(block_indices(chunk_sizes=zarray.write_chunk_sizes)) + # 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 + ) + chunk_sizes = tuple( + declared_chunk_data_sizes(d, n) for d, n in zip(declared, zarray.shape, strict=True) + ) + assert zarray.write_chunk_sizes == chunk_sizes + + block_indexer, array_indexer = data.draw(block_indices(chunk_sizes=chunk_sizes)) expected = nparray[array_indexer] # sync get, via both the .blocks interface and the dedicated method @@ -570,12 +587,29 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N assert asdict_dict["fill_value"] == -9223372036854775808 +@given(data=st.data()) +def test_rectilinear_chunks_declares_edges_per_dimension(data: st.DataObject) -> None: + """`rectilinear_chunks` draws an explicit edge list for every dimension, summing + to the extent (any positive edges, for a zero extent); a 0-d shape has none.""" + shape = data.draw( + npst.array_shapes(min_dims=0, max_dims=3, min_side=0, max_side=20), label="shape" + ) + chunks = data.draw(rectilinear_chunks(shape=shape), label="chunks") + assert len(chunks) == len(shape) + for edges, extent in zip(chunks, shape, strict=True): + assert isinstance(edges, list) + assert edges + assert all(type(edge) is int and edge >= 1 for edge in edges) + assert extent == 0 or sum(edges) == extent + + def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: """Bare-int dims in rectilinear metadata (the spec's step-size shorthand, produced by a scalar dimension of a mixed chunk spec) must pass through the `chunks=` conversion unchanged. Wrapping one in a single-element list turns "repeat to cover the axis" into "exactly one - chunk" and re-creation fails the sum-to-span check.""" + chunk" and re-creation fails the sum-to-span check. Edge tuples become lists, + as zarr 3.4.0 returned them.""" from zarr.core.metadata.v3 import RectilinearChunkGridMetadata from zarr.storage import MemoryStore from zarr.testing.strategies import chunks_param_from_rectilinear @@ -585,10 +619,114 @@ def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: grid = src.metadata.chunk_grid # type: ignore[union-attr] assert isinstance(grid, RectilinearChunkGridMetadata) assert grid.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] dst = zarr.create_array( MemoryStore(), shape=src.shape, - chunks=chunks_param_from_rectilinear(grid), + chunks=chunks, dtype="uint8", ) assert dst.metadata.chunk_grid == grid # type: ignore[union-attr] + + +@given(data=st.data()) +def test_rectilinear_chunk_grid_declarations(data: st.DataObject) -> None: + """Every `chunk_shapes` declaration the rectilinear spec allows — bare-int + steps, edge lists written in full or run-length encoded in any grouping, + edges overhanging the extent — parses to its expanded edges, and the + re-serialized form parses back to the same grid.""" + shape = data.draw(npst.array_shapes(max_dims=3, min_side=0, max_side=20), label="shape") + declaration, chunk_shapes = data.draw( + rectilinear_chunk_shape_declarations(shape=shape), label="declaration" + ) + stored: RectilinearChunkGridMetadataJSON = { + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": declaration}, + } + meta = RectilinearChunkGridMetadata.from_dict(stored) + assert meta.chunk_shapes == chunk_shapes + + serialized = json.loads(json.dumps(meta.to_dict())) + assert serialized["name"] == "rectilinear" + assert RectilinearChunkGridMetadata.from_dict(serialized) == meta + + # The declaration is read correctly inside a whole stored metadata document. + document = { + "zarr_format": 3, + "node_type": "array", + "shape": list(shape), + "data_type": "uint8", + "chunk_grid": stored, + "chunk_key_encoding": {"name": "default", "configuration": {"separator": "/"}}, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + assert ArrayV3Metadata.from_dict(document).chunk_grid == meta # type: ignore[arg-type] + + +def _rle_expand(dim: list[Any]) -> list[int]: + """Expand one stored run-length encoded dimension: bare edges and + `[size, count]` pairs.""" + edges: list[int] = [] + for item in dim: + if type(item) is int: + edges.append(item) + else: + size, count = item + edges.extend([size] * count) + return edges + + +@given(data=st.data()) +def test_create_array_stores_declared_rectilinear_chunks(data: st.DataObject) -> None: + """A `chunks=` specification mixing bare ints and edge lists in any + arrangement is stored as a rectilinear grid whose `chunk_shapes` are + exactly the specification: bare ints stay bare ints, and edge lists keep + their edges (gh-4374, gh-4272). Checked on the stored JSON.""" + shape = data.draw(npst.array_shapes(max_dims=3, min_side=0, max_side=20), label="shape") + chunks = data.draw(_rectilinear_chunks(shape=shape), label="chunks") + arr = zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, dtype="uint8") + + zarr_json = sync(arr.store.get(ZARR_JSON, prototype=default_buffer_prototype())) + assert zarr_json is not None + stored = json.loads(zarr_json.to_bytes())["chunk_grid"] + assert stored["name"] == "rectilinear" + assert stored["configuration"]["kind"] == "inline" + stored_dims = stored["configuration"]["chunk_shapes"] + assert [dim if type(dim) is int else _rle_expand(dim) for dim in stored_dims] == chunks + + +@given(data=st.data()) +def test_rectilinear_zero_length_axis_round_trip(data: st.DataObject) -> None: + """An array declared with rectilinear `chunks=` over a zero-length axis + holds the data written after that axis grows, by `append` or by `resize`, + also when it grows past the declared edges.""" + shape = list(data.draw(npst.array_shapes(max_dims=3, min_side=0, max_side=6), label="shape")) + axis = data.draw(st.integers(0, len(shape) - 1), label="zero-length axis") + shape[axis] = 0 + chunks = data.draw(_rectilinear_chunks(shape=tuple(shape)), label="chunks") + store = MemoryStore() + arr = zarr.create_array(store, shape=tuple(shape), chunks=chunks, dtype="int16", fill_value=-1) + assert_array_equal(arr[...], np.full(shape, -1, dtype="int16")) + + declared = chunks[axis] + declared_span = sum(declared) if isinstance(declared, list) else declared + rows = data.draw(st.integers(1, 2 * declared_span + 2), label="rows") + if isinstance(declared, list): + event("grown axis declared as an edge list") + if rows > declared_span: + event("grown past the declared edges") + grown = [*shape] + grown[axis] = rows + values = data.draw(npst.arrays(np.dtype("int16"), tuple(grown)), label="values") + if data.draw(st.booleans(), label="append"): + arr.append(values, axis=axis) + else: + arr.resize(tuple(grown)) + arr[...] = values + assert arr.shape == tuple(grown) + assert_array_equal(arr[...], values) + assert_array_equal(zarr.open_array(store, mode="r")[...], values) diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 6a43cb4f3b..c20bca9dab 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -3079,65 +3079,43 @@ def test_rectilinear_roundtrip(json_input: RectilinearChunkGridMetadataJSON) -> import hypothesis.strategies as st from hypothesis import event, given, settings - -@st.composite -def rectilinear_chunks_st(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[int]]: - """Generate valid rectilinear chunk shapes for a given array shape.""" - chunk_shapes: list[list[int]] = [] - for size in shape: - assert size > 0 - max_chunks = min(size, 10) - nchunks = draw(st.integers(min_value=1, max_value=max_chunks)) - if nchunks == 1: - chunk_shapes.append([size]) - else: - 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)] - ) - return chunk_shapes +from tests.conftest import declared_chunk_data_sizes +from zarr.testing.strategies import _rectilinear_chunks @st.composite -def rectilinear_arrays_st(draw: st.DrawFn) -> tuple[zarr.Array[Any], np.ndarray[Any, Any]]: - """Generate a rectilinear zarr array with random data, shape, and chunks.""" +def rectilinear_arrays_st( + draw: st.DrawFn, +) -> tuple[zarr.Array[Any], np.ndarray[Any, Any], list[int | list[int]]]: + """Generate a rectilinear zarr array with random data, shape, and chunks, + with the `chunks=` it was created with.""" from zarr.storage import MemoryStore ndim = draw(st.integers(min_value=1, max_value=3)) shape = draw(st.tuples(*[st.integers(min_value=2, max_value=20) for _ in range(ndim)])) - chunk_shapes = draw(rectilinear_chunks_st(shape=shape)) + chunk_shapes = draw(_rectilinear_chunks(shape=shape)) event(f"ndim={ndim}, shape={shape}") a = np.arange(int(np.prod(shape)), dtype="int32").reshape(shape) store = MemoryStore() z = zarr.create_array(store=store, shape=shape, chunks=chunk_shapes, dtype="int32") z[:] = a - return z, a + return z, a, chunk_shapes @settings(deadline=None, max_examples=50) @given(data=st.data()) def test_property_block_indexing_rectilinear(data: st.DataObject) -> None: """Property test: block indexing on rectilinear arrays matches numpy.""" - z, a = data.draw(rectilinear_arrays_st()) - grid = ChunkGrid.from_metadata(z.metadata) + z, a, chunks = data.draw(rectilinear_arrays_st()) for dim in range(a.ndim): - dim_grid = grid._dimensions[dim] - block_ix = data.draw(st.integers(min_value=0, max_value=dim_grid.nchunks - 1)) + # The block extents come from the declaration, not from zarr's grid code. + sizes = declared_chunk_data_sizes(chunks[dim], a.shape[dim]) + block_ix = data.draw(st.integers(min_value=0, max_value=len(sizes) - 1)) sel = [slice(None)] * a.ndim - start = dim_grid.chunk_offset(block_ix) - stop = start + dim_grid.data_size(block_ix) - sel[dim] = slice(start, stop) + start = sum(sizes[:block_ix]) + sel[dim] = slice(start, start + sizes[block_ix]) block_sel: list[slice | int] = [slice(None)] * a.ndim block_sel[dim] = block_ix np.testing.assert_array_equal(