From 44fe43419acb0805c7dc7b59fc8bdb5ddfacb06c Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 6 Nov 2026 09:16:45 +0100 Subject: [PATCH 1/2] Normalize negative indices in orthogonal selections get_oselection_numpy(), set_oselection_numpy() and NDArray.oindex passed their index lists to c-blosc2 unchanged. Negative indices, which NumPy counts from the end of the axis, selected wrong data or made c-blosc2 write outside its chunk buffers, and an index equal to the axis length returned the chunk padding instead of raising. Now negative indices count from the end like in NumPy, and indices outside the axis raise IndexError before anything is read or written. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01DiPFzsfE5i74cBSxZN13M8 --- src/blosc2/ndarray.py | 28 ++++++++++++-- tests/ndarray/test_ndarray.py | 69 +++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 3 deletions(-) diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index bf739cb7d..f56186315 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -4579,23 +4579,45 @@ def _get_set_findex_default(self, _slice, out=None, value=None): out = super().set_slice((start, stop), chunk) return out + def _normalize_oselection(self, key: list | np.ndarray) -> list: + """Count negative indices from the end of their axis, like NumPy; IndexError outside the axis.""" + shape = self.shape + if len(key) != len(shape): + raise ValueError(f"Key must have {len(shape)} dimensions, got {len(key)}.") + normalized = [] + for indices, size in zip(key, shape, strict=True): + if isinstance(indices, np.ndarray): + indices = indices.tolist() # min(), max() and the extension's copy are faster on lists + if len(indices) > 0: + low, high = builtins.min(indices), builtins.max(indices) + if low < -size or high >= size: + bad = next(i for i in indices if not -size <= i < size) # the first, as in NumPy + axis = len(normalized) + raise IndexError(f"index {bad} is out of bounds for axis {axis} with size {size}") + if low < 0: + indices = [i + size if i < 0 else i for i in indices] + normalized.append(indices) + return normalized + def get_oselection_numpy(self, key: list | np.ndarray) -> np.ndarray: """ Select independently from self along axes specified in key. Key must be same length as self shape. + Negative indices count from the end of their axis; out-of-bounds indices raise IndexError. See Zarr https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#orthogonal-indexing. """ - shape = tuple(len(k) for k in key) + self.shape[len(key) :] + key = self._normalize_oselection(key) # Create the array to store the result - arr = np.empty(shape, dtype=self.dtype) + arr = np.empty(tuple(len(k) for k in key), dtype=self.dtype) return super().get_oindex_numpy(arr, key) def set_oselection_numpy(self, key: list | np.ndarray, arr: NDArray) -> np.ndarray: """ Select independently from self along axes specified in key and set to entries in arr. Key must be same length as self shape. + Negative indices count from the end of their axis; out-of-bounds indices raise IndexError. See Zarr https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#orthogonal-indexing. """ - return super().set_oindex_numpy(key, arr) + return super().set_oindex_numpy(self._normalize_oselection(key), arr) def _try_subsample_gather(self, start, stop, step, shape): """Fast path for a coarse strided read via the sparse-gather primitive. diff --git a/tests/ndarray/test_ndarray.py b/tests/ndarray/test_ndarray.py index 9a06fdf19..fd5d6f1d8 100644 --- a/tests/ndarray/test_ndarray.py +++ b/tests/ndarray/test_ndarray.py @@ -808,6 +808,75 @@ def test_oindex(): np.testing.assert_allclose(arr[:], nparr) +@pytest.mark.parametrize( + "key", + [ + ([-1], [0, 1]), + ([0, 1], [-3]), + ([-10, 9, 0], [-1, 0]), # -shape and the last index + ([3, 3, -7], [5, -5, 5]), # repeated, also through negative aliases + ([], [0, 1]), + ([], []), + (np.array([-1, 2]), np.array([1, -2], dtype=np.int8)), + (np.array([9], dtype=np.uint64), range(-3, 3)), + ], +) +def test_oindex_normalizes_like_numpy(key): + # oindex goes through get_oselection_numpy() and set_oselection_numpy() + x = np.arange(100, dtype=np.int32).reshape(10, 10) + a = blosc2.asarray(x, chunks=(4, 4), blocks=(2, 2)) + key_copy = [np.array(k) for k in key] + expected = x[np.ix_(*key)] + np.testing.assert_array_equal(a.oindex[key], expected) + + values = np.arange(1000, 1000 + expected.size, dtype=np.int32).reshape(expected.shape) + a.oindex[key] = values + x[np.ix_(*key)] = values + np.testing.assert_array_equal(a[:], x) + for k, k_copy in zip(key, key_copy, strict=True): + np.testing.assert_array_equal(k, k_copy) # the caller's indices are not changed + + +@pytest.mark.parametrize( + ("key", "match"), + [ + (([10], [0, 1]), "index 10 is out of bounds for axis 0 with size 10"), + (([0, 1], [-11]), "index -11 is out of bounds for axis 1 with size 10"), + (([0, 11, -12, 10], [0]), "index 11 is out of bounds for axis 0 with size 10"), + (([0], [2**63]), "out of bounds"), + ((np.array([0, 1]), np.array([1, 2**63], dtype=np.uint64)), "out of bounds"), + ], +) +def test_oindex_rejects_bad_indices(key, match): + x = np.arange(100, dtype=np.int32).reshape(10, 10) + a = blosc2.asarray(x, chunks=(4, 4), blocks=(2, 2)) + with pytest.raises(IndexError, match=match): + a.oindex[key] + with pytest.raises(IndexError, match=match): + a.oindex[key] = np.zeros((len(key[0]), len(key[1])), dtype=np.int32) + np.testing.assert_array_equal(a[:], x) + + +def test_oindex_random_negative_indices(): + rng = np.random.default_rng(8) + for shape, chunks, blocks in [ + ((13,), (5,), (2,)), + ((9, 11), (4, 5), (2, 3)), + ((6, 7, 8), (3, 4, 5), (2, 2, 2)), + ]: + x = rng.integers(0, 1000, size=shape, dtype=np.int64) + a = blosc2.asarray(x, chunks=chunks, blocks=blocks) + for _ in range(10): + # Indices in [-n, n): negatives, repeats and the edges of every axis + key = [rng.integers(-n, n, size=rng.integers(0, 2 * n)).tolist() for n in shape] + expected = x[np.ix_(*key)] + np.testing.assert_array_equal(a.oindex[key], expected) + values = rng.integers(0, 1000, size=expected.shape, dtype=np.int64) + a.oindex[key] = values + x[np.ix_(*key)] = values + np.testing.assert_array_equal(a[:], x) + + @pytest.mark.parametrize("c", [None, 3]) def test_fancy_index(c): # Test 1d From 4c419a6f061ea52cc7c779b62ee11fd993486ef2 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 6 Nov 2026 12:24:43 +0100 Subject: [PATCH 2/2] Check the value of orthogonal selection writes set_oselection_numpy() and NDArray.oindex[...] = value handed the value's buffer to c-blosc2 with a size computed from the selection, not from the value. A value with fewer items was read past its end (the extra items came from whatever memory followed), a value of another dtype was copied as raw bytes, a value of the right size but another shape was taken in C order, and scalars and lists raised TypeError. Now a value that is not already a C-contiguous array of the selection's shape and dtype goes through NumPy's assignment into an array of that shape and dtype: it is converted, broadcast as in NumPy (leading length-1 dimensions ignored), and a shape that does not broadcast raises NumPy's ValueError before anything is written. Values that already fit are passed on without a copy. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01DiPFzsfE5i74cBSxZN13M8 --- src/blosc2/ndarray.py | 23 ++++++++++++++-- tests/ndarray/test_ndarray.py | 49 +++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index f56186315..b9a3ee2f7 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -4610,14 +4610,33 @@ def get_oselection_numpy(self, key: list | np.ndarray) -> np.ndarray: arr = np.empty(tuple(len(k) for k in key), dtype=self.dtype) return super().get_oindex_numpy(arr, key) - def set_oselection_numpy(self, key: list | np.ndarray, arr: NDArray) -> np.ndarray: + def set_oselection_numpy(self, key: list | np.ndarray, arr: Any) -> np.ndarray: """ Select independently from self along axes specified in key and set to entries in arr. Key must be same length as self shape. Negative indices count from the end of their axis; out-of-bounds indices raise IndexError. + ``arr`` is converted to the array's dtype and broadcast to the selection's shape as in NumPy + (a scalar, a list, an array or an NDArray). See Zarr https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#orthogonal-indexing. """ - return super().set_oindex_numpy(self._normalize_oselection(key), arr) + key = self._normalize_oselection(key) + shape = tuple(map(len, key)) + # The extension reads the selection's bytes straight from arr's buffer: only a C-contiguous + # array of the selection's shape and dtype skips NumPy's assignment + if not ( + isinstance(arr, np.ndarray) + and arr.shape == shape + and arr.dtype == self.dtype + and arr.flags.c_contiguous + ): + if isinstance(arr, blosc2.Operand): + arr = arr[()] + elif isinstance(arr, list | tuple): + arr = np.asarray(arr) # nested lists broadcast like arrays, as in NumPy's indexing + value = np.empty(shape, dtype=self.dtype) + value[...] = arr + arr = value + return super().set_oindex_numpy(key, arr) def _try_subsample_gather(self, start, stop, step, shape): """Fast path for a coarse strided read via the sparse-gather primitive. diff --git a/tests/ndarray/test_ndarray.py b/tests/ndarray/test_ndarray.py index fd5d6f1d8..c37016001 100644 --- a/tests/ndarray/test_ndarray.py +++ b/tests/ndarray/test_ndarray.py @@ -877,6 +877,55 @@ def test_oindex_random_negative_indices(): np.testing.assert_array_equal(a[:], x) +@pytest.mark.parametrize( + "value", + [ + np.arange(6, dtype=np.int32).reshape(2, 3), # the selection's shape and dtype + np.full((2, 3), 1.7), + np.arange(6, dtype=">i4").reshape(2, 3), + np.arange(12, dtype=np.int32).reshape(2, 6)[:, ::2], # not contiguous + np.asfortranarray(np.arange(6, dtype=np.int32).reshape(2, 3)), + np.arange(3), # broadcast rows + np.arange(2).reshape(2, 1), # broadcast columns + np.ones((1, 2, 3)), # a leading length-1 dimension + 7, + [1, 2, 3], + [[1, 2, 3], [4, 5, 6]], + [[[1, 2, 3], [4, 5, 6]]], + blosc2.asarray(np.arange(6, dtype=np.int16).reshape(2, 3)), + ], +) +def test_oindex_set_value_like_numpy(value): + x = np.arange(100, dtype=np.int32).reshape(10, 10) + a = blosc2.asarray(x, chunks=(4, 4), blocks=(2, 2)) + key = ([1, 8], [0, 5, 2]) + a.oindex[key] = value + x[np.ix_(*key)] = value[()] if isinstance(value, blosc2.NDArray) else value + np.testing.assert_array_equal(a[:], x) + + +@pytest.mark.parametrize( + "value", + [ + np.arange(6, dtype=np.int32).reshape(3, 2), # the selection's size and dtype, not its shape + np.arange(2, dtype=np.int32), + np.arange(6.0).reshape(3, 2), + np.arange(7), + np.ones((2, 2, 3)), + [[1, 2]], + ], +) +def test_oindex_set_value_wrong_shape(value): + x = np.arange(100, dtype=np.int32).reshape(10, 10) + a = blosc2.asarray(x, chunks=(4, 4), blocks=(2, 2)) + key = ([1, 8], [0, 5, 2]) + with pytest.raises(ValueError): + x.copy()[np.ix_(*key)] = value # NumPy rejects it too + with pytest.raises(ValueError, match="could not broadcast"): + a.oindex[key] = value + np.testing.assert_array_equal(a[:], x) + + @pytest.mark.parametrize("c", [None, 3]) def test_fancy_index(c): # Test 1d