diff --git a/changes/4377.misc.md b/changes/4377.misc.md new file mode 100644 index 0000000000..096c1b14c9 --- /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: `rectilinear_chunks` mixes bare-int steps and explicit edge lists in any arrangement for `chunks=`, 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. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index db01697f1e..5017587b83 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -26,7 +26,7 @@ from zarr.codecs.zstd import ZstdCodec from zarr.core.array import Array, CompressorsLike, SerializerLike from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding -from zarr.core.common import JSON, AccessModeLiteral, ZarrFormat +from zarr.core.common import JSON, AccessModeLiteral, ZarrFormat, compress_rle from zarr.core.dtype import get_data_type_from_native_dtype from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata @@ -337,14 +337,29 @@ def arrays( # - 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) + # The per-dimension list form of `chunks=` can only express a grid with + # at least one edge list (bare ints alone declare a regular grid) whose + # edge lists sum exactly to the extent. Other grids are passed as the + # metadata object, which `create_array` also accepts. + expressible_as_lists = any( + isinstance(dim, tuple) for dim in chunk_grid_meta.chunk_shapes + ) and all( + isinstance(dim, int) or sum(dim) == extent + for dim, extent in zip(chunk_grid_meta.chunk_shapes, nparray.shape, strict=True) + ) + if expressible_as_lists and draw(st.booleans(), label="chunks as lists"): + event("rectilinear chunks=: per-dimension lists") + chunks_param = chunks_param_from_rectilinear(chunk_grid_meta) + else: + event("rectilinear chunks=: metadata object") + chunks_param = chunk_grid_meta elif isinstance(chunk_grid_meta, RegularChunkGridMetadata): chunks_param = chunk_grid_meta.chunk_shape else: @@ -394,7 +409,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 (gh-4374, gh-4272). + assert a.metadata.chunk_grid == chunk_grid_meta assert shard_shape is None assert a.basename == name, (a.basename, name) @@ -436,58 +453,156 @@ def chunks_param_from_rectilinear( @st.composite -def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[int]]: - """Generate valid rectilinear chunk shapes for a given array shape. +def rectilinear_dim_edges(draw: st.DrawFn, *, extent: int) -> list[int]: + """Explicit chunk edge lengths summing exactly to `extent`. - 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. + 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 20 chunks per dimension keeps property + tests fast. """ - 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)] + assert extent > 0 + if extent == 1: + return [1] + if draw(st.booleans(), label="uneven edges"): + event("rectilinear edges: uneven") + nchunks = draw(st.integers(min_value=1, max_value=min(extent, 20))) + dividers = sorted( + draw( + st.lists( + st.integers(min_value=1, max_value=extent - 1), + min_size=nchunks - 1, + max_size=nchunks - 1, + unique=True, ) + ) + ) + return [b - a for a, b in zip([0, *dividers], [*dividers, extent], strict=True)] + size = draw(st.integers(min_value=math.ceil(extent / 20), max_value=extent)) + edges = [size] * (extent // size) + if extent % size: + edges.append(extent % size) + if draw(st.booleans(), label="shuffle uniform edges"): + event("rectilinear edges: uniform, shuffled") + return list(draw(st.permutations(edges))) + event("rectilinear edges: uniform") + return edges + + +def _rectilinear_step(draw: st.DrawFn, *, extent: 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=math.ceil(extent / 20), max_value=extent + 3)) + event("rectilinear dim: bare int" + (", larger than extent" if step > extent else "")) + return step + + +@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. 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`. + """ + if not shape: + return [] + forced_list = draw(st.integers(min_value=0, max_value=len(shape) - 1)) + chunks: list[int | list[int]] = [] + for i, extent in enumerate(shape): + assert extent > 0 + if i != forced_list and draw(st.booleans(), label="bare int"): + chunks.append(_rectilinear_step(draw, extent=extent)) + else: + event("rectilinear dim: edge list") + chunks.append(draw(rectilinear_dim_edges(extent=extent))) + event( + "rectilinear chunks: mixed bare ints and edge lists" + if any(isinstance(c, int) for c in chunks) + else "rectilinear chunks: edge lists only" + ) + return chunks + + +RectilinearDimDeclaration = int | list[int | list[int]] + + +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 or an arbitrary + grouping, which may split a run across pairs and use `count == 1`.""" + if draw(st.booleans(), label="canonical rle"): + event("rectilinear rle: canonical") + return compress_rle(edges) + 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 + count = draw(st.integers(min_value=1, max_value=run)) + if count == 1 and draw(st.booleans(), label="bare edge"): + encoded.append(edges[i]) + else: + encoded.append([edges[i], count]) + i += count + return encoded + + +@st.composite +def rectilinear_chunk_shape_declarations( + draw: st.DrawFn, *, shape: tuple[int, ...], overhang: bool = True +) -> tuple[list[RectilinearDimDeclaration], 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). With `overhang`, edge lists + may sum beyond the extent, which the spec allows and a shrinking resize + produces. + + Returns `(declaration, chunk_shapes)`: the JSON value to store, and the + `chunk_shapes` that parsing it must produce. + """ + declaration: list[RectilinearDimDeclaration] = [] + chunk_shapes: list[int | tuple[int, ...]] = [] + for extent in shape: + assert extent > 0 + if draw(st.booleans(), label="bare int"): + step = _rectilinear_step(draw, extent=extent) + declaration.append(step) + chunk_shapes.append(step) + continue + edges = draw(rectilinear_dim_edges(extent=extent)) + if overhang and draw(st.booleans(), label="overhang"): + if draw(st.booleans(), label="trailing edge"): + event("rectilinear edges: overhang, trailing edge beyond extent") + edges = [*edges, draw(st.integers(min_value=1, max_value=5))] 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)) + event("rectilinear edges: overhang, last edge past extent") + edges[-1] += draw(st.integers(min_value=1, max_value=5)) + if draw(st.booleans(), label="write edges in full"): + event("rectilinear edges: written in full") + declaration.append(list(edges)) else: - chunk_shapes.append([1]) - return chunk_shapes + declaration.append(_rle_encode(draw, edges)) + chunk_shapes.append(tuple(edges)) + return declaration, tuple(chunk_shapes) + + +@st.composite +def rectilinear_chunk_grids( + draw: st.DrawFn, *, shape: tuple[int, ...], overhang: bool = True +) -> RectilinearChunkGridMetadata: + """A `RectilinearChunkGridMetadata` parsed from a drawn stored declaration.""" + declaration, _ = draw(rectilinear_chunk_shape_declarations(shape=shape, overhang=overhang)) + return RectilinearChunkGridMetadata.from_dict( + {"name": "rectilinear", "configuration": {"kind": "inline", "chunk_shapes": declaration}} + ) @st.composite @@ -509,9 +624,8 @@ def chunk_grids( 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))) @@ -832,7 +946,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( diff --git a/tests/test_properties.py b/tests/test_properties.py index 2794ad3cb0..4776d1ca0e 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -20,7 +20,9 @@ from zarr.abc.store import Store from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata.v3 import RectilinearChunkGridMetadata from zarr.core.sync import sync +from zarr.storage import MemoryStore from zarr.testing.strategies import ( array_metadata, arrays, @@ -31,6 +33,8 @@ numpy_arrays, orthogonal_indices, rectilinear_arrays, + rectilinear_chunk_shape_declarations, + rectilinear_chunks, sharded_arrays, simple_arrays, stores, @@ -496,3 +500,60 @@ def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: 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=1, max_side=20), label="shape") + declaration, chunk_shapes = data.draw( + rectilinear_chunk_shape_declarations(shape=shape), label="declaration" + ) + stored = { + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": declaration}, + } + meta = RectilinearChunkGridMetadata.from_dict(stored) # type: ignore[arg-type] + 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] + + +@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. Checked on the stored JSON, not only the + in-memory metadata: zarr 3.2.x stored such grids as "regular" (gh-4374).""" + shape = data.draw(npst.array_shapes(max_dims=3, min_side=1, 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" + declared = RectilinearChunkGridMetadata( + chunk_shapes=tuple(tuple(c) if isinstance(c, list) else c for c in chunks) + ) + assert RectilinearChunkGridMetadata.from_dict(stored) == declared + assert arr.metadata.chunk_grid == declared diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index b8289d2135..5b011f175c 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -3066,32 +3066,7 @@ 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 zarr.testing.strategies import rectilinear_chunks @st.composite @@ -3101,7 +3076,7 @@ def rectilinear_arrays_st(draw: st.DrawFn) -> tuple[zarr.Array[Any], np.ndarray[ 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)