Check the value's shape in NDArray.__setitem__ and broadcast it as NumPy does

src/blosc2/blosc2_ext.pyx

@@ -6474,9 +6474,10 @@ cdef class NDArray:
 
         For keys made of integers (Python or NumPy, not bools), slices with unit
         step and Ellipsis, follow shape changes made by another handle (SWMR) and
-        return ``(start, stop, int_mask)``: the region as lists and which dimensions
-        are indexed with an integer.  Any other key, or one that does not fit the
-        shape, gives None (the general path then raises the appropriate exception).
+        return ``(start, stop, int_mask, shape)``: the region as lists, which
+        dimensions are indexed with an integer, and NumPy's shape of the selection
+        (a tuple without those dimensions).  Any other key, or one that does not fit
+        the shape, gives None (the general path then raises the appropriate exception).
         """
         cdef tuple items = key if type(key) is tuple else (key,)
         cdef Py_ssize_t nkey = len(items)
@@ -6496,7 +6497,7 @@ cdef class NDArray:
                 return None
         cdef int ndim = self.array.ndim, j
         return ([start_[j] for j in range(ndim)], [stop_[j] for j in range(ndim)],
-                tuple([int_dims[j] for j in range(ndim)]))
+                tuple([int_dims[j] for j in range(ndim)]), tuple(out_shape))
 
     def get_basic_numpy(self, key):
         """Return ``self[key]`` as a NumPy array if *key* is a basic index, else None.

src/blosc2/ctable.py

@@ -2939,7 +2939,8 @@ def null():
         pred = dc.codes != code if negate else dc.codes == code
         if len(dc.codes) != n_phys:
             physical = blosc2.zeros(n_phys, dtype=np.bool_)
-            physical[: len(dc.codes)] = pred
+            n = min(len(dc.codes), n_phys)  # the codes can have more slots than the table
+            physical[:n] = pred[:n]
             pred = physical
         pred = pred & valid
         return pred if other is None else kleene(pred, exclude_nulls=negate)
@@ -2949,7 +2950,8 @@ def _dictionary_null_pred(self, dc, spec, n_phys: int):
         null = dc.codes == np.int32(spec.null_code)
         if len(dc.codes) != n_phys:
             physical = blosc2.zeros(n_phys, dtype=np.bool_)
-            physical[: len(dc.codes)] = null
+            n = min(len(dc.codes), n_phys)
+            physical[:n] = null[:n]
             null = physical
         return null
 

src/blosc2/ndarray.py

@@ -343,6 +343,52 @@ def get_ndarray_start_stop(ndim, key, shape):
     return start, stop, tuple(step), none_mask
 
 
+def _selection_shape(start, stop, step, int_mask, none_mask) -> tuple:
+    """NumPy's shape of the selection of a key resolved by process_key() and
+    get_ndarray_start_stop(): ``len(range(start, stop, step))`` per sliced dimension, nothing for
+    an integer index and 1 for every None (`int_mask` and `none_mask` have one entry per key item).
+    """
+    dims = zip(start, stop, step, strict=True)
+    shape = []
+    for is_int, is_none in zip(int_mask, none_mask, strict=True):
+        if is_none:
+            shape.append(1)
+            continue
+        st, sp, stp = next(dims)
+        if not is_int:
+            shape.append(len(range(st, sp, stp)))
+    return tuple(shape)
+
+
+def _check_broadcast(vshape: tuple, shape: tuple) -> bool:
+    """Check that a value of shape `vshape` can be assigned to a selection of shape `shape` as in
+    NumPy, which ignores leading dimensions of length 1 of the value and broadcasts the rest;
+    raise NumPy's ValueError if not.  Return whether the value must be broadcast, i.e. it does not
+    hold the items of the selection (in C order: dimensions of length 1 do not matter there).
+    """
+    extra = builtins.max(len(vshape) - len(shape), 0)
+    if builtins.any(n != 1 for n in vshape[:extra]) or builtins.any(
+        n not in (1, s) for n, s in zip(reversed(vshape), reversed(shape), strict=False)
+    ):
+        raise ValueError(f"could not broadcast input array from shape {vshape} into shape {shape}")
+    return math.prod(vshape) != math.prod(shape)
+
+
+def _broadcast_value(value: np.ndarray, shape: tuple) -> np.ndarray:
+    """`value` broadcast to `shape` as checked by _check_broadcast() (a read-only view)."""
+    extra = value.ndim - len(shape)
+    return np.broadcast_to(value.reshape(value.shape[extra:]) if extra > 0 else value, shape)
+
+
+def _setitem_value(value, shape: tuple):
+    """`value` (not a scalar) for a selection of shape `shape` that is written piece by piece
+    (``value[idx]``): as it is if it has that shape, else broadcast to it (a NumPy view)."""
+    if value.shape == shape:
+        return value
+    _check_broadcast(value.shape, shape)
+    return _broadcast_value(value[()] if isinstance(value, blosc2.Operand) else value, shape)
+
+
 def are_partitions_aligned(shape, chunks, blocks):
     """
     Check if the partitions defined by chunks and blocks are aligned with the shape.
@@ -5066,12 +5112,12 @@ def __setitem__(  # noqa: C901
         Parameters
         ----------
         key: int, slice or sequence of slices
-            The index or indices specifying the slice(s) to be updated. Note that the step parameter
-            is not yet supported.
-        value: Py_Object Supporting the Buffer Protocol
-            An object supporting the
-            `Buffer Protocol <https://docs.python.org/3/c-api/buffer.html>`_
-            which will be used to overwrite the specified slice(s).
+            The index or indices specifying the slice(s) to be updated.
+        value: scalar or array-like
+            The value(s) to write: a scalar, or an array (NumPy array, :ref:`NDArray`, or any object
+            supporting the `Buffer Protocol <https://docs.python.org/3/c-api/buffer.html>`_) that
+            is broadcast to the shape of the selection as in NumPy. A value that cannot be
+            broadcast raises ValueError.
 
         Examples
         --------
@@ -5107,13 +5153,16 @@ def __setitem__(  # noqa: C901
             key_, mask = process_key(key, self.shape)  # internally handles key an integer
         if hasattr(value, "shape") and value.shape == () and not getattr(value.dtype, "fields", None):
             value = value.item()
+        # 0-d values are scalars now and nothing below makes new ones: check once (it is not cheap)
+        scalar = np.isscalar(value)
         # Convert to SimpleProxy for e.g. JAX, Tensorflow, PyTorch (not NumPy arrays: that costs more
         # than a small write, for the same result)
-        if not (np.isscalar(value) or type(value) is np.ndarray):
+        if not (scalar or type(value) is np.ndarray):
             value = blosc2.as_simpleproxy(value)
 
+        # `shape` is NumPy's shape of the selection, which the value must broadcast to
         if basic is not None:
-            start, stop, mask = basic
+            start, stop, mask, shape = basic
             none_mask = ()
         elif builtins.any(isinstance(k, list | np.ndarray) for k in key_):  # fancy indexing
             _slice = ndindex.ndindex(key_).expand(
@@ -5127,38 +5176,52 @@ def __setitem__(  # noqa: C901
                     _slice = ndindex.ndindex(()).expand(self.shape)  # just get whole array
                 else:  # do nothing
                     return self
+            if not scalar:
+                # The shape of a stride-0 stand-in for self with an extra empty dimension (no memory)
+                stand_in = np.broadcast_to(np.empty((), dtype=np.uint8), (*self.shape, 0))
+                value = _setitem_value(value, stand_in[(*_slice.raw, slice(None))].shape[:-1])
             return self._mark_indexes_stale(self._get_set_findex_default(_slice, value=value))
         else:
             start, stop, step, none_mask = get_ndarray_start_stop(self.ndim, key_, self.shape)
+            shape = _selection_shape(start, stop, step, mask, none_mask)
             if step != (1,) * self.ndim:  # handle non-unit or negative steps
                 if np.any(none_mask):
                     raise ValueError("Cannot mix non-unit steps and None indexing for __setitem__.")
+                if not scalar:
+                    value = _setitem_value(value, shape)
                 return self._mark_indexes_stale(
                     self._get_set_nonunit_steps((start, stop, step, mask), value=value)
                 )
 
-        shape = [sp - st for sp, st in zip(stop, start, strict=False)]
+        # set_slice() checks only the size of the buffer: check the shape before writing anything
+        broadcast = not scalar and (vshape := value.shape) != shape and _check_broadcast(vshape, shape)
         if (
             isinstance(value, NDArray)
+            and not broadcast
             and not builtins.any(none_mask)
             and self._set_array_chunkwise(start, stop, mask, value)
         ):
             return self._mark_indexes_stale(self)
         if isinstance(value, blosc2.Operand):  # handles SimpleProxy, NDArray, LazyExpr etc.
             value = value[()]  # convert to numpy
-        if np.isscalar(value) or value.shape == ():
+        if broadcast and value.size == 1:
+            value = value.reshape(())  # fill as with a scalar (chunk-wise in large regions)
+        if scalar or value.shape == ():
             if self._fill_chunkwise(start, stop, value):
                 return self._mark_indexes_stale(self)
             value = np.full(shape, value, dtype=self.dtype)
-        if value.dtype != self.dtype:  # handles decompressed NDArray too
-            try:
-                value = value.astype(self.dtype)
-            except ComplexWarning:
-                # Expression type inference can promote complex functions (e.g.
-                # conj) applied to real arrays; preserve the destination dtype.
-                value = value.real.astype(self.dtype)
-        # set_slice() takes one C-contiguous buffer: copy views with gaps or negative strides and
-        # Fortran-ordered arrays (a no-op for C-contiguous ones)
+        else:
+            if value.dtype != self.dtype:  # handles decompressed NDArray too
+                try:
+                    value = value.astype(self.dtype, order="C")
+                except ComplexWarning:
+                    # Expression type inference can promote complex functions (e.g.
+                    # conj) applied to real arrays; preserve the destination dtype.
+                    value = value.real.astype(self.dtype, order="C")
+            if broadcast:  # after the conversion, which then converts only the given items
+                value = _broadcast_value(value, shape)
+        # set_slice() takes one C-contiguous buffer: copy views with gaps, negative or zero strides
+        # (broadcast values) and Fortran-ordered arrays (a no-op for C-contiguous ones)
         value = np.ascontiguousarray(value)
 
         return self._mark_indexes_stale(super().set_slice((start, stop), value))

tests/ndarray/test_setitem.py

@@ -170,6 +170,90 @@ def test_structured_scalar_assignment():
         np.testing.assert_array_equal(array[:], expected)
 
 
+def setitem_values(sel):
+    """Values for a selection of shape `sel`: right shapes, wrong ones, broadcastable ones."""
+    n = int(np.prod(sel))
+    items = np.arange(100, 100 + n, dtype=np.int32)
+    values = [
+        items.reshape(sel),
+        items.reshape(sel[::-1]),  # transposed: wrong unless sel is symmetric
+        items,  # flat: wrong unless sel is 1-d
+        items.reshape((*sel, 1)),  # trailing 1: wrong
+        np.arange(n + 1, dtype=np.int32),  # wrong size
+        np.zeros(0, dtype=np.int32),  # wrong size
+        items.reshape((1, *sel)),  # leading 1s are ignored
+        items.reshape((1, 1, *sel)),
+        np.array([7], dtype=np.int32),  # broadcast
+        items.astype(np.float64).reshape(sel) + 0.5,  # converted
+    ]
+    if sel:
+        values.append(items.reshape((sel[0], 1, *sel[1:])))  # inner 1: wrong
+        values.append(np.arange(sel[-1], dtype=np.int32) - 9)  # broadcast along the last axis
+        values.append(np.arange(sel[-1], dtype=np.float32) + 0.5)  # broadcast and converted
+    if len(sel) > 1:
+        values.append(np.arange(sel[-2], dtype=np.int32).reshape(-1, 1) - 7)  # along the others
+    return values
+
+
+# set_slice() checks only the size of the value's buffer: values of a wrong shape and the right size
+# were written in C order, and broadcastable values raised RuntimeError
+@pytest.mark.parametrize(
+    ("shape", "chunks", "key"),
+    [
+        ((4, 6), (2, 3), (slice(0, 2), slice(0, 3))),
+        ((4, 6), (2, 3), (slice(1, 4), slice(1, 6))),
+        ((4, 6), (2, 3), Ellipsis),
+        ((4, 6), (2, 3), (0, slice(0, 3))),
+        ((4, 6), (2, 3), (slice(0, 3), 2)),
+        ((4, 6), (2, 3), (slice(0, 2), None, slice(0, 3))),
+        ((4, 6), (2, 3), (slice(0, 0), slice(0, 3))),
+        ((4, 6), (2, 3), (slice(0, 4, 2), slice(0, 3))),
+        ((4, 6), (2, 3), (slice(None, None, -1), slice(0, 6, 2))),
+        ((4, 6), (2, 3), (1, slice(5, 0, -2))),
+        ((4, 6), (2, 3), np.array([0, 3])),
+        ((4, 6), (2, 3), np.array([True, False, True, True])),
+        ((6, 8, 10), (3, 4, 5), (slice(1, 5), 2, slice(0, 10))),
+        ((40, 50), (10, 10), (slice(None), slice(None))),  # chunk-wise fill and copy
+        ((40, 50), (10, 10), (slice(5, 37), 7)),
+        ((), (), Ellipsis),
+    ],
+)
+@pytest.mark.parametrize("kind", ["numpy", "list", "NDArray"])
+def test_setitem_value_shape_like_numpy(shape, chunks, key, kind):
+    base = np.arange(int(np.prod(shape)), dtype=np.int32).reshape(shape)
+    for value in setitem_values(base[key].shape):
+        if kind == "list":
+            v = value.tolist()
+        elif kind == "NDArray" and value.size > 0:
+            v = blosc2.asarray(value)
+        else:
+            v = value
+        a = blosc2.asarray(base, chunks=chunks)
+        expected = base.copy()
+        try:
+            # Lists are taken as np.asarray() makes them (NumPy limits their dimensions)
+            expected[key] = np.asarray(v) if kind == "list" else value
+        except ValueError:
+            with pytest.raises(ValueError, match="could not broadcast"):
+                a[key] = v
+            expected = base  # nothing written
+        else:
+            a[key] = v
+        np.testing.assert_array_equal(a[()], expected, err_msg=f"value of shape {value.shape}")
+
+
+def test_setitem_lazyexpr_value_shape():
+    a = blosc2.zeros((4, 6), dtype=np.float64, chunks=(2, 3))
+    row = blosc2.asarray(np.arange(3, dtype=np.float64))
+    a[1:3, 0:3] = row + 1  # broadcast
+    expected = np.zeros((4, 6))
+    expected[1:3, 0:3] = np.arange(3) + 1
+    np.testing.assert_array_equal(a[:], expected)
+    with pytest.raises(ValueError, match="could not broadcast"):
+        a[0:3, 0:2] = blosc2.asarray(np.ones((2, 3))) * 2  # 6 items, but the wrong shape
+    np.testing.assert_array_equal(a[:], expected)
+
+
 def test_setitem_fancy_index():
     out = blosc2.zeros(10)
     idx = np.array([1, 6, 7])