Add NDArray.read_into(): decompress a selection into an existing array

CMakeLists.txt

@@ -15,9 +15,9 @@ endif()
 project(python-blosc2)
 
 # Update to the latest minimum C-Blosc2 version, and to the actual bundled one
-set(BLOSC2_MIN_VERSION 3.5.0)
-set(BLOSC2_BUNDLED_VERSION v3.5.0)
-#set(BLOSC2_BUNDLED_VERSION a54e259aeb1c0e62d6bbea28934369d4831e4046)
+set(BLOSC2_MIN_VERSION 3.5.1)
+#set(BLOSC2_BUNDLED_VERSION v3.5.0)
+set(BLOSC2_BUNDLED_VERSION e3ff64f261522803c6b9af7be913ef1c530cbc04)
 
 if(WIN32 AND NOT CMAKE_C_COMPILER_ID STREQUAL "Clang")
     message(FATAL_ERROR "Windows builds require clang-cl. Set CC/CXX to clang-cl or configure CMake with -T ClangCL.")

doc/reference/ndarray.rst

@@ -25,6 +25,7 @@ It is a direct alias for ``array.vlmeta`` and uses the same persistent storage.
         __len__
         __getitem__
         __setitem__
+        read_into
         take
 
     Utility Methods
@@ -34,6 +35,7 @@ It is a direct alias for ``array.vlmeta`` and uses the same persistent storage.
     .. automethod:: __len__
     .. automethod:: __getitem__
     .. automethod:: __setitem__
+    .. automethod:: read_into
     .. automethod:: take
 
 Constructors

src/blosc2/blosc2_ext.pyx

@@ -648,6 +648,9 @@ cdef extern from "b2nd.h":
     int b2nd_get_slice_cbuffer(b2nd_array_t *array,
                                int64_t *start, int64_t *stop,
                                void *buffer, int64_t *buffershape, int64_t buffersize) nogil
+    int b2nd_get_slice_cbuffer_into(b2nd_array_t *array,
+                                    int64_t *start, int64_t *stop,
+                                    void *buffer, int64_t *buffershape, int64_t buffersize) nogil
     int b2nd_set_slice_cbuffer(void *buffer, int64_t *buffershape, int64_t buffersize,
                                int64_t *start, int64_t *stop, b2nd_array_t *array)
     int b2nd_get_slice(b2nd_context_t *ctx, b2nd_array_t **array, b2nd_array_t *src, const int64_t *start,
@@ -4386,6 +4389,73 @@ cdef class NDArray:
         _check_rc(rc, "Error while getting the buffer")
         return arr
 
+    def get_basic_into(self, np.ndarray out, key):
+        """NDArray.read_into() for a basic index: decompress ``self[key]`` into the NumPy array
+        `out` and return it, or return None if *key* is not a basic index.
+
+        The data goes straight into `out` if it has the array's dtype and is a box of a C-order
+        buffer (see _box_layout()); otherwise into a new C-contiguous array, which is returned
+        instead for the caller to copy into `out`.  Keys are as in get_basic_numpy(), but without
+        None; any other key, and any key that is out of range, gives None, so that the caller can
+        use the general indexing path, which also raises the appropriate exception.  A selection
+        of another shape than `out` raises ValueError.
+        """
+        cdef tuple items = key if type(key) is tuple else (key,)
+        cdef Py_ssize_t nkey = len(items), i
+        cdef char[_MAX_BASIC_KEY] kinds
+        cdef int64_t[_MAX_BASIC_KEY] ints
+        cdef Py_ssize_t[_MAX_BASIC_KEY] starts, stops
+        cdef Py_ssize_t nidx = self._classify_basic_key(items, kinds, ints, starts, stops)
+        if nidx < 0:
+            return None
+        for i in range(nkey):
+            if kinds[i] == _KEY_NEWAXIS:
+                return None
+
+        cdef int64_t[B2ND_MAX_DIM] start_, stop_, buffershape_
+        cdef c_bool[B2ND_MAX_DIM] int_dims
+        cdef int64_t tick, buffersize = 0
+        cdef int rc = 0, attempt
+        cdef c_bool direct = False
+        target = None
+        with self._handle_lock:
+            # As in get_basic_numpy(): a key resolved against the cached shape is right if the read
+            # does not find shape changes made by another handle (SWMR); otherwise follow them and
+            # try again.  Also follow them first if the selection does not have out's shape (they
+            # may be why) or is empty (there is no read then).
+            for attempt in range(_BASIC_READ_ATTEMPTS):
+                if attempt > 0:
+                    _check_rc(b2nd_refresh(self.array), "Error while refreshing the array")
+                tick = self.array.last_tick
+                out_shape = self._resolve_basic_key(nkey, kinds, ints, starts, stops, nidx, start_, stop_,
+                                                    int_dims)
+                if out_shape is None:
+                    if attempt > 0:
+                        return None
+                    continue
+                if tuple(out_shape) != (<object> out).shape:
+                    if attempt > 0:
+                        raise ValueError(f"out has shape {(<object> out).shape}, but the selection has "
+                                         f"shape {tuple(out_shape)}")
+                    continue
+                if np.PyArray_SIZE(out) == 0:
+                    if attempt > 0:
+                        return out
+                    continue
+                if target is None:
+                    direct = out.dtype == self.dtype and self._box_layout(out, int_dims, buffershape_,
+                                                                          &buffersize)
+                    target = out if direct else np.empty(out_shape, dtype=self.dtype)
+                if direct:
+                    rc = self._read_slice(start_, stop_, np.PyArray_DATA(out), buffershape_, buffersize, True)
+                else:
+                    rc = self._get_slice_cbuffer(target, start_, stop_)
+                if rc >= 0 and self.array.last_tick == tick:
+                    return target
+        # Out of attempts: as the general path, return what the last read gave
+        _check_rc(rc, "Error while getting the slice")
+        return target
+
     cdef _resolve_basic_key(self, Py_ssize_t nkey, const char *kinds, const int64_t *ints,
                             const Py_ssize_t *starts, const Py_ssize_t *stops, Py_ssize_t nidx,
                             int64_t *start_, int64_t *stop_, c_bool *int_dims=NULL):
@@ -4449,7 +4519,6 @@ cdef class NDArray:
 
     cdef int _get_slice_cbuffer(self, arr, int64_t *start_, int64_t *stop_) except? -9999:
         # The caller holds self._handle_lock; returns the rc of b2nd_get_slice_cbuffer()
-        global _nogil_reads
         cdef int64_t[B2ND_MAX_DIM] buffershape_
         cdef int i
         for i in range(self.array.ndim):
@@ -4459,22 +4528,73 @@ cdef class NDArray:
         cdef int rc = 0
         PyObject_GetBuffer(arr, &view, PyBUF_SIMPLE)
         try:
-            # Decompress without the GIL, so that threads reading other arrays (or other handles
-            # of the same file) run in parallel; everything else that uses this super-chunk
-            # holds the lock meanwhile.
-            if _can_release_gil(self.array.sc):
-                _nogil_reads += 1
-                with nogil:
-                    rc = b2nd_get_slice_cbuffer(self.array, start_, stop_,
-                                                <void *> view.buf, buffershape_, view.len)
-                _nogil_reads -= 1
-            else:
-                rc = b2nd_get_slice_cbuffer(self.array, start_, stop_,
-                                            <void *> view.buf, buffershape_, view.len)
+            rc = self._read_slice(start_, stop_, view.buf, buffershape_, view.len, False)
         finally:
             PyBuffer_Release(&view)
         return rc
 
+    cdef int _read_slice(self, int64_t *start_, int64_t *stop_, void *buf, int64_t *buffershape_,
+                         int64_t buffersize, c_bool into) except? -9999:
+        """b2nd_get_slice_cbuffer(), or b2nd_get_slice_cbuffer_into() if `into`; returns its rc.
+        The caller holds self._handle_lock. The decompression runs without the GIL when no Python
+        callback can run on the way, so that threads reading other arrays (or other handles of the
+        same file) run in parallel; everything else that uses this super-chunk holds the lock
+        meanwhile."""
+        global _nogil_reads
+        cdef int rc
+        if _can_release_gil(self.array.sc):
+            _nogil_reads += 1
+            with nogil:
+                if into:
+                    rc = b2nd_get_slice_cbuffer_into(self.array, start_, stop_, buf, buffershape_, buffersize)
+                else:
+                    rc = b2nd_get_slice_cbuffer(self.array, start_, stop_, buf, buffershape_, buffersize)
+            _nogil_reads -= 1
+        elif into:
+            rc = b2nd_get_slice_cbuffer_into(self.array, start_, stop_, buf, buffershape_, buffersize)
+        else:
+            rc = b2nd_get_slice_cbuffer(self.array, start_, stop_, buf, buffershape_, buffersize)
+        return rc
+
+    cdef c_bool _box_layout(self, np.ndarray out, const c_bool *int_dims, int64_t *buffershape_,
+                            int64_t *buffersize):
+        """Whether `out`, of the shape of a selection whose integer-indexed dimensions are
+        `int_dims`, is a box of a C-order buffer that starts at out's first item; if so, fill that
+        buffer's shape and its size in bytes from there (for b2nd_get_slice_cbuffer_into()).
+        Integer-indexed dimensions and the others of length 1 are size-1 axes of the box, whose
+        strides do not matter. Of the longer ones, the innermost must be contiguous, and every
+        other one's stride a whole multiple of the next one's that spans at least that one's
+        length."""
+        cdef np.npy_intp *shape = np.PyArray_DIMS(out)
+        cdef np.npy_intp *strides = np.PyArray_STRIDES(out)
+        cdef int64_t itemsize = np.PyArray_ITEMSIZE(out)
+        cdef int64_t[B2ND_MAX_DIM] lens, steps
+        cdef int64_t size = itemsize
+        cdef int ndim = self.array.ndim, d, j = 0, inner = -1
+        for d in range(ndim):
+            if int_dims[d]:
+                lens[d] = 1
+                steps[d] = 0
+            else:
+                lens[d] = shape[j]
+                steps[d] = strides[j]
+                j += 1
+            buffershape_[d] = lens[d]
+        for d in range(ndim - 1, -1, -1):  # the long dimensions, innermost first
+            if lens[d] <= 1:
+                continue
+            if inner < 0:
+                if steps[d] != itemsize or itemsize <= 0:
+                    return False
+            elif steps[d] % steps[inner] != 0 or steps[d] // steps[inner] < lens[inner]:
+                return False
+            else:
+                buffershape_[inner] = steps[d] // steps[inner]
+            size += (lens[d] - 1) * steps[d]
+            inner = d
+        buffersize[0] = size
+        return True
+
     @_handle_locked
     def get_1d_span_numpy(self, arr, int64_t nchunk, int32_t start, int32_t nitems):
         if self.ndim != 1:

src/blosc2/ndarray.py

@@ -4967,6 +4967,78 @@ def _remember_read(self, key, nparr: np.ndarray) -> None:
         self._last_read.clear()
         self._last_read[make_key_hashable(key)] = nparr
 
+    def read_into(self, out: np.ndarray, key: Any = None, casting: str = "same_kind") -> np.ndarray:
+        """Decompress ``self[key]`` into the existing NumPy array ``out``.
+
+        This is ``out[...] = self[key]`` without the temporary array: the data is decompressed
+        straight to its place in ``out``, which can be a view into a larger array, e.g. the
+        interior of a padded array, one item of a batch or one channel of a multi-channel array.
+        Nothing of the larger array outside ``out`` is touched. Reusing ``out`` across calls also
+        saves the page faults of a fresh result array every time.
+
+        Parameters
+        ----------
+        out: np.ndarray
+            Writeable NumPy array with the shape of ``self[key]``. Contiguous or not: any layout
+            works. Data goes straight into ``out`` (no temporary) when ``out`` has ``self.dtype`` and
+            its last axis is contiguous, with any strides in the other axes that are whole
+            multiples of the next axis' stride. That covers C-contiguous arrays and every box of one,
+            including views with steps in all but the last axis. Other layouts (e.g. Fortran order,
+            transposed views, steps along the last axis) and other dtypes go through one temporary.
+            For a CPU ``torch.Tensor`` ``t``, pass ``t.numpy()``.
+        key: int, slice, Ellipsis or tuple of them, or any key that :meth:`__getitem__` accepts
+            The selection to read, as in ``self[key]``; None (default) reads the whole array.
+            Integers, slices with step 1 and an Ellipsis decompress straight into ``out``; other
+            keys (steps, integer or boolean arrays, ``None``) go through ``self[key]``.
+        casting: str
+            How ``self.dtype`` may be converted to ``out.dtype``, as in :func:`numpy.copyto`.
+            Default ``"same_kind"``.
+
+        Returns
+        -------
+        out: np.ndarray
+            The ``out`` array.
+
+        Examples
+        --------
+        >>> import numpy as np
+        >>> import blosc2
+        >>> a = blosc2.arange(12, dtype=np.int32, shape=(3, 4))
+        >>> padded = np.full((5, 6), -1, dtype=np.int32)
+        >>> _ = a.read_into(padded[1:4, 1:5])  # into the interior of a padded array
+        >>> padded
+        array([[-1, -1, -1, -1, -1, -1],
+               [-1,  0,  1,  2,  3, -1],
+               [-1,  4,  5,  6,  7, -1],
+               [-1,  8,  9, 10, 11, -1],
+               [-1, -1, -1, -1, -1, -1]], dtype=int32)
+        >>> batch = np.empty((2, 2, 4), dtype=np.int32)
+        >>> _ = a.read_into(batch[1], (slice(1, 3), slice(None)))  # into one item of a batch
+        >>> batch[1]
+        array([[ 4,  5,  6,  7],
+               [ 8,  9, 10, 11]], dtype=int32)
+        """
+        if not isinstance(out, np.ndarray):
+            raise TypeError(f"out must be a NumPy array, not {type(out).__name__}")
+        if not out.flags.writeable:
+            raise ValueError("out is read-only")
+        if key is None:
+            key = ()
+        # Integers, slices with step 1 and an Ellipsis: straight into out, or into a temporary for
+        # another dtype or layout (None: another key, see get_basic_into)
+        value = self.get_basic_into(out, key)
+        if value is out:
+            return out
+        if value is None:  # steps, arrays, newaxis, masks, fields...: through __getitem__
+            value = self[key]
+            if isinstance(value, (blosc2.LazyExpr, NDField)):
+                value = value[:]
+            value = np.asarray(value)
+            if out.shape != value.shape:
+                raise ValueError(f"out has shape {out.shape}, but the selection has shape {value.shape}")
+        np.copyto(out, value, casting=casting)
+        return out
+
     def __setitem__(  # noqa: C901
         self,
         key: int | slice | Sequence[slice | int | np.bool_ | np.ndarray[int | np.bool_] | None] | None,