From 85f2455f4a54346ebcf1f55da23be62e33efce1e Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 6 Nov 2026 09:32:54 +0100 Subject: [PATCH 1/2] Write only the table's slots of a dictionary column's code predicate _dictionary_eq() and _dictionary_null_pred() assign the predicate over all of the column's codes to a mask with one item per physical row. The codes have one item per allocated slot, which can be more than the rows (4096 vs. 16 in the tests), and this worked only because NDArray.__setitem__ wrote the first items of any value at least as large as the selection. Write min(len(codes), n_phys) items explicitly (the same result, and only those are evaluated), so that __setitem__ can check the value's shape as NumPy does. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01DiPFzsfE5i74cBSxZN13M8 --- src/blosc2/ctable.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/blosc2/ctable.py b/src/blosc2/ctable.py index 1b69d372a..c9b01ef0a 100644 --- a/src/blosc2/ctable.py +++ b/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 From a50b9db5c6b3620cdc5721f0d596dd7b899a15f0 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 6 Nov 2026 09:32:54 +0100 Subject: [PATCH 2/2] Check the value's shape in NDArray.__setitem__ and broadcast it as NumPy set_slice() takes the value as one buffer and checks only that it is large enough, so a[0:2, 0:3] = np.arange(6).reshape(3, 2) wrote the items in C order, a too large value had its first items written, and a value NumPy broadcasts, such as a[0:2, 0:3] = np.arange(3), raised "RuntimeError: Error while setting the slice". The paths for steps and fancy keys indexed the value with the selection's coordinates, so there too large values were cut and broadcastable ones raised IndexError. Compute NumPy's shape of the selection on every path (no dimension for an integer index, 1 for each None; get_basic_start_stop() already has it for basic keys and now returns it) and check the value's shape against it before anything is written: leading dimensions of length 1 are ignored and the rest must broadcast, as in NumPy; anything else raises NumPy's "could not broadcast" ValueError. Values that hold the items of the selection in C order are written as before, without a copy. Values that need broadcasting are converted to the array's dtype first, then broadcast and copied once; a size-1 value fills like a scalar. The step and fancy paths get the value broadcast to the selection's shape. np.isscalar() (~250 ns for an array) is now called once per call, and the basic path reuses the shape the extension computes, so same-shape and scalar writes into small arrays got 2-8% faster instead of slower. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01DiPFzsfE5i74cBSxZN13M8 --- src/blosc2/blosc2_ext.pyx | 9 +-- src/blosc2/ndarray.py | 101 +++++++++++++++++++++++++++------- tests/ndarray/test_setitem.py | 84 ++++++++++++++++++++++++++++ 3 files changed, 171 insertions(+), 23 deletions(-) diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index de90e254c..9889dbea0 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/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. diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index bf739cb7d..3117dbc39 100644 --- a/src/blosc2/ndarray.py +++ b/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 `_ - 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 `_) 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)) diff --git a/tests/ndarray/test_setitem.py b/tests/ndarray/test_setitem.py index 858c5b765..4092dfa54 100644 --- a/tests/ndarray/test_setitem.py +++ b/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])