diff --git a/changes/4375.bugfix.md b/changes/4375.bugfix.md new file mode 100644 index 0000000000..689d8e3be7 --- /dev/null +++ b/changes/4375.bugfix.md @@ -0,0 +1,3 @@ +Arrays whose stored `regular` chunk grid mixes chunk sizes with lists of chunk edge lengths, such as `"chunk_shape": [2, [5, 10, 5]]` (written by zarr 3.2.0 and 3.2.1 for `chunks=(2, (5, 10, 5))`, and as `[2, [5.0, 10.0, 5.0]]` for float edges), can be read again, without enabling `array.rectilinear_chunks`. The grid is read as the rectilinear chunk grid it describes, with a `ZarrUserWarning`; re-saving the array's metadata (or writing chunks to it) stores that rectilinear grid, which requires the flag. A group's consolidated metadata keeps such an array's metadata as it was stored. + +The `array.rectilinear_chunks` flag now gates reading and storing array metadata that declares a rectilinear chunk grid, including in a group's consolidated metadata, instead of constructing `RectilinearChunkGridMetadata`. An operation that would store such metadata with the flag off (creating an array, `Array.resize`, deleting a member of a group with consolidated metadata, `create_hierarchy`) raises before it deletes or writes anything, and the error names the array, also when it is a member of a group's consolidated metadata. diff --git a/changes/4376.bugfix.md b/changes/4376.bugfix.md new file mode 100644 index 0000000000..510a9e28a0 --- /dev/null +++ b/changes/4376.bugfix.md @@ -0,0 +1 @@ +Chunk specifications are judged only by the chunk normalizer: the separate duck-typed check for rectilinear input is gone, so the Zarr format 2 and sharding restrictions apply to what the normalized grid says was declared. A chunk size is an integer by Python's integer protocol: numpy integer scalars and 0-d integer arrays are accepted wherever a chunk or shard shape is. The legacy `zarr.create(..., zarr_format=2)` no longer tests the truth value of a numpy array given as `chunks`, so numpy arrays work there. A non-integer scalar chunk specification raises a `TypeError` naming the problem. An explicit list of chunk edges declares a rectilinear dimension everywhere, including a stored rectilinear chunk grid passed as `chunks=` whose edges happen to be uniform, so the Zarr format 2 and sharding restrictions treat it the same whether it is given as lists or as metadata, and `zarr.create` stores it as a rectilinear chunk grid, as `create_array` does. A `RectilinearChunkGridMetadata` made only of bare integers declares no explicit edge list, so it is read as the regular grid it describes: it is accepted for Zarr format 2 and as the chunk shape of a sharded array. 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/changes/4431.removal.md b/changes/4431.removal.md new file mode 100644 index 0000000000..598e634b82 --- /dev/null +++ b/changes/4431.removal.md @@ -0,0 +1,7 @@ +Metadata built in code is strict about chunk edge lengths: `ArrayV2Metadata(chunks=...)`, `RegularChunkGridMetadata`, `RectilinearChunkGridMetadata` and `ShardingCodec(chunk_shape=...)` take them as Python `int`s of at least 1. A chunk size of 0 now raises a `ValueError`, as in `ArrayV2Metadata(chunks=(0,))` for an empty array. Code that passes an array's shape as its chunk shape, such as `ArrayV2Metadata(chunks=arr.shape)` or `RegularChunkGridMetadata(chunk_shape=arr.shape)`, fails for an array with a zero-length axis: use `tuple(max(1, s) for s in arr.shape)`. A `bool` or a NumPy integer, including a scalar `chunks=np.int64(5)`, raises a `TypeError`, as a float already did, and so does a chunk shape that is not a list or tuple. This reverses what 3.4.1 said of these classes: `ArrayV2Metadata(chunks=(0,))` is no longer accepted, and a `bool` chunk edge length is no longer read as the `int` it equals. (This entry assumes 3.4.1 is released before this release.) The array creation functions still accept NumPy integers in `chunks=` and `shards=`. + +This applies to stored metadata too wherever the stored-document repairs do not read it: a stored chunk shape that is a bare integer instead of a list, as in a Zarr format 2 `"chunks": 4` or a sharding codec's `"chunk_shape": 2`, which earlier releases read as a one-dimensional chunk shape, now raises a `TypeError` when the array is opened. A group whose consolidated metadata includes such an array does not open from its consolidated metadata at all, and the error names that member: open the group with `use_consolidated=False` to reach its other members. A stored `true` as a rectilinear chunk grid's bare step (`"chunk_shapes": [true]`) or as the count of a run-length pair (`[[4, true]]`), which no known writer stored, raises a `TypeError` too. The chunk sizes the repairs read open as before: 0 or `false` as a regular chunk size (in Zarr format 2 `chunks` or a `regular` grid's `chunk_shape`); `true` as a regular chunk size, as an edge in a rectilinear edge list, as the size of a run-length pair, or as a sharding codec's inner chunk size; and integral float edges in a rectilinear chunk grid (`[[4.0, 2]]`). + +Other than `chunks=False` or `shards=False` as the whole specification (one chunk spanning every axis), a boolean anywhere in a chunk specification given to the array creation functions (`chunks=True`, `np.True_`, `np.False_`, a boolean array, `chunks=(True, 5)`, `chunks=(False, 5)`, or a boolean edge in a list) raises one `TypeError`, instead of being read as a chunk size of 1 or 0 or raising a different error for each spelling. `chunks=True` and `create_array(chunks=None)` raised a `ValueError` before; both errors say how to ask for automatic chunking. The legacy `zarr.create(..., zarr_format=2)` no longer reads a falsy `chunks` as not given: `chunks=0`, `np.int64(0)` and `chunks=[]` raise a `ValueError`, and `chunks=False` means one chunk spanning the array, as it does for Zarr format 3, instead of automatic chunking; `chunks=None` still chunks automatically. + +`zarr.testing.strategies.chunks_param_from_rectilinear` is removed; nothing in zarr used it after `arrays()` began drawing its `chunks=` argument from `rectilinear_chunks`. `rectilinear_chunks` now returns `list[int | list[int]]` instead of `list[list[int]]`: a dimension may be a bare int step, which can exceed the extent, so code that treats every dimension as a list of edges (`len(dim)`, `sum(dim)`) must handle ints. It also requires `shape` to have at least one dimension. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 85d445752d..cf6886774f 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -45,9 +45,9 @@ merge_and_encode_chunk, ) from zarr.core.common import ( - ShapeLike, + ChunkShape, + parse_chunk_shape, parse_named_configuration, - parse_shapelike, product, ) from zarr.core.config import config as zarr_config @@ -459,7 +459,7 @@ class ShardingCodec( def __init__( self, *, - chunk_shape: ShapeLike, + chunk_shape: ChunkShape, codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(endian="little"),), index_codecs: Iterable[Codec | dict[str, JSON]] = ( BytesCodec(endian="little"), @@ -468,7 +468,7 @@ def __init__( index_location: ShardingCodecIndexLocation | IndexLocation = "end", subchunk_write_order: SubchunkWriteOrder = "morton", ) -> None: - chunk_shape_parsed = parse_shapelike(chunk_shape) + chunk_shape_parsed = parse_chunk_shape(chunk_shape) codecs_parsed = parse_codecs(codecs) index_codecs_parsed = parse_codecs(index_codecs) _check_index_codecs_fixed_size(index_codecs_parsed) @@ -507,7 +507,7 @@ def __getstate__(self) -> dict[str, Any]: def __setstate__(self, state: dict[str, Any]) -> None: config = state["configuration"] - object.__setattr__(self, "chunk_shape", parse_shapelike(config["chunk_shape"])) + object.__setattr__(self, "chunk_shape", parse_chunk_shape(config["chunk_shape"])) object.__setattr__(self, "codecs", parse_codecs(config["codecs"])) object.__setattr__(self, "index_codecs", parse_codecs(config["index_codecs"])) object.__setattr__(self, "index_location", _parse_index_location(config["index_location"])) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 9788f60743..50800d4e51 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -45,7 +45,6 @@ ChunkGrid, _is_auto, _is_keep, - _is_rectilinear_chunks, guess_chunks, normalize_chunks_nd, resolve_outer_and_inner_chunks, @@ -120,10 +119,12 @@ ) from zarr.core.metadata.io import ( ARRAY_DOCUMENTS, + encode_documents, parse_stored_array, read_documents, save_metadata, save_new_metadata, + store_documents, upsert_metadata, ) from zarr.core.metadata.v2 import ( @@ -210,14 +211,8 @@ def _chunk_sizes_from_shape( def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: """Array metadata from a metadata object or a metadata document, naming the array at - `path` in warnings about how an invalid document was read. - - `ArrayV2Metadata` accepts a chunk size of 0, as it always has, though only an - invalid document holds one: such metadata is read as the documents it would store - are (see `zarr.core.metadata.repair`), so an array can be built from it. No data - was read or written under that chunk size, so the reading is silent.""" - if isinstance(data, ArrayV2Metadata) and 0 in data.chunks: - return parse_stored_array(data.to_buffer_dict(default_buffer_prototype()), 2) + `path` in warnings about how an invalid document was read. A metadata object is + valid as built: the metadata constructors are strict.""" if isinstance(data, ArrayMetadata): return data if isinstance(data, dict): @@ -463,6 +458,11 @@ async def _create( # Unify the v2 (chunks) and v3 (chunk_shape) parameter names _raw_chunks = chunks if chunks is not None else chunk_shape + if _raw_chunks is None: + item_size = dtype_parsed.item_size if isinstance(dtype_parsed, HasItemSize) else 1 + outer_chunks = guess_chunks(shape, item_size) + else: + outer_chunks = normalize_chunks_nd(_raw_chunks, shape) config_parsed = parse_array_config(config) @@ -484,13 +484,6 @@ async def _create( if order is not None: _warn_order_kwarg() - item_size = 1 - if isinstance(dtype_parsed, HasItemSize): - item_size = dtype_parsed.item_size - if _raw_chunks is None: - outer_chunks = guess_chunks(shape, item_size) - else: - outer_chunks = normalize_chunks_nd(_raw_chunks, shape) chunk_grid = create_chunk_grid_metadata(outer_chunks) result = await cls._create_v3( store_path, @@ -516,17 +509,8 @@ async def _create( ) if dimension_names is not None: raise ValueError("dimension_names cannot be used for arrays with zarr_format 2.") - if _is_rectilinear_chunks(_raw_chunks): + if not outer_chunks.is_regular: raise ValueError("Zarr format 2 does not support rectilinear chunk grids.") - - item_size = 1 - if isinstance(dtype_parsed, HasItemSize): - item_size = dtype_parsed.item_size - _raw = chunks or chunk_shape - if _raw is None: - outer_chunks = guess_chunks(shape, item_size) - else: - outer_chunks = normalize_chunks_nd(_raw, shape) _chunks = outer_chunks.chunk_shape if order is None: @@ -646,7 +630,6 @@ async def _create_v3( dimension_names=dimension_names, attributes=attributes, ) - array = cls(metadata=metadata, store_path=store_path, config=config) await save_new_metadata(store_path, metadata, overwrite=overwrite, ensure_parents=True) return array @@ -727,7 +710,6 @@ async def _create_v2( compressor=compressor_parsed, attributes=attributes, ) - array = cls(metadata=metadata, store_path=store_path, config=config) await save_new_metadata(store_path, metadata, overwrite=overwrite, ensure_parents=True) return array @@ -1607,7 +1589,9 @@ async def get_coordinate_selection( async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None: """Store `metadata` as this array's own documents, then clear the - `_stored_document` mark (see `_stored_document_replaced`).""" + `_stored_document` mark (see `_stored_document_replaced`). `_resize` stores the + documents it encoded before deleting chunks directly, then clears the mark the + same way.""" await save_metadata(self.store_path, metadata, ensure_parents=ensure_parents) self._stored_document_replaced() @@ -1642,14 +1626,14 @@ async def _store_repaired_document(self) -> None: zarr_format = self.metadata.zarr_format documents = await read_documents(self.store_path, ARRAY_DOCUMENTS[zarr_format]) try: - current = parse_stored_array(documents, zarr_format) + current = parse_stored_array(documents, zarr_format, str(self.store_path)) except ArrayNotFoundError: pass else: if _chunk_layout(current) != _chunk_layout(self.metadata): raise ValueError( - f"The metadata stored for the array at {str(self.store_path)!r} has " - "changed since this array was opened: reopen the array to write to it." + f"Array {str(self.store_path)!r}: the metadata stored has changed since " + "this array was opened; reopen the array to write to it. Nothing was stored." ) if current._stored_document is not None: await upsert_metadata(self.store_path, current, documents) @@ -4568,8 +4552,17 @@ async def init_array( chunk_key_encoding, zarr_format=zarr_format ) - # Validate rectilinear chunks constraints - if _is_rectilinear_chunks(chunks): + # Normalize the user's chunks into a canonical ChunkGrid + + if _is_auto(chunks): + max_bytes = None if shards is None else SHARDED_INNER_CHUNK_MAX_BYTES + chunks_normalized = guess_chunks(shape_parsed, item_size, max_bytes=max_bytes) + else: + chunks_normalized = normalize_chunks_nd(chunks, shape_parsed) + + # Validate rectilinear chunks constraints. The normalized grid is the one + # judge of what the user declared. + if not chunks_normalized.is_regular: if zarr_format == 2: raise ValueError("Zarr format 2 does not support rectilinear chunk grids.") if shards is not None: @@ -4579,14 +4572,6 @@ async def init_array( "chunks=(inner_size, ...), shards=[[shard_sizes], ...]" ) - # Normalize the user's chunks into a canonical ChunkGrid - - if _is_auto(chunks): - max_bytes = None if shards is None else SHARDED_INNER_CHUNK_MAX_BYTES - chunks_normalized = guess_chunks(shape_parsed, item_size, max_bytes=max_bytes) - else: - chunks_normalized = normalize_chunks_nd(chunks, shape_parsed) - # Resolve chunks + shards into outer_chunks (grid metadata) and # inner (sub-chunk structure for ShardingCodec, None if no sharding) outer_chunks, inner = resolve_outer_and_inner_chunks( @@ -4908,10 +4893,9 @@ def _stored_rectilinear_grid_or_none( """The *stored* rectilinear chunk grid, or None if the stored grid is regular (in which case `.chunks` and `.shards` are defined). - Dispatches on the stored metadata, not the runtime ``ChunkGrid``: the - runtime grid collapses a rectilinear dimension whose edges happen to be - uniform to a ``FixedDimension`` as an optimization, so it can report regular - for an array whose stored metadata — and therefore `.chunks` — is + Dispatches on the stored metadata, not the runtime `ChunkGrid`: a + rectilinear grid whose dimensions are all bare-int step sizes is regular at + runtime, while its stored metadata — and therefore `.chunks` — is rectilinear. Zarr format 2 grids are always regular. """ if isinstance(metadata, ArrayV3Metadata) and isinstance( @@ -5904,6 +5888,10 @@ async def _resize( # ensure deletion is only run if array is shrinking as the delete_outside_chunks path is unbounded in memory only_growing = all(new >= old for new, old in zip(new_shape, array.metadata.shape, strict=True)) + # Encode the new metadata before deleting any chunk: metadata that cannot be stored + # then fails with the store untouched. + documents = encode_documents(array.store_path, new_metadata) + if delete_outside_chunks and not only_growing: # Remove all chunks outside of the new shape old_chunk_coords = set(array._chunk_grid.all_chunk_coords()) @@ -5922,7 +5910,8 @@ async def _delete_key(key: str) -> None: ) # Write new metadata - await array._save_metadata(new_metadata) + await store_documents(array.store_path, documents) + array._stored_document_replaced() # Update metadata and chunk_grid (in place) object.__setattr__(array, "metadata", new_metadata) diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index f7dbffadf0..efc118ecd0 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -3,7 +3,6 @@ import bisect import itertools import math -import numbers import operator import warnings from dataclasses import dataclass, field @@ -11,9 +10,11 @@ from typing import ( TYPE_CHECKING, Any, + Final, Literal, NamedTuple, Protocol, + SupportsIndex, cast, runtime_checkable, ) @@ -34,6 +35,7 @@ from collections.abc import Iterable, Iterator, Sequence from zarr.core.array import ShardsLike + from zarr.core.common import ChunksLike from zarr.core.metadata import ArrayMetadata SHARDED_INNER_CHUNK_MAX_BYTES: int = 1048576 @@ -325,28 +327,49 @@ def _is_keep(spec: object) -> TypeIs[Literal["keep"]]: return isinstance(spec, str) and spec == "keep" -def _is_rectilinear_chunks(chunks: Any) -> bool: - """Check if chunks specifies a rectilinear grid along any dimension. +_AUTOMATIC_CHUNKING: Final = ( + ' For automatic chunking, pass "auto" to create_array (as chunks= or shards=), ' + "or chunks=None to zarr.create." +) +"""How to ask for automatic chunking, for the errors of specifications that once did.""" + + +def _chunk_int(value: object) -> int | None: + """Return `value` as an `int` if it is an integer chunk size, else `None`. - Returns True for nested sequences like [[10, 20], [5, 5]], for mixed - per-dimension specs like (5, [10, 20]) regardless of which dimension - carries the sequence, and for stored rectilinear grid metadata. - Returns False for flat sequences like (10, 10) or [10, 10]. + This is the chunk normalizer's one integer test. An integer is anything + Python's integer protocol (`__index__`) accepts: `int`, numpy integer + scalars and 0-d integer arrays. Floats and arrays with dimensions are not + integers. + + A boolean (`bool`, `np.bool_`, a boolean array) is a flag, not a size: it + raises a `TypeError` here, the one place every spelling of a chunk size + passes through, rather than being read as a size of 0 or 1. """ - from zarr.core.metadata.v3 import RectilinearChunkGridMetadata + if isinstance(value, bool | np.bool_) or ( + isinstance(value, np.ndarray) and value.dtype == np.bool_ + ): + raise TypeError(f"A bool is not a chunk size; got {value!r}.") + try: + return operator.index(cast("SupportsIndex", value)) + except TypeError: + # Not an integer; numpy arrays define `__index__` but only 0-d integer arrays honour it. + return None - if isinstance(chunks, RectilinearChunkGridMetadata): - return True - if isinstance(chunks, (str, int, ChunkGrid)): - return False - if not hasattr(chunks, "__iter__"): - return False + +def _chunk_list(chunks: object) -> list[Any]: + """Return the elements of a chunk specification that is not an integer, or raise + a `TypeError` naming it (a float, `None`, a 0-d non-integer array).""" try: - return any( - hasattr(elem, "__iter__") and not isinstance(elem, (str, bytes, int)) for elem in chunks - ) + # `list` also takes a sequence with only `__getitem__` and `__len__`, which + # `Iterable` does not recognize. + return list(cast("Iterable[Any]", chunks)) except TypeError: - return False + # Not iterable, or a 0-d numpy array, whose iteration raises. + raise TypeError( + f"Chunk specification must be an integer or an iterable of integers; got " + f"{chunks!r} of type {type(chunks).__name__}." + ) from None def is_regular_1d(dim_chunks: Sequence[int]) -> bool: @@ -434,9 +457,9 @@ 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. If all - edges are identical and cover the extent, the dimension is - stored as ``FixedDimension``; otherwise as ``VaryingDimension``. + - A ``Sequence[int]`` — 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`. """ extents = parse_shapelike(array_shape) if len(extents) != len(chunk_sizes): @@ -449,20 +472,7 @@ def from_sizes( if isinstance(dim_spec, int): dims.append(FixedDimension(size=dim_spec, extent=extent)) else: - edges_list = list(dim_spec) - if not edges_list: - raise ValueError("Each dimension must have at least one chunk") - edge_sum = sum(edges_list) - if ( - edges_list[0] > 0 - and all(e == edges_list[0] for e in edges_list) - and ( - extent == edge_sum or len(edges_list) == ceildiv_int(extent, edges_list[0]) - ) - ): - dims.append(FixedDimension(size=edges_list[0], extent=extent)) - else: - dims.append(VaryingDimension(edges_list, extent=extent)) + dims.append(VaryingDimension(dim_spec, extent=extent)) return cls(dimensions=tuple(dims)) # -- Properties -- @@ -758,29 +768,22 @@ def normalize_chunks_1d(chunks: int | Iterable[object], span: int, unit: int = 1 the last chunk may overhang the span. On a zero-length span any non-empty list of positive sizes is kept: the chunks the axis grows into. """ - # `numbers.Integral` rather than `int` so that numpy integer scalars (which are not - # `int` subclasses) take the uniform-chunk path instead of being treated as a sequence. - # The `-1` sentinel check lives inside this branch so that numpy-array chunk - # specifications never hit an ambiguous-truth-value error on `chunks == -1`. - if isinstance(chunks, numbers.Integral): - chunk_size = int(chunks) + chunk_size = _chunk_int(chunks) + if chunk_size is not None: if chunk_size < -1 or chunk_size == 0: raise ValueError(f"Chunk size must be positive or -1, got {chunk_size}") if chunk_size == -1: return FixedDimension(size=full_span_chunk_size(span, unit), extent=span) return FixedDimension(size=chunk_size, extent=span) else: - try: - chunk_list = list(chunks) # type: ignore[arg-type] - except TypeError: - raise TypeError( - f"Chunk specification must be an integer or an iterable of integers; got " - f"{chunks!r} of type {type(chunks).__name__}." - ) from None + chunk_list = _chunk_list(chunks) if not chunk_list: raise ValueError("Chunk specification must not be empty") + as_ints = [_chunk_int(c) for c in chunk_list] non_int = [ - (idx, c) for idx, c in enumerate(chunk_list) if not isinstance(c, numbers.Integral) + (idx, c) + for idx, (c, i) in enumerate(zip(chunk_list, as_ints, strict=True)) + if i is None ] if non_int: non_int_idxs, non_int_vals = [*zip(*non_int, strict=False)] @@ -789,7 +792,7 @@ def normalize_chunks_1d(chunks: int | Iterable[object], span: int, unit: int = 1 f"at indices {non_int_idxs!r}. Chunk sizes must be declared as a flat sequence of " f"positive integers (e.g. [3, 3, 1])." ) - ints: list[int] = [int(c) for c in chunk_list] # type: ignore[call-overload] + ints = [i for i in as_ints if i is not None] if any(c <= 0 for c in ints): raise ValueError(f"All chunk sizes must be positive, got {ints}") if span > 0 and sum(ints) != span: @@ -798,7 +801,7 @@ def normalize_chunks_1d(chunks: int | Iterable[object], span: int, unit: int = 1 def normalize_chunks_nd( - chunks: Any, + chunks: ChunksLike, shape: tuple[int, ...], unit: tuple[int, ...] | None = None, ) -> ChunkGrid: @@ -832,18 +835,19 @@ def normalize_chunks_nd( if isinstance(chunks, RectilinearChunkGridMetadata): return ChunkGrid.from_sizes(shape, chunks.chunk_shapes) - if chunks is None or chunks is True: - raise ValueError( - f'{chunks!r} is not a valid chunk input. Use chunks=None or chunks="auto" from the top-level API for auto-chunking, or pass an int / tuple of ints.' - ) - - # handle no chunking: one chunk covering every axis. + # A bool as the whole specification: `False` is one chunk covering every axis. + # `True` meant automatic chunking in zarr 2, and `None` means it in `zarr.create`, + # so their errors name the spellings that chunk automatically here. if chunks is False: chunks = -1 + elif chunks is True: + raise TypeError(f"A bool is not a chunk size; got True.{_AUTOMATIC_CHUNKING}") + elif chunks is None: + raise TypeError(f"None is not a chunk specification.{_AUTOMATIC_CHUNKING}") - # handle 1D convenience form. bool is excluded above so this only catches actual ints. - if isinstance(chunks, numbers.Integral): - chunks = tuple(int(chunks) for _ in shape) + # handle 1D convenience form: one integer applies to every dimension. + chunk_size = _chunk_int(chunks) + chunks = (chunk_size,) * len(shape) if chunk_size is not None else _chunk_list(chunks) # handle bad dimensionality if len(chunks) != len(shape): @@ -965,15 +969,10 @@ def resolve_outer_and_inner_chunks( if shard_shape is None: return ChunkLayout(outer_chunks=chunks) - # Rectilinear shards: normalize the nested sequence directly. - if _is_rectilinear_chunks(shard_shape): - outer = normalize_chunks_nd(shard_shape, array_shape) - return ChunkLayout(outer_chunks=outer, inner=ChunkLayout(outer_chunks=chunks)) - - # Extract the flat chunk shape (uniform size per dimension) for arithmetic. - chunk_shape_flat = chunks.chunk_shape - + shard_spec: ChunksLike if _is_auto(shard_shape): + # Extract the flat chunk shape (uniform size per dimension) for arithmetic. + chunk_shape_flat = chunks.chunk_shape warnings.warn( "Automatic shard shape inference is experimental and may change without notice.", ZarrUserWarning, @@ -997,11 +996,13 @@ def resolve_outer_and_inner_chunks( _shards_out += (c_shape * num_chunks_per_shard_axis,) else: _shards_out += (c_shape,) - shard_flat = _shards_out + shard_spec = _shards_out elif isinstance(shard_shape, dict): - shard_flat = tuple(shard_shape["shape"]) + shard_spec = shard_shape["shape"] else: - shard_flat = cast("tuple[int, ...]", shard_shape) + shard_spec = shard_shape - outer = normalize_chunks_nd(shard_flat, array_shape, unit=chunk_shape_flat) + # Regular and rectilinear shard specifications go through the one normalizer; + # a full-span shard is a multiple of the inner chunk. + outer = normalize_chunks_nd(shard_spec, array_shape, unit=chunks.chunk_shape) return ChunkLayout(outer_chunks=outer, inner=ChunkLayout(outer_chunks=chunks)) diff --git a/src/zarr/core/common.py b/src/zarr/core/common.py index de009d0b50..39d192dfe0 100644 --- a/src/zarr/core/common.py +++ b/src/zarr/core/common.py @@ -45,6 +45,9 @@ type ChunksLike = ShapeLike | Iterable[int | Iterable[int]] | ChunkGridMetadata # For backwards compatibility ChunkCoords = tuple[int, ...] +type ChunkShape = tuple[int, ...] | list[int] +"""A regular chunk shape as the metadata constructors take it: a list or tuple of one +chunk edge length (an `int` of at least 1) per axis (see `parse_chunk_shape`).""" ZarrFormat = Literal[2, 3] NodeType = Literal["array", "group"] JSON = str | int | float | bool | Mapping[str, "JSON"] | Sequence["JSON"] | None @@ -288,20 +291,19 @@ def _subject(name: str, axis: int | None) -> str: def _parse_positive_int(value: object, name: str, axis: int | None) -> int: - """`value` as an `int` of at least 1. A `bool` is read as the `int` it equals; any - other type, a NumPy integer or a float (even an integral one: stored documents with - integral floats are read by `zarr.core.metadata.repair`), is rejected.""" + """`value` as an `int` of at least 1. A `bool`, a NumPy integer or a float (even an + integral one: stored documents with integral floats are read by + `zarr.core.metadata.repair`) is rejected.""" subject = _subject(name, axis) - if not isinstance(value, int): + if isinstance(value, bool) or not isinstance(value, int): raise TypeError(f"{subject} must be an int, got {value!r}") if value < 1: raise ValueError(f"{subject} must be >= 1, got {value!r}") - return int(value) + return value def parse_chunk_edge(size: object, axis: int | None = None) -> int: - """Check that `size` is a chunk edge length: an `int` of at least 1 (a `bool` is - read as the `int` it equals). + """Check that `size` is a chunk edge length: an `int` (not a `bool`) of at least 1. This is the one rule for chunk edge lengths in metadata: bare chunk sizes, explicit edges and run-length encoded sizes. `axis`, when given, is named in the error. @@ -310,14 +312,12 @@ def parse_chunk_edge(size: object, axis: int | None = None) -> int: def parse_chunk_shape(data: object) -> tuple[int, ...]: - """Check a regular chunk shape: an iterable, other than a string or a mapping, of one - chunk edge length per axis (see `parse_chunk_edge`).""" + """Check a regular chunk shape: a list or tuple of one chunk edge length per axis + (see `parse_chunk_edge`).""" match data: - case str() | Mapping(): - pass - case Iterable(): + case list() | tuple(): return tuple(parse_chunk_edge(size, axis) for axis, size in enumerate(data)) - raise TypeError(f"A chunk shape must be an iterable of chunk edge lengths, got {data!r}") + raise TypeError(f"A chunk shape must be a list or tuple of chunk edge lengths, got {data!r}") def expand_rle(data: Sequence[object], axis: int | None = None) -> list[int]: diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 7d4e26a1dd..9f552d2752 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -49,8 +49,13 @@ from zarr.core.dtype import parse_data_type from zarr.core.json_parse import parse_field from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata -from zarr.core.metadata.io import save_metadata, save_new_metadata -from zarr.core.metadata.v3 import AllowedExtraField, parse_extra_fields +from zarr.core.metadata.io import ( + encode_documents, + save_metadata, + save_new_metadata, + store_documents, +) +from zarr.core.metadata.v3 import AllowedExtraField, check_storable, parse_extra_fields from zarr.core.sync import SyncMixin, sync from zarr.errors import ( ArrayNotFoundError, @@ -192,29 +197,36 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Consoli f"Invalid value for metadata items. key='{k}', type='{type(v).__name__}'" ) - # zarr_format is present in v2 and v3. - zarr_format = parse_zarr_format(v["zarr_format"]) - - if zarr_format == 3: - node_type = parse_node_type(v.get("node_type", None)) - if node_type == "group": - metadata[k] = GroupMetadata.from_dict(v, path=member) - elif node_type == "array": - metadata[k] = ArrayV3Metadata.from_dict(v, path=member) - else: - assert_never(node_type) - elif zarr_format == 2: - if "shape" in v: - metadata[k] = ArrayV2Metadata.from_dict(v, path=member) - else: - metadata[k] = GroupMetadata.from_dict(v, path=member) - else: - assert_never(zarr_format) + try: + metadata[k] = cls._member_from_dict(v, member) + except (TypeError, ValueError) as e: + e.add_note(f"Member {member!r} of the consolidated metadata.") + raise cls._flat_to_nested(metadata) return cls(metadata=metadata) + @staticmethod + def _member_from_dict( + data: dict[str, JSON], path: str + ) -> ArrayV2Metadata | ArrayV3Metadata | GroupMetadata: + """Read one member of consolidated metadata, at `path`.""" + # zarr_format is present in v2 and v3. + zarr_format = parse_zarr_format(data["zarr_format"]) + if zarr_format == 3: + node_type = parse_node_type(data.get("node_type")) + if node_type == "group": + return GroupMetadata.from_dict(data, path=path) + if node_type == "array": + return ArrayV3Metadata.from_dict(data, path=path) + assert_never(node_type) + if zarr_format == 2: + if "shape" in data: + return ArrayV2Metadata.from_dict(data, path=path) + return GroupMetadata.from_dict(data, path=path) + assert_never(zarr_format) + @staticmethod def _flat_to_nested( metadata: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata], @@ -380,6 +392,16 @@ class GroupMetadata(Metadata): extra_fields: dict[str, AllowedExtraField] = field(default_factory=dict) def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: + if self.consolidated_metadata is not None: + for path, member in self.consolidated_metadata.flattened_metadata.items(): + # A member read from a document that had to be repaired is stored as it + # was stored (see `ConsolidatedMetadata.to_dict`). + if isinstance(member, ArrayV3Metadata) and member._stored_document is None: + try: + check_storable(member) + except ValueError as e: + e.add_note(f"Array {path!r} in the consolidated metadata.") + raise indent = config.get("json_indent") if self.zarr_format == 3: return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)} @@ -834,11 +856,24 @@ async def delitem(self, key: str) -> None: Array or group name """ store_path = self.store_path / key - + consolidated = self.metadata.consolidated_metadata + if consolidated is None: + await store_path.delete_dir() + return + # Encode the group metadata without the member before deleting it: metadata that + # cannot be stored then fails with the store and this group untouched. What is + # stored is encoded after the deletion, from the metadata as it then is, so + # concurrent deletions each store the deletions made before them. + members = {name: node for name, node in consolidated.metadata.items() if name != key} + encode_documents( + self.store_path, + replace(self.metadata, consolidated_metadata=replace(consolidated, metadata=members)), + ) await store_path.delete_dir() - if self.metadata.consolidated_metadata: - self.metadata.consolidated_metadata.metadata.pop(key, None) - await self._save_metadata() + # In place, so every handle sharing this consolidated metadata (a parent's or a + # subgroup's) sees the deletion. + consolidated.metadata.pop(key, None) + await store_documents(self.store_path, encode_documents(self.store_path, self.metadata)) async def get[DefaultT]( self, key: str, default: DefaultT | None = None @@ -3307,11 +3342,11 @@ async def create_hierarchy( else: nodes_explicit[k] = v - # Build every node before deleting or storing anything: metadata that no array or - # group can be built from then fails with the store untouched. - built = _build_nodes(store, nodes_explicit) + # Build and encode every node before deleting anything: a node that cannot be built + # or whose metadata cannot be stored then fails with the store untouched. + built, documents = _prepare_nodes(store, nodes_explicit) await asyncio.gather(*(store.delete_dir(key) for key in to_delete_keys)) - async for key, node in _store_nodes(store, nodes_explicit, built): + async for key, node in _store_nodes(store, nodes_explicit, built, documents): yield key, node @@ -3341,34 +3376,42 @@ async def create_nodes( AsyncGroup | AsyncArray The created nodes in the order they are created. """ - async for key, node in _store_nodes(store, nodes, _build_nodes(store, nodes)): + async for key, node in _store_nodes(store, nodes, *_prepare_nodes(store, nodes)): yield key, node -def _build_nodes( +def _prepare_nodes( store: Store, nodes: Mapping[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata] -) -> dict[str, AsyncGroup | AnyAsyncArray]: - """The array or group each of `nodes` describes, at its path in `store`.""" - return { +) -> tuple[dict[str, AsyncGroup | AnyAsyncArray], dict[str, Buffer]]: + """The array or group each of `nodes` describes, at its path in `store`, and the + metadata documents of `nodes`, by their keys in the store: every node is built and + encoded before anything is stored.""" + built = { path: _build_node(store=store, path=path, metadata=meta) for path, meta in nodes.items() } + documents = { + _join_paths([path, key]): value + for path, metadata in nodes.items() + for key, value in encode_documents(StorePath(store, path), metadata).items() + } + return built, documents async def _store_nodes( store: Store, nodes: Mapping[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], built: Mapping[str, AsyncGroup | AnyAsyncArray], + documents: Mapping[str, Buffer], ) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]: - """Store the metadata of `nodes` and yield the nodes `_build_nodes` built from them - (see `create_nodes`).""" + """Store the `documents` encoded from `nodes` and yield the nodes built from them, as + `_prepare_nodes` returns them (see `create_nodes`).""" # Note: the only way to alter this value is via the config. If that's undesirable for some reason, # then we should consider adding a keyword argument to this function semaphore = asyncio.Semaphore(config.get("async.concurrency")) - create_tasks: list[Coroutine[None, None, str]] = [] - - for key, value in nodes.items(): - # make the key absolute - create_tasks.extend(_persist_metadata(store, key, value, semaphore=semaphore)) + create_tasks = [ + _set_return_key(store=store, key=key, value=value, semaphore=semaphore) + for key, value in documents.items() + ] created_object_keys = [] @@ -3881,23 +3924,6 @@ async def _set_return_key( return key -def _persist_metadata( - store: Store, - path: str, - metadata: ArrayV2Metadata | ArrayV3Metadata | GroupMetadata, - semaphore: asyncio.Semaphore | None = None, -) -> tuple[Coroutine[None, None, str], ...]: - """ - Prepare to save a metadata document to storage, returning a tuple of coroutines that must be awaited. - """ - - to_save = metadata.to_buffer_dict(default_buffer_prototype()) - return tuple( - _set_return_key(store=store, key=_join_paths([path, key]), value=value, semaphore=semaphore) - for key, value in to_save.items() - ) - - async def create_rooted_hierarchy( *, store: Store, diff --git a/src/zarr/core/metadata/io.py b/src/zarr/core/metadata/io.py index 5a9237327d..c1cd927b94 100644 --- a/src/zarr/core/metadata/io.py +++ b/src/zarr/core/metadata/io.py @@ -70,8 +70,28 @@ def _diff( yield DocumentChange(path, stored, new) +def encode_documents( + store_path: StorePath, metadata: ArrayMetadata | GroupMetadata +) -> dict[str, Buffer]: + """The metadata documents `metadata` stores under `store_path`, by key (see + `to_buffer_dict`). + + An operation that deletes or writes store content encodes its documents first, so + metadata that cannot be stored fails with the store untouched; the error then names + the node at `store_path`, as the warnings about stored documents do. + """ + from zarr.core.group import GroupMetadata + + try: + return metadata.to_buffer_dict(default_buffer_prototype()) + except ValueError as e: + node = "Group" if isinstance(metadata, GroupMetadata) else "Array" + e.add_note(f"{node} {str(store_path)!r}: nothing was stored.") + raise + + async def store_documents(store_path: StorePath, documents: Mapping[str, Buffer]) -> None: - """Store metadata documents encoded by `to_buffer_dict` under `store_path`.""" + """Store metadata documents encoded by `encode_documents` under `store_path`.""" await asyncio.gather( *(set_or_delete(store_path / key, value) for key, value in documents.items()) ) @@ -97,7 +117,7 @@ async def upsert_metadata( The documents are encoded before any is stored, so metadata that cannot be stored fails with the store untouched. """ - documents = metadata.to_buffer_dict(default_buffer_prototype()) + documents = encode_documents(store_path, metadata) changes = diff_documents( {key: buffer_to_json_object(buf) for key, buf in stored.items() if key in documents}, {key: buffer_to_json_object(buf) for key, buf in documents.items()}, @@ -114,11 +134,17 @@ async def upsert_metadata( """The store keys of the metadata documents of an array of each Zarr format.""" -def parse_stored_array(documents: Mapping[str, Buffer], zarr_format: ZarrFormat) -> ArrayMetadata: +def parse_stored_array( + documents: Mapping[str, Buffer], zarr_format: ZarrFormat, path: str | None = None +) -> ArrayMetadata: """The metadata of an array from its documents (by store key, see `ARRAY_DOCUMENTS`), read with the repairs but without their warnings (whoever asks has warned, or reads metadata built in code), and marked (see `mark_repaired`) if they had to be - repaired. Raises `ArrayNotFoundError` if there is no array document among them.""" + repaired. Raises `ArrayNotFoundError` if there is no array document among them. + + Only operations that store metadata read documents this way, so documents read as a + rectilinear chunk grid require the rectilinear chunks flag, as storing it does; the + error names the array at `path`.""" from zarr.core.array import ( _array_metadata_dict_v2, _array_metadata_dict_v3, @@ -132,7 +158,8 @@ def parse_stored_array(documents: Mapping[str, Buffer], zarr_format: ZarrFormat) else: raise ArrayNotFoundError(f"No Zarr format {zarr_format} array metadata document.") repaired, readings = repair_array_document(stored, zarr_format) - return mark_repaired(parse_array_metadata(dict(repaired)), stored, readings, None, warn=False) + metadata = parse_array_metadata(dict(repaired), path) + return mark_repaired(metadata, stored, readings, None, warn=False) def _build_parents(store_path: StorePath, zarr_format: ZarrFormat) -> dict[str, GroupMetadata]: @@ -172,7 +199,7 @@ async def save_metadata( ------ ValueError """ - to_save = metadata.to_buffer_dict(default_buffer_prototype()) + to_save = encode_documents(store_path, metadata) await _write_metadata(store_path, to_save, metadata.zarr_format, ensure_parents=ensure_parents) @@ -185,8 +212,8 @@ async def save_new_metadata( ) -> None: """Save the metadata of a new array or group, replacing any existing node if requested. - The metadata is encoded before the store is modified, so metadata that cannot be - encoded raises without deleting an existing node. + The metadata is encoded first (see `encode_documents`), so metadata that cannot be + stored raises, naming the node, without deleting an existing node. Parameters ---------- @@ -201,7 +228,7 @@ async def save_new_metadata( ensure_parents : bool, optional Create any missing parent groups, and check no existing parents are arrays. """ - to_save = metadata.to_buffer_dict(default_buffer_prototype()) + to_save = encode_documents(store_path, metadata) if overwrite and store_path.store.supports_deletes: await store_path.delete_dir() else: diff --git a/src/zarr/core/metadata/repair.py b/src/zarr/core/metadata/repair.py index 06ce97dce7..bd0651fb26 100644 --- a/src/zarr/core/metadata/repair.py +++ b/src/zarr/core/metadata/repair.py @@ -1,9 +1,10 @@ """Repairs that read invalid stored array metadata documents written by older software. -This is the only place invalid metadata is read leniently. A repair maps a stored -array metadata document (parsed JSON) to a valid one. `ArrayV2Metadata.from_dict` and -`ArrayV3Metadata.from_dict` apply the repairs for their Zarr format, so every path -that parses a stored document, including consolidated metadata, goes through them. +This is the only place invalid metadata is read leniently; the metadata constructors +are strict. A repair maps a stored array metadata document (parsed JSON) to a valid +one. `ArrayV2Metadata.from_dict` and `ArrayV3Metadata.from_dict` apply the repairs for +their Zarr format, so every path that parses a stored document, including consolidated +metadata, goes through them. A reading warns only where the user must act on it; a reading that gives what zarr read from the same document before these repairs existed is silent, so a document @@ -50,6 +51,13 @@ class Reading(NamedTuple): """Returns `None` if the document needs no repair, else the repaired document and how it was read.""" +RESAVE_HINT: Final = ( + "To store valid metadata, open the array writable and call `array.update_attributes({})`; " + "if a group holds consolidated metadata for the array, then also call " + "`zarr.consolidate_metadata` on that group." +) +"""How to store the repair of a document whose array holds data.""" + RECREATE_HINT: Final = ( "The array holds only its fill value, so recreate it with the chunk shape you want; " "nothing is lost: `zarr.from_array(array.store, name=array.path, data=array, " @@ -95,13 +103,20 @@ def _is_int_list(value: object) -> TypeGuard[list[int]]: return isinstance(value, list) and all(isinstance(v, int) for v in value) +def _abbreviate(value: JSON, limit: int = 60) -> str: + """`value` as JSON, cut to at most `limit` characters.""" + text = json.dumps(value) + return text if len(text) <= limit else f"{text[: limit - 3]}..." + + def _read_chunk_size( size: JSON, span: int | None, unit: int -) -> tuple[int, bool, str | None] | None: - """Read one entry of a stored regular chunk shape as a chunk edge length. +) -> tuple[int | list[JSON], bool, str | None] | None: + """Read one entry of a stored regular chunk shape as a chunk edge length, or as the + chunk edge lengths of its axis. - Returns the edge length, whether it moves chunks (it is not the size a reader of - the stored entry uses), and, where the user must act on how it was read, how it was + Returns the reading, whether it moves chunks (it is not the size a reader of the + stored entry uses), and, where the user must act on how it was read, how it was read; `None` if the entry cannot be read, which leaves it for the metadata constructors to check. A JSON int >= 1 is kept, JSON `true` is read as 1, and 0 or JSON `false` is read as `unit`, the smallest chunk edge length the axis can have (1, @@ -110,7 +125,10 @@ def _read_chunk_size( does not depend on how far the axis has grown since. On an axis of positive length no chunk can have been stored under a chunk size of 0, so the array holds only its fill value. `span` is `None` where no stored 0 is known, as in the inner chunk shape of a - sharding codec: 0 is then left as stored. + sharding codec: 0 is then left as stored. A flat JSON list is kept as the chunk edge + lengths of its axis, which only a rectilinear chunk grid can declare (see + `_invalid_chunk_sizes_v3`); its edges are read as those of a stored rectilinear chunk + grid (see `_invalid_edge_lengths_v3`). """ match size: case True: @@ -130,23 +148,25 @@ def _read_chunk_size( "chunk size of 0" ), ) + case list() if not any(isinstance(edge, list) for edge in size): + return size, False, None return None def _read_chunk_shape( stored: JSON, spans: Sequence[int | None], units: Iterable[int] = () -) -> tuple[list[int], Reading] | None: +) -> tuple[list[int | list[JSON]], Reading | None] | None: """Read a stored regular chunk shape, entry by entry (see `_read_chunk_size`), for axes of lengths `spans` whose chunks are multiples of `units` (1 where not given). - Returns the chunk shape and how it was read, if any entry was read as another value + Returns the chunk shape and how it was read if any entry was read as another value (the warning says how the chunk shape was read and, as the array then holds only its - fill value, recommends recreating it, see `RECREATE_HINT`); `None` if it cannot be - read or needs no repair. + fill value, recommends recreating it, see `RECREATE_HINT`), else `None`; `None` if + it cannot be read. """ if not (isinstance(stored, list) and len(stored) == len(spans)): return None - edges: list[int] = [] + edges: list[int | list[JSON]] = [] changed = moves_chunks = False readings: list[str] = [] axes = zip(stored, spans, chain(units, repeat(1)), strict=False) @@ -161,11 +181,11 @@ def _read_chunk_shape( if how is not None: readings.append(f"{json.dumps(size)} in dimension {axis} as {how}") if not changed: - return None + return edges, None warning = ( - f"The stored chunk shape {json.dumps(stored)} is invalid: chunk sizes must be " - f"integers of at least 1. It is read as {edges}, reading {'; '.join(readings)}. " - f"{RECREATE_HINT}" + f"The stored chunk shape {_abbreviate(stored)} is invalid: chunk sizes must be " + f"integers of at least 1. It is read as {_abbreviate(edges)}, reading " + f"{'; '.join(readings)}. {RECREATE_HINT}" if readings else None ) @@ -177,7 +197,7 @@ def _invalid_chunk_sizes_v2(doc: ArrayDocument) -> tuple[ArrayDocument, Reading] if not _is_int_list(shape): return None match _read_chunk_shape(doc.get("chunks"), shape): - case chunks, reading: + case chunks, Reading() as reading: return {**doc, "chunks": chunks}, reading return None @@ -192,10 +212,13 @@ def _read_codec(codec: JSON) -> JSON | None: case {"name": "sharding_indexed", "configuration": Mapping() as configuration}: repaired: dict[str, JSON] = {} stored = configuration.get("chunk_shape") - if isinstance(stored, list) and ( - read := _read_chunk_shape(stored, [None] * len(stored)) + match ( + _read_chunk_shape(stored, [None] * len(stored)) + if isinstance(stored, list) + else None ): - repaired["chunk_shape"] = read[0] + case chunk_shape, Reading(): + repaired["chunk_shape"] = chunk_shape if isinstance(codecs := configuration.get("codecs"), list) and ( inner := _read_codecs(codecs) ): @@ -249,11 +272,35 @@ def _invalid_chunk_sizes_v3(doc: ArrayDocument) -> tuple[ArrayDocument, Reading] configuration = grid.get("configuration") if not isinstance(configuration, Mapping): return None - match _read_chunk_shape(configuration.get("chunk_shape"), shape, units): - case chunk_shape, reading: - repaired = {**configuration, "chunk_shape": chunk_shape} - return {**doc, "chunk_grid": {**grid, "configuration": repaired}}, reading - return None + read = _read_chunk_shape(configuration.get("chunk_shape"), shape, units) + if read is None: + return None + chunk_shape, reading = read + edge_axes = [axis for axis, size in enumerate(chunk_shape) if isinstance(size, list)] + if not edge_axes: + if reading is None: + return None + repaired = {**configuration, "chunk_shape": chunk_shape} + return {**doc, "chunk_grid": {**grid, "configuration": repaired}}, reading + if len(edge_axes) == len(chunk_shape): + # Only a mix of chunk sizes and edge lists was ever stored in a regular grid. + return None + # The user must act: re-saving this grid needs the rectilinear chunks flag. + as_rectilinear = ( + f"The stored chunk grid is named 'regular', but its chunk shape lists chunk edge " + f"lengths in dimensions {edge_axes}, which only a rectilinear chunk grid can declare. " + "It is read as that rectilinear chunk grid. Re-saving the metadata stores that " + "rectilinear chunk grid, so each step that follows requires " + f"`zarr.config.set({{'array.rectilinear_chunks': True}})`. {RESAVE_HINT}" + ) + rectilinear: JSON = { + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": chunk_shape}, + } + warning = " ".join(filter(None, (reading and reading.warning, as_rectilinear))) + # Only zarr 3.2.x reads the stored grid, so the array stores the rectilinear grid + # before it writes chunks, for every other reader to find them. + return {**doc, "chunk_grid": rectilinear}, Reading(moves_chunks=True, warning=warning) def _read_edge_length(edge: JSON) -> JSON: diff --git a/src/zarr/core/metadata/v2.py b/src/zarr/core/metadata/v2.py index 136267fc58..02d2ee9be9 100644 --- a/src/zarr/core/metadata/v2.py +++ b/src/zarr/core/metadata/v2.py @@ -38,8 +38,9 @@ JSON, ZARRAY_JSON, ZATTRS_JSON, + ChunkShape, MemoryOrder, - ShapeLike, + parse_chunk_shape, parse_shapelike, ) from zarr.core.config import config, parse_indexing_order @@ -82,7 +83,7 @@ def __init__( *, shape: tuple[int, ...], dtype: ZDType[TBaseDType, TBaseScalar], - chunks: tuple[int, ...], + chunks: ChunkShape, fill_value: Any, order: MemoryOrder, dimension_separator: Literal[".", "/"] = ".", @@ -330,10 +331,9 @@ def parse_compressor(data: object) -> Numcodec | None: raise ValueError(msg) -def parse_chunks(chunks: ShapeLike, shape: tuple[int, ...]) -> tuple[int, ...]: - """Check a chunk shape: one non-negative integer per array axis (see - `parse_shapelike`). Stored chunk sizes of 0 are read by `zarr.core.metadata.repair`.""" - chunks_parsed = parse_shapelike(chunks) +def parse_chunks(chunks: object, shape: tuple[int, ...]) -> tuple[int, ...]: + """Check a chunk shape: one chunk edge length (an `int` >= 1) per array axis.""" + chunks_parsed = parse_chunk_shape(chunks) if len(chunks_parsed) != len(shape): raise ValueError( f"The `shape` and `chunks` attributes must have the same length. " diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index e58082aeb9..2d5d169e6b 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -296,10 +296,6 @@ def from_dict(cls, data: RegularChunkGridMetadataJSON) -> Self: # type: ignore[ return cls(chunk_shape=parse_chunk_shape(configuration["chunk_shape"])) -class RectilinearChunksDisabledError(ValueError): - """Rectilinear chunk grids are used while the `array.rectilinear_chunks` flag is off.""" - - @dataclass(frozen=True, kw_only=True) class RectilinearChunkGridMetadata(Metadata): """Metadata-only description of a rectilinear chunk grid. @@ -319,12 +315,6 @@ class RectilinearChunkGridMetadata(Metadata): chunk_shapes: tuple[int | tuple[int, ...], ...] def __post_init__(self) -> None: - if not config.get("array.rectilinear_chunks"): - raise RectilinearChunksDisabledError( - "Rectilinear chunk grids are experimental and disabled by default. " - "Enable them with: zarr.config.set({'array.rectilinear_chunks': True}) " - "or set the environment variable ZARR_ARRAY__RECTILINEAR_CHUNKS=True" - ) object.__setattr__(self, "chunk_shapes", _validate_chunk_shapes(self.chunk_shapes)) @property @@ -392,6 +382,37 @@ def from_dict(cls, data: RectilinearChunkGridMetadataJSON) -> Self: # type: ign ChunkGridMetadata = RegularChunkGridMetadata | RectilinearChunkGridMetadata +class RectilinearChunksDisabledError(ValueError): + """Rectilinear chunk grids are used while the `array.rectilinear_chunks` flag is off.""" + + +def _check_rectilinear_chunks_enabled() -> None: + """Raise unless rectilinear chunks are enabled. + + The flag gates storing and reading array metadata documents that declare a + rectilinear chunk grid; the chunk grid metadata classes themselves are not gated. + """ + if not config.get("array.rectilinear_chunks"): + raise RectilinearChunksDisabledError( + "Rectilinear chunk grids are experimental and disabled by default. " + "Enable them with: zarr.config.set({'array.rectilinear_chunks': True}) " + "or set the environment variable ZARR_ARRAY__RECTILINEAR_CHUNKS=True" + ) + + +def check_storable(metadata: ArrayV3Metadata) -> None: + """Raise if `metadata` may not be stored: a rectilinear chunk grid requires the + rectilinear chunks flag. `zarr.core.metadata.io.encode_documents` names the node in + the error. + + Every serialization of array metadata for a store calls this, before the store is + touched: `ArrayV3Metadata.to_buffer_dict` and, for the arrays in a group's + consolidated metadata, `GroupMetadata.to_buffer_dict`. + """ + if isinstance(metadata.chunk_grid, RectilinearChunkGridMetadata): + _check_rectilinear_chunks_enabled() + + def create_chunk_grid_metadata( chunks: ChunkGrid, ) -> ChunkGridMetadata: @@ -639,6 +660,7 @@ def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: return self.chunk_key_encoding.encode_chunk_key(chunk_coords) def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: + check_storable(self) indent = config.get("json_indent") return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)} @@ -646,7 +668,17 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: """Read a stored `zarr.json` array document. An invalid document that `zarr.core.metadata.repair` can read is read as repaired; a reading the user - must act on warns, naming the array at `path`.""" + must act on warns, and a document the rectilinear chunks flag refuses raises, + naming the array at `path`.""" + # The flag gates what the document declares, so it is checked before repairs. + chunk_grid = data.get("chunk_grid") + if isinstance(chunk_grid, Mapping) and chunk_grid.get("name") == "rectilinear": + try: + _check_rectilinear_chunks_enabled() + except ValueError as e: + if path is not None: + e.add_note(f"Array {path!r}: nothing was read.") + raise repaired, readings = repair_array_document(data, 3) # a new dict, because we are modifying it _data = dict(repaired) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 00932c96ce..4d3613d6b6 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) @@ -533,72 +548,178 @@ def simple_arrays( ) -def chunks_param_from_rectilinear( - meta: RectilinearChunkGridMetadata, -) -> list[int | list[int]]: - """Convert rectilinear chunk grid metadata into a `chunks=` argument. +# 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 - Explicit edge tuples 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] + +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_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, 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 - 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: + 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: - chunk_shapes.append([1]) - return chunk_shapes + 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 +734,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))) @@ -943,7 +1057,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/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.py b/tests/test_array.py index 9d3ab32bfa..c0bb96b1cd 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -1168,7 +1168,7 @@ def test_auto_partition_auto_shards_with_auto_chunks_should_be_close_to_1MiB() - [(10, 10), [10, 10], np.array([10, 10]), (np.int64(10), np.int64(10))], ids=["tuple", "list", "array", "numpy-scalars"], ) -def test_chunks_and_shards(chunks: ChunksLike, shards: ShardsLike) -> None: +def test_chunks_and_shards(chunks: ChunksLike, shards: Any) -> None: store = StorePath(MemoryStore()) shape = (100, 100) expected_chunks = normalize_chunks_nd(chunks, shape).chunk_shape diff --git a/tests/test_array_stateful.py b/tests/test_array_stateful.py index a25b056c7d..7e47645118 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_chunk_grids.py b/tests/test_chunk_grids.py index 207f121d11..f6b90534e5 100644 --- a/tests/test_chunk_grids.py +++ b/tests/test_chunk_grids.py @@ -1,4 +1,5 @@ import contextlib +import re from typing import Any, Literal, cast import numpy as np @@ -62,6 +63,20 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None: assert all(0 < c <= max(s, 1) for c, s in zip(chunks, shape, strict=False)) +class _IndexedSequence: + """A sequence with only `__getitem__` and `__len__`: `list` takes it, but + `collections.abc.Iterable` does not recognize it.""" + + def __init__(self, *items: int) -> None: + self._items = items + + def __getitem__(self, index: int) -> int: + return self._items[index] + + def __len__(self) -> int: + return len(self._items) + + @pytest.mark.parametrize( ("chunks", "shape", "expected"), [ @@ -104,8 +119,14 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None: ((1, 3, np.int64(16), np.int64(16)), (1, 3, 32, 32), (1, 3, 16, 16)), ((np.int32(30), np.int64(-1)), (100, 20), (30, 20)), (np.array([10, 10]), (100, 100), (10, 10)), + # 0-d integer arrays are integers, as the scalar form and per dimension + (np.array(10), (100, 100), (10, 10)), + ((np.array(5), np.int64(-1)), (10, 6), (5, 6)), # rectilinear chunks given as numpy arrays ((np.array([60, 40]), np.array([50, 50])), (100, 100), ((60, 40), (50, 50))), + # a sequence with only `__getitem__` and `__len__` is a list of chunk sizes + (_IndexedSequence(5, 5), (10, 10), (5, 5)), + ((_IndexedSequence(6, 4), 5), (10, 10), ((6, 4), 5)), ], ) def test_normalize_chunks( @@ -198,6 +219,13 @@ def test_chunk_layout_nested() -> None: msg="must be an integer or an iterable of integers; got 2.5 of type float", escape=True, ), + # an integral float is still not an integer + ExpectFail( + input=(10.0, 100), + exception=TypeError, + id="integral-float-scalar", + msg="got 10.0 of type float", + ), ExpectFail( input=([10, -1, 10], 100), exception=ValueError, @@ -251,21 +279,14 @@ def test_normalize_chunks_1d_errors(case: ExpectFail[tuple[Any, int]]) -> None: @pytest.mark.parametrize( "case", [ + ExpectFail(input=("foo", (100,)), exception=ValueError, id="string", msg="dimensions"), + # A 0-d array is an integer only if its dtype is. ExpectFail( - input=(None, (100,)), - exception=ValueError, - id="none", - msg="None is not a valid chunk input", - ), - # `True` is rejected explicitly because bool is a subclass of int — without - # this guard, `chunks=True` would silently produce size-1 chunks. - ExpectFail( - input=(True, (100,)), - exception=ValueError, - id="true", - msg="True is not a valid chunk input", + input=(np.array(2.0), (100,)), + exception=TypeError, + id="0-d-float-array", + msg="must be an integer or an iterable of integers", ), - ExpectFail(input=("foo", (100,)), exception=ValueError, id="string", msg="dimensions"), ExpectFail( input=((100, 10), (100,)), exception=ValueError, id="too-many-dims", msg="dimensions" ), @@ -290,6 +311,46 @@ def test_normalize_chunks_nd_errors(case: ExpectFail[tuple[Any, tuple[int, ...]] normalize_chunks_nd(chunks, shape) +@pytest.mark.parametrize( + "chunks", + [ + True, + np.True_, + np.array(True), + np.array([True, True]), + (True, 5), + (np.True_, 5), + (np.array(True), 5), + [[True, 9], 10], + [[np.True_, 9], 10], + np.False_, + np.array(False), + np.array([False, False]), + (False, 5), + (np.False_, 5), + (np.array(False), 5), + [[False, 9], 10], + [[np.False_, 9], 10], + ], + ids=repr, +) +def test_normalize_chunks_nd_rejects_bool(chunks: Any) -> None: + """A boolean is a flag, not a chunk size: every spelling of one, as the whole + specification, as a dimension's size, or as an edge in a dimension's list, is + rejected with the same error rather than read as a size of 0 or 1. The one + exception is `False` as the whole specification (one chunk covering every axis).""" + with pytest.raises(TypeError, match="A bool is not a chunk size"): + normalize_chunks_nd(chunks, (10, 10)) + + +@pytest.mark.parametrize("chunks", [True, None]) +def test_normalize_chunks_nd_names_automatic_chunking(chunks: Any) -> None: + """`chunks=True` meant automatic chunking in zarr 2, and `chunks=None` means it in + `zarr.create`: their errors say how to ask for it here.""" + with pytest.raises(TypeError, match=re.escape('For automatic chunking, pass "auto"')): + normalize_chunks_nd(chunks, (10, 10)) + + @pytest.mark.parametrize( "case", [ diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 188d6d85c1..f027617a7c 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -765,15 +765,13 @@ def test_structured_dtype_fill_value() -> None: [ ShardingCodec(chunk_shape=(8, 8)), ShardingCodec(chunk_shape=(8, 8), subchunk_write_order="lexicographic"), - ShardingCodec(chunk_shape=(0,)), ], - ids=["default", "lexicographic", "zero-chunk-size"], + ids=["default", "lexicographic"], ) def test_pickle(codec: ShardingCodec) -> None: """ShardingCodec round-trips through pickle, including the non-serialized ``subchunk_write_order`` (which ``to_dict`` omits and which must not silently - revert to the ``morton`` default), and an inner chunk size of 0, which the - constructor accepts.""" + revert to the ``morton`` default).""" restored = pickle.loads(pickle.dumps(codec)) assert restored == codec assert restored.chunk_shape == codec.chunk_shape diff --git a/tests/test_group.py b/tests/test_group.py index 2fbd4d2650..6976ef536c 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import contextlib import inspect import json @@ -1733,11 +1734,13 @@ async def test_group_getitem_consolidated(self, store: Store) -> None: rg2 = await rg1.get_group("g2") assert rg2.metadata.consolidated_metadata == ConsolidatedMetadata(metadata={}) - async def test_group_delitem_consolidated(self, store: Store) -> None: + async def test_group_delitem_consolidated(self, store: Store, zarr_format: ZarrFormat) -> None: + """Deleting a member removes it from the consolidated metadata in memory and in + every document that stores it, so the group reopens without it.""" if isinstance(store, ZipStore): raise pytest.skip("Not implemented") - root = await AsyncGroup.from_store(store=store) + root = await AsyncGroup.from_store(store=store, zarr_format=zarr_format) # Set up the test structure with # / # g0/ # group /g0 @@ -1759,23 +1762,80 @@ async def test_group_delitem_consolidated(self, store: Store) -> None: x2 = await x1.create_group("x2") await x2.create_array("data", shape=(1,), dtype="uint8") - with pytest.warns( # noqa: PT031 - ZarrUserWarning, - match="Consolidated metadata is currently not part in the Zarr format 3 specification.", - ): - if isinstance(store, ZipStore): - with pytest.warns(UserWarning, match="Duplicate name"): - await zarr.api.asynchronous.consolidate_metadata(store) - else: - await zarr.api.asynchronous.consolidate_metadata(store) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "Consolidated metadata is currently not part", ZarrUserWarning + ) + await zarr.api.asynchronous.consolidate_metadata(store) - group = await zarr.api.asynchronous.open_consolidated(store=store) - assert len(group.metadata.consolidated_metadata.metadata) == 2 - assert "g0" in group.metadata.consolidated_metadata.metadata + group = await zarr.api.asynchronous.open_consolidated(store=store, zarr_format=zarr_format) + assert group.metadata.consolidated_metadata is not None + assert sorted(group.metadata.consolidated_metadata.metadata) == ["g0", "x0"] await group.delitem("g0") - assert len(group.metadata.consolidated_metadata.metadata) == 1 - assert "g0" not in group.metadata.consolidated_metadata.metadata + assert sorted(group.metadata.consolidated_metadata.metadata) == ["x0"] + + reopened = await zarr.api.asynchronous.open_consolidated( + store=store, zarr_format=zarr_format + ) + assert reopened.metadata.consolidated_metadata is not None + assert sorted(reopened.metadata.consolidated_metadata.metadata) == ["x0"] + + def test_group_delitem_consolidated_aliased(self, store: Store) -> None: + """A subgroup read from its parent's consolidated metadata shares it, so a member + deleted through the subgroup is gone through the parent too.""" + if isinstance(store, ZipStore): + raise pytest.skip("Not implemented") + + root = zarr.create_group(store) + root.create_group("sub").create_array("b", shape=(4,), chunks=(2,), dtype="i4") + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "Consolidated metadata is currently not part", ZarrUserWarning + ) + zarr.consolidate_metadata(store) + + group = zarr.open_group(store, mode="r+", use_consolidated=True) + sub = group["sub"] + assert isinstance(sub, Group) + del sub["b"] + assert "b" not in sub + assert "b" not in group["sub"] + with pytest.raises(KeyError): + group["sub/b"] + + async def test_group_delitem_consolidated_concurrent(self, zarr_format: ZarrFormat) -> None: + """Concurrent deletions through one handle each store the deletions made before + them, so the stored consolidated metadata lists what the handle lists. Here the + deletions finish in the reverse of the order they started in.""" + + delays = [0.03, 0.02, 0.01] + + class SlowDeletes(LatencyStore): + async def delete_dir(self, prefix: str) -> None: + await asyncio.sleep(delays.pop(0)) + await super().delete_dir(prefix) + + store = SlowDeletes(MemoryStore()) + root = await AsyncGroup.from_store(store=store, zarr_format=zarr_format) + for name in "abcd": + await root.create_array(name, shape=(2,), dtype="i4") + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "Consolidated metadata is currently not part", ZarrUserWarning + ) + await zarr.api.asynchronous.consolidate_metadata(store) + group = await zarr.api.asynchronous.open_consolidated(store=store, zarr_format=zarr_format) + + await asyncio.gather(*(group.delitem(name) for name in "abc")) + + reopened = await zarr.api.asynchronous.open_consolidated( + store=store, zarr_format=zarr_format + ) + assert group.metadata.consolidated_metadata is not None + assert reopened.metadata.consolidated_metadata is not None + assert list(group.metadata.consolidated_metadata.metadata) == ["d"] + assert list(reopened.metadata.consolidated_metadata.metadata) == ["d"] def test_open_consolidated_raises(self, store: Store) -> None: if isinstance(store, ZipStore): diff --git a/tests/test_metadata/test_io.py b/tests/test_metadata/test_io.py index 999e5a15a6..d9e35f2c16 100644 --- a/tests/test_metadata/test_io.py +++ b/tests/test_metadata/test_io.py @@ -26,6 +26,7 @@ from zarr.abc.store import Store from zarr.core.buffer import Buffer from zarr.core.common import JSON + from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata V3_DOC: dict[str, JSON] = { "zarr_format": 3, @@ -191,17 +192,21 @@ def test_upsert_metadata_identical_stores_nothing(zarr_format: Literal[2, 3]) -> assert store.sets == 0 -def test_upsert_metadata_unstorable_leaves_store_untouched(monkeypatch: pytest.MonkeyPatch) -> None: - """Metadata that cannot be encoded fails before the store is written.""" - store_path, metadata = _legacy(3) +def test_upsert_metadata_unstorable_leaves_store_untouched() -> None: + """Metadata that may not be stored (a rectilinear chunk grid without the rectilinear + chunks flag) fails before the store is read or written, naming the array.""" + store_path, _ = _legacy(3) before = _documents(store_path.store) - - def refuse(*args: object) -> None: - raise ValueError("cannot be stored") - - monkeypatch.setattr(ArrayV3Metadata, "to_buffer_dict", refuse) - with pytest.raises(ValueError, match="cannot be stored"): + with zarr.config.set({"array.rectilinear_chunks": True}): + metadata = zarr.create_array( + MemoryStore(), shape=(3,), chunks=[[1, 2]], dtype="int16" + ).metadata + with ( + zarr.config.set({"array.rectilinear_chunks": False}), + pytest.raises(ValueError, match="experimental and disabled") as info, + ): _upsert(store_path, metadata) + assert info.value.__notes__ == [f"Array {str(store_path)!r}: nothing was stored."] assert _documents(store_path.store) == before diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index da35922da7..1211732e7f 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -21,6 +21,7 @@ from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.repair import ( RECREATE_HINT, + RESAVE_HINT, repair_array_document, ) from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata @@ -253,7 +254,7 @@ def _open_strictly(path: Path, mode: Literal["r", "a", "r+"] = "r") -> AnyArray: def test_stored_negative_chunk_size_rejected() -> None: """No known writer stored a negative chunk size: it is rejected, not repaired.""" - with pytest.raises(ValueError, match="^Expected all values to be non-negative"): + with pytest.raises(ValueError, match="^Dimension 0: chunk edge length must be >= 1, got -1$"): _read_strictly(_v2_doc([4], [-1])) @@ -264,6 +265,13 @@ def test_stored_chunk_shape_ndim_mismatch_rejected() -> None: _read_strictly(_v3_doc([4, 4], [0])) +def test_stored_zero_inner_chunk_size_rejected() -> None: + """No known writer stored an inner chunk size of 0, and no span defines one: it is + rejected, not repaired.""" + with pytest.raises(ValueError, match="^Dimension 0: chunk edge length must be >= 1, got 0$"): + _read_strictly(_v3_doc([4], [4], inner=[0])) + + @pytest.mark.parametrize("inner", [[0], [False]]) def test_stored_zero_chunk_size_of_shard_with_invalid_inner_chunk_shape_rejected( inner: list[Any], @@ -330,9 +338,9 @@ def test_read_invalid_edges_in_rectilinear_grid( (_v3_doc([4], [4.5]), "Dimension 0: chunk edge length must be an int, got 4.5"), ( _v3_doc([8], [4], inner=[2.0]), - "Expected an iterable of integers. Got [2.0] instead.", + "Dimension 0: chunk edge length must be an int, got 2.0", ), - (_v2_doc([20], [10.0]), "Expected an iterable of integers. Got [10.0] instead."), + (_v2_doc([20], [10.0]), "Dimension 0: chunk edge length must be an int, got 10.0"), ( _rectilinear_doc([8], [[[4, 2.0]]]), "Dimension 0: RLE repeat count must be an int, got 2.0", @@ -427,51 +435,47 @@ def _second_edge(grid: RectilinearChunkGridMetadata) -> int: CHUNK_EDGE_SITES: dict[str, Callable[[Any], object]] = { "regular": lambda size: RegularChunkGridMetadata(chunk_shape=(size,)).chunk_shape[0], + "v2": lambda size: _v2_metadata((size,)).chunks[0], "rectilinear-bare": lambda size: _rectilinear((size,)).chunk_shapes[0], "rectilinear-edge": lambda size: _second_edge(_rectilinear(((4, size),))), "rectilinear-bare-json": lambda size: _rectilinear_from_dict([size]).chunk_shapes[0], "rectilinear-edge-json": lambda size: _second_edge(_rectilinear_from_dict([[4, size]])), "rectilinear-rle-json": lambda size: _second_edge(_rectilinear_from_dict([[[size, 2]]])), + "sharding-inner": lambda size: ShardingCodec(chunk_shape=(size,)).chunk_shape[0], } -"""Each place chunk grid metadata reads a chunk edge length, returning the edge it read.""" +"""Each place metadata built in code takes a chunk edge length, returning the edge it +took.""" @pytest.mark.parametrize("site", CHUNK_EDGE_SITES) -@pytest.mark.parametrize( - ("size", "expected"), - [(4, 4), (True, 1)], - ids=["int", "bool"], -) -def test_metadata_reads_integer_chunk_edge(site: str, size: object, expected: int) -> None: - """Chunk grid metadata reads an `int` or a `bool` as the `int` chunk edge length it - equals.""" - edge = CHUNK_EDGE_SITES[site](size) +def test_metadata_takes_int_chunk_edge(site: str) -> None: + """Metadata built in code takes an `int` chunk edge length of at least 1 as is.""" + edge = CHUNK_EDGE_SITES[site](4) assert type(edge) is int - assert edge == expected + assert edge == 4 @pytest.mark.parametrize("site", CHUNK_EDGE_SITES) @pytest.mark.parametrize( "size", - [4.0, np.float64(4.0), 4.5, float("inf"), "4", None, np.int64(4)], - ids=["float", "numpy-float", "fractional", "inf", "str", "none", "numpy-int"], + [True, False, 4.0, np.float64(4.0), 4.5, float("inf"), "4", None, np.int64(4)], + ids=["true", "false", "float", "numpy-float", "fractional", "inf", "str", "none", "numpy-int"], ) -def test_metadata_rejects_non_integer_chunk_edge(site: str, size: object) -> None: - """A chunk edge length in metadata built in code is an `int`: a float is rejected, - even an integral one (stored documents with integral floats are read by the - repairs), and so is a NumPy integer, as zarr 3.4.0 rejected one.""" +def test_metadata_rejects_non_int_chunk_edge(site: str, size: object) -> None: + """Metadata built in code takes chunk edge lengths as `int`s only, everywhere: not a + `bool`, a float (stored documents with integral floats are read by the repairs) or + a NumPy integer.""" with pytest.raises( - TypeError, - match=re.escape(f"Dimension 0: chunk edge length must be an int, got {size!r}"), + TypeError, match=re.escape(f"Dimension 0: chunk edge length must be an int, got {size!r}") ): CHUNK_EDGE_SITES[site](size) @pytest.mark.parametrize("site", CHUNK_EDGE_SITES) -@pytest.mark.parametrize("size", [0, False, -1]) +@pytest.mark.parametrize("size", [0, -1]) def test_metadata_rejects_chunk_edge_below_one(site: str, size: int) -> None: - """Chunk grid metadata built in code is strict: a chunk edge length below 1 is - rejected, without a warning.""" + """Metadata built in code is strict: a chunk edge length below 1 is rejected, + without a warning.""" with warnings.catch_warnings(): warnings.simplefilter("error", ZarrUserWarning) with pytest.raises( @@ -480,73 +484,37 @@ def test_metadata_rejects_chunk_edge_below_one(site: str, size: int) -> None: CHUNK_EDGE_SITES[site](size) -@pytest.mark.parametrize("chunk_shape", [4, np.int64(4), None, "44", {"4": 4}]) -def test_regular_chunk_grid_rejects_chunk_shape_not_a_sequence(chunk_shape: Any) -> None: - """A regular chunk shape is an iterable of chunk edge lengths, but not a string or a - mapping, which is rejected as a whole, not entry by entry.""" +CHUNK_SHAPE_SITES: dict[str, Callable[[Any], object]] = { + "regular": lambda chunk_shape: RegularChunkGridMetadata(chunk_shape=chunk_shape), + "v2": _v2_metadata, + "sharding-inner": lambda chunk_shape: ShardingCodec(chunk_shape=chunk_shape), +} + + +@pytest.mark.parametrize("site", CHUNK_SHAPE_SITES) +@pytest.mark.parametrize("chunk_shape", [4, np.int64(4), "10", {"a": 1}, range(1, 2)]) +def test_metadata_rejects_chunk_shape_not_list_or_tuple(site: str, chunk_shape: object) -> None: + """A regular chunk shape is a list or tuple; anything else is rejected as a whole, + not iterated as if its elements were chunk edge lengths.""" with pytest.raises( TypeError, match=re.escape( - f"A chunk shape must be an iterable of chunk edge lengths, got {chunk_shape!r}" + f"A chunk shape must be a list or tuple of chunk edge lengths, got {chunk_shape!r}" ), ): - RegularChunkGridMetadata(chunk_shape=chunk_shape) - - -def _sharding_chunk_shape(chunks: Any) -> tuple[tuple[int, ...], object]: - codec = ShardingCodec(chunk_shape=chunks) - configuration = cast("dict[str, JSON]", codec.to_dict()["configuration"]) - return codec.chunk_shape, configuration["chunk_shape"] + CHUNK_SHAPE_SITES[site](chunk_shape) -CHUNK_SHAPE_SITES: dict[str, Callable[[Any], tuple[tuple[int, ...], object]]] = { - "v2": lambda chunks: ((md := _v2_metadata(chunks)).chunks, md.to_dict()["chunks"]), - "sharding-inner": _sharding_chunk_shape, -} -"""`ArrayV2Metadata` and `ShardingCodec` read a chunk shape as an array shape, returning -the chunk shape and the value `to_dict` writes for it.""" - - -@pytest.mark.parametrize("site", CHUNK_SHAPE_SITES) -@pytest.mark.parametrize( - ("chunks", "expected"), - [ - ((4,), (4,)), - ([4], (4,)), - (4, (4,)), - (np.int64(4), (4,)), - ((np.int64(4),), (4,)), - (np.array([4]), (4,)), - ((True,), (1,)), - ((0,), (0,)), - ((False,), (0,)), - (range(4, 5), (4,)), - ], -) -def test_chunk_shape_read_as_array_shape( - site: str, chunks: object, expected: tuple[int, ...] -) -> None: - """`ArrayV2Metadata` and `ShardingCodec` read their chunk shape as `parse_shapelike` - reads an array shape: an integer or an iterable of non-negative integers, including - NumPy integers and bools. A chunk size of 0 is written back as given; reading a - stored 0 is `zarr.core.metadata.repair`' business.""" - parsed, written = CHUNK_SHAPE_SITES[site](chunks) - assert parsed == expected - assert all(type(size) is int for size in parsed) - assert written == expected - - -@pytest.mark.parametrize("site", CHUNK_SHAPE_SITES) -def test_chunk_shape_read_as_array_shape_rejects_negative(site: str) -> None: - with pytest.raises(ValueError, match="Expected all values to be non-negative"): - CHUNK_SHAPE_SITES[site]((-1,)) - - -@pytest.mark.parametrize("site", CHUNK_SHAPE_SITES) -@pytest.mark.parametrize("chunks", [(4.0,), "4", None]) -def test_chunk_shape_read_as_array_shape_rejects_non_integer(site: str, chunks: object) -> None: - with pytest.raises(TypeError, match="Expected an"): - CHUNK_SHAPE_SITES[site](chunks) +def test_consolidated_member_error_names_member() -> None: + """A group whose consolidated metadata holds a member that cannot be read does not + open from it; the error names the member.""" + member = {**_v2_doc([5], [4]), "chunks": 4} + with pytest.raises(TypeError, match="A chunk shape must be a list or tuple") as info: + ConsolidatedMetadata.from_dict( + {"kind": "inline", "must_understand": False, "metadata": {"sub/a": member}}, + path="group", + ) + assert info.value.__notes__ == ["Member 'group/sub/a' of the consolidated metadata."] def _rewrite_doc(path: Path, zarr_format: Literal[2, 3], edit: Any) -> None: @@ -833,7 +801,7 @@ def test_stale_handle_write_after_chunk_grid_change_raises( _rewrite_doc(path, zarr_format, change) documents = {p.name: p.read_bytes() for p in path.iterdir()} - with pytest.raises(ValueError, match="has changed since this array was opened: reopen"): + with pytest.raises(ValueError, match="has changed since this array was opened; reopen"): stale[0:3] = [7, 8, 9] assert stale.metadata._stored_document is not None @@ -965,39 +933,13 @@ def test_respelled_document_write_stores_only_chunks( ], ids=["codec-instance", "numpy-integer"], ) -def test_from_dict_keeps_values_that_are_not_json(codec: Any) -> None: +def test_repair_keep_values_that_are_not_json(codec: Any) -> None: """`from_dict` takes metadata built in code as well as stored documents, so the repairs read values that are not JSON (codec instances, NumPy integers) as they are, - without encoding them, and the constructors read them as before.""" + without encoding them, and leave them for the constructors to check.""" doc: dict[str, Any] = _v3_doc([4], [4]) doc["codecs"] = [codec] - metadata = ArrayV3Metadata.from_dict(doc) - assert metadata._stored_document is None - assert isinstance(metadata.codecs[0], ShardingCodec) - - -@pytest.mark.parametrize(("shape", "expected"), [((0,), (1,)), ((3,), (1,))]) -def test_array_from_metadata_with_chunk_size_zero(shape: tuple[int], expected: tuple[int]) -> None: - """`ArrayV2Metadata` accepts a chunk size of 0, as a stored document may hold it. An - array built from such metadata reads it as the repairs read that document, silently - (no data was read or written under it): `create_hierarchy` stores the metadata as - given and yields such an array, which stores the repair before its first write.""" - metadata = ArrayV2Metadata(shape=shape, chunks=(0,), dtype=Int16(), fill_value=0, order="C") - store = MemoryStore() - with warnings.catch_warnings(): - warnings.simplefilter("error", ZarrUserWarning) - nodes = dict(zarr.create_hierarchy(store=store, nodes={"a": metadata})) - array = nodes["a"] - assert isinstance(array, zarr.Array) - assert array.chunks == expected - assert json.loads(store._store_dict["a/.zarray"].to_bytes())["chunks"] == [0] - - array[...] = 1 - - # Writing an empty selection stores no chunks, so it stores no metadata either. - resaved = list(expected) if array.size else [0] - assert json.loads(store._store_dict["a/.zarray"].to_bytes())["chunks"] == resaved - np.testing.assert_array_equal(zarr.open_array(store, path="a")[...], np.ones(shape)) + assert repair_array_document(doc, 3) == (doc, []) def test_array_from_metadata_with_numpy_scalar_codec_configuration() -> None: @@ -1180,7 +1122,7 @@ def test_consolidated_repaired_member_write_after_chunk_grid_change_raises( _rewrite_doc(path / "a", zarr_format, _store_chunk_size_3) documents = {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} - with pytest.raises(ValueError, match="has changed since this array was opened: reopen"): + with pytest.raises(ValueError, match="has changed since this array was opened; reopen"): array[0:3] = [7, 8, 9] assert {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} == documents @@ -1309,3 +1251,340 @@ def test_legacy_chunk_size_consolidated(tmp_path: Path, zarr_format: Literal[2, array = reopened[name] assert isinstance(array, zarr.Array) assert array.chunks == (1,) + + +# A document copied verbatim from a store that zarr 3.2.1 wrote for +# `create_array(shape=(6, 20), chunks=(2, (5, 10, 5)), dtype="float32")`. +MIXED_REGULAR_GRID_DOC = """{ + "shape": [6, 20], + "data_type": "float32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [2, [5, 10, 5]]}}, + "chunk_key_encoding": {"name": "default", "configuration": {"separator": "/"}}, + "fill_value": 0.0, + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "zstd", "configuration": {"level": 0, "checksum": false}} + ], + "attributes": {}, + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +}""" + + +def _mixed_doc(shape: list[int], chunk_shape: list[Any]) -> dict[str, JSON]: + doc: dict[str, JSON] = json.loads(MIXED_REGULAR_GRID_DOC) + doc["shape"] = shape + doc["chunk_grid"] = {"name": "regular", "configuration": {"chunk_shape": chunk_shape}} + return doc + + +@pytest.mark.parametrize( + ("doc", "expected", "warning"), + [ + (json.loads(MIXED_REGULAR_GRID_DOC), (2, (5, 10, 5)), r"^The stored chunk grid .* \[1\]"), + ( + _mixed_doc([6, 20, 4], [2, [5, 10, 5], [1, 3]]), + (2, (5, 10, 5), (1, 3)), + r"^The stored chunk grid .* in dimensions \[1, 2\]", + ), + (_mixed_doc([6, 20], [2, [20]]), (2, (20,)), r"^The stored chunk grid .* \[1\]"), + (_mixed_doc([6, 12], [2, [5, 10, 5]]), (2, (5, 10, 5)), r"^The stored chunk grid"), + (_mixed_doc([0, 20], [2, [5, 10, 5]]), (2, (5, 10, 5)), r"^The stored chunk grid"), + (_mixed_doc([6, 0], [2, [5, 10, 5]]), (2, (5, 10, 5)), r"^The stored chunk grid"), + # JSON true is read as 1 silently; only the grid reading warns. + (_mixed_doc([6, 20], [True, [5, 10, 5]]), (1, (5, 10, 5)), r"^The stored chunk grid"), + (_mixed_doc([6, 20], [2, [5.0, 10.0, 5.0]]), (2, (5, 10, 5)), r"^The stored chunk grid"), + ( + _mixed_doc([4, 10_000], [0, [10] * 1000]), + (1, (10,) * 1000), + r"^The stored chunk shape \[0, \[10, 10, .*\.\.\. is invalid: .* read as \[1, \[10, .*\.\.\.,", + ), + ], + ids=[ + "written", + "3d", + "one-edge", + "shrunk", + "empty-int-axis", + "empty-edge-axis", + "true", + "float-edges", + "long", + ], +) +def test_read_edge_lists_in_regular_grid( + doc: dict[str, JSON], expected: tuple[int | tuple[int, ...], ...], warning: str +) -> None: + """A `regular` chunk grid whose chunk shape mixes chunk sizes with lists of chunk + edge lengths is read as the rectilinear chunk grid it describes, without the + rectilinear chunks flag. `from_dict` warns once, naming the array and the axes, + quoting a bounded part of the chunk shape, and saying that re-saving requires the + flag.""" + with ( + zarr.config.set({"array.rectilinear_chunks": False}), + warnings.catch_warnings(record=True) as record, + ): + warnings.simplefilter("always") + metadata = ArrayV3Metadata.from_dict(doc, path="group/array") + assert metadata.chunk_grid == RectilinearChunkGridMetadata(chunk_shapes=expected) + [message] = [str(w.message) for w in record] + assert re.search(warning, message.removeprefix("Array 'group/array': ")) + assert message.startswith("Array 'group/array': ") + assert ( + "Re-saving the metadata stores that rectilinear chunk grid, so each step that " + "follows requires `zarr.config.set({'array.rectilinear_chunks': True})`. " + ) in message + assert message.endswith(RESAVE_HINT) + # The 1000-edge list is abbreviated: the message is two sentences and two hints, not + # a dump of the edges. + assert len(message) < 1600 + + +def _rejected_without_warning(doc: dict[str, JSON]) -> pytest.ExceptionInfo[Exception]: + with warnings.catch_warnings(): + warnings.simplefilter("error", ZarrUserWarning) + with pytest.raises((TypeError, ValueError)) as info: + ArrayV3Metadata.from_dict(doc) + return info + + +def test_regular_grid_of_only_edge_lists_rejected() -> None: + """A regular chunk shape made only of edge lists was never stored (a rectilinear + chunk grid was), so it is not read as rectilinear.""" + info = _rejected_without_warning(_mixed_doc([6, 20], [[1, 5], [5, 10, 5]])) + assert info.match(re.escape("Dimension 0: chunk edge length must be an int, got [1, 5]")) + + +def test_run_length_encoded_edges_in_regular_grid_rejected() -> None: + """Run-length encoded edges were never stored in a regular chunk shape.""" + info = _rejected_without_warning(_mixed_doc([6, 20], [2, [[5, 2], 10]])) + assert info.match(re.escape("Dimension 1: chunk edge length must be an int, got [[5, 2], 10]")) + + +@pytest.mark.parametrize("edge", [5.5, "5"], ids=["fractional", "string"]) +def test_non_int_edge_in_regular_grid_rejected(edge: object) -> None: + """An edge that is not an int is reported as such, not blamed on its list.""" + info = _rejected_without_warning(_mixed_doc([6, 20], [2, [edge, 15]])) + assert info.match(re.escape(f"Dimension 1: chunk edge length must be an int, got {edge!r}")) + + +def test_edge_below_one_in_regular_grid_rejected() -> None: + info = _rejected_without_warning(_mixed_doc([6, 20], [2, [0, 20]])) + assert info.match("Dimension 1: chunk edge length must be >= 1, got 0") + + +def test_short_edges_in_regular_grid_rejected() -> None: + info = _rejected_without_warning(_mixed_doc([6, 20], [2, [5, 10]])) + assert info.match("sum to 15 but array shape extent is 20") + + +def _store_mixed_array(path: Path) -> np.ndarray[Any, np.dtype[np.float32]]: + """Write the chunks of the verbatim document and the document itself at `path`.""" + data = np.arange(120, dtype="float32").reshape(6, 20) + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = zarr.create_array(path, shape=data.shape, chunks=(2, (5, 10, 5)), dtype="float32") + arr[...] = data + (path / "zarr.json").write_text(MIXED_REGULAR_GRID_DOC) + return data + + +def _update_attributes(arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]: + arr.update_attributes({}) + return data + + +def _write(arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]: + arr[:2] = -data[:2] + return np.concatenate([-data[:2], data[2:]]) + + +@pytest.mark.parametrize("store_metadata", [_update_attributes, _write], ids=["re-save", "write"]) +def test_edge_lists_in_regular_grid_round_trip( + tmp_path: Path, + store_metadata: Callable[[zarr.Array[Any], np.ndarray[Any, Any]], np.ndarray[Any, Any]], +) -> None: + """A store holding the verbatim document opens without the rectilinear chunks flag + and reads its data. With the flag, re-saving the metadata, or writing chunks (which + stores the metadata first), stores the rectilinear chunk grid, which then opens + cleanly.""" + path = tmp_path / "mixed.zarr" + data = _store_mixed_array(path) + + with ( + zarr.config.set({"array.rectilinear_chunks": False}), + pytest.warns(ZarrUserWarning, match="read as that rectilinear chunk grid"), + ): + arr = zarr.open_array(path, mode="a") + np.testing.assert_array_equal(arr[...], data) + + with zarr.config.set({"array.rectilinear_chunks": True}): + expected = store_metadata(arr, data) + assert json.loads((path / "zarr.json").read_text())["chunk_grid"] == { + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": [2, [5, 10, 5]]}, + } + with warnings.catch_warnings(): + warnings.simplefilter("error", ZarrUserWarning) + reopened = zarr.open_array(path) + np.testing.assert_array_equal(reopened[...], expected) + + +def _store_mixed_group(path: Path) -> None: + """A group at `path` holding the verbatim document at `mixed` and a regular array + `n`, with consolidated metadata that quotes the verbatim document.""" + group = zarr.open_group(path, mode="w") + _store_mixed_array(path / "mixed") + group.create_array("n", data=np.arange(4), chunks=(2,)) + with zarr.config.set({"array.rectilinear_chunks": True}): + zarr.consolidate_metadata(path) + group_doc = json.loads((path / "zarr.json").read_text()) + group_doc["consolidated_metadata"]["metadata"]["mixed"] = json.loads(MIXED_REGULAR_GRID_DOC) + (path / "zarr.json").write_text(json.dumps(group_doc)) + + +def _resize(path: Path) -> None: + zarr.open_array(path / "mixed", mode="a").resize((6, 5)) + + +def _write_chunks(path: Path) -> None: + zarr.open_array(path / "mixed", mode="a")[...] = 1 + + +def _delete_member(path: Path) -> None: + del zarr.open_group(path, mode="a")["n"] + + +def _overwrite_hierarchy(path: Path) -> None: + mixed = zarr.open_array(path / "mixed") + list(zarr.create_hierarchy(store=LocalStore(path), nodes={"n": mixed.metadata}, overwrite=True)) + + +def _overwrite_with_create(path: Path) -> None: + zarr.create(shape=(4,), chunks=[[2, 2]], dtype="int64", store=path / "n", overwrite=True) # type: ignore[arg-type] + + +def _overwrite_with_create_array(path: Path) -> None: + zarr.create_array(path / "n", shape=(4,), chunks=[[2, 2]], dtype="int64", overwrite=True) + + +def _set_group_attribute(path: Path) -> None: + zarr.open_group(path, mode="a").attrs["x"] = 1 + + +@pytest.mark.filterwarnings( + "ignore:.*read as that rectilinear chunk grid:zarr.errors.ZarrUserWarning" +) +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.parametrize( + "action", + [ + _resize, + _write_chunks, + _overwrite_hierarchy, + _overwrite_with_create, + _overwrite_with_create_array, + ], + ids=["resize", "write", "overwrite-hierarchy", "overwrite-create", "overwrite-create-array"], +) +def test_store_untouched_without_flag(tmp_path: Path, action: Callable[[Path], None]) -> None: + """An operation that would store the rectilinear chunk grid read from the verbatim + document fails without the flag before it deletes or writes anything.""" + path = tmp_path / "group.zarr" + _store_mixed_group(path) + stored = {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} + with ( + zarr.config.set({"array.rectilinear_chunks": False}), + pytest.raises(ValueError, match="experimental and disabled by default"), + ): + action(path) + assert {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} == stored + + +RECTILINEAR_GRID: dict[str, Any] = { + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": [2, [5, 10, 5]]}, +} + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.parametrize("member", ["mixed", "sub/mixed"]) +def test_consolidate_edge_lists_in_regular_grid(tmp_path: Path, member: str) -> None: + """Consolidating a group holding the verbatim document copies it as stored, with or + without the flag, and leaves the member's document as it is, so the group opens and + reads the member without the flag. Once the member's metadata is re-saved (which + needs the flag), consolidating copies its rectilinear chunk grid.""" + path = tmp_path / "group.zarr" + zarr.open_group(path, mode="w").create_group("sub") + data = _store_mixed_array(path / member) + for flag in (False, True): + with ( + zarr.config.set({"array.rectilinear_chunks": flag}), + pytest.warns(ZarrUserWarning, match="read as that rectilinear chunk grid"), + ): + zarr.consolidate_metadata(path) + assert (path / member / "zarr.json").read_text() == MIXED_REGULAR_GRID_DOC + assert _consolidated_member(path, 3, member) == json.loads(MIXED_REGULAR_GRID_DOC) + with ( + zarr.config.set({"array.rectilinear_chunks": False}), + pytest.warns(ZarrUserWarning, match="read as that rectilinear chunk grid"), + ): + mixed = zarr.open_group(path, mode="r")[member] + assert isinstance(mixed, zarr.Array) + np.testing.assert_array_equal(mixed[...], data) + + with zarr.config.set({"array.rectilinear_chunks": True}): + with pytest.warns(ZarrUserWarning, match="read as that rectilinear chunk grid"): + zarr.open_array(path / member, mode="a").update_attributes({}) + zarr.consolidate_metadata(path) + assert _consolidated_member(path, 3, member)["chunk_grid"] == RECTILINEAR_GRID + with warnings.catch_warnings(): + warnings.simplefilter("error", ZarrUserWarning) + reopened = zarr.open_group(path, mode="r")[member] + assert isinstance(reopened, zarr.Array) + np.testing.assert_array_equal(reopened[...], data) + + +@pytest.mark.filterwarnings( + "ignore:.*read as that rectilinear chunk grid:zarr.errors.ZarrUserWarning" +) +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.parametrize("action", [_delete_member, _set_group_attribute], ids=["delete", "attrs"]) +@pytest.mark.parametrize("flag", [False, True]) +def test_group_write_stores_mixed_member_as_stored( + tmp_path: Path, action: Callable[[Path], None], flag: bool +) -> None: + """A group write, with or without the flag, stores the consolidated copy of the + verbatim document as stored and leaves the member's own document as it is: the + group opens and reads the member without the flag.""" + path = tmp_path / "group.zarr" + _store_mixed_group(path) + with zarr.config.set({"array.rectilinear_chunks": flag}): + action(path) + assert _consolidated_member(path, 3, "mixed") == json.loads(MIXED_REGULAR_GRID_DOC) + assert (path / "mixed" / "zarr.json").read_text() == MIXED_REGULAR_GRID_DOC + with zarr.config.set({"array.rectilinear_chunks": False}): + mixed = zarr.open_group(path, mode="r", use_consolidated=True)["mixed"] + assert isinstance(mixed, zarr.Array) + assert mixed.read_chunk_sizes == ((2, 2, 2), (5, 10, 5)) + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +def test_delete_member_without_flag_keeps_group(tmp_path: Path) -> None: + """A deletion that would store a rectilinear chunk grid in the consolidated metadata + fails without the flag before it deletes anything, and leaves the group listing the + member.""" + path = tmp_path / "group.zarr" + with zarr.config.set({"array.rectilinear_chunks": True}): + group = zarr.open_group(path, mode="w") + group.create_array("r", shape=(6, 20), chunks=(2, (5, 10, 5)), dtype="float32") + group.create_array("n", shape=(2,), chunks=(1,), dtype="int8") + zarr.consolidate_metadata(path) + group = zarr.open_group(path, mode="a", use_consolidated=True) + stored = {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} + with zarr.config.set({"array.rectilinear_chunks": False}): + with pytest.raises(ValueError, match="experimental and disabled"): + del group["n"] + assert group.metadata.consolidated_metadata is not None + assert sorted(group.metadata.consolidated_metadata.metadata) == ["n", "r"] + assert {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} == stored diff --git a/tests/test_metadata/test_v3.py b/tests/test_metadata/test_v3.py index be2564a570..94240c36b9 100644 --- a/tests/test_metadata/test_v3.py +++ b/tests/test_metadata/test_v3.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import re from typing import TYPE_CHECKING import pytest @@ -19,6 +20,7 @@ ARRAY_METADATA_KEYS, ArrayMetadataJSON_V3, ArrayV3Metadata, + RegularChunkGridMetadata, create_chunk_grid_metadata, parse_codecs, parse_dimension_names, @@ -154,6 +156,15 @@ def test_create_chunk_grid_metadata_unknown_dimension_type() -> None: create_chunk_grid_metadata(grid) +def test_regular_chunk_grid_rejects_edge_lists() -> None: + """A regular chunk grid only accepts integer chunk edge lengths.""" + with pytest.raises( + TypeError, + match=re.escape("Dimension 1: chunk edge length must be an int, got (5, 10, 5)"), + ): + RegularChunkGridMetadata(chunk_shape=(2, (5, 10, 5))) # type: ignore[arg-type] + + # --------------------------------------------------------------------------- # Types # --------------------------------------------------------------------------- diff --git a/tests/test_properties.py b/tests/test_properties.py index fcb1076441..9b38fd4f2b 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,8 +27,10 @@ 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 ( array_metadata, arrays, @@ -38,6 +41,8 @@ numpy_arrays, orthogonal_indices, rectilinear_arrays, + rectilinear_chunk_shape_declarations, + rectilinear_chunks, sharded_arrays, simple_arrays, stores, @@ -297,7 +302,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,25 +586,102 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N assert asdict_dict["fill_value"] == -9223372036854775808 -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.""" - from zarr.core.metadata.v3 import RectilinearChunkGridMetadata - from zarr.storage import MemoryStore - from zarr.testing.strategies import chunks_param_from_rectilinear +@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 - with zarr.config.set({"array.rectilinear_chunks": True}): - 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) - dst = zarr.create_array( - MemoryStore(), - shape=src.shape, - chunks=chunks_param_from_rectilinear(grid), - dtype="uint8", - ) - assert dst.metadata.chunk_grid == grid # type: ignore[union-attr] + +@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_store/test_zip.py b/tests/test_store/test_zip.py index 32b18c5273..8931875125 100644 --- a/tests/test_store/test_zip.py +++ b/tests/test_store/test_zip.py @@ -6,7 +6,7 @@ import shutil import tempfile import zipfile -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import numpy as np import pytest @@ -407,3 +407,23 @@ def test_zipstore_close_lifecycle(tmp_path: Path) -> None: lambda: ZipStoreLifecycleMachine(tmp_path), settings=settings(max_examples=50, deadline=None), ) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_resize_needing_deletes_leaves_array_unchanged( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A zip store cannot delete, so a resize that must delete chunks fails before it + stores the new metadata.""" + path = tmp_path / "a.zip" + store = ZipStore(path, mode="w") + arr = create_array(store, shape=(20,), chunks=(5,), dtype="i4", zarr_format=zarr_format) + arr[:] = np.arange(20) + with pytest.raises(NotImplementedError): + arr.resize((5,)) + assert arr.shape == (20,) + store.close() + names = zipfile.ZipFile(path).namelist() + assert len(names) == len(set(names)) + reopened = zarr.open_array(ZipStore(path, mode="r"), mode="r") + np.testing.assert_array_equal(reopened[:], np.arange(20)) diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 6a43cb4f3b..261c1594d0 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -9,21 +9,26 @@ from __future__ import annotations import re +from functools import partial from typing import TYPE_CHECKING, Any import numpy as np import pytest import zarr +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, ChunkSpec, FixedDimension, VaryingDimension, - _is_rectilinear_chunks, + guess_chunks, ) from zarr.core.common import compress_rle, expand_rle +from zarr.core.dtype import UInt8 from zarr.core.metadata.v3 import ( + ArrayV3Metadata, RectilinearChunkGridMetadata, RectilinearChunkGridMetadataJSON, RegularChunkGridMetadata, @@ -102,30 +107,57 @@ def test_dimension_index_to_chunk_last_valid( # --------------------------------------------------------------------------- +_RECTILINEAR_DOC = minimal_metadata_dict_v3( + shape=(30, 50), + chunk_grid={ + "name": "rectilinear", + "configuration": {"kind": "inline", "chunk_shapes": [[10, 20], [25, 25]]}, + }, +) + + @pytest.mark.parametrize( "action", [ - lambda: RectilinearChunkGridMetadata(chunk_shapes=((10, 20), (25, 25))), - lambda: RectilinearChunkGridMetadata.from_dict( - { - "name": "rectilinear", - "configuration": {"kind": "inline", "chunk_shapes": [[10, 20, 30], [50, 50]]}, - } - ), + lambda: ArrayV3Metadata.from_dict(dict(_RECTILINEAR_DOC)), # type: ignore[arg-type] + lambda: ArrayV3Metadata( + shape=(30, 50), + data_type=UInt8(), + chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((10, 20), (25, 25))), + chunk_key_encoding={"name": "default"}, + fill_value=0, + codecs=[{"name": "bytes"}], + attributes=None, + dimension_names=None, + ).to_buffer_dict(default_buffer_prototype()), lambda: zarr.create_array(MemoryStore(), shape=(30,), chunks=[[10, 20]], dtype="int32"), ], - ids=["constructor", "from_dict", "create_array"], + ids=["read", "store", "create_array"], ) def test_rectilinear_feature_flag_blocked(action: Any) -> None: - """Rectilinear chunk operations raise ValueError when the feature flag is disabled""" + """Reading or storing an array metadata document that declares a rectilinear chunk + grid raises ValueError when the feature flag is disabled.""" with zarr.config.set({"array.rectilinear_chunks": False}): with pytest.raises(ValueError, match="experimental and disabled by default"): action() -def test_rectilinear_feature_flag_enabled() -> None: - """Rectilinear chunk grid construction succeeds when the feature flag is enabled""" +def test_rectilinear_feature_flag_names_stored_array() -> None: + """Opening a stored array whose document the flag refuses names the array.""" + store = MemoryStore() with zarr.config.set({"array.rectilinear_chunks": True}): + zarr.create_array(store, name="r", shape=(30,), chunks=[[10, 20]], dtype="int32") + with ( + zarr.config.set({"array.rectilinear_chunks": False}), + pytest.raises(ValueError, match="experimental and disabled by default") as info, + ): + zarr.open_array(store, path="r") + assert info.value.__notes__ == [f"Array {str(store) + '/r'!r}: nothing was read."] + + +def test_rectilinear_metadata_classes_not_gated() -> None: + """The flag gates stored documents, not the chunk grid metadata classes.""" + with zarr.config.set({"array.rectilinear_chunks": False}): grid = RectilinearChunkGridMetadata(chunk_shapes=((10, 20), (25, 25))) assert grid.ndim == 2 @@ -307,9 +339,12 @@ def test_chunk_spec( ((100, 200), (10, 20), True, 2, (10, 20)), ((), (), True, 0, ()), ((60, 100), [[10, 20, 30], [25, 25, 25, 25]], False, 2, None), - ((30, 50), [[10, 10, 10], [25, 25]], True, 2, (10, 25)), # uniform edges → regular + # explicit edge lists stay rectilinear even when uniform (gh-4272), including + # uniform edges overhanging the extent after a shrinking resize + ((30, 50), [[10, 10, 10], [25, 25]], False, 2, None), + ((35,), [[10, 10, 10, 10]], False, 1, None), ], - ids=["regular", "zero-dim", "rectilinear", "uniform-becomes-regular"], + ids=["regular", "zero-dim", "rectilinear", "uniform-edges", "uniform-edges-overhang"], ) def test_chunk_grid_construction( array_shape: tuple[int, ...], @@ -329,13 +364,6 @@ def test_chunk_grid_construction( _ = g.chunk_shape -def test_chunk_grid_rectilinear_uniform_dim_is_fixed() -> None: - """A rectilinear grid with all-same sizes in one dim stores it as Fixed.""" - g = ChunkGrid.from_sizes((60, 100), [[10, 20, 30], [25, 25, 25, 25]]) - assert isinstance(g._dimensions[0], VaryingDimension) - assert isinstance(g._dimensions[1], FixedDimension) - - # --------------------------------------------------------------------------- # ChunkGrid query tests # --------------------------------------------------------------------------- @@ -495,7 +523,6 @@ def test_chunk_grid_iter() -> None: [ ([[10, 3]], [10, 10, 10]), ([[10, 2], [20, 1]], [10, 10, 20]), - ([[True, 2], [3, 1]], [1, 1, 3]), ], ) def test_rle_expand(compressed: list[Any], expected: list[int]) -> None: @@ -533,6 +560,7 @@ def test_rle_roundtrip() -> None: ([[-10, 2]], "Chunk edge length must be >= 1"), ([[5, 0]], "RLE repeat count must be >= 1"), ([[5, -1]], "RLE repeat count must be >= 1"), + ([[5, 2, 1]], r"RLE entries must be an integer or \[size, count\], got \[5, 2, 1\]"), ], ids=[ "zero-edge", @@ -541,6 +569,7 @@ def test_rle_roundtrip() -> None: "negative-rle-size", "zero-rle-count", "negative-rle-count", + "rle-entry-of-three", ], ) def test_rle_expand_rejects_invalid(rle_input: list[Any], match: str) -> None: @@ -556,24 +585,29 @@ def test_rle_expand_rejects_invalid(rle_input: list[Any], match: str) -> None: ([10.0], "Chunk edge length must be an int, got 10.0"), ([[10.0, 3]], "Chunk edge length must be an int, got 10.0"), (["10"], "Chunk edge length must be an int, got '10'"), + ([True], "Chunk edge length must be an int, got True"), ([[10, 3.5]], "RLE repeat count must be an int, got 3.5"), ([[10, 3.0]], "RLE repeat count must be an int, got 3.0"), ([np.int64(10)], "Chunk edge length must be an int, got np.int64(10)"), ([[10, np.int64(3)]], "RLE repeat count must be an int, got np.int64(3)"), + ([[10, True]], "RLE repeat count must be an int, got True"), ], ids=[ "fractional-edge", "float-edge", "float-rle-size", "string-edge", + "bool-edge", "fractional-count", "float-count", "numpy-int-edge", "numpy-int-count", + "bool-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.""" + """expand_rle takes `int`s only, not `bool`s, floats or NumPy integers; stored + integral floats and JSON `true` are read by the repairs.""" with pytest.raises(TypeError, match=re.escape(match)): expand_rle(rle_input) @@ -586,7 +620,7 @@ def test_rle_expand_rejects_non_int(rle_input: list[Any], match: str) -> None: ([[5, 0]], "RLE repeat count must be >= 1"), ([[5, 2, 1]], r"RLE entries must be an integer or \[size, count\]"), ], - ids=["zero-edge", "fractional-edge", "zero-rle-count", "rle-entry-of-three"], + ids=["zero-edge", "float-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.""" @@ -594,55 +628,6 @@ def test_rle_expand_names_dimension(rle_input: list[Any], match: str) -> None: expand_rle(rle_input, axis=2) -# --------------------------------------------------------------------------- -# _is_rectilinear_chunks tests -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - ("value", "expected"), - [ - ([[10, 20], [5, 5]], True), - (((10, 20), (5, 5)), True), - ((10, 20), False), - ([10, 20], False), - (10, False), - ("auto", False), - ([], False), - ([[]], True), - (ChunkGrid.from_sizes((10,), (5,)), False), - (None, False), - (3.14, False), - ], - ids=[ - "nested-lists", - "nested-tuples", - "flat-tuple", - "flat-list", - "single-int", - "string", - "empty-list", - "empty-nested-list", - "chunk-grid-instance", - "none", - "float", - ], -) -def test_is_rectilinear_chunks(value: Any, expected: bool) -> None: - """_is_rectilinear_chunks correctly identifies nested sequences as rectilinear""" - assert _is_rectilinear_chunks(value) is expected - - -def test_is_rectilinear_chunks_handles_broken_iterable() -> None: - """_is_rectilinear_chunks returns False for objects that raise on iteration.""" - - class BrokenIter: - def __iter__(self) -> Any: - raise TypeError("cannot iterate") - - assert _is_rectilinear_chunks(BrokenIter()) is False - - # --------------------------------------------------------------------------- # Serialization tests # --------------------------------------------------------------------------- @@ -681,7 +666,7 @@ def test_from_metadata_unknown_chunk_grid_type() -> None: def test_from_sizes_rejects_empty_edge_list() -> None: """ChunkGrid.from_sizes raises ValueError when a dimension has an empty edge list.""" - with pytest.raises(ValueError, match="at least one chunk"): + with pytest.raises(ValueError, match="edges must not be empty"): ChunkGrid.from_sizes((10,), ([],)) @@ -1236,16 +1221,81 @@ def test_iter_shard_regions_bounds_check() -> None: @pytest.mark.parametrize( - "chunks", [(5, [10, 20, 70]), ([10, 20, 70], 5)], ids=["int-first", "list-first"] + ("shape", "chunks"), + [ + ((30, 100), (5, [10, 20, 70])), + ((100, 30), ([10, 20, 70], 5)), + # uniform edges are still rectilinear, whether given as lists or as + # stored metadata (gh-4272) + ((30, 100), ([10, 10, 10], 50)), + # (metadata is built lazily: constructing it needs the rectilinear flag) + ((30, 100), partial(RectilinearChunkGridMetadata, chunk_shapes=((10, 10, 10), 50))), + ], + ids=["int-first", "list-first", "uniform-list", "uniform-metadata"], ) -def test_mixed_chunks_gates_are_order_independent(chunks: Any) -> None: - """The v2 and chunks+shards gates must recognize a mixed per-dimension spec - as rectilinear regardless of which dimension carries the sequence.""" - shape = (30, 100) if isinstance(chunks[0], int) else (100, 30) +def test_rectilinear_chunks_gates(shape: tuple[int, ...], chunks: Any) -> None: + """The v2 and chunks+shards gates recognize every rectilinear spec as + rectilinear, regardless of which dimension carries the sequence and + whether its edges happen to be uniform.""" + chunks = chunks() if callable(chunks) else chunks with pytest.raises(ValueError, match="Rectilinear chunks with sharding is not supported"): zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, shards=(10, 10), dtype="uint8") with pytest.raises(ValueError, match="Zarr format 2 does not support rectilinear chunk grids"): zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, zarr_format=2, dtype="uint8") + with pytest.raises(ValueError, match="Zarr format 2 does not support rectilinear chunk grids"): + zarr.create(store=MemoryStore(), shape=shape, chunks=chunks, zarr_format=2, dtype="uint8") + + +@pytest.mark.parametrize( + ("chunks", "expected"), + [ + (None, None), + (False, (4096, 4096)), + (-1, (4096, 4096)), + (5, (5, 5)), + ([5, 3], (5, 3)), + (np.int64(5), (5, 5)), + (np.array(7), (7, 7)), + (np.array([5, 3]), (5, 3)), + ], + ids=repr, +) +def test_legacy_create_v2_chunks(chunks: Any, expected: tuple[int, ...] | None) -> None: + """The legacy `zarr.create` Zarr format 2 path chunks automatically only when + `chunks` is `None` (`expected` is `None`); any other `chunks`, falsy or not, is a + chunk specification, read as `create_array` reads it. It never tests the truth + value of a numpy array, which has none.""" + shape = (4096, 4096) + arr = zarr.create(store=MemoryStore(), shape=shape, chunks=chunks, dtype="uint8", zarr_format=2) + assert arr.chunks == (expected or guess_chunks(shape, 1).chunk_shape) + assert all(type(c) is int for c in arr.chunks) + + +def test_legacy_create_v2_chunks_zero_rejected() -> None: + """A chunk size of 0 is a chunk specification, not a request for automatic + chunking, and a chunk size must be at least 1.""" + with pytest.raises(ValueError, match="Chunk size must be positive or -1, got 0"): + zarr.create(store=MemoryStore(), shape=(10,), chunks=0, dtype="uint8", zarr_format=2) + + +def test_legacy_create_v2_chunks_empty_rejected() -> None: + """An empty chunk specification is not a request for automatic chunking: it has no + dimension, where the array has one.""" + chunks: Any = [] # outside the annotated type, which has no list + with pytest.raises(ValueError, match="chunks has 0 dimensions but shape has 1 dimensions"): + zarr.create(store=MemoryStore(), shape=(10,), chunks=chunks, dtype="uint8", zarr_format=2) + + +def test_legacy_create_v2_chunks_all_zero_array_raises() -> None: + """A numpy `chunks` array with more than one element is given even when every element + is 0, so the legacy Zarr format 2 path rejects its 0 sizes instead of chunking + automatically.""" + # `zarr.create` does not declare numpy arrays for `chunks`; the legacy path reads them. + chunks: Any = np.array([0, 0]) + with pytest.raises(ValueError, match="Chunk size must be positive or -1, got 0"): + zarr.create( + store=MemoryStore(), shape=(2**12, 2**12), chunks=chunks, dtype="int32", zarr_format=2 + ) def test_from_array_keep_preserves_all_bare_int_rectilinear_grid() -> None: @@ -1311,15 +1361,22 @@ def test_from_array_keep_is_o1_in_chunk_count() -> None: assert grid.chunk_shapes == (3, (10, 20)) -def test_resize_uniform_rectilinear_appends_edge() -> None: +@pytest.mark.parametrize( + "rect_chunks", + [[[10, 10, 10]], partial(RectilinearChunkGridMetadata, chunk_shapes=((10, 10, 10),))], + ids=["list", "metadata"], +) +def test_resize_uniform_rectilinear_appends_edge(rect_chunks: Any) -> None: """Growing an explicitly rectilinear array whose edges look regular appends a new edge chunk, while the same sizes declared as a scalar extend the uniform pattern instead (gh-4272), so an append-only workload writing the grown region touches exactly one new chunk.""" from zarr.core.metadata.v3 import ArrayV3Metadata + rect_chunks = rect_chunks() if callable(rect_chunks) else rect_chunks rect_store: dict[str, Any] = {} - rect = zarr.create_array(store=rect_store, shape=(30,), chunks=[[10, 10, 10]], dtype="int32") + rect = zarr.create_array(store=rect_store, shape=(30,), chunks=rect_chunks, dtype="int32") + assert isinstance(rect.metadata.chunk_grid, RectilinearChunkGridMetadata) rect.resize((45,)) assert isinstance(rect.metadata, ArrayV3Metadata) rect_grid = rect.metadata.chunk_grid @@ -2218,15 +2275,6 @@ def test_overflow_multidim() -> None: assert spec.codec_shape == (30, 40) -def test_overflow_uniform_edges_collapses_to_fixed() -> None: - """Uniform edges where len == ceildiv(extent, edge) collapse to FixedDimension.""" - g = ChunkGrid.from_sizes((35,), [[10, 10, 10, 10]]) - assert isinstance(g._dimensions[0], FixedDimension) - assert g.is_regular - assert g.chunk_sizes == ((10, 10, 10, 5),) - assert g._dimensions[0].nchunks == 4 - - def test_overflow_index_to_chunk_near_extent() -> None: """Index lookup near and at the extent boundary.""" d = VaryingDimension([10, 20, 30, 40], extent=50) @@ -3079,65 +3127,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(