Check the indices and values of orthogonal selections

src/blosc2/ndarray.py

@@ -4579,22 +4579,63 @@ 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:
+    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.
         """
+        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):

tests/ndarray/test_ndarray.py

@@ -808,6 +808,124 @@ 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(
+    "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