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/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 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])