diff --git a/changes/4376.bugfix.md b/changes/4376.bugfix.md new file mode 100644 index 0000000000..97a5a58d0a --- /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 v2 and sharding restrictions apply to what the normalized grid says was declared. numpy values are accepted everywhere a chunk or shard shape is: 0-d arrays unwrap to their scalar, and the legacy `zarr.create(chunks=np.array(...), zarr_format=2)` no longer fails on the truth value of an array. A non-integer scalar chunk specification raises a `TypeError` naming the problem. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 047e7bb3b3..b31e6f04e8 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, @@ -514,17 +513,15 @@ 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): - 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: + if _raw_chunks is None: outer_chunks = guess_chunks(shape, item_size) else: - outer_chunks = normalize_chunks_nd(_raw, shape) + outer_chunks = normalize_chunks_nd(_raw_chunks, shape) + if not outer_chunks.is_regular: + raise ValueError("Zarr format 2 does not support rectilinear chunk grids.") _chunks = outer_chunks.chunk_shape if order is None: @@ -4530,8 +4527,18 @@ async def init_array( await _prepare_overwrite(store_path, zarr_format=zarr_format, overwrite=overwrite) - # 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; stored rectilinear metadata counts as + # rectilinear even when its edges happen to be uniform. + if isinstance(chunks, RectilinearChunkGridMetadata) or 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: @@ -4541,14 +4548,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( diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index cc27366027..9477e6e5bf 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -14,7 +14,6 @@ Literal, NamedTuple, Protocol, - cast, runtime_checkable, ) @@ -334,30 +333,6 @@ 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. - - 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]. - """ - from zarr.core.metadata.v3 import RectilinearChunkGridMetadata - - if isinstance(chunks, RectilinearChunkGridMetadata): - return True - if isinstance(chunks, (str, int, ChunkGrid)): - return False - if not hasattr(chunks, "__iter__"): - return False - try: - return any( - hasattr(elem, "__iter__") and not isinstance(elem, (str, bytes, int)) for elem in chunks - ) - except TypeError: - return False - - def is_regular_1d(dim_chunks: Sequence[int]) -> bool: """Check if a single dimension's chunk sizes represent a regular grid. @@ -756,6 +731,8 @@ def normalize_chunks_1d(chunks: int | Iterable[object], span: int) -> DimensionG # `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, np.ndarray) and chunks.ndim == 0: + chunks = chunks[()] if isinstance(chunks, numbers.Integral): chunk_size = int(chunks) if chunk_size < -1 or chunk_size == 0: @@ -834,9 +811,16 @@ def normalize_chunks_nd( chunks = -1 # handle 1D convenience form. bool is excluded above so this only catches actual ints. + if isinstance(chunks, np.ndarray) and chunks.ndim == 0: + chunks = chunks[()] if isinstance(chunks, numbers.Integral): chunks = tuple(int(chunks) for _ in shape) + if not hasattr(chunks, "__len__"): + raise TypeError( + f"Chunk specification must be an integer or an iterable of integers; got " + f"{chunks!r} of type {type(chunks).__name__}." + ) # handle bad dimensionality if len(chunks) != len(shape): raise ValueError( @@ -952,15 +936,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: Any 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, @@ -984,11 +963,12 @@ 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) + # Regular and rectilinear shard specifications go through the one normalizer. + outer = normalize_chunks_nd(shard_spec, array_shape) return ChunkLayout(outer_chunks=outer, inner=ChunkLayout(outer_chunks=chunks)) diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index b8289d2135..104407f2b2 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -19,7 +19,6 @@ ChunkSpec, FixedDimension, VaryingDimension, - _is_rectilinear_chunks, ) from zarr.core.common import compress_rle, expand_rle from zarr.core.metadata.v3 import ( @@ -557,55 +556,6 @@ def test_expand_rle_pair_with_float_count() -> None: assert result == [10, 10, 10] -# --------------------------------------------------------------------------- -# _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 # --------------------------------------------------------------------------- @@ -1211,6 +1161,73 @@ def test_mixed_chunks_gates_are_order_independent(chunks: Any) -> None: zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, zarr_format=2, dtype="uint8") +@pytest.mark.parametrize( + ("chunks", "expected"), + [ + pytest.param(np.int64(2), (2, 2), id="numpy-scalar"), + pytest.param(np.array(2), (2, 2), id="0-d-array"), + pytest.param((np.int64(5), np.int32(3)), (5, 3), id="tuple-of-numpy-ints"), + pytest.param(np.array([5, 3], dtype=np.uint8), (5, 3), id="1-d-array"), + pytest.param([np.int64(5), np.int64(3)], (5, 3), id="list-of-numpy-ints"), + pytest.param((np.array(5), 3), (5, 3), id="0-d-array-element"), + pytest.param((np.int64(-1), 3), (10, 3), id="numpy-minus-one"), + ], +) +@pytest.mark.parametrize("creator", ["create_array-v3", "create_array-v2", "create-v2"]) +def test_numpy_chunk_specs_normalize(chunks: Any, expected: tuple[int, ...], creator: str) -> None: + """numpy scalars, 0-d arrays, and integer arrays are accepted wherever a + chunk shape is, in every zarr format and creation API, and are stored as + plain ints. The legacy `zarr.create` v2 path used to test the spec's truth + value, which a numpy array does not have.""" + store = MemoryStore() + if creator == "create_array-v3": + arr = zarr.create_array(store, shape=(10, 6), chunks=chunks, dtype="uint8") + elif creator == "create_array-v2": + arr = zarr.create_array(store, shape=(10, 6), chunks=chunks, dtype="uint8", zarr_format=2) + else: + arr = zarr.create(store=store, shape=(10, 6), chunks=chunks, dtype="uint8", zarr_format=2) + assert arr.chunks == expected + assert all(type(c) is int for c in arr.chunks) + + +@pytest.mark.parametrize( + ("chunks", "expected"), + [ + pytest.param((2, np.array([3, 3])), (2, (3, 3)), id="int-and-array-edges"), + pytest.param((np.int64(2), [3, 3]), (2, (3, 3)), id="numpy-int-and-list"), + pytest.param(np.array([[5, 5], [3, 3]]), ((5, 5), (3, 3)), id="2-d-array"), + pytest.param(([np.int64(5), np.int64(5)], 3), ((5, 5), 3), id="edges-of-numpy-ints"), + ], +) +def test_numpy_rectilinear_chunk_specs_normalize( + chunks: Any, expected: tuple[int | tuple[int, ...], ...] +) -> None: + """numpy values inside a rectilinear `chunks=` spec are stored as plain ints + and the grid kind is decided by the normalized grid, not by duck typing.""" + arr = zarr.create_array(MemoryStore(), shape=(10, 6), chunks=chunks, dtype="uint8") + assert arr.metadata.chunk_grid == RectilinearChunkGridMetadata(chunk_shapes=expected) + + +@pytest.mark.parametrize( + "shards", + [np.array([10, 6]), (np.int64(10), np.int32(6)), [np.int64(10), np.int64(6)]], + ids=["array", "tuple-of-numpy-ints", "list-of-numpy-ints"], +) +def test_numpy_shard_specs_normalize(shards: Any) -> None: + """numpy values in `shards=` normalize like those in `chunks=`.""" + arr = zarr.create_array( + MemoryStore(), shape=(10, 6), chunks=(2, 3), shards=shards, dtype="uint8" + ) + assert arr.shards == (10, 6) + assert arr.chunks == (2, 3) + + +def test_zero_dim_float_array_chunks_rejected() -> None: + """A 0-d array unwraps to its scalar, which must still be an integer.""" + with pytest.raises(TypeError, match="must be an integer or an iterable of integers"): + zarr.create_array(MemoryStore(), shape=(10,), chunks=np.array(2.0), dtype="uint8") + + def test_from_array_keep_preserves_all_bare_int_rectilinear_grid() -> None: """A rectilinear grid using the bare-int shorthand on every dimension can only arrive from externally written metadata — `create_array` never