diff --git a/changes/4257.bugfix.md b/changes/4257.bugfix.md index 6f2740ccb4..f643cf31f0 100644 --- a/changes/4257.bugfix.md +++ b/changes/4257.bugfix.md @@ -1 +1 @@ -Numpy integers are accepted as chunk sizes again. Since 3.3.0 a per-dimension chunk size that was a numpy integer (e.g. `chunks=(np.int64(2), np.int64(2))`, as produced by any computed chunk shape) raised `TypeError: 'numpy.int64' object is not iterable`, because the scalar chunk path narrowed on `int` while its caller dispatched on `numbers.Integral`. Numpy arrays are now also accepted as chunk specifications, and a chunk specification that is neither an integer nor iterable now reports the offending value instead of failing with an opaque iteration error. +Numpy integers are accepted as chunk sizes again. Since 3.3.0 a per-dimension chunk size that was a numpy integer (e.g. `chunks=(np.int64(2), np.int64(2))`, as produced by any computed chunk shape) raised `TypeError: 'numpy.int64' object is not iterable`, because the scalar chunk path narrowed on `int` while its caller dispatched on `numbers.Integral`. The same regression had broken numpy arrays as chunk specifications through the legacy `zarr.create` / `zarr.array` / `zarr.open_array` functions, which accepted them in 2.x and 3.2.x; those work again. (`zarr.create_array` and the functions built on it gain numpy-array support separately, in #4329.) A chunk specification that is neither an integer nor iterable now reports the offending value instead of failing with an opaque iteration error. diff --git a/changes/4329.feature.md b/changes/4329.feature.md new file mode 100644 index 0000000000..65b0bc598f --- /dev/null +++ b/changes/4329.feature.md @@ -0,0 +1 @@ +`zarr.create_array`, `Group.create_array`, `zarr.from_array`, and the entry points built on them now accept a numpy array as the `chunks` or `shards` specification, alongside ints, tuples, and numpy integer scalars. This is new for that API: it has never accepted numpy arrays in any 3.x release, because each entry point compared the specification to the `"auto"` or `"keep"` sentinel string before normalizing it, and for a numpy array that comparison raised numpy's ambiguous-truth-value `ValueError`. Those sentinel checks are now guarded so array-like specifications reach the normalizer, bringing this API in line with the legacy `zarr.create` / `zarr.array` / `zarr.open_array` functions, which have accepted numpy arrays since 2.x. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 3f67eb7c75..1061415216 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -43,6 +43,8 @@ from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, + _is_auto, + _is_keep, _is_rectilinear_chunks, guess_chunks, normalize_chunks_nd, @@ -4537,7 +4539,7 @@ async def init_array( # Normalize the user's chunks into a canonical ChunkGrid - if chunks == "auto": + 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: @@ -4893,10 +4895,12 @@ def _parse_keep_array_attr( DimensionNamesLike, dict[str, JSON] | None, ]: + # chunks / shards may be numpy arrays, so their sentinel checks go through the + # isinstance-guarded helpers rather than a bare ``==`` against the string. if isinstance(data, Array): rectilinear_grid = _stored_rectilinear_grid_or_none(data.metadata) sharded = _sharding_codec(data.metadata) is not None - if chunks == "keep": + if _is_keep(chunks): if rectilinear_grid is None or sharded: # `.chunks` is the inner chunk shape when sharding is used, and # inner chunks are regular whatever the shape of the shard grid. @@ -4908,7 +4912,7 @@ def _parse_keep_array_attr( # and trailing edges beyond the extent left behind by a # shrinking resize. chunks = rectilinear_grid - if shards == "keep": + if _is_keep(shards): if rectilinear_grid is None: shards = data.shards elif sharded: @@ -4960,9 +4964,9 @@ def _parse_keep_array_attr( # array's in-memory metadata and the new array's. attributes = copy.deepcopy(dict(data.attrs)) else: - if chunks == "keep": + if _is_keep(chunks): chunks = "auto" - if shards == "keep": + if _is_keep(shards): shards = None if zarr_format is None: zarr_format = 3 diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 6242994fdf..cc27366027 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -11,6 +11,7 @@ from typing import ( TYPE_CHECKING, Any, + Literal, NamedTuple, Protocol, cast, @@ -19,6 +20,7 @@ import numpy as np import numpy.typing as npt +from typing_extensions import TypeIs import zarr from zarr.core.common import ( @@ -314,6 +316,24 @@ def is_boundary(self) -> bool: # list of ints (explicit edges), or mixed RLE (e.g. [[10, 3], 5]). +def _is_auto(spec: object) -> TypeIs[Literal["auto"]]: + """Check whether a chunk or shard specification is the ``"auto"`` sentinel. + + Specifications may be numpy arrays, whose ``==`` against a string is + elementwise and cannot be used in a boolean context, so the string check + must be guarded by ``isinstance``. + """ + return isinstance(spec, str) and spec == "auto" + + +def _is_keep(spec: object) -> TypeIs[Literal["keep"]]: + """Check whether a chunk or shard specification is the ``"keep"`` sentinel. + + See `_is_auto` for why this is not a bare ``==`` comparison. + """ + return isinstance(spec, str) and spec == "keep" + + def _is_rectilinear_chunks(chunks: Any) -> bool: """Check if chunks specifies a rectilinear grid along any dimension. @@ -940,7 +960,7 @@ def resolve_outer_and_inner_chunks( # Extract the flat chunk shape (uniform size per dimension) for arithmetic. chunk_shape_flat = chunks.chunk_shape - if shard_shape == "auto": + if _is_auto(shard_shape): warnings.warn( "Automatic shard shape inference is experimental and may change without notice.", ZarrUserWarning, diff --git a/tests/test_array.py b/tests/test_array.py index a1cd687ab9..4fcba12207 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -32,6 +32,7 @@ AsyncArray, CompressorsLike, FiltersLike, + ShardsLike, _iter_chunk_coords, _iter_chunk_regions, _iter_shard_coords, @@ -53,7 +54,7 @@ resolve_outer_and_inner_chunks, ) from zarr.core.chunk_key_encodings import ChunkKeyEncodingParams -from zarr.core.common import JSON, ZarrFormat, ceildiv +from zarr.core.common import JSON, ChunksLike, ZarrFormat, ceildiv from zarr.core.dtype import ( DateTime64, Float32, @@ -1128,14 +1129,24 @@ def test_auto_partition_auto_shards_with_auto_chunks_should_be_close_to_1MiB() - assert all(s % c == 0 for s, c in zip(shard_shape, chunk_shape, strict=True)) -def test_chunks_and_shards() -> None: +@pytest.mark.parametrize( + "chunks", + [(5, 5), [5, 5], np.array([5, 5]), (np.int64(5), np.int64(5))], + ids=["tuple", "list", "array", "numpy-scalars"], +) +@pytest.mark.parametrize( + "shards", + [(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: store = StorePath(MemoryStore()) shape = (100, 100) - chunks = (5, 5) - shards = (10, 10) + expected_chunks = normalize_chunks_nd(chunks, shape).chunk_shape + expected_shards = normalize_chunks_nd(shards, shape).chunk_shape arr_v3 = zarr.create_array(store=store / "v3", shape=shape, chunks=chunks, dtype="i4") - assert arr_v3.chunks == chunks + assert arr_v3.chunks == expected_chunks assert arr_v3.shards is None arr_v3_sharding = zarr.create_array( @@ -1145,13 +1156,13 @@ def test_chunks_and_shards() -> None: shards=shards, dtype="i4", ) - assert arr_v3_sharding.chunks == chunks - assert arr_v3_sharding.shards == shards + assert arr_v3_sharding.chunks == expected_chunks + assert arr_v3_sharding.shards == expected_shards arr_v2 = zarr.create_array( store=store / "v2", shape=shape, chunks=chunks, zarr_format=2, dtype="i4" ) - assert arr_v2.chunks == chunks + assert arr_v2.chunks == expected_chunks assert arr_v2.shards is None @@ -1757,14 +1768,26 @@ async def test_creation_from_other_zarr_format( @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True) @pytest.mark.parametrize("store2", ["local", "memory", "zip"], indirect=["store2"]) @pytest.mark.parametrize("src_chunks", [(40, 10), (11, 50)]) -@pytest.mark.parametrize("new_chunks", [(40, 10), (11, 50)]) +@pytest.mark.parametrize( + "new_chunks", [(40, 10), (11, 50), [40, 10], np.array([11, 50]), (np.int64(40), np.int64(10))] +) +@pytest.mark.parametrize( + "new_shards", + [None, (440, 100), [440, 100], np.array([440, 100]), (np.int64(440), np.int64(100))], + ids=["none", "tuple", "list", "array", "numpy-scalars"], +) +@pytest.mark.parametrize("source_as_numpy", [False, True], ids=["zarr", "numpy"]) async def test_from_array( store: Store, store2: Store, src_chunks: tuple[int, int], - new_chunks: tuple[int, int], + new_chunks: ChunksLike, + new_shards: ShardsLike | None, + source_as_numpy: bool, zarr_format: ZarrFormat, ) -> None: + if zarr_format == 2 and new_shards is not None: + pytest.skip("Zarr format 2 does not support sharding") src_fill_value = 2 src_dtype = np.dtype("uint8") src_attributes = None @@ -1776,6 +1799,7 @@ async def test_from_array( store=store, fill_value=src_fill_value, attributes=src_attributes, + zarr_format=zarr_format, ) src[:] = np.arange(1000).reshape((100, 10)) @@ -1783,18 +1807,21 @@ async def test_from_array( new_attributes: dict[str, JSON] = {"foo": "bar"} result = zarr.from_array( - data=src, + data=np.asarray(src) if source_as_numpy else src, store=store2, chunks=new_chunks, + shards=new_shards, fill_value=new_fill_value, attributes=new_attributes, + zarr_format=zarr_format, ) np.testing.assert_array_equal(result[:], src[:]) assert result.fill_value == new_fill_value assert result.dtype == src_dtype assert result.attrs == new_attributes - assert result.chunks == new_chunks + np.testing.assert_array_equal(result.chunks, new_chunks) + np.testing.assert_array_equal(result.shards, new_shards) @pytest.mark.parametrize("store", ["local"], indirect=True) diff --git a/tests/test_group.py b/tests/test_group.py index f7f2333ef5..4ce326621f 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -60,8 +60,9 @@ import pathlib from collections.abc import Callable + from zarr.core.array import ShardsLike from zarr.core.buffer.core import Buffer - from zarr.core.common import JSON, ZarrFormat + from zarr.core.common import JSON, ChunksLike, ZarrFormat from zarr.core.dtype import ZDType, ZDTypeLike @@ -783,21 +784,35 @@ async def test_group_update_attributes_async(store: Store, zarr_format: ZarrForm @pytest.mark.parametrize("name", ["a", "/a"]) +@pytest.mark.parametrize( + "chunks", + [(2, 2), [2, 2], np.array([2, 2]), (np.int64(2), np.int64(2))], + ids=["tuple", "list", "array", "numpy-scalars"], +) +@pytest.mark.parametrize( + "shards", + [None, (4, 4), [4, 4], np.array([4, 4]), (np.int64(4), np.int64(4))], + ids=["none", "tuple", "list", "array", "numpy-scalars"], +) def test_group_create_array( store: Store, zarr_format: ZarrFormat, overwrite: bool, name: str, + chunks: ChunksLike, + shards: ShardsLike | None, ) -> None: """ - Test `Group.from_store` + Test `Group.create_array` """ + if zarr_format == 2 and shards is not None: + pytest.skip("Zarr format 2 does not support sharding") group = Group.from_store(store, zarr_format=zarr_format) shape = (10, 10) dtype = "uint8" data = np.arange(np.prod(shape)).reshape(shape).astype(dtype) - array = group.create_array(name=name, shape=shape, dtype=dtype) + array = group.create_array(name=name, shape=shape, dtype=dtype, chunks=chunks, shards=shards) array[:] = data if not overwrite: @@ -807,6 +822,8 @@ def test_group_create_array( assert array.path == normalize_path(name) assert array.name == f"/{array.path}" + assert array.chunks == (2, 2) + np.testing.assert_array_equal(array.shards, shards) assert array.shape == shape assert array.dtype == np.dtype(dtype) assert np.array_equal(array[:], data)