From f50daa1c0112b825f4b11e593fa28e05b211177d Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Thu, 21 Nov 2024 20:38:46 -0700 Subject: [PATCH] Error when attempt to set with an array of incompatible shape. Closes #2469 --- src/zarr/core/array.py | 18 ++++++++++++++---- tests/test_indexing.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 249168723..b72657a3c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2842,14 +2842,24 @@ def set_coordinate_selection( try: from numcodecs.compat import ensure_ndarray_like - value = ensure_ndarray_like(value) # TODO replace with agnostic + value_array = ensure_ndarray_like(value) # TODO replace with agnostic except TypeError: # Handle types like `list` or `tuple` - value = np.array(value) # TODO replace with agnostic + value_array = np.array(value) # TODO replace with agnostic if hasattr(value, "shape") and len(value.shape) > 1: - value = np.array(value).reshape(-1) + value_array = np.array(value).reshape(-1) - sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + if indexer.sel_shape != value_array.shape: + raise ValueError( + f"Attempting to set a selection of {indexer.sel_shape[0]} " + f"elements with an array of {value_array.shape[0]} elements." + ) + + sync( + self._async_array._set_selection( + indexer, value_array, fields=fields, prototype=prototype + ) + ) @_deprecate_positional_args def get_block_selection( diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 479ce201b..ab14bae73 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1936,3 +1936,21 @@ def test_zero_sized_chunks(store: StorePath, shape: list[int]) -> None: z = Array.create(store=store, shape=shape, chunk_shape=shape, zarr_format=3, dtype="f8") z[...] = 42 assert_array_equal(z[...], np.zeros(shape, dtype="f8")) + + +@pytest.mark.parametrize("store", ["memory"], indirect=["store"]) +def test_vectorized_indexing_incompatible_shape(store) -> None: + # GH2469 + shape = (4, 4) + chunks = (2, 2) + fill_value = 32767 + arr = zarr.create( + shape, + store=store, + chunks=chunks, + dtype=np.int16, + fill_value=fill_value, + codecs=[zarr.codecs.BytesCodec(), zarr.codecs.BloscCodec()], + ) + with pytest.raises(ValueError, match="Attempting to set"): + arr[np.array([1, 2]), np.array([1, 2])] = np.array([[-1, -2], [-3, -4]])