From 09e7a47f1a1f9eff9622390d9d1fbdd8255e2fff Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Wed, 28 Oct 2026 21:32:42 +0100 Subject: [PATCH 01/12] Release the GIL in NDArray slice reads NDArray.get_slice_numpy() decompressed with the GIL held, so threads reading separate arrays, or separate handles of one file, ran one at a time; on NFS this included the wait for the server. Release the GIL around b2nd_get_slice_cbuffer(). The GIL was also what kept two threads from using one super-chunk at the same time. Replace the per-NDArray read_lock with one reentrant lock per super-chunk, shared by the SChunk wrapper, the NDArray, its views and vlmeta, and hold it in every entry point into C. Keep the GIL when a Python postfilter, codec or filter could run during the read, and on wasm32. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj7Wi7Sc8HtQpELyNXB --- src/blosc2/blosc2_ext.pyx | 173 ++++++++++++++++++++---- src/blosc2/core.py | 4 + tests/ndarray/test_concurrent_access.py | 161 ++++++++++++++++++++++ 3 files changed, 310 insertions(+), 28 deletions(-) create mode 100644 tests/ndarray/test_concurrent_access.py diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 64affa623..47740ecbe 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -14,8 +14,10 @@ import ast import atexit import pathlib import sys +import functools import time import warnings +from threading import RLock import _ctypes @@ -365,6 +367,8 @@ cdef extern from "blosc2.h": blosc2_context* blosc2_create_dctx(blosc2_dparams dparams) nogil + int blosc2_ctx_get_dparams(blosc2_context *ctx, blosc2_dparams *dparams) + void blosc2_free_ctx(blosc2_context * context) nogil int blosc2_set_maskout(blosc2_context *ctx, c_bool *maskout, int nblocks) @@ -622,7 +626,7 @@ cdef extern from "b2nd.h": int b2nd_free(b2nd_array_t *array) int b2nd_get_slice_cbuffer(b2nd_array_t *array, int64_t *start, int64_t *stop, - void *buffer, int64_t *buffershape, int64_t buffersize) + 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, @@ -1707,11 +1711,30 @@ cdef get_chunk_repeatval(blosc2_cparams cparams, const int32_t nbytes, raise RuntimeError("Problems when creating the repeated values chunk") +def _handle_locked(method): + """Run `method` holding the handle lock of its first argument (an SChunk, NDArray or vlmeta). + + There is one such lock per C super-chunk, shared by the SChunk wrapper, the NDArray built on it + and that array's views. NDArray.get_slice_numpy() decompresses without the GIL, so every other + operation that uses the super-chunk's C state (its contexts, frame, metalayers) must hold this + lock instead of relying on the GIL. It is reentrant because Python callbacks (prefilters, + postfilters, UDFs) can call back into the same array. + """ + @functools.wraps(method) + def locked(self, *args, **kwargs): + with self._handle_lock: + return method(self, *args, **kwargs) + return locked + + cdef class SChunk: cdef blosc2_schunk *schunk cdef c_bool _is_view + cdef readonly object _handle_lock def __init__(self, _schunk=None, chunksize=2 ** 24, data=None, **kwargs): + # Replaced by the NDArray's lock when this is an NDArray's SChunk (see NDArray.__init__) + self._handle_lock = RLock() # hold on to a bytestring of urlpath for the lifetime of the instance # because its value is referenced via a C-pointer urlpath = kwargs.get("urlpath", None) @@ -1881,6 +1904,7 @@ cdef class SChunk: def contiguous(self): return self.schunk.storage.contiguous + @_handle_locked def get_cparams(self): if self.schunk.storage.cparams.compcode in blosc2.Codec._value2member_map_: codec = blosc2.Codec(self.schunk.storage.cparams.compcode) @@ -1914,6 +1938,7 @@ cdef class SChunk: return cparams + @_handle_locked def update_cparams(self, new_cparams): cdef blosc2_cparams* cparams = self.schunk.storage.cparams codec = new_cparams.codec @@ -1957,9 +1982,11 @@ cdef class SChunk: self.schunk.filters = self.schunk.storage.cparams.filters self.schunk.filters_meta = self.schunk.storage.cparams.filters_meta + @_handle_locked def get_dparams(self): return blosc2.DParams(nthreads=self.schunk.storage.dparams.nthreads) + @_handle_locked def update_dparams(self, new_dparams): cdef blosc2_dparams* dparams = self.schunk.storage.dparams dparams.nthreads = new_dparams.nthreads @@ -1971,6 +1998,7 @@ cdef class SChunk: if self.schunk.dctx == NULL: raise RuntimeError("Could not create decompression context") + @_handle_locked def append_data(self, data): cdef Py_buffer buf PyObject_GetBuffer(data, &buf, PyBUF_SIMPLE) @@ -1998,6 +2026,7 @@ cdef class SChunk: raise RuntimeError("Could not append the chunk") return rc + @_handle_locked def fill_special(self, nitems, special_value, value): if value is None: return blosc2_schunk_fill_special(self.schunk, nitems, special_value, self.chunksize) @@ -2046,6 +2075,7 @@ cdef class SChunk: return self.nchunks + @_handle_locked def decompress_chunk(self, nchunk, dst=None): cdef uint8_t *chunk cdef c_bool needs_free @@ -2079,6 +2109,7 @@ cdef class SChunk: if size < 0: raise RuntimeError(f"Error while decompressing the specified chunk, error code: {size}") + @_handle_locked def get_chunk(self, nchunk): cdef uint8_t *chunk cdef c_bool needs_free @@ -2090,6 +2121,7 @@ cdef class SChunk: free(chunk) return ret_chunk + @_handle_locked def get_lazychunk(self, nchunk): cdef uint8_t *chunk cdef c_bool needs_free @@ -2113,6 +2145,7 @@ cdef class SChunk: free(chunk) return ret_chunk + @_handle_locked def get_vlblock(self, nchunk, nblock): cdef uint8_t *block cdef int32_t destsize @@ -2123,12 +2156,14 @@ cdef class SChunk: free(block) return ret_block + @_handle_locked def delete_chunk(self, nchunk): rc = blosc2_schunk_delete_chunk(self.schunk, nchunk) if rc < 0: raise RuntimeError("Could not delete the desired chunk") return rc + @_handle_locked def append_chunk(self, chunk): cdef const uint8_t[:] typed_view_chunk mem_view_chunk = memoryview(chunk) @@ -2139,6 +2174,7 @@ cdef class SChunk: raise RuntimeError("Could not append the desired chunk") return rc + @_handle_locked def insert_chunk(self, nchunk, chunk): cdef const uint8_t[:] typed_view_chunk mem_view_chunk = memoryview(chunk) @@ -2149,6 +2185,7 @@ cdef class SChunk: raise RuntimeError("Could not insert the desired chunk") return rc + @_handle_locked def insert_data(self, nchunk, data, copy): cdef blosc2_context *cctx cdef Py_buffer buf @@ -2179,6 +2216,7 @@ cdef class SChunk: raise RuntimeError("Could not insert the desired chunk") return rc + @_handle_locked def update_chunk(self, nchunk, chunk): cdef const uint8_t[:] typed_view_chunk mem_view_chunk = memoryview(chunk) @@ -2189,6 +2227,7 @@ cdef class SChunk: raise RuntimeError("Could not update the desired chunk") return rc + @_handle_locked def reorder_offsets(self, order): cdef np.ndarray[np.int64_t, ndim=1] offsets_order = np.ascontiguousarray(order, dtype=np.int64) rc = blosc2_schunk_reorder_offsets(self.schunk, offsets_order.data) @@ -2196,6 +2235,7 @@ cdef class SChunk: raise RuntimeError("Could not reorder the chunk offsets") return None + @_handle_locked def update_data(self, nchunk, data, copy): cdef Py_buffer buf PyObject_GetBuffer(data, &buf, PyBUF_SIMPLE) @@ -2226,6 +2266,7 @@ cdef class SChunk: return rc # This is used internally for prefiltering + @_handle_locked def _prefilter_data(self, nchunk, data, chunk_data): cdef Py_buffer buf PyObject_GetBuffer(data, &buf, PyBUF_SIMPLE) @@ -2241,6 +2282,7 @@ cdef class SChunk: raise RuntimeError("The result could not fit ") return size + @_handle_locked def get_slice(self, start=0, stop=None, out=None): cdef int64_t nitems = self.schunk.nbytes // self.schunk.typesize start, stop, _ = slice(start, stop, 1).indices(nitems) @@ -2265,6 +2307,7 @@ cdef class SChunk: if rc < 0: raise RuntimeError("Error while getting the slice") + @_handle_locked def set_slice(self, value, start=0, stop=None): cdef int64_t nitems = self.schunk.nbytes // self.schunk.typesize start, stop = self._massage_key(start, stop, nitems) @@ -2370,6 +2413,7 @@ cdef class SChunk: if rc < 0: raise RuntimeError("Error while setting the slice") + @_handle_locked def to_cframe(self): cdef c_bool needs_free cdef uint8_t *cframe @@ -2387,9 +2431,11 @@ cdef class SChunk: return out + @_handle_locked def _avoid_cframe_free(self, avoid_cframe_free): blosc2_schunk_avoid_cframe_free(self.schunk, avoid_cframe_free) + @_handle_locked def lock(self): """Take the exclusive frame lock and hold it across several operations. @@ -2402,12 +2448,14 @@ cdef class SChunk: if rc < 0: raise RuntimeError("Could not lock the super-chunk frame") + @_handle_locked def unlock(self): """Release the frame lock taken with `lock()`.""" cdef int rc = blosc2_schunk_unlock(self.schunk) if rc < 0: raise RuntimeError("Could not unlock the super-chunk frame") + @_handle_locked def refresh(self): cdef int rc = blosc2_schunk_refresh(self.schunk) _check_rc(rc, "Error while refreshing the schunk") @@ -2433,6 +2481,7 @@ cdef class SChunk: return start, stop + @_handle_locked def _set_postfilter(self, func, dtype_input, dtype_output=None): # Get user data func_id = func.__name__ @@ -2467,7 +2516,11 @@ cdef class SChunk: if self.schunk.dctx == NULL: raise RuntimeError("Could not create decompression context") - cpdef remove_postfilter(self, func_name, _new_ctx=True): + @_handle_locked + def remove_postfilter(self, func_name, _new_ctx=True): + self._remove_postfilter(func_name, _new_ctx) + + cdef _remove_postfilter(self, func_name, _new_ctx): if func_name is not None: del blosc2.postfilter_funcs[func_name] @@ -2487,6 +2540,7 @@ cdef class SChunk: # Avoid creating new dctx when calling this from the __dealloc__ self.schunk.dctx = NULL + @_handle_locked def _set_filler(self, func, inputs_id, dtype_output): if self.schunk.storage.cparams.nthreads > 1: raise AttributeError("compress `nthreads` must be 1 when assigning a prefilter") @@ -2516,6 +2570,7 @@ cdef class SChunk: if self.schunk.cctx == NULL: raise RuntimeError("Could not create compression context") + @_handle_locked def _set_prefilter(self, func, dtype_input, dtype_output=None): if self.schunk.storage.cparams.nthreads > 1: raise AttributeError("compress `nthreads` must be 1 when assigning a prefilter") @@ -2551,7 +2606,11 @@ cdef class SChunk: if self.schunk.cctx == NULL: raise RuntimeError("Could not create compression context") - cpdef remove_prefilter(self, func_name, _new_ctx=True): + @_handle_locked + def remove_prefilter(self, func_name, _new_ctx=True): + self._remove_prefilter(func_name, _new_ctx) + + cdef _remove_prefilter(self, func_name, _new_ctx): cdef udf_udata* udf_data cdef user_filters_udata* udata cdef mm_udata* mm_data @@ -2630,9 +2689,9 @@ cdef class SChunk: if self.schunk != NULL and not self._is_view: # Free prefilters and postfilters params if self.schunk.storage.cparams.prefilter != NULL: - self.remove_prefilter(func_name=None, _new_ctx=False) + self._remove_prefilter(None, False) if self.schunk.storage.dparams.postfilter != NULL: - self.remove_postfilter(func_name=None, _new_ctx=False) + self._remove_postfilter(None, False) # Free the C-Blosc2 super-chunk with the GIL held so threadpool # teardown cannot race with active miniexpr workers. @@ -3320,6 +3379,13 @@ cdef class vlmeta: def __init__(self, schunk): self.schunk = schunk + @property + def _handle_lock(self): + # Looked up on every call: an NDArray swaps its own lock into the owning SChunk after + # that SChunk (and this object) were created + return self._owner._handle_lock + + @_handle_locked def set_vlmeta(self, name, content, **cparams): cdef blosc2_cparams ccparams create_cparams_from_kwargs(&ccparams, cparams) @@ -3335,6 +3401,7 @@ cdef class vlmeta: if rc < 0: raise RuntimeError + @_handle_locked def get_vlmeta(self, name): name = name.encode("utf-8") if isinstance(name, str) else name rc = blosc2_vlmeta_exists(self.schunk, name) @@ -3348,6 +3415,7 @@ cdef class vlmeta: raise RuntimeError return content[:content_len] + @_handle_locked def del_vlmeta(self, name): name = name.encode("utf-8") if isinstance(name, str) else name rc = blosc2_vlmeta_delete(self.schunk, name) @@ -3357,6 +3425,7 @@ cdef class vlmeta: def nvlmetalayers(self): return self.schunk.nvlmetalayers + @_handle_locked def get_names(self): cdef char** names = malloc(self.schunk.nvlmetalayers * sizeof (char *)) rc = blosc2_vlmeta_get_names(self.schunk, names) @@ -3365,6 +3434,7 @@ cdef class vlmeta: res = [names[i].decode("utf-8") for i in range(rc)] return res + @_handle_locked def to_dict(self): cdef char** names = malloc(self.schunk.nvlmetalayers * sizeof (char*)) rc = blosc2_vlmeta_get_names(self.schunk, names) @@ -3376,12 +3446,14 @@ cdef class vlmeta: return res +@_handle_locked def meta__contains__(self, name): cdef blosc2_schunk *schunk = self.c_schunk name = name.encode("utf-8") if isinstance(name, str) else name n = blosc2_meta_exists(schunk, name) return False if n < 0 else True +@_handle_locked def meta__getitem__(self, name): cdef blosc2_schunk *schunk = self.c_schunk name = name.encode("utf-8") if isinstance(name, str) else name @@ -3393,6 +3465,7 @@ def meta__getitem__(self, name): return res +@_handle_locked def meta__setitem__(self, name, content): cdef blosc2_schunk *schunk = self.c_schunk name = name.encode("utf-8") if isinstance(name, str) else name @@ -3405,6 +3478,7 @@ def meta__len__(self): cdef blosc2_schunk *schunk = self.c_schunk return schunk.nmetalayers +@_handle_locked def meta_keys(self): cdef blosc2_schunk *schunk = self.c_schunk keys = [] @@ -3814,24 +3888,39 @@ cdef class slice_flatter: self.done = 1 +cdef c_bool _can_release_gil(blosc2_schunk *sc) except? False: + """Whether decompressing from `sc` may run without the GIL, i.e. no Python callback can run on + the way: no postfilter on its decompression context and no codec or filter registered from + Python (any chunk could use one, whatever the super-chunk's own cparams say). Never on wasm32, + where releasing the GIL is known to give wrong results (see blosc2.set_releasegil).""" + if blosc2.IS_WASM or blosc2.ucodecs_registry or blosc2.ufilters_registry: + return False + cdef blosc2_dparams dparams + if sc.dctx == NULL or blosc2_ctx_get_dparams(sc.dctx, &dparams) < 0: + return False + return dparams.postfilter == NULL + + cdef class NDArray: cdef b2nd_array_t* array - cdef PyThread_type_lock read_lock - cdef c_bool owns_read_lock + cdef readonly object _handle_lock def __init__(self, array, base=None): self._dtype = None self.array = PyCapsule_GetPointer(array, "b2nd_array_t*") self.base = base # add reference to base if NDArray is a view - if base is None: - self.read_lock = PyThread_allocate_lock() - if self.read_lock == NULL: - raise MemoryError("Could not allocate NDArray read lock") - self.owns_read_lock = True + # One handle lock per C super-chunk (see _handle_locked): expand_dims/squeeze views share + # their base's, and the SChunk wrapper over the same super-chunk (blosc2.NDArray._schunk, + # created before this runs) is switched to it. + schunk = getattr(self, "_schunk", None) + if base is not None: + self._handle_lock = (base)._handle_lock + elif schunk is not None: + self._handle_lock = (schunk)._handle_lock else: - # expand_dims/squeeze views share the base SChunk, so reads through - # all aliases must be protected by the same lock. - self.read_lock = (base).read_lock + self._handle_lock = RLock() + if schunk is not None: + (schunk)._handle_lock = self._handle_lock @property def c_array(self): @@ -3899,22 +3988,29 @@ cdef class NDArray: buffershape_[i] = stop_[i] - start_[i] cdef Py_buffer view - cdef int rc + cdef int rc = 0 PyObject_GetBuffer(arr, &view, PyBUF_SIMPLE) - # Waiting for a reader already using this SChunk must not retain the - # GIL: that reader can need the GIL again from a Python postfilter. - with nogil: - PyThread_acquire_lock(self.read_lock, 1) try: - rc = b2nd_get_slice_cbuffer(self.array, start_, stop_, - view.buf, buffershape_, view.len) + # Waiting for the lock releases the GIL (RLock does that), so a holder that needs the + # GIL back, e.g. for a Python postfilter, can finish. + with self._handle_lock: + # 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): + with nogil: + rc = b2nd_get_slice_cbuffer(self.array, start_, stop_, + view.buf, buffershape_, view.len) + else: + rc = b2nd_get_slice_cbuffer(self.array, start_, stop_, + view.buf, buffershape_, view.len) finally: - PyThread_release_lock(self.read_lock) - PyBuffer_Release(&view) + PyBuffer_Release(&view) _check_rc(rc, "Error while getting the buffer") return arr + @_handle_locked def get_1d_span_numpy(self, arr, int64_t nchunk, int32_t start, int32_t nitems): if self.ndim != 1: raise ValueError("get_1d_span_numpy is only supported for 1-D arrays") @@ -3997,6 +4093,7 @@ cdef class NDArray: return arr + @_handle_locked def get_sparse_numpy(self, arr, coords): cdef np.ndarray[np.int64_t, ndim=1, mode="c"] coords_ = np.ascontiguousarray(coords, dtype=np.int64) cdef Py_buffer view @@ -4014,6 +4111,7 @@ cdef class NDArray: _check_rc(rc, "Error while getting the sparse selection") return arr + @_handle_locked def get_oindex_numpy(self, arr, key): """ Orthogonal indexing. Key is a tuple of lists of integer indices. @@ -4046,6 +4144,7 @@ cdef class NDArray: free(key_) return arr + @_handle_locked def set_oindex_numpy(self, key, arr): """ Orthogonal indexing. Set elements of self with arr using key. @@ -4079,6 +4178,7 @@ cdef class NDArray: return arr + @_handle_locked def get_slice(self, key, mask, **kwargs): start, stop = key shape = tuple(sp - st for sp, st in zip(stop, start)) @@ -4120,6 +4220,7 @@ cdef class NDArray: return ndarray + @_handle_locked def set_slice(self, key, ndarray): ndim = self.ndim start, stop = key @@ -4138,6 +4239,7 @@ cdef class NDArray: return self + @_handle_locked def tobytes(self): buffersize = self.size * self.array.sc.typesize buffer = bytes(buffersize) @@ -4146,6 +4248,7 @@ cdef class NDArray: return buffer + @_handle_locked def to_cframe(self): cdef c_bool needs_free cdef uint8_t *cframe @@ -4160,6 +4263,7 @@ cdef class NDArray: return out + @_handle_locked def copy(self, dtype, **kwargs): chunks = kwargs.pop("chunks", self.chunks) blocks = kwargs.pop("blocks", self.blocks) @@ -4180,6 +4284,7 @@ cdef class NDArray: return ndarray + @_handle_locked def resize(self, new_shape): cdef int64_t new_shape_[B2ND_MAX_DIM] for i, s in enumerate(new_shape): @@ -4187,6 +4292,7 @@ cdef class NDArray: _check_rc(b2nd_resize(self.array, new_shape_, NULL), "Error while resizing the array") + @_handle_locked def refresh(self): cdef int rc = b2nd_refresh(self.array) _check_rc(rc, "Error while refreshing the array") @@ -4356,6 +4462,7 @@ cdef class NDArray: return udata + @_handle_locked def _set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, jit=None, candidate_blocks=None): # Set prefilter for miniexpr @@ -4538,6 +4645,7 @@ cdef class NDArray: "status": _me_compile_status_name(rc), } + @_handle_locked def _set_pref_matmul(self, inputs, fp_accuracy): # Set prefilter for miniexpr cdef blosc2_cparams* cparams = self.array.sc.storage.cparams @@ -4558,6 +4666,7 @@ cdef class NDArray: if self.array.sc.cctx == NULL: raise RuntimeError("Could not create compression context") + @_handle_locked def _set_pref_udf(self, func, inputs_id): if self.array.sc.storage.cparams.nthreads > 1: raise AttributeError("compress `nthreads` must be 1 when assigning a prefilter") @@ -4580,6 +4689,7 @@ cdef class NDArray: if self.array.sc.cctx == NULL: raise RuntimeError("Could not create compression context") + @_handle_locked def _set_postf_udf(self, func, inputs_id): if self.array.sc.storage.dparams.nthreads > 1: raise AttributeError("decompress `nthreads` must be 1 when assigning a postfilter") @@ -4608,8 +4718,6 @@ cdef class NDArray: def __dealloc__(self): if self.array != NULL: _check_rc(b2nd_free(self.array), "Error while freeing the array") - if self.owns_read_lock and self.read_lock != NULL: - PyThread_free_lock(self.read_lock) cdef b2nd_context_t* create_b2nd_context(shape, chunks, blocks, dtype, kwargs): @@ -4859,6 +4967,7 @@ def ndarray_from_cframe(cframe, copy=False): return ndarray +@_handle_locked def array_get_slice_nchunks(array: NDArray, key): start, stop = key cdef int64_t[B2ND_MAX_DIM] start_, stop_ @@ -4875,6 +4984,7 @@ def array_get_slice_nchunks(array: NDArray, key): return res +@_handle_locked def schunk_get_slice_nchunks(schunk: SChunk, key): start, stop = key nitems = schunk.nbytes // schunk.typesize @@ -4904,8 +5014,13 @@ def concat(arr1: NDArray, arr2: NDArray, axis: int, **kwargs): raise RuntimeError("Error while creating the context for concatenation") cdef b2nd_array_t *array - _check_rc(b2nd_concatenate(ctx, arr1.array, arr2.array, axis, copy, &array), - "Error while concatenating the arrays") + cdef int rc = 0 + # Both handle locks, always in the same order, so that concat(a, b) and concat(b, a) in two + # threads cannot deadlock (they are the same lock for views of one array) + first, second = sorted((arr1._handle_lock, arr2._handle_lock), key=id) + with first, second: + rc = b2nd_concatenate(ctx, arr1.array, arr2.array, axis, copy, &array) + _check_rc(rc, "Error while concatenating the arrays") _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") if copy: @@ -4916,6 +5031,7 @@ def concat(arr1: NDArray, arr2: NDArray, axis: int, **kwargs): # Return the first array, which now contains the concatenated data return arr1 +@_handle_locked def expand_dims(arr1: NDArray, axis_mask: list[bool], final_dims: int) -> blosc2.NDArray: """ Add new dummy axis to NDArray object at specified dimension. @@ -4933,6 +5049,7 @@ def expand_dims(arr1: NDArray, axis_mask: list[bool], final_dims: int) -> blosc2 return blosc2.NDArray(_schunk=PyCapsule_New(view.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(view, "b2nd_array_t*", NULL), _base=new_base) +@_handle_locked def squeeze(arr1: NDArray, axis_mask: list[bool]) -> blosc2.NDArray: """ Remove axis from NDArray object at specified dimensions. diff --git a/src/blosc2/core.py b/src/blosc2/core.py index b2ebfb75e..21a4eafbc 100644 --- a/src/blosc2/core.py +++ b/src/blosc2/core.py @@ -1449,6 +1449,10 @@ def set_releasegil(gilstate: bool) -> bool: small performance penalty with releasing the GIL that will more harshly penalize small block sizes. + Reading slices of an :ref:`NDArray` releases the GIL while decompressing + regardless of this setting, unless a Python postfilter, codec or filter + could run during the read. + Examples -------- >>> oldReleaseState = blosc2.set_releasegil(True) diff --git a/tests/ndarray/test_concurrent_access.py b/tests/ndarray/test_concurrent_access.py new file mode 100644 index 000000000..e41418584 --- /dev/null +++ b/tests/ndarray/test_concurrent_access.py @@ -0,0 +1,161 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +# NDArray.get_slice_numpy() decompresses without the GIL. These tests check that other threads +# really run meanwhile, and that every other operation on the same array still waits for it. + +import threading +import time + +import numpy as np +import pytest + +import blosc2 + +SHAPE, CHUNKS, BLOCKS = (64, 256, 256), (8, 256, 256), (2, 64, 256) + + +def run_for(seconds, *targets): + """Run each target in a loop in its own thread for `seconds`; return the errors they raised.""" + stop = time.monotonic() + seconds + errors = [] + + def loop(target): + try: + while time.monotonic() < stop and not errors: + target() + except Exception as e: # report anything a racing thread hits + errors.append(e) + + threads = [threading.Thread(target=loop, args=(t,)) for t in targets] + for t in threads: + t.start() + for t in threads: + t.join(timeout=seconds + 60) + assert not any(t.is_alive() for t in threads), "threads did not finish (deadlock?)" + return errors + + +@pytest.fixture +def array_and_values(tmp_path, request): + values = np.arange(np.prod(SHAPE), dtype=np.float32).reshape(SHAPE) + kwargs = {"urlpath": tmp_path / "a.b2nd", "mode": "w"} if request.param == "disk" else {} + return blosc2.asarray(values, chunks=CHUNKS, blocks=BLOCKS, **kwargs), values + + +def check_region(array, values): + """Read [8:24, 10:200]; a concurrent writer only ever writes values or values + 1 into + [8:24, 0:100], one chunk (8 rows) at a time.""" + got, want = array[8:24, 10:200], values[8:24, 10:200] + for c in (0, 8): + g, w = got[c : c + 8], want[c : c + 8] + np.testing.assert_array_equal(g[:, 90:], w[:, 90:]) + assert np.array_equal(g[:, :90], w[:, :90]) or np.array_equal(g[:, :90], w[:, :90] + 1) + + +def test_handle_lock_is_shared(): + a = blosc2.asarray(np.arange(100).reshape(10, 10), chunks=(5, 5)) + view = blosc2.expand_dims(a, axis=0) + for other in (a.schunk, a.schunk.vlmeta, view, view.schunk): + assert other._handle_lock is a._handle_lock + + +def test_get_slice_releases_the_gil(monkeypatch): + # Python-registered codecs or filters (from other tests in this process) keep the GIL + monkeypatch.setattr(blosc2, "ucodecs_registry", {}) + monkeypatch.setattr(blosc2, "ufilters_registry", {}) + a = blosc2.linspace( + 0, + 1, + 32 * 1024 * 1024, + shape=(4096, 8192), + dtype=np.float32, + cparams={"codec": blosc2.Codec.ZSTD, "clevel": 1}, + dparams={"nthreads": 1}, + ) + stamps, stop = [], threading.Event() + + def ticker(): + while not stop.is_set(): + stamps.append(time.perf_counter()) + time.sleep(0) + + t = threading.Thread(target=ticker) + t.start() + try: + t0 = time.perf_counter() + _ = a[:] + t1 = time.perf_counter() + finally: + stop.set() + t.join() + inside = [s for s in stamps if t0 <= s <= t1] + gaps = np.diff([t0, *inside, t1]) + # Holding the GIL while decompressing leaves one gap as long as the whole read + assert gaps.max() < 0.5 * (t1 - t0), f"longest gap {gaps.max():.3f} s of a {t1 - t0:.3f} s read" + + +@pytest.mark.parametrize("array_and_values", ["memory", "disk"], indirect=True) +def test_read_while_writing_same_handle(array_and_values): + a, values = array_and_values + k = [0] + + def write(): + k[0] += 1 + a[8:24, 0:100] = values[8:24, 0:100] + (k[0] % 2) + + assert run_for(1.0, lambda: check_region(a, values), write) == [] + + +@pytest.mark.parametrize("array_and_values", ["memory"], indirect=True) +@pytest.mark.parametrize("other", ["schunk", "vlmeta", "dparams", "resize"]) +def test_read_while_other_access_same_handle(array_and_values, other): + a, values = array_and_values + k = [0] + + def access(): + k[0] += 1 + if other == "schunk": + n = 1 + k[0] % 6 + a.schunk.update_chunk(n, a.schunk.get_chunk(n)) + a.schunk.decompress_chunk(k[0] % 8) + elif other == "vlmeta": + a.schunk.vlmeta[f"k{k[0] % 5}"] = list(range(k[0] % 50)) + elif other == "dparams": + a.schunk.dparams = blosc2.DParams(nthreads=1 + k[0] % 2) + elif other == "resize": + # grow and shrink past the region the reader checks + a.resize((SHAPE[0] + 8 * (k[0] % 2), *SHAPE[1:])) + + assert run_for(1.0, lambda: check_region(a, values), access) == [] + + +def test_concurrent_reads_with_python_postfilter(): + values = np.arange(np.prod(SHAPE), dtype=np.float32).reshape(SHAPE) + a = blosc2.asarray(values, chunks=CHUNKS, blocks=BLOCKS, dparams={"nthreads": 1}) + + @a.schunk.postfilter(np.float32, np.float32) + def add_one(inp, out, offset): + out[:] = inp + 1 + + def read(): + np.testing.assert_array_equal(a[8:24, 10:200], values[8:24, 10:200] + 1) + + assert run_for(1.0, read, read) == [] + + +def test_concat_both_orders_does_not_deadlock(): + a = blosc2.asarray(np.zeros((16, 64), dtype=np.int32), chunks=(8, 64)) + b = blosc2.asarray(np.ones((16, 64), dtype=np.int32), chunks=(8, 64)) + errors = run_for( + 1.0, + lambda: blosc2.concat([a, b], axis=0), + lambda: blosc2.concat([b, a], axis=0), + lambda: a[:], + lambda: b[:], + ) + assert errors == [] From ddd70de032f3095352cce05b51ecedba9fe8b8c6 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Wed, 28 Oct 2026 22:26:22 +0100 Subject: [PATCH 02/12] Make GIL-free slice reads safe around fork, exit and lazy expressions Hold the operands' handle locks while the miniexpr prefilter of a lazy expression reads their super-chunks from C. Count the reads that run without the GIL, so that exit and the registration of Python codecs and filters wait for them, and free handle locks that other threads held at a fork in the child. Lock the metalayer counters too. Skip the new tests on wasm32, shorten them, retry the GIL check on busy machines, and add tests for lazy expressions and fork. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj7Wi7Sc8HtQpELyNXB --- src/blosc2/blosc2_ext.pyx | 139 ++++++++++++++--- src/blosc2/lazyexpr.py | 10 +- tests/ndarray/test_concurrent_access.py | 192 ++++++++++++++++-------- 3 files changed, 257 insertions(+), 84 deletions(-) diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 47740ecbe..1ccf65d5b 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -17,6 +17,7 @@ import sys import functools import time import warnings +import weakref from threading import RLock import _ctypes @@ -876,7 +877,13 @@ blosc2_init() @atexit.register def destroy(): - blosc2_destroy() + # Daemon threads can still be inside a slice read without the GIL (see + # NDArray.get_slice_numpy), and blosc2_destroy() frees thread pools and plugins such a + # read may be using. If they do not finish soon, leave the cleanup to the OS. + global _nogil_reads_blocked + _nogil_reads_blocked += 1 + if _wait_for_nogil_reads(timeout=1.0): + blosc2_destroy() def register_wasm_jit_helpers(uintptr_t instantiate_ptr, uintptr_t free_ptr): @@ -1727,6 +1734,81 @@ def _handle_locked(method): return locked +# Every handle lock, so that a forked child can free those that other threads held at the fork +_all_handle_locks = weakref.WeakSet() + + +def _new_handle_lock(): + lock = RLock() + _all_handle_locks.add(lock) + return lock + + +cdef class HandleLocks: + """Reusable context manager that holds the handle locks of `arrays` (objects without one are + skipped), always in the same order, so that threads locking overlapping sets cannot deadlock. + For C code that reads other arrays' super-chunks outside their methods: concat() and the + miniexpr prefilter of lazy expressions. + """ + cdef list locks + + def __init__(self, arrays): + locks = {id(lock): lock for lock in (getattr(a, "_handle_lock", None) for a in arrays) + if lock is not None} + self.locks = [locks[key] for key in sorted(locks)] + + def __enter__(self): + cdef Py_ssize_t n = 0 + try: + for lock in self.locks: + lock.acquire() + n += 1 + except BaseException: + for lock in reversed(self.locks[:n]): + lock.release() + raise + return self + + def __exit__(self, *exc_info): + for lock in reversed(self.locks): + lock.release() + + +# Slice reads running without the GIL (see NDArray.get_slice_numpy), and whether new ones must +# keep it (> 0). Both only change with the GIL held. +cdef int _nogil_reads = 0 +cdef int _nogil_reads_blocked = 0 + + +def _wait_for_nogil_reads(timeout=None): + """Wait until no slice read runs without the GIL; return False if `timeout` (s) ran out. + + For changes to C state that such reads use (registering a Python codec or filter, exit). + The caller increments `_nogil_reads_blocked` first, so that no new ones start meanwhile. + """ + deadline = None if timeout is None else time.monotonic() + timeout + while _nogil_reads: + if deadline is not None and time.monotonic() > deadline: + return False + time.sleep(0.0005) # releases the GIL, so the readers can finish + return True + + +def _reset_after_fork(): + # The threads that held handle locks or read without the GIL at the fork do not exist in + # the child, so nothing would ever release those locks there + global _nogil_reads, _nogil_reads_blocked + _nogil_reads = 0 + _nogil_reads_blocked = 0 + for lock in list(_all_handle_locks): + if not lock._is_owned(): + lock._at_fork_reinit() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_after_fork) + + cdef class SChunk: cdef blosc2_schunk *schunk cdef c_bool _is_view @@ -1734,7 +1816,7 @@ cdef class SChunk: def __init__(self, _schunk=None, chunksize=2 ** 24, data=None, **kwargs): # Replaced by the NDArray's lock when this is an NDArray's SChunk (see NDArray.__init__) - self._handle_lock = RLock() + self._handle_lock = _new_handle_lock() # hold on to a bytestring of urlpath for the lifetime of the instance # because its value is referenced via a C-pointer urlpath = kwargs.get("urlpath", None) @@ -3422,6 +3504,7 @@ cdef class vlmeta: if rc < 0: raise RuntimeError("Could not delete the vlmeta") + @_handle_locked def nvlmetalayers(self): return self.schunk.nvlmetalayers @@ -3474,6 +3557,7 @@ def meta__setitem__(self, name, content): raise ValueError("The length of the content in a metalayer cannot change.") blosc2_meta_update(schunk, name, content, len(content)) +@_handle_locked def meta__len__(self): cdef blosc2_schunk *schunk = self.c_schunk return schunk.nmetalayers @@ -3705,12 +3789,21 @@ def register_codec(codec_name, id, encoder=None, decoder=None, version=1): else: codec.decoder = general_decoder - rc = blosc2_register_codec(&codec) - if rc < 0: - raise RuntimeError("Error while registering codec") + # A slice read that found no Python codec could otherwise meet `decoder` without the GIL: keep + # new reads from releasing it and wait for those that did. Nothing below releases the GIL + # before the registry is updated. + global _nogil_reads_blocked + _nogil_reads_blocked += 1 + try: + _wait_for_nogil_reads() + rc = blosc2_register_codec(&codec) + if rc < 0: + raise RuntimeError("Error while registering codec") - if encoder and decoder: - blosc2.ucodecs_registry[id] = (codec_name, encoder, decoder) + if encoder and decoder: + blosc2.ucodecs_registry[id] = (codec_name, encoder, decoder) + finally: + _nogil_reads_blocked -= 1 cdef int general_forward(const uint8_t* input_buffer, uint8_t* output_buffer, int32_t size, @@ -3772,11 +3865,18 @@ def register_filter(id, forward, backward, filter_name): filter.name = malloc(strlen(filter_name_) + 1) strcpy(filter.name, filter_name_) - rc = blosc2_register_filter(&filter) - if rc < 0: - raise RuntimeError("Error while registering filter") - if forward and backward: - blosc2.ufilters_registry[id] = (forward, backward) + # As in register_codec() + global _nogil_reads_blocked + _nogil_reads_blocked += 1 + try: + _wait_for_nogil_reads() + rc = blosc2_register_filter(&filter) + if rc < 0: + raise RuntimeError("Error while registering filter") + if forward and backward: + blosc2.ufilters_registry[id] = (forward, backward) + finally: + _nogil_reads_blocked -= 1 cdef _check_rc(rc, message): if rc < 0: @@ -3892,8 +3992,9 @@ cdef c_bool _can_release_gil(blosc2_schunk *sc) except? False: """Whether decompressing from `sc` may run without the GIL, i.e. no Python callback can run on the way: no postfilter on its decompression context and no codec or filter registered from Python (any chunk could use one, whatever the super-chunk's own cparams say). Never on wasm32, - where releasing the GIL is known to give wrong results (see blosc2.set_releasegil).""" - if blosc2.IS_WASM or blosc2.ucodecs_registry or blosc2.ufilters_registry: + where releasing the GIL is known to give wrong results (see blosc2.set_releasegil), nor while + _wait_for_nogil_reads() callers need the reads to stop.""" + if _nogil_reads_blocked or blosc2.IS_WASM or blosc2.ucodecs_registry or blosc2.ufilters_registry: return False cdef blosc2_dparams dparams if sc.dctx == NULL or blosc2_ctx_get_dparams(sc.dctx, &dparams) < 0: @@ -3918,7 +4019,7 @@ cdef class NDArray: elif schunk is not None: self._handle_lock = (schunk)._handle_lock else: - self._handle_lock = RLock() + self._handle_lock = _new_handle_lock() if schunk is not None: (schunk)._handle_lock = self._handle_lock @@ -3978,6 +4079,7 @@ cdef class NDArray: return dtype def get_slice_numpy(self, arr, key): + global _nogil_reads start, stop = key cdef int64_t[B2ND_MAX_DIM] start_, stop_ @@ -3998,9 +4100,11 @@ cdef class NDArray: # 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_, view.buf, buffershape_, view.len) + _nogil_reads -= 1 else: rc = b2nd_get_slice_cbuffer(self.array, start_, stop_, view.buf, buffershape_, view.len) @@ -5015,10 +5119,7 @@ def concat(arr1: NDArray, arr2: NDArray, axis: int, **kwargs): cdef b2nd_array_t *array cdef int rc = 0 - # Both handle locks, always in the same order, so that concat(a, b) and concat(b, a) in two - # threads cannot deadlock (they are the same lock for views of one array) - first, second = sorted((arr1._handle_lock, arr2._handle_lock), key=id) - with first, second: + with HandleLocks((arr1, arr2)): rc = b2nd_concatenate(ctx, arr1.array, arr2.array, axis, copy, &array) _check_rc(rc, "Error while concatenating the arrays") _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 6cf988169..28af9bbb8 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -1983,10 +1983,13 @@ def _miniexpr_eligible_operand(op): nct = res_eval.schunk.nchunks if bpc > 0 and _cb_anchor.shape[0] >= nct * bpc: chunk_has_candidate = _cb_anchor[: nct * bpc].reshape(nct, bpc).any(axis=1) + # The prefilter reads the operands' super-chunks from C, outside their methods + operand_locks = blosc2.blosc2_ext.HandleLocks(_pref_ops.values()) for nchunk in range(res_eval.schunk.nchunks): if chunk_has_candidate is not None and not chunk_has_candidate[nchunk]: continue - res_eval.schunk.update_data(nchunk, data, copy=False) + with operand_locks: + res_eval.schunk.update_data(nchunk, data, copy=False) # _cb_anchor (if any) stays referenced until fast_eval returns, which # is after the finally below removes the prefilter — so the buffer # safely outlives every prefilter dereference. @@ -2924,8 +2927,11 @@ def reduce_slices( # noqa: C901 data = np.empty(res_eval.schunk.chunksize, dtype=np.uint8) chunk_data = np.empty(res_eval.schunk.chunksize + blosc2.MAX_OVERHEAD, dtype=np.uint8) # Exercise prefilter for each chunk + # The prefilter reads the operands' super-chunks from C, outside their methods + operand_locks = blosc2.blosc2_ext.HandleLocks(operands.values()) for nchunk in range(res_eval.schunk.nchunks): - res_eval.schunk._prefilter_data(nchunk, data, chunk_data) + with operand_locks: + res_eval.schunk._prefilter_data(nchunk, data, chunk_data) except Exception as e: use_miniexpr = False if callable(expression) and _is_dsl_kernel_expression(expression): diff --git a/tests/ndarray/test_concurrent_access.py b/tests/ndarray/test_concurrent_access.py index e41418584..186870e50 100644 --- a/tests/ndarray/test_concurrent_access.py +++ b/tests/ndarray/test_concurrent_access.py @@ -8,6 +8,9 @@ # NDArray.get_slice_numpy() decompresses without the GIL. These tests check that other threads # really run meanwhile, and that every other operation on the same array still waits for it. +import itertools +import os +import signal import threading import time @@ -16,7 +19,16 @@ import blosc2 +pytestmark = pytest.mark.skipif(blosc2.IS_WASM, reason="Pyodide cannot start threads") + SHAPE, CHUNKS, BLOCKS = (64, 256, 256), (8, 256, 256), (2, 64, 256) +VALUES = np.arange(np.prod(SHAPE), dtype=np.float32).reshape(SHAPE) +# Per stress test; without the locking, the races show up within the first few reads +SECONDS = 0.5 + + +def make_array(**kwargs): + return blosc2.asarray(VALUES, chunks=CHUNKS, blocks=BLOCKS, **kwargs) def run_for(seconds, *targets): @@ -40,21 +52,20 @@ def loop(target): return errors -@pytest.fixture -def array_and_values(tmp_path, request): - values = np.arange(np.prod(SHAPE), dtype=np.float32).reshape(SHAPE) - kwargs = {"urlpath": tmp_path / "a.b2nd", "mode": "w"} if request.param == "disk" else {} - return blosc2.asarray(values, chunks=CHUNKS, blocks=BLOCKS, **kwargs), values +def check_region(array): + """Read [8:24, 10:200]; a concurrent write_region() only ever writes VALUES or VALUES + 1 into + [8:24, 0:100], one chunk (8 rows) at a time.""" + got, want = array[8:24, 10:200], VALUES[8:24, 10:200] + for chunk in (slice(0, 8), slice(8, 16)): + got_rows, want_rows = got[chunk], want[chunk] + np.testing.assert_array_equal(got_rows[:, 90:], want_rows[:, 90:]) + assert np.array_equal(got_rows[:, :90], want_rows[:, :90]) or np.array_equal( + got_rows[:, :90], want_rows[:, :90] + 1 + ) -def check_region(array, values): - """Read [8:24, 10:200]; a concurrent writer only ever writes values or values + 1 into - [8:24, 0:100], one chunk (8 rows) at a time.""" - got, want = array[8:24, 10:200], values[8:24, 10:200] - for c in (0, 8): - g, w = got[c : c + 8], want[c : c + 8] - np.testing.assert_array_equal(g[:, 90:], w[:, 90:]) - assert np.array_equal(g[:, :90], w[:, :90]) or np.array_equal(g[:, :90], w[:, :90] + 1) +def write_region(array, k): + array[8:24, 0:100] = VALUES[8:24, 0:100] + (k % 2) def test_handle_lock_is_shared(): @@ -77,85 +88,140 @@ def test_get_slice_releases_the_gil(monkeypatch): cparams={"codec": blosc2.Codec.ZSTD, "clevel": 1}, dparams={"nthreads": 1}, ) - stamps, stop = [], threading.Event() - def ticker(): - while not stop.is_set(): - stamps.append(time.perf_counter()) - time.sleep(0) + def longest_gap(): + """Longest time a ticker thread did not run during one read, as a fraction of the read.""" + stamps, stop = [], threading.Event() - t = threading.Thread(target=ticker) - t.start() - try: - t0 = time.perf_counter() - _ = a[:] - t1 = time.perf_counter() - finally: - stop.set() - t.join() - inside = [s for s in stamps if t0 <= s <= t1] - gaps = np.diff([t0, *inside, t1]) - # Holding the GIL while decompressing leaves one gap as long as the whole read - assert gaps.max() < 0.5 * (t1 - t0), f"longest gap {gaps.max():.3f} s of a {t1 - t0:.3f} s read" + def ticker(): + while not stop.is_set(): + stamps.append(time.perf_counter()) + time.sleep(0) + + t = threading.Thread(target=ticker) + t.start() + try: + t0 = time.perf_counter() + _ = a[:] + t1 = time.perf_counter() + finally: + stop.set() + t.join() + inside = [s for s in stamps if t0 <= s <= t1] + return np.diff([t0, *inside, t1]).max() / (t1 - t0) + + # Holding the GIL while decompressing leaves one gap as long as nearly the whole read, every + # time; a busy machine can delay the ticker now and then, hence a few attempts + gaps = [] + while len(gaps) < 5 and (not gaps or gaps[-1] >= 0.5): + gaps.append(longest_gap()) + assert gaps[-1] < 0.5, f"longest gaps, as fractions of the read: {gaps}" + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_read_while_writing_same_handle(tmp_path, on_disk): + a = make_array(urlpath=tmp_path / "a.b2nd", mode="w") if on_disk else make_array() + k = itertools.count() + assert run_for(SECONDS, lambda: check_region(a), lambda: write_region(a, next(k))) == [] -@pytest.mark.parametrize("array_and_values", ["memory", "disk"], indirect=True) -def test_read_while_writing_same_handle(array_and_values): - a, values = array_and_values - k = [0] +def update_chunks(a, k): + n = 1 + k % 6 + a.schunk.update_chunk(n, a.schunk.get_chunk(n)) + a.schunk.decompress_chunk(k % 8) - def write(): - k[0] += 1 - a[8:24, 0:100] = values[8:24, 0:100] + (k[0] % 2) - assert run_for(1.0, lambda: check_region(a, values), write) == [] +def write_vlmeta(a, k): + a.schunk.vlmeta[f"k{k % 5}"] = list(range(k % 50)) -@pytest.mark.parametrize("array_and_values", ["memory"], indirect=True) -@pytest.mark.parametrize("other", ["schunk", "vlmeta", "dparams", "resize"]) -def test_read_while_other_access_same_handle(array_and_values, other): - a, values = array_and_values - k = [0] +def swap_dparams(a, k): + a.schunk.dparams = blosc2.DParams(nthreads=1 + k % 2) - def access(): - k[0] += 1 - if other == "schunk": - n = 1 + k[0] % 6 - a.schunk.update_chunk(n, a.schunk.get_chunk(n)) - a.schunk.decompress_chunk(k[0] % 8) - elif other == "vlmeta": - a.schunk.vlmeta[f"k{k[0] % 5}"] = list(range(k[0] % 50)) - elif other == "dparams": - a.schunk.dparams = blosc2.DParams(nthreads=1 + k[0] % 2) - elif other == "resize": - # grow and shrink past the region the reader checks - a.resize((SHAPE[0] + 8 * (k[0] % 2), *SHAPE[1:])) - assert run_for(1.0, lambda: check_region(a, values), access) == [] +def resize(a, k): + # grow and shrink past the region check_region() reads + a.resize((SHAPE[0] + 8 * (k % 2), *SHAPE[1:])) + + +@pytest.mark.parametrize( + "access", [update_chunks, write_vlmeta, swap_dparams, resize], ids=lambda f: f.__name__ +) +def test_read_while_other_access_same_handle(access): + a = make_array() + k = itertools.count() + assert run_for(SECONDS, lambda: check_region(a), lambda: access(a, next(k))) == [] def test_concurrent_reads_with_python_postfilter(): - values = np.arange(np.prod(SHAPE), dtype=np.float32).reshape(SHAPE) - a = blosc2.asarray(values, chunks=CHUNKS, blocks=BLOCKS, dparams={"nthreads": 1}) + a = make_array(dparams={"nthreads": 1}) @a.schunk.postfilter(np.float32, np.float32) def add_one(inp, out, offset): out[:] = inp + 1 def read(): - np.testing.assert_array_equal(a[8:24, 10:200], values[8:24, 10:200] + 1) + np.testing.assert_array_equal(a[8:24, 10:200], VALUES[8:24, 10:200] + 1) + + assert run_for(SECONDS, read, read) == [] + - assert run_for(1.0, read, read) == [] +def test_lazy_expression_while_reading_operands(tmp_path): + # The miniexpr prefilter reads on-disk operands from C, so it must hold their locks too + a = make_array(urlpath=tmp_path / "a.b2nd", mode="w") + b = make_array(urlpath=tmp_path / "b.b2nd", mode="w") + + def compute(): + np.testing.assert_array_equal((a + b).compute()[8:24], 2 * VALUES[8:24]) + + assert run_for(SECONDS, compute, compute, lambda: check_region(a), lambda: check_region(b)) == [] def test_concat_both_orders_does_not_deadlock(): a = blosc2.asarray(np.zeros((16, 64), dtype=np.int32), chunks=(8, 64)) b = blosc2.asarray(np.ones((16, 64), dtype=np.int32), chunks=(8, 64)) errors = run_for( - 1.0, + SECONDS, lambda: blosc2.concat([a, b], axis=0), lambda: blosc2.concat([b, a], axis=0), lambda: a[:], lambda: b[:], ) assert errors == [] + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +@pytest.mark.filterwarnings("ignore::DeprecationWarning") # fork() with threads running +def test_fork_while_another_thread_holds_the_lock(): + # As when forking while another thread reads without the GIL: the child must not inherit + # a lock that no thread of its own will ever release + a = make_array(dparams={"nthreads": 1}) + held, release = threading.Event(), threading.Event() + + def hold(): + with a._handle_lock: + held.set() + release.wait() + + t = threading.Thread(target=hold) + t.start() + try: + assert held.wait(timeout=10) + pid = os.fork() + if pid == 0: + ok = False + try: + ok = np.array_equal(a[8:24, 10:200], VALUES[8:24, 10:200]) + finally: + os._exit(0 if ok else 1) + deadline = time.monotonic() + 30 + while (status := os.waitpid(pid, os.WNOHANG))[0] == 0: + if time.monotonic() > deadline: + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + pytest.fail("the child hung on a lock held by a thread of the parent") + time.sleep(0.01) + assert os.waitstatus_to_exitcode(status[1]) == 0 + finally: + release.set() + t.join() From 8fee90badd1d224feaab539e483f2db263c705b2 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:39:51 +0100 Subject: [PATCH 03/12] Fix crashes with frame-backed operands in miniexpr and matmul The miniexpr prefilter read the chunks of every in-memory operand from sc->data, and the fast matmul prefilter did so for every operand. Frames have no chunk table there: in-memory arrays created with contiguous=True, arrays from ndarray_from_cframe() (copy=False, the default) and, for matmul, any on-disk array. Any lazy expression or matmul with such an operand dereferenced a NULL table and crashed the interpreter. Read frame-backed in-memory operands through blosc2_schunk_get_chunk(), as on-disk ones already are, and keep frame-backed operands off the fast matmul path: in its Python checks, and in _set_pref_matmul() itself, which now raises ValueError for them before changing any state. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/blosc2_ext.pyx | 11 ++++++++++- src/blosc2/linalg.py | 5 +++++ tests/ndarray/test_lazyexpr.py | 16 ++++++++++++++++ tests/ndarray/test_linalg.py | 28 ++++++++++++++++++++++++++++ 4 files changed, 59 insertions(+), 1 deletion(-) diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 1ccf65d5b..d43eab991 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -2961,7 +2961,9 @@ cdef int aux_miniexpr(me_udata *udata, int64_t nchunk, int32_t nblock, continue ndarr = udata.inputs[i] - if ndarr.sc.storage.urlpath == NULL: + # sc.data only holds the chunks of in-memory super-chunks that are not frames + # (contiguous=True or from a cframe); read frames through blosc2_schunk_get_chunk() + if ndarr.sc.storage.urlpath == NULL and ndarr.sc.frame == NULL: src = ndarr.sc.data[nchunk] else: input_cache = &udata.input_chunk_caches[i] @@ -4751,6 +4753,13 @@ cdef class NDArray: @_handle_locked def _set_pref_matmul(self, inputs, fp_accuracy): + cdef b2nd_array_t* operand + for name in ("x1", "x2"): + operand = inputs[name].c_array + # aux_matmul() reads the chunks from sc.data, which frames (on disk, or in memory + # with contiguous=True) do not fill + if operand.sc.frame != NULL: + raise ValueError("The fast matmul path does not support frame-backed operands") # Set prefilter for miniexpr cdef blosc2_cparams* cparams = self.array.sc.storage.cparams cparams.prefilter = matmul_prefilter diff --git a/src/blosc2/linalg.py b/src/blosc2/linalg.py index f6b864395..1dd4fd14a 100644 --- a/src/blosc2/linalg.py +++ b/src/blosc2/linalg.py @@ -77,6 +77,11 @@ def _matmul_can_use_fast_path( if not all_ndarray: return False + # The prefilter reads the operands' chunks straight from memory, which frames (on disk, + # or in memory with contiguous=True) do not keep + if any(op.schunk.urlpath is not None or op.schunk.contiguous for op in (x1, x2)): + return False + # The current prefilter-backed implementation is only supported for 2-D layouts. if result.ndim != 2 or x1.ndim != 2 or x2.ndim != 2: return False diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 5a099cbac..763d4de15 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -2406,3 +2406,19 @@ def test_cumulative_block_layout(monkeypatch, blocks): expected = np.cumulative_sum(data, axis=0) np.testing.assert_array_equal(result[:], expected) np.testing.assert_array_equal(out[:], expected) + + +@pytest.mark.parametrize("source", ["contiguous", "cframe"]) +def test_frame_backed_inmemory_operands(source): + # In-memory frames (contiguous=True, or views of a cframe) keep no chunk table, so the + # miniexpr prefilter must fetch their chunks like it does for on-disk operands + na = np.arange(270, dtype=np.int16).reshape(10, 27) + if source == "contiguous": + a = blosc2.asarray(na, chunks=(3, 4), blocks=(1, 2), contiguous=True) + else: + a = blosc2.ndarray_from_cframe(blosc2.asarray(na, chunks=(3, 4), blocks=(1, 2)).to_cframe()) + assert a.schunk.contiguous + assert a.schunk.urlpath is None + b = blosc2.asarray(na * 2, chunks=(3, 4), blocks=(1, 2)) + np.testing.assert_array_equal((a + 1)[:], na + 1) + np.testing.assert_array_equal((a + b).compute()[:], na * 3) diff --git a/tests/ndarray/test_linalg.py b/tests/ndarray/test_linalg.py index 2827f8a23..a93b0b1c2 100644 --- a/tests/ndarray/test_linalg.py +++ b/tests/ndarray/test_linalg.py @@ -1283,3 +1283,31 @@ def test_matmul_broadcast_batch_chunks(): blosc2.asarray(b, chunks=(1, 2, 2)), blosc2.asarray(a, chunks=(1, 2, 2)), chunks=(1, 2, 2) ) np.testing.assert_allclose(res[:], np.matmul(b, a)) + + +@pytest.mark.parametrize("storage", ["contiguous", "disk"]) +def test_matmul_frame_backed_operands(storage, tmp_path): + # The fast matmul prefilter reads in-memory chunk tables; frames must take the general path + x = np.arange(64 * 64, dtype=np.float32).reshape(64, 64) / 100 + y = x.T.copy() + if storage == "contiguous": + a, b = blosc2.asarray(x, contiguous=True), blosc2.asarray(y, contiguous=True) + else: + a = blosc2.asarray(x, urlpath=str(tmp_path / "x.b2nd"), mode="w") + b = blosc2.asarray(y, urlpath=str(tmp_path / "y.b2nd"), mode="w") + np.testing.assert_allclose(blosc2.matmul(a, b)[:], x @ y, rtol=1e-5) + + +@pytest.mark.parametrize("storage", ["contiguous", "disk"]) +def test_set_pref_matmul_rejects_frame_backed_operands(storage, tmp_path): + # The prefilter itself refuses frames too, rather than leaving that to matmul's Python checks + x = np.ones((64, 64), dtype=np.float32) + if storage == "contiguous": + a = blosc2.asarray(x, contiguous=True) + else: + a = blosc2.asarray(x, urlpath=str(tmp_path / "x.b2nd"), mode="w") + b = blosc2.asarray(x) + out = blosc2.zeros((64, 64), dtype=np.float32) + for inputs in ({"x1": a, "x2": b}, {"x1": b, "x2": a}): + with pytest.raises(ValueError, match="frame-backed"): + out._set_pref_matmul(inputs, fp_accuracy=blosc2.FPAccuracy.DEFAULT) From 0677313b76e2123af90b92c852b1bc1847a67160 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:40:36 +0100 Subject: [PATCH 04/12] Make the general NDArray.__getitem__ path accept the keys NumPy accepts The general indexing path (ndindex) raised, or returned a wrong shape, for keys that NumPy indexes fine: 554 of 18000 random keys in a comparison with NumPy. Fix the classes found: - slice bounds beyond the index range, and bool bounds or steps, next to ..., None or NumPy integers raised IndexError or TypeError; normalize them as NumPy reads them (clamped, bools as ints) before ndindex sees them; - 0-d integer arrays index as integers (AttributeError before); - scalar bools are 0-d boolean indices, as in NumPy: their dimension goes among adjacent advanced indices (integers included), else in front, and next to array indices they broadcast (True: no effect, False: an empty result or IndexError). The result shape comes from NumPy on a stride-0 stand-in for the array, so it costs no memory; - None next to (adjacent) array indices is supported; - 0-d arrays take NumPy scalar bools (AttributeError before) and raise IndexError for integer lists (ValueError before). Advanced indices separated by slices, ... or None still raise NotImplementedError (NumPy moves their dimensions to the front there). After this, none of the 18000 keys gives another result than NumPy; the remaining differences are that limitation and which error a key that is invalid for several reasons reports first. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/ndarray.py | 134 +++++++++++++++++----- tests/ndarray/test_getitem.py | 203 ++++++++++++++++++++++++++++++++++ 2 files changed, 309 insertions(+), 28 deletions(-) diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 98ab50b0a..e2f9a7343 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -11,6 +11,8 @@ import inspect import itertools import math +import operator +import sys import weakref from abc import abstractmethod from collections import OrderedDict, namedtuple @@ -245,6 +247,42 @@ def is_inside_new_expr() -> bool: return builtins.any(frame_info.function in {"_new_expr", "open_lazyarray"} for frame_info in stack) +def _normalize_slice(s): + """Return slice s with bool or out-of-range integer bounds and step as NumPy reads them. + + NumPy (like slice.indices()) accepts any integers there, clamping them to the index range, + and bools as integers; ndindex rejects both. Slices without such values, or with a value + that is not an integer at all, are returned unchanged. + """ + values, changed = [], False + for v, lo in ((s.start, -sys.maxsize - 1), (s.stop, -sys.maxsize - 1), (s.step, -sys.maxsize)): + if v is not None: + try: + i = operator.index(v) + except TypeError: + return s + if isinstance(v, bool) or not lo <= i <= sys.maxsize: + v, changed = builtins.min(builtins.max(int(i), lo), sys.maxsize), True + values.append(v) + return slice(*values) if changed else s + + +def _normalize_key_item(k): + if isinstance(k, slice): + return _normalize_slice(k) + # NumPy indexes with a 0-d integer array as with an integer + if isinstance(k, np.ndarray) and k.ndim == 0 and np.issubdtype(k.dtype, np.integer): + return int(k) + return k + + +def _normalize_key(key): + """Turn the parts of an indexing key that NumPy accepts but ndindex does not into equivalents.""" + if isinstance(key, tuple): + return tuple(_normalize_key_item(k) for k in key) + return _normalize_key_item(key) + + def _is_scalar_bool(k): """True for Python/numpy scalar bools and 0-d bool arrays (numpy's scalar-bool indexing).""" return isinstance(k, (bool, np.bool_)) or ( @@ -252,6 +290,23 @@ def _is_scalar_bool(k): ) +def _check_adjacent_advanced(key, count_bools=False): + """Raise NotImplementedError unless the advanced indices in the tuple `key` (arrays, integers + and, with `count_bools`, scalar bools) are adjacent: NumPy moves the broadcast dimensions to + the front otherwise.""" + adv = [ + i + for i, k in enumerate(key) + if (count_bools and _is_scalar_bool(k)) + or isinstance(k, list | np.ndarray) + or (isinstance(k, int | np.integer) and not isinstance(k, bool)) + ] + if adv[-1] - adv[0] + 1 != len(adv): + raise NotImplementedError( + "Array indices separated by slices, ellipses (...), or newaxes (None) are not supported" + ) + + def make_key_hashable(key): if isinstance(key, slice): return (key.start, key.stop, key.step) @@ -4545,33 +4600,52 @@ def _try_sparse_fancy_index(self, key) -> np.ndarray | None: if not isinstance(key, (list, np.ndarray)): return None key_arr = np.asarray(key) - if not (np.issubdtype(key_arr.dtype, np.integer) and key_arr.ndim >= 1): + if not (np.issubdtype(key_arr.dtype, np.integer) and key_arr.ndim >= 1) or self.ndim == 0: return None # 1-D: axis=None (flat); ndim>1: axis=0 (row selection) return self._take_numpy(key_arr, axis=None if self.ndim == 1 else 0) - def _scalar_bool_key_to_newaxis(self, key): - """Rewrite a tuple key's scalar bools to numpy's single-newaxis semantics. + def _drop_scalar_bools(self, key): + """Handle the scalar bools of a tuple key like NumPy. - Scalar bools act as one np.newaxis-like dimension of length 1 (all True) - or 0 (any False), placed at the first bool's position. Returns - ``(key, None)`` unchanged when there are no scalar bools, ``(new_key, - None)`` to continue processing, or ``(None, empty_shape)`` when any bool - is False and the result is an empty array of that shape. + NumPy treats them as 0-d boolean (advanced) indices: together they add one + dimension of length 1 (all True) or 0 (any False), which goes where the + advanced indices (integers included) are if those are adjacent, else in + front. Returns ``(key, None)`` unchanged when there are no scalar bools, + ``(rest, shape)`` when the result is ``self[rest]`` (the key without the + bools) reshaped to ``shape``, ``(None, shape)`` when the result is an empty + array of that shape, or ``(rest, None)`` to continue processing with rest. """ if not builtins.any(_is_scalar_bool(k) for k in key): return key, None - first = next(i for i, k in enumerate(key) if _is_scalar_bool(k)) - new_key = key[:first] + (None,) + tuple(k for k in key[first + 1 :] if not _is_scalar_bool(k)) - if builtins.all(k for k in key if _is_scalar_bool(k)): # all True -> plain newaxis - return new_key, None - # any False -> the newaxis dim is empty; shape from the all-True form - key_idx = ndindex.ndindex(new_key).expand(self.shape) - shape = list(key_idx.newshape(self.shape)) - raw = key_idx.raw - pos = builtins.sum(not isinstance(k, int) for k in raw[: raw.index(None)]) - shape[pos] = 0 - return None, tuple(shape) + rest = tuple(k for k in key if not _is_scalar_bool(k)) + all_true = builtins.all(k for k in key if _is_scalar_bool(k)) + # NumPy works out the shape (and validates the key) on a stride-0 stand-in for self; that + # costs no memory for basic indices, nor for an empty result + stand_in = np.broadcast_to(np.empty((), dtype=np.uint8), self.shape) + if not builtins.any(isinstance(k, list | np.ndarray) for k in rest): + shape = stand_in[key].shape + return (rest, shape) if all_true else (None, shape) + # Together with array indices, a False empties the broadcast result (or does not broadcast) + if not all_true: + return None, stand_in[key].shape + # and True bools broadcast away, unless they separate the advanced indices + _check_adjacent_advanced(key, count_bools=True) + return rest, None + + def _getitem_fancy_newaxis(self, key): + """Fancy indexing with np.newaxis in a tuple key. + + Index without the Nones and insert their dimensions where NumPy puts them. + The advanced indices (arrays and integers) must be adjacent, as without + Nones. + """ + _check_adjacent_advanced(key) + nparr = self[tuple(k for k in key if k is not None)] + # NumPy places the new axes on a stride-0 stand-in for self; an extra trailing dimension of + # length 0 keeps its result empty, so this costs no memory + stand_in = np.broadcast_to(np.empty((), dtype=np.uint8), (*self.shape, 0)) + return nparr.reshape(stand_in[(*key, slice(None))].shape[:-1]) def _getitem_bool_mask(self, key): """Handle boolean array key with optional sparse-gather fast path. @@ -4581,6 +4655,8 @@ def _getitem_bool_mask(self, key): """ if not (hasattr(key, "dtype") and np.issubdtype(key.dtype, np.bool_) and key.shape == self.shape): return None + if isinstance(key, np.generic): # a scalar bool, whose shape matches a 0-d array's + return None # For sparse boolean masks, converting to flat indices and using the # sparse-gather path is faster than decompressing every data chunk. try: @@ -4663,17 +4739,17 @@ def __getitem__( # noqa: C901 key = key[()] if isinstance(key, NDArray) else key # key not iterable key = tuple(k[()] if isinstance(k, NDArray) else k for k in key) if isinstance(key, tuple) else key + key = _normalize_key(key) - # Scalar bools inside a tuple act like a single np.newaxis dimension of - # length 1 (all True) or 0 (any False), placed at the first bool's - # position (numpy semantics); rewrite to the equivalent None form so the - # regular machinery handles them. Bare scalar bools are handled below in - # the fancy branch. (numpy rejects scalar-bool tuples in __setitem__, so - # this stays getitem-only.) + # Scalar bools inside a tuple are 0-d boolean indices, as in NumPy (see + # _drop_scalar_bools). Bare scalar bools are handled below in the fancy + # branch. (__setitem__ does not support them.) if isinstance(key, tuple): - key, empty_shape = self._scalar_bool_key_to_newaxis(key) - if empty_shape is not None: - return np.empty(empty_shape, dtype=self.dtype) + key, bool_shape = self._drop_scalar_bools(key) + if key is None: + return np.empty(bool_shape, dtype=self.dtype) + if bool_shape is not None: + return self[key].reshape(bool_shape) # Check boolean array key early to avoid expensive process_key / nonzero result = self._getitem_bool_mask(key) @@ -4703,6 +4779,8 @@ def __getitem__( # noqa: C901 return np.expand_dims(self._get_set_findex_default(_slice, out=out), 0) else: # do nothing return np.empty((0,) + self.shape, dtype=self.dtype) + if isinstance(key, tuple) and builtins.any(k is None for k in key): + return self._getitem_fancy_newaxis(key) return self.get_fselection_numpy(key) # fancy index default, can be quite slow start, stop, step, none_mask = get_ndarray_start_stop(self.ndim, key_, self.shape) diff --git a/tests/ndarray/test_getitem.py b/tests/ndarray/test_getitem.py index 437cab4f5..567d56ed9 100644 --- a/tests/ndarray/test_getitem.py +++ b/tests/ndarray/test_getitem.py @@ -912,3 +912,206 @@ def test_scalar_bool_key_matches_numpy(shape, key): assert got_exc == want_exc if want_exc is None: np.testing.assert_array_equal(got, want) + + +def _assert_getitem_matches_numpy(npa, a, key): + try: + want = npa[key] + except Exception as e: + with pytest.raises(type(e)): + a[key] + return + got = a[key] + assert got.shape == np.shape(want), key + np.testing.assert_array_equal(got, want) + + +# NumPy clamps slice bounds beyond the index range and reads bool bounds/steps as ints; next to ..., +# None or NumPy integers (the ndindex path) these used to raise +@pytest.mark.parametrize( + "key", + [ + (slice(-(2**70), 2**64), np.int64(1)), + (Ellipsis, slice(2**64, None, -1)), + (None, slice(None, 2**64, 2)), + (slice(3, -(2**70), -1), np.uint8(0)), + (slice(1, None, True), Ellipsis), + (None, slice(True, 3)), + slice(-(2**70), np.int32(-1), 3), + (slice(None, None, 2**70), None), + (slice(-(2**70), -2, 0), np.int64(0)), # step 0: ValueError, as in NumPy + (slice(1.5, 3), None), # not an integer: TypeError, as in NumPy + ], +) +def test_getitem_huge_slice_bounds_and_bool_steps(key): + npa = np.arange(20).reshape(5, 4) + _assert_getitem_matches_numpy(npa, blosc2.asarray(npa), key) + + +# NumPy indexes with 0-d integer arrays as with integers; they used to raise AttributeError +@pytest.mark.parametrize( + "key", + [ + np.array(2), + np.array(-1, dtype=np.int8), + (None, slice(None), np.array(2), np.int64(1)), + (np.array(1), Ellipsis), + (slice(1, 3), np.array(np.uint16(0)), None), + (np.array(5),), # out of range: IndexError, as in NumPy + ], +) +def test_getitem_0d_integer_array(key): + npa = np.arange(60).reshape(3, 4, 5) + _assert_getitem_matches_numpy(npa, blosc2.asarray(npa), key) + + +# NumPy places the dimension of scalar bools among the other advanced indices (integers count as +# such) when they are adjacent, else in front; also after an explicit None. Shapes used to differ. +@pytest.mark.parametrize( + "key", + [ + (None, False, slice(-1, 0)), + (None, np.False_, None), + (None, True, 1, None), + (0, Ellipsis, False), + (0, Ellipsis, True), + (1, slice(None), True), + (True, slice(None), 0), + (slice(None), True, slice(None), np.True_), + (Ellipsis, 0, True), + (True, 7), # out of range: IndexError, as in NumPy + ], +) +def test_getitem_scalar_bool_placement(key): + npa = np.arange(24).reshape(2, 3, 4) + _assert_getitem_matches_numpy(npa, blosc2.asarray(npa), key) + + +# Scalar bools next to array indices broadcast with them in NumPy (True: no effect, False: empty +# result or IndexError); they used to give AttributeError or a wrong shape +@pytest.mark.parametrize( + "key", + [ + (True, [0, 1]), + ([0, 1], np.True_), + ([[0], [1]], True, [1, 2]), + (False, [0, 1]), # shapes (0,) and (2,) do not broadcast: IndexError, as in NumPy + (False, [0]), + (slice(None), [0], np.False_), + (np.array([True, False]), False), + (False, slice(None), [0, 1]), + ], +) +def test_getitem_scalar_bool_with_array_index(key): + npa = np.arange(24).reshape(2, 3, 4) + _assert_getitem_matches_numpy(npa, blosc2.asarray(npa), key) + + +def test_getitem_scalar_bool_separating_array_index(): + # NumPy moves the broadcast dimensions to the front here: a documented limitation + a = blosc2.asarray(np.arange(24).reshape(2, 3, 4)) + with pytest.raises(NotImplementedError): + a[True, :, [0, 1]] + + +# np.newaxis next to array indices used to raise AttributeError +@pytest.mark.parametrize( + "key", + [ + (None, [1, 0]), + ([1, 0], None), + (None, slice(None), [1, 0]), + (slice(None), [1, 0], None, slice(1, 3)), + (np.array([0, 1]), None), + (None, [0, 1], [1, 2], None), + (None, 1, [0, 2]), + (Ellipsis, None, [[0], [1]]), + (None, np.array([True, False]), Ellipsis, None), + (None, [0, 5]), # out of range: IndexError, as in NumPy + ], +) +def test_getitem_newaxis_with_array_index(key): + npa = np.arange(24).reshape(2, 3, 4) + _assert_getitem_matches_numpy(npa, blosc2.asarray(npa), key) + + +def test_getitem_newaxis_separating_array_index(): + # NumPy moves the broadcast dimensions to the front here: a documented limitation + a = blosc2.asarray(np.arange(24).reshape(2, 3, 4)) + with pytest.raises(NotImplementedError): + a[[0, 1], None, [1, 2]] + + +# Keys for 0-d arrays: NumPy scalar bools used to raise AttributeError, integer lists ValueError +@pytest.mark.parametrize( + "key", [np.True_, np.False_, True, np.array(False), (np.True_, None), [0], [1, 0], (Ellipsis, np.False_)] +) +def test_getitem_0d_array_scalar_bools_and_lists(key): + npa = np.array(3.5) + _assert_getitem_matches_numpy(npa, blosc2.asarray(npa), key) + + +def _random_numpy_key(rng, ndim, n): + """A random key from the classes the general path used to get wrong (see the tests above): + huge slice bounds, bool steps, 0-d integer arrays, scalar bools, and None next to integers + and array indices. All integers are valid for a dimension of size n.""" + items, nidx, ellipsis = [], 0, False + bounds = [None, None, -n - 1, -1, 0, 1, n, 2**64, -(2**64), 2**70] + for _ in range(int(rng.integers(0, ndim + 3))): + r = rng.random() + if r < 0.25 and nidx < ndim: + step = [None, 1, True, np.int64(1), 2, -1][int(rng.integers(6))] + items.append( + slice(bounds[int(rng.integers(len(bounds)))], bounds[int(rng.integers(len(bounds)))], step) + ) + nidx += 1 + elif r < 0.45 and nidx < ndim: + i = int(rng.integers(-n, n)) + items.append([i, np.int64(i), np.array(i)][int(rng.integers(3))]) + nidx += 1 + elif r < 0.6: + items.append(None) + elif r < 0.75: + items.append([True, False, np.True_, np.False_][int(rng.integers(4))]) + elif r < 0.85 and not ellipsis: + items.append(Ellipsis) + ellipsis = True + elif nidx < ndim: + items.append(rng.integers(-n, n, size=int(rng.integers(1, 3))).tolist()) + nidx += 1 + return tuple(items) + + +def _separates_advanced_indices(key): + # Array indices (with integers and scalar bools) not adjacent: a documented limitation + adv = [i for i, k in enumerate(key) if k is not None and k is not Ellipsis and not isinstance(k, slice)] + has_array = any(isinstance(k, list) for k in key) + return has_array and adv[-1] - adv[0] + 1 != len(adv) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_getitem_random_keys_match_numpy(tmp_path, on_disk): + rng = np.random.default_rng(3) + for ishape, shape in enumerate([(), (5,), (3, 3), (3, 3, 3), (2, 2, 2, 2)]): + npa = np.arange(math.prod(shape), dtype=np.float32).reshape(shape) + kwargs = {"urlpath": tmp_path / f"rnd{ishape}.b2nd", "mode": "w"} if on_disk else {} + a = blosc2.asarray(npa, **kwargs) + for _ in range(300): + key = _random_numpy_key(rng, len(shape), shape[0] if shape else 1) + try: + want = npa[key] + except Exception as e: + want_exc = type(e) + else: + want_exc = None + try: + got, got_exc = a[key], None + except Exception as e: + got, got_exc = None, type(e) + if got_exc is NotImplementedError: + assert _separates_advanced_indices(key), key + continue + assert got_exc is want_exc, key + if want_exc is None: + assert got.shape == np.shape(want), key + np.testing.assert_array_equal(got, want) From 4301191fd0832d064a4010d14ce88c316cc893e1 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:43:56 +0100 Subject: [PATCH 05/12] Fix SChunk.update_special(), failed data writes and decompress_chunk() - update_special() built its special chunk in a temporary SChunk with default cparams, so on an NDArray's super-chunk the chunk header had another block size than the array (e.g. 20000 instead of 800 bytes) and every later read of the array failed: b2nd reads blocks selectively and checks their number per chunk. The new SChunk._special_chunk() makes zero, NaN, uninitialized and repeated-value chunks with the super-chunk's own cparams. - update_data() and insert_data() with copy=False did not free their compressed chunk when the C call failed (nor on a compression error), so it leaked wherever c-blosc2 did not free it itself. They now free it whenever the call fails. This needs the next c-blosc2 release (or the companion c-blosc2 PR) for its caller-keeps-the-chunk-on-failure rule; with an older c-blosc2, a failed frame update can free the chunk twice. - decompress_chunk() without dst fetched the whole compressed chunk only to read its size from the header, then read it again to decompress it; it now gets the lazy chunk (header and block offsets) for the size. astype(float64) of a 300x512x512 float32 array on disk: 385 -> 366 ms. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/blosc2_ext.pyx | 93 ++++++++++++++++++++++++++----------- src/blosc2/schunk.py | 5 +- tests/test_schunk.py | 21 +++++++++ tests/test_schunk_update.py | 36 ++++++++++++++ 4 files changed, 126 insertions(+), 29 deletions(-) diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index d43eab991..477b378bf 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -189,6 +189,7 @@ cdef extern from "blosc2.h": ctypedef enum: BLOSC2_SPECIAL_ZERO BLOSC2_SPECIAL_NAN + BLOSC2_SPECIAL_VALUE BLOSC2_SPECIAL_UNINIT ctypedef enum: @@ -516,6 +517,9 @@ cdef extern from "blosc2.h": int blosc2_schunk_reorder_offsets(blosc2_schunk *schunk, int64_t *offsets_order) int64_t blosc2_schunk_frame_len(blosc2_schunk* schunk) + int blosc2_chunk_zeros(blosc2_cparams cparams, const int32_t nbytes, void *dest, int32_t destsize) + int blosc2_chunk_nans(blosc2_cparams cparams, const int32_t nbytes, void *dest, int32_t destsize) + int blosc2_chunk_uninit(blosc2_cparams cparams, const int32_t nbytes, void *dest, int32_t destsize) int blosc2_chunk_repeatval(blosc2_cparams cparams, const int32_t nbytes, void *dest, int32_t destsize, const void *repeatval) @@ -1710,6 +1714,20 @@ cdef create_storage(blosc2_storage *storage, kwargs): storage.io = NULL +cdef _repeatval_array(value, int32_t typesize): + """`value` as an array of one item of `typesize` bytes, for repeated-value chunks.""" + array = np.array([value]) + if array.dtype.itemsize != typesize: + if isinstance(value, int): + dtype = np.dtype('i'+ str(typesize)) + elif isinstance(value, float): + dtype = np.dtype('f' + str(typesize)) + else: + raise ValueError("value size in bytes must match with typesize") + array = np.array([value], dtype=dtype) + return array + + cdef get_chunk_repeatval(blosc2_cparams cparams, const int32_t nbytes, void *dest, int32_t destsize, Py_buffer *repeatval): if blosc2_chunk_repeatval(cparams, nbytes, dest, destsize, repeatval.buf) < 0: @@ -2120,15 +2138,7 @@ cdef class SChunk: if self.nbytes > 0 or self.cbytes > 0: raise RuntimeError("Filling with special values only works on empty SChunks") # Get a void pointer to the value - array = np.array([value]) - if array.dtype.itemsize != self.typesize: - if isinstance(value, int): - dtype = np.dtype('i'+ str(self.typesize)) - elif isinstance(value, float): - dtype = np.dtype('f' + str(self.typesize)) - else: - raise ValueError("value size in bytes must match with typesize") - array = np.array([value], dtype=dtype) + array = _repeatval_array(value, self.typesize) cdef Py_buffer buf PyObject_GetBuffer(array, &buf, PyBUF_SIMPLE) # Create chunk with repeated values @@ -2157,22 +2167,44 @@ cdef class SChunk: return self.nchunks + @_handle_locked + def _special_chunk(self, int32_t nbytes, int special_value, value=None): + """Return a special chunk of `nbytes` bytes (`special_value` is a SpecialValue value; VALUE + repeats `value`, as in fill_special()). It is made with the cparams of this super-chunk, + so that its header has the same block size as the other chunks: b2nd reads blocks + selectively and checks the number of blocks per chunk.""" + cdef blosc2_cparams *cparams = self.schunk.storage.cparams + cdef int32_t destsize = BLOSC_EXTENDED_HEADER_LENGTH + cparams.typesize + cdef uint8_t *dest = malloc(destsize) + cdef Py_buffer buf + cdef int rc + try: + if special_value == BLOSC2_SPECIAL_ZERO: + rc = blosc2_chunk_zeros(cparams[0], nbytes, dest, destsize) + elif special_value == BLOSC2_SPECIAL_NAN: + rc = blosc2_chunk_nans(cparams[0], nbytes, dest, destsize) + elif special_value == BLOSC2_SPECIAL_UNINIT: + rc = blosc2_chunk_uninit(cparams[0], nbytes, dest, destsize) + elif special_value == BLOSC2_SPECIAL_VALUE: + array = _repeatval_array(value, self.typesize) + PyObject_GetBuffer(array, &buf, PyBUF_SIMPLE) + rc = blosc2_chunk_repeatval(cparams[0], nbytes, dest, destsize, buf.buf) + PyBuffer_Release(&buf) + else: + raise ValueError(f"Not a special value: {special_value}") + if rc < 0: + raise RuntimeError("Could not create the special chunk") + return PyBytes_FromStringAndSize( dest, rc) + finally: + free(dest) + @_handle_locked def decompress_chunk(self, nchunk, dst=None): cdef uint8_t *chunk cdef c_bool needs_free - rc = blosc2_schunk_get_chunk(self.schunk, nchunk, &chunk, &needs_free) - - if rc < 0: - raise RuntimeError("Error while getting the chunk") - cdef int32_t nbytes cdef int32_t cbytes cdef int32_t blocksize - blosc2_cbuffer_sizes(chunk, &nbytes, &cbytes, &blocksize) - if needs_free: - free(chunk) - cdef Py_buffer buf if dst is not None: PyObject_GetBuffer(dst, &buf, PyBUF_SIMPLE) @@ -2181,6 +2213,13 @@ cdef class SChunk: size = blosc2_schunk_decompress_chunk(self.schunk, nchunk, buf.buf, buf.len) PyBuffer_Release(&buf) else: + # Only the header is needed: a lazy chunk does not read the whole chunk from a frame twice + rc = blosc2_schunk_get_lazychunk(self.schunk, nchunk, &chunk, &needs_free) + if rc < 0: + raise RuntimeError("Error while getting the chunk") + blosc2_cbuffer_sizes(chunk, &nbytes, &cbytes, &blocksize) + if needs_free: + free(chunk) dst = PyBytes_FromStringAndSize(NULL, nbytes) if dst is None: raise RuntimeError("Could not get a bytes object") @@ -2283,16 +2322,17 @@ cdef class SChunk: else: size = blosc2_compress_ctx(self.schunk.cctx, buf.buf, buf.len, chunk, len_chunk) PyBuffer_Release(&buf) - if size < 0: - raise RuntimeError("Could not compress the data") - elif size == 0: + if size <= 0: free(chunk) + if size < 0: + raise RuntimeError("Could not compress the data") raise RuntimeError("The result could not fit ") chunk = realloc(chunk, size) _check_comp_length('chunk', size) rc = blosc2_schunk_insert_chunk(self.schunk, nchunk, chunk, copy) - if copy: + # Without a copy, the schunk owns the chunk only if the insert succeeded + if copy or rc < 0: free(chunk) if rc < 0: raise RuntimeError("Could not insert the desired chunk") @@ -2332,16 +2372,17 @@ cdef class SChunk: size = blosc2_compress_ctx(self.schunk.cctx, buf.buf, buf.len, chunk, len_chunk) PyBuffer_Release(&buf) - if size < 0: - raise RuntimeError("Could not compress the data") - elif size == 0: + if size <= 0: free(chunk) + if size < 0: + raise RuntimeError("Could not compress the data") raise RuntimeError("The result could not fit ") chunk = realloc(chunk, size) _check_comp_length('chunk', size) rc = blosc2_schunk_update_chunk(self.schunk, nchunk, chunk, copy) - if copy: + # Without a copy, the schunk owns the chunk only if the update succeeded + if copy or rc < 0: free(chunk) if rc < 0: raise RuntimeError("Could not update the desired chunk") diff --git a/src/blosc2/schunk.py b/src/blosc2/schunk.py index aa058e044..ff69826a5 100644 --- a/src/blosc2/schunk.py +++ b/src/blosc2/schunk.py @@ -815,13 +815,12 @@ def update_special( total_items = self.nbytes // self.typesize nitems = min(chunk_items, total_items - nchunk * chunk_items) - tmp = SChunk(chunksize=self.chunksize, cparams=blosc2.CParams(typesize=self.typesize)) - tmp.fill_special(nitems, special_value, value) + chunk = super()._special_chunk(nitems * self.typesize, special_value.value, value) # A special chunk is indistinguishable from one that was never written, so # a reader that keeps its own record of what it has (:ref:`Proxy` does) # cannot see this happen. Counting it lets such a reader notice for free. self.nspecialized += 1 - nchunks = self.update_chunk(nchunk, tmp.get_chunk(0)) + nchunks = self.update_chunk(nchunk, chunk) self._forget_fetched(nchunk) return nchunks diff --git a/tests/test_schunk.py b/tests/test_schunk.py index aed61f499..03aa802d7 100644 --- a/tests/test_schunk.py +++ b/tests/test_schunk.py @@ -349,6 +349,27 @@ def test_schunk_update_special(): schunk2.update_special(5, blosc2.SpecialValue.UNINIT) +@pytest.mark.parametrize( + ("special", "value", "expected"), + [ + (blosc2.SpecialValue.ZERO, None, 0.0), + (blosc2.SpecialValue.NAN, None, np.nan), + (blosc2.SpecialValue.VALUE, np.float64(7.5).tobytes(), 7.5), + (blosc2.SpecialValue.VALUE, 7.5, 7.5), + ], +) +def test_schunk_update_special_ndarray(special, value, expected): + # The special chunk must have the block size of the array (b2nd reads blocks selectively, and + # checks their number per chunk), not the one a plain SChunk would choose + x = np.arange(100 * 100, dtype=np.float64).reshape(100, 100) + a = blosc2.asarray(x, chunks=(50, 50), blocks=(10, 10)) + a.schunk.update_special(1, special, value) + assert blosc2.get_cbuffer_sizes(a.schunk.get_chunk(1))[2] == 10 * 10 * 8 + x[:50, 50:] = expected + np.testing.assert_array_equal(a[:], x) + np.testing.assert_array_equal(a[2:5, 48:53], x[2:5, 48:53]) + + def test_get_lazychunk_sections(tmp_path): """A lazy chunk keeps its block offsets and trailer; other chunks stay capped.""" urlpath = tmp_path / "lazy.b2nd" diff --git a/tests/test_schunk_update.py b/tests/test_schunk_update.py index 4115c6b0b..7391a6bb1 100644 --- a/tests/test_schunk_update.py +++ b/tests/test_schunk_update.py @@ -5,7 +5,9 @@ # SPDX-License-Identifier: BSD-3-Clause ####################################################################### +import os import random +import sys import numpy as np import pytest @@ -143,3 +145,37 @@ def test_append_chunk_variable_sizes(contiguous, urlpath): assert [reopened.decompress_chunk(i) for i in range(reopened.nchunks)] == expected blosc2.remove_urlpath(urlpath) + + +def _rss_bytes(): + with open("/proc/self/statm") as f: + return int(f.read().split()[1]) * os.sysconf("SC_PAGE_SIZE") + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="needs /proc and unlink of open files") +@pytest.mark.parametrize("method", ["update_data", "insert_data"]) +def test_failed_data_write_without_copy_frees_chunk(tmp_path, method): + # A failed update_chunk/insert_chunk without a copy leaves the compressed + # chunk with the caller, so update_data/insert_data must free it + urlpath = str(tmp_path / "failed_write.b2frame") + nbytes = 2 * 2**20 + rng = np.random.default_rng(0) + data = rng.integers(0, 256, nbytes, dtype=np.uint8) # incompressible, so chunks stay large + schunk = blosc2.SChunk( + chunksize=nbytes, + data=data, + urlpath=urlpath, + contiguous=True, + mode="w", + cparams={"typesize": 1, "clevel": 1}, + ) + os.remove(urlpath) # the frame keeps working through its cached reader, but cannot write + write = getattr(schunk, method) + with pytest.raises(RuntimeError): + write(0, data, copy=False) # warm-up + rss0 = _rss_bytes() + for _ in range(40): + with pytest.raises(RuntimeError): + write(0, data, copy=False) + # 40 leaked chunks would be 80 MB + assert _rss_bytes() - rss0 < 30 * 2**20 From 6abf3f6eabc2ad1536b23e32e9b6a1681dc8c00a Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:50:41 +0100 Subject: [PATCH 06/12] Clip the last chunk when collecting the True positions of a bool mask LazyExpr._collect_flat_indices_from_bool_ndarray(), used for 1-D masks, took every chunk as chunk_len items, so the padding of the last chunk could yield positions past the end of the array: blosc2.full(1000, True, chunks=(300,)) gave 1200 positions. Clip each chunk to the items inside the array. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/lazyexpr.py | 16 ++++++++-------- tests/ndarray/test_getitem.py | 12 ++++++++++++ 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 28af9bbb8..188d4c612 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -4489,29 +4489,29 @@ def _collect_flat_indices_from_bool_ndarray(bool_ndarray): Flat indices of True positions (int64). """ chunk_len = bool_ndarray.chunks[0] + nitems = bool_ndarray.shape[0] all_indices = [] for info in bool_ndarray.iterchunks_info(): + offset = info.nchunk * chunk_len + # Items of this chunk inside the array: the last chunk's padding is not data, and it + # can hold anything (e.g. the fill value of a blosc2.full() array) + nvalid = builtins.min(chunk_len, nitems - offset) # Skip special-value chunks that are entirely False if info.special == blosc2.SpecialValue.ZERO: continue if info.special == blosc2.SpecialValue.VALUE: - if not info.repeated_value: # repeated_value is False/0 + if not bool(info.repeated_value): # repeated_value is False/0 continue # repeated_value is True: all elements in this chunk are True - offset = info.nchunk * chunk_len - all_indices.append(np.arange(offset, offset + chunk_len, dtype=np.int64)) + all_indices.append(np.arange(offset, offset + nvalid, dtype=np.int64)) continue # Normal chunk: decompress and scan for True positions raw = bool_ndarray.schunk.decompress_chunk(info.nchunk) - arr = np.frombuffer(raw, dtype=np.bool_) - # Truncate to the logical chunk size (buffer may include padding) - if len(arr) > chunk_len: - arr = arr[:chunk_len] + arr = np.frombuffer(raw, dtype=np.bool_)[:nvalid] idx = np.flatnonzero(arr) if len(idx) > 0: - offset = info.nchunk * chunk_len all_indices.append(idx + offset) if not all_indices: diff --git a/tests/ndarray/test_getitem.py b/tests/ndarray/test_getitem.py index 567d56ed9..66f7e2262 100644 --- a/tests/ndarray/test_getitem.py +++ b/tests/ndarray/test_getitem.py @@ -1115,3 +1115,15 @@ def test_getitem_random_keys_match_numpy(tmp_path, on_disk): if want_exc is None: assert got.shape == np.shape(want), key np.testing.assert_array_equal(got, want) + + +@pytest.mark.parametrize("fill", [True, False]) +@pytest.mark.parametrize("modified", [False, True]) +@pytest.mark.parametrize("nitems", [900, 1000]) +def test_collect_flat_indices_clips_last_chunk(fill, modified, nitems): + # The last chunk's padding holds the fill value, and must not yield indices past the end + mask = blosc2.full(nitems, fill, dtype=np.bool_, chunks=(300,), blocks=(100,)) + if modified: + mask[nitems - 50 : nitems - 40] = not fill # makes the last chunk a regular one + indices = blosc2.LazyExpr._collect_flat_indices_from_bool_ndarray(mask) + np.testing.assert_array_equal(indices, np.flatnonzero(mask[:])) From 8ed4a4837eaa7173051010c2a18118bcd23c1324 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:51:45 +0100 Subject: [PATCH 07/12] Start blosc2.nthreads from BLOSC_NTHREADS and keep NUMEXPR_NUM_THREADS C-Blosc2 uses BLOSC_NTHREADS, when set to a number it accepts, for every compression and decompression context, whatever nthreads the parameters ask for. blosc2.nthreads, and so the CParams/DParams defaults and numexpr's thread count, ignored it: with BLOSC_NTHREADS=4, Python reported 28 threads and gave numexpr 28 while C used 4. Start blosc2.nthreads from what C-Blosc2 gives a compression context (the new blosc2_ext.context_nthreads(), via blosc2_ctx_get_cparams()), and say in set_nthreads()'s docs that the variable wins at the C level. Importing blosc2 also overrode NUMEXPR_NUM_THREADS, which set_nthreads() recommends setting; it now leaves numexpr alone then, and otherwise gives it at most numexpr.MAX_THREADS (BLOSC_NTHREADS can go beyond that, which numexpr refuses with an error). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/__init__.py | 16 +++++++++-- src/blosc2/blosc2_ext.pyx | 14 ++++++++++ src/blosc2/core.py | 5 +++- tests/test_numexpr_threads.py | 52 +++++++++++++++++++++++++++++++++++ 4 files changed, 83 insertions(+), 4 deletions(-) diff --git a/src/blosc2/__init__.py b/src/blosc2/__init__.py index 39008fd0b..2cdb8e733 100644 --- a/src/blosc2/__init__.py +++ b/src/blosc2/__init__.py @@ -503,6 +503,8 @@ def _raise(exc): cpu_info = get_cpu_info() nthreads = ncores = cpu_info.get("count", 1) """Number of threads to be used in compression/decompression. + +It starts from the ``BLOSC_NTHREADS`` environment variable if that is set. """ # Protection against too many threads nthreads = min(nthreads, 64) @@ -515,14 +517,22 @@ def _raise(exc): # Experiments say that, when using a large number of threads, it is better to not use them all if nthreads > 16: nthreads -= nthreads // 8 - # Only call set_num_threads if within NUMEXPR_MAX_THREADS limit to avoid warning + # C-Blosc2 uses BLOSC_NTHREADS, when set to a valid number, for every compression and + # decompression, whatever the nthreads in their parameters; start from it (as C-Blosc2 reads + # it), so that nthreads and the CParams/DParams defaults say how many threads are used + from .blosc2_ext import context_nthreads + + nthreads = context_nthreads(nthreads) + # Only call set_num_threads if within NUMEXPR_MAX_THREADS limit to avoid warning, and not + # over the thread count that the user gave numexpr with NUMEXPR_NUM_THREADS numexpr_max_env = os.environ.get("NUMEXPR_MAX_THREADS") numexpr_max: int | None = None if numexpr_max_env is not None: with contextlib.suppress(ValueError): numexpr_max = int(numexpr_max_env) - if numexpr_max is None or nthreads <= numexpr_max: - numexpr.set_num_threads(nthreads) + if not os.environ.get("NUMEXPR_NUM_THREADS") and (numexpr_max is None or nthreads <= numexpr_max): + # BLOSC_NTHREADS can go beyond numexpr's own maximum, which it refuses with an error + numexpr.set_num_threads(min(nthreads, numexpr.MAX_THREADS)) # This import must be before ndarray and schunk from .storage import ( # noqa: I001 diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 477b378bf..06975399d 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -369,6 +369,8 @@ cdef extern from "blosc2.h": blosc2_context* blosc2_create_dctx(blosc2_dparams dparams) nogil + int blosc2_ctx_get_cparams(blosc2_context *ctx, blosc2_cparams *cparams) + int blosc2_ctx_get_dparams(blosc2_context *ctx, blosc2_dparams *dparams) void blosc2_free_ctx(blosc2_context * context) nogil @@ -1270,6 +1272,18 @@ def set_nthreads(nthreads): else: return rc +def context_nthreads(int16_t nthreads): + """The number of threads that C-Blosc2 gives a compression context asking for `nthreads`: the + BLOSC_NTHREADS environment variable instead, when that is set to a number it accepts.""" + cdef blosc2_cparams cparams = BLOSC2_CPARAMS_DEFAULTS + cparams.nthreads = nthreads + cdef blosc2_context *cctx = blosc2_create_cctx(cparams) + if cctx == NULL: + raise MemoryError() + blosc2_ctx_get_cparams(cctx, &cparams) + blosc2_free_ctx(cctx) + return cparams.nthreads + def set_blocksize(size_t blocksize=0): blosc1_set_blocksize(blocksize) diff --git a/src/blosc2/core.py b/src/blosc2/core.py index 21a4eafbc..1e3dea7bd 100644 --- a/src/blosc2/core.py +++ b/src/blosc2/core.py @@ -1286,7 +1286,10 @@ def set_nthreads(nthreads: int) -> int: Notes ----- The number of threads can also be set via the ``BLOSC_NTHREADS`` environment - variable (e.g., ``export BLOSC_NTHREADS=1``). Additionally, you may want to set + variable (e.g., ``export BLOSC_NTHREADS=1``), which :attr:`~blosc2.nthreads` + then starts from. While it is set, C-Blosc2 uses it for every compression and + decompression, over this function and the ``nthreads`` of :class:`CParams` and + :class:`DParams`. Additionally, you may want to set ``NUMEXPR_NUM_THREADS`` (e.g., ``export NUMEXPR_NUM_THREADS=1``) as well since numexpr is used under the hood when performing some operations. Note that this function only sets the number of threads used by Blosc, not the number diff --git a/tests/test_numexpr_threads.py b/tests/test_numexpr_threads.py index 9d162ff3c..2a95e4bca 100644 --- a/tests/test_numexpr_threads.py +++ b/tests/test_numexpr_threads.py @@ -43,3 +43,55 @@ def test_numexpr_max_threads_no_warning(): assert "nthreads cannot be larger" not in result.stderr, ( f"Unexpected numexpr warning in stderr: {result.stderr}" ) + + +def _import_blosc2_and_print(env_update, code): + env = os.environ.copy() + env.pop("BLOSC_NTHREADS", None) + env.update(env_update) + result = subprocess.run( + [sys.executable, "-c", f"import blosc2; {code}"], capture_output=True, text=True, env=env, check=True + ) + return result.stdout.split() + + +@pytest.mark.skipif(blosc2.IS_WASM, reason="subprocess is not supported on emscripten/wasm32") +def test_blosc_nthreads_env(): + # C-Blosc2 uses BLOSC_NTHREADS for every compression and decompression context; nthreads and + # the parameter defaults must say so + code = ( + "a = blosc2.zeros(10); " + "print(blosc2.nthreads, blosc2.CParams().nthreads, blosc2.DParams().nthreads, a.cparams.nthreads)" + ) + assert _import_blosc2_and_print({"BLOSC_NTHREADS": "3"}, code) == ["3"] * 4 + # C-Blosc2 reads the leading integer (strtol) + assert _import_blosc2_and_print({"BLOSC_NTHREADS": "4abc"}, code) == ["4"] * 4 + assert _import_blosc2_and_print({"BLOSC_NTHREADS": "2.5"}, code) == ["2"] * 4 + default = _import_blosc2_and_print({}, "print(blosc2.nthreads)") + for invalid in ["0", "-2", "abc", ""]: + assert _import_blosc2_and_print({"BLOSC_NTHREADS": invalid}, "print(blosc2.nthreads)") == default + + +@pytest.mark.skipif(blosc2.IS_WASM, reason="subprocess is not supported on emscripten/wasm32") +def test_numexpr_num_threads_env_is_kept(): + # set_nthreads() recommends NUMEXPR_NUM_THREADS for numexpr; importing blosc2 must not override it + code = "import numexpr; print(numexpr.get_num_threads())" + assert _import_blosc2_and_print({"NUMEXPR_NUM_THREADS": "2"}, code) == ["2"] + assert _import_blosc2_and_print({"BLOSC_NTHREADS": "3"}, code) == ["3"] + + +@pytest.mark.skipif(blosc2.IS_WASM, reason="subprocess is not supported on emscripten/wasm32") +def test_blosc_nthreads_env_beyond_numexpr_max(): + # nthreads says what C-Blosc2 uses; numexpr gets as many threads as it takes, without errors + env = os.environ.copy() + for name in ("NUMEXPR_NUM_THREADS", "NUMEXPR_MAX_THREADS"): + env.pop(name, None) + env["BLOSC_NTHREADS"] = "100" + code = "import blosc2, numexpr; print(blosc2.nthreads, numexpr.get_num_threads(), numexpr.MAX_THREADS)" + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, env=env, check=True + ) + assert "nthreads cannot be larger" not in result.stdout + result.stderr + nthreads, numexpr_nthreads, numexpr_max = map(int, result.stdout.split()) + assert nthreads == 100 + assert numexpr_nthreads == min(100, numexpr_max) From 7d7561737f170c94fc24ced4f246e6c2936c0cec Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:41:47 +0100 Subject: [PATCH 08/12] Resolve basic keys of NDArray.__getitem__/__setitem__ in Cython Keys made of integers (Python or NumPy, not bools), unit-step slices, one Ellipsis and, for reads, None went through the general path: 8-30 us of Python per call (ndindex as soon as a NumPy integer, ... or None appears), more than decompressing a small selection. NDArray.get_basic_numpy() classifies such keys in Cython, resolves them against the shape and reads the result under one acquisition of the handle lock, through _get_slice_cbuffer(), the GIL-releasing read that get_slice_numpy() now shares with it. get_basic_start_stop() resolves the same keys (without None) for __setitem__. Any other key, or one that does not fit the shape, gives None and takes the unchanged general path, which raises the same exceptions as before. NumPy integers without __index__ (timedelta64) are left to it too, where they raise IndexError as in NumPy. SWMR: a read resolves the key against the cached shape and relies on the refresh inside the C read; if b2nd_array_t.last_tick shows that this refresh changed the shape (or the key did not fit, or the read failed), it refreshes, resolves and reads again, so a key is never resolved against a stale shape, at one file check per on-disk read instead of two. A write refreshes under the handle lock before resolving. Keys with None now return arrays that own their data (they were views of a temporary). Measured (Ryzen 9 7950X): a[3, 5, 7] of a 300x512x512 float32 array 34 -> 25 us in memory and 48 -> 35 us on disk, a[500] of a 1000x1000 array 21.5 -> 13.9 us, a[np.int64(7), np.int32(9)] 37.5 -> 7.3 us, iterating over 2000 rows 64 -> 44 ms; a[np.int64(i), j] = v 22 us less. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- doc/reference/ndarray.rst | 2 +- src/blosc2/blosc2_ext.pyx | 252 +++++++++++++++++++++++++++++++--- src/blosc2/ndarray.py | 62 ++++++--- tests/ndarray/test_getitem.py | 68 +++++++++ tests/ndarray/test_setitem.py | 53 +++++++ tests/test_swmr.py | 26 ++++ 6 files changed, 422 insertions(+), 41 deletions(-) diff --git a/doc/reference/ndarray.rst b/doc/reference/ndarray.rst index 30d07cd5c..8c492fce7 100644 --- a/doc/reference/ndarray.rst +++ b/doc/reference/ndarray.rst @@ -14,7 +14,7 @@ It is a direct alias for ``array.vlmeta`` and uses the same persistent storage. .. autoclass:: NDArray :members: :inherited-members: - :exclude-members: get_slice, set_slice, get_slice_numpy, get_oindex_numpy, set_oindex_numpy + :exclude-members: get_slice, set_slice, get_slice_numpy, get_basic_numpy, get_basic_start_stop, get_oindex_numpy, set_oindex_numpy :member-order: groupwise :Special Methods: diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 06975399d..4e7fcbd7c 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -30,8 +30,13 @@ from cpython cimport ( PyBytes_FromStringAndSize, PyObject_GetBuffer, ) +from cpython.bool cimport PyBool_Check +from cpython.long cimport PyLong_Check +from cpython.number cimport PyIndex_Check +from cpython.pyport cimport PY_SSIZE_T_MAX, PY_SSIZE_T_MIN from cpython.ref cimport Py_INCREF, Py_DECREF from cpython.pycapsule cimport PyCapsule_GetPointer, PyCapsule_New +from cpython.slice cimport PySlice_AdjustIndices, PySlice_Check, PySlice_Unpack from cython.operator cimport dereference from libc.stdint cimport uintptr_t from libc.stdlib cimport free, malloc, realloc, calloc @@ -614,6 +619,7 @@ cdef extern from "b2nd.h": int64_t chunk_array_strides[B2ND_MAX_DIM] char *dtype int8_t dtype_format + int64_t last_tick ctypedef struct b2nd_context_t: pass @@ -4059,6 +4065,18 @@ cdef c_bool _can_release_gil(blosc2_schunk *sc) except? False: return dparams.postfilter == NULL +# Kinds of key items in NDArray.get_basic_numpy() and get_basic_start_stop() +cdef enum: + _KEY_INT + _KEY_SLICE + _KEY_NEWAXIS + _KEY_ELLIPSIS + _MAX_BASIC_KEY = 64 # longer keys (only possible with many None) take the general path + _BASIC_READ_ATTEMPTS = 3 + +cdef object _np_integer = np.integer + + cdef class NDArray: cdef b2nd_array_t* array cdef readonly object _handle_lock @@ -4136,40 +4154,238 @@ cdef class NDArray: return dtype def get_slice_numpy(self, arr, key): - global _nogil_reads start, stop = key cdef int64_t[B2ND_MAX_DIM] start_, stop_ - cdef int64_t[B2ND_MAX_DIM] buffershape_ + cdef int rc for i in range(self.ndim): start_[i] = start[i] stop_[i] = stop[i] + # Waiting for the lock releases the GIL (RLock does that), so a holder that needs the + # GIL back, e.g. for a Python postfilter, can finish. + with self._handle_lock: + rc = self._get_slice_cbuffer(arr, start_, stop_) + _check_rc(rc, "Error while getting the buffer") + + return arr + + cdef Py_ssize_t _classify_basic_key(self, tuple items, char *kinds, int64_t *ints, Py_ssize_t *starts, + Py_ssize_t *stops, c_bool setitem=False) except -2: + """First pass of get_basic_numpy() and get_basic_start_stop() (`setitem`): classify the + key items without looking at the shape. + + Fill kinds, ints, starts and stops (_MAX_BASIC_KEY entries) and return the number of + items that index a dimension, or -1 if the key is not a basic one (or has too many). + """ + cdef Py_ssize_t nkey = len(items) + if nkey > _MAX_BASIC_KEY: + return -1 + cdef Py_ssize_t i, step, nidx = 0 + cdef c_bool ellipsis = False + for i in range(nkey): + k = items[i] + if k is None: + if setitem: + return -1 + kinds[i] = _KEY_NEWAXIS + elif PySlice_Check(k): + try: + PySlice_Unpack(k, &starts[i], &stops[i], &step) + except Exception: + return -1 + if step != 1: + return -1 + # __setitem__'s general path rejects some bool steps and bounds clamped by + # PySlice_Unpack (beyond the index range): leave them all to it + if setitem and ( + PyBool_Check(k.step) + or starts[i] in (PY_SSIZE_T_MIN, PY_SSIZE_T_MAX) + or stops[i] == PY_SSIZE_T_MIN + or (stops[i] == PY_SSIZE_T_MAX and k.stop is not None) + ): + return -1 + kinds[i] = _KEY_SLICE + nidx += 1 + elif k is Ellipsis: + if ellipsis: + return -1 + ellipsis = True + kinds[i] = _KEY_ELLIPSIS + # NumPy integers without __index__ (timedelta64) are no index for NumPy either + elif (PyLong_Check(k) and not PyBool_Check(k)) or ( + isinstance(k, _np_integer) and PyIndex_Check(k) + ): + try: + ints[i] = k + except Exception: # e.g. beyond int64 + return -1 + kinds[i] = _KEY_INT + nidx += 1 + else: + return -1 + return nidx if nidx <= self.array.ndim else -1 + + def get_basic_start_stop(self, key): + """Resolve a basic index key without None for __setitem__, else return None. + + 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). + """ + cdef tuple items = key if type(key) is tuple else (key,) + cdef Py_ssize_t nkey = len(items) + 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, True) + if nidx < 0: + return None + cdef int64_t[B2ND_MAX_DIM] start_, stop_ + cdef c_bool[B2ND_MAX_DIM] int_dims + with self._handle_lock: + _check_rc(b2nd_refresh(self.array), "Error while refreshing the array") + out_shape = self._resolve_basic_key(nkey, kinds, ints, starts, stops, nidx, start_, stop_, + int_dims) + if out_shape is None: + 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)])) + + def get_basic_numpy(self, key): + """Return ``self[key]`` as a NumPy array if *key* is a basic index, else None. + + Only keys made of integers (Python or NumPy, not bools), slices with unit + step, Ellipsis and None are handled here, after following shape changes + made by another handle (SWMR). Any other key, and any key that is out of + range or otherwise invalid, gives None, so that the caller can use the + general indexing path, which also raises the appropriate exception. + """ + cdef tuple items = key if type(key) is tuple else (key,) + cdef Py_ssize_t nkey = len(items) + 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 + + cdef int64_t[B2ND_MAX_DIM] start_, stop_ + cdef int64_t tick + cdef int rc, attempt + with self._handle_lock: + # The read itself follows shape changes made by another handle (SWMR) before reading, + # so a key resolved against the cached shape is right if the read did not find any. + # This saves a second check of the file per read. Otherwise (or if the key does not + # fit the cached shape, or the read failed) follow the changes first and try again; + # another handle can change the shape once more meanwhile, hence a few attempts. + 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_) + if out_shape is None: + if attempt > 0: + return None + continue + arr = np.empty(out_shape, dtype=self.dtype) + rc = self._get_slice_cbuffer(arr, start_, stop_) + if rc >= 0 and self.array.last_tick == tick: + return arr + # Out of attempts: as the general path, return what the last read gave + _check_rc(rc, "Error while getting the buffer") + return arr + + 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): + """Resolve a key classified by _classify_basic_key() against the cached shape. + + Fill start_ and stop_ (and int_dims, which dimensions are indexed with an integer, if + given) and return the shape of the result, or None if the key does not fit the shape. + """ + cdef int64_t n, idx + cdef Py_ssize_t i, start, stop + cdef int dim = 0 + out_shape = [] + for i in range(nkey): + if kinds[i] == _KEY_NEWAXIS: + out_shape.append(1) + continue + if kinds[i] == _KEY_ELLIPSIS: + dim = self._whole_dims(dim, self.array.ndim - nidx, start_, stop_, int_dims, out_shape) + continue + n = self.array.shape[dim] + if kinds[i] == _KEY_INT: + idx = ints[i] + if idx < 0: + idx += n + if idx < 0 or idx >= n: + return None + start_[dim] = idx + stop_[dim] = idx + 1 + else: + if n > PY_SSIZE_T_MAX: + return None + start = starts[i] + stop = stops[i] + PySlice_AdjustIndices( n, &start, &stop, 1) + if start >= stop: + start = stop = 0 + start_[dim] = start + stop_[dim] = stop + out_shape.append(stop - start) + if int_dims != NULL: + int_dims[dim] = kinds[i] == _KEY_INT + dim += 1 + # The dimensions after the key (none if it has an ellipsis) + self._whole_dims(dim, self.array.ndim - dim, start_, stop_, int_dims, out_shape) + return out_shape + + cdef int _whole_dims(self, int dim, int count, int64_t *start_, int64_t *stop_, c_bool *int_dims, + list out_shape): + """Select `count` whole dimensions from `dim` on for _resolve_basic_key(); return the next.""" + cdef int64_t n + cdef int j + for j in range(count): + n = self.array.shape[dim] + start_[dim] = 0 + stop_[dim] = n + if int_dims != NULL: + int_dims[dim] = False + out_shape.append(n) + dim += 1 + return dim + + 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): buffershape_[i] = stop_[i] - start_[i] cdef Py_buffer view cdef int rc = 0 PyObject_GetBuffer(arr, &view, PyBUF_SIMPLE) try: - # Waiting for the lock releases the GIL (RLock does that), so a holder that needs the - # GIL back, e.g. for a Python postfilter, can finish. - with self._handle_lock: - # 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_, - view.buf, buffershape_, view.len) - _nogil_reads -= 1 - else: + # 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_, view.buf, buffershape_, view.len) + _nogil_reads -= 1 + else: + rc = b2nd_get_slice_cbuffer(self.array, start_, stop_, + view.buf, buffershape_, view.len) finally: PyBuffer_Release(&view) - _check_rc(rc, "Error while getting the buffer") - - return arr + return rc @_handle_locked def get_1d_span_numpy(self, arr, int64_t nchunk, int32_t start, int32_t nitems): diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index e2f9a7343..87ad6e9a4 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -4721,6 +4721,14 @@ def __getitem__( # noqa: C901 [3.3333, 3.3333, 3.3333, 3.3333, 3.3333], [3.3333, 3.3333, 3.3333, 3.3333, 3.3333]]) """ + # Fast path for the common basic indexing (integers, unit-step slices, Ellipsis, None); + # None means that the key needs the general path below + nparr = self.get_basic_numpy(key) + if nparr is not None: + if self._keep_last_read: + self._remember_read(key, nparr) + return nparr + # Follow shape changes made by another handle (SWMR) before validating the key self.refresh() # The more general case (this is quite slow) @@ -4805,12 +4813,14 @@ def __getitem__( # noqa: C901 nparr = np.expand_dims(nparr, axis=[i for i, n in enumerate(none_mask) if n]) if self._keep_last_read: - self._last_read.clear() - inmutable_key = make_key_hashable(key) - self._last_read[inmutable_key] = nparr - + self._remember_read(key, nparr) return nparr + def _remember_read(self, key, nparr: np.ndarray) -> None: + """Make `nparr`, read with `key`, the content of the last-read cache (see keep_last_read).""" + self._last_read.clear() + self._last_read[make_key_hashable(key)] = nparr + def __setitem__( self, key: int | slice | Sequence[slice | int | np.bool_ | np.ndarray[int | np.bool_] | None] | None, @@ -4846,21 +4856,30 @@ def __setitem__( [3.3333, 3.3333, 3.3333, 3.3333, 3.3333, 3.3333, 3.3333, 3.3333]]) """ blosc2_ext.check_access_mode(self.schunk.urlpath, self.schunk.mode) - # Follow shape changes made by another handle (SWMR) before validating the key - self.refresh() + # Fast path for basic keys (integers, unit-step slices, Ellipsis), which also follows shape + # changes made by another handle (SWMR); None means that the key needs the general path + basic = self.get_basic_start_stop(key) + if basic is None: + # Follow shape changes made by another handle (SWMR) before validating the key + self.refresh() - # key not iterable - key = key[()] if isinstance(key, NDArray) else key - key = tuple(k[()] if isinstance(k, NDArray) else k for k in key) if isinstance(key, tuple) else key + # key not iterable + key = key[()] if isinstance(key, NDArray) else key + key = ( + tuple(k[()] if isinstance(k, NDArray) else k for k in key) if isinstance(key, tuple) else key + ) - key_, mask = process_key(key, self.shape) # internally handles key an integer + 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() value = ( value if np.isscalar(value) else blosc2.as_simpleproxy(value) ) # convert to SimpleProxy for e.g. JAX, Tensorflow, PyTorch - if builtins.any(isinstance(k, list | np.ndarray) for k in key_): # fancy indexing + if basic is not None: + start, stop, mask = basic + none_mask = () + elif builtins.any(isinstance(k, list | np.ndarray) for k in key_): # fancy indexing _slice = ndindex.ndindex(key_).expand( self.shape ) # handles negative indices -> positive internally @@ -4877,17 +4896,16 @@ def __setitem__( indexing.mark_indexes_stale(self) return result - - start, stop, step, none_mask = get_ndarray_start_stop(self.ndim, key_, self.shape) - - 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__.") - result = self._get_set_nonunit_steps((start, stop, step, mask), value=value) - from . import indexing - - indexing.mark_indexes_stale(self) - return result + else: + start, stop, step, none_mask = get_ndarray_start_stop(self.ndim, key_, self.shape) + 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__.") + result = self._get_set_nonunit_steps((start, stop, step, mask), value=value) + from . import indexing + + indexing.mark_indexes_stale(self) + return result shape = [sp - st for sp, st in zip(stop, start, strict=False)] if isinstance(value, blosc2.Operand): # handles SimpleProxy, NDArray, LazyExpr etc. diff --git a/tests/ndarray/test_getitem.py b/tests/ndarray/test_getitem.py index 66f7e2262..472f483fe 100644 --- a/tests/ndarray/test_getitem.py +++ b/tests/ndarray/test_getitem.py @@ -914,6 +914,59 @@ def test_scalar_bool_key_matches_numpy(shape, key): np.testing.assert_array_equal(got, want) +# Keys for NDArray.get_basic_numpy(), the fast path of __getitem__ for integers, unit-step slices, +# Ellipsis and None. Mixed with keys it must leave to the general path (invalid, out of range, steps, +# bools, huge bounds, fancy), as any key must give the same result or exception either way. +def _random_key_item(rng, n): + r = rng.random() + if r < 0.3: + i = int(rng.integers(-n - 2, n + 2)) + return i if rng.random() < 0.5 else np.dtype(rng.choice(["i1", "i8", "u2"])).type(abs(i) % 100) + if r < 0.7: + bounds = [None, *rng.integers(-n - 3, n + 3, size=2).tolist()] + start, stop = rng.choice(bounds), rng.choice(bounds) + step = rng.choice([None, None, 1, np.int64(1), 2, -1, True]) + return slice(start, stop, step) + if r < 0.8: + return Ellipsis + if r < 0.9: + return None + return [True, 1.0, 2**64, slice(2**64), [0], np.array(1)][int(rng.integers(6))] + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_basic_index_fast_path_matches_general_path(tmp_path, monkeypatch, on_disk): + rng = np.random.default_rng(1) + shapes = [(), (0,), (7,), (5, 0, 3), (6, 5), (4, 3, 5), (3, 2, 4, 3)] + for ishape, shape in enumerate(shapes): + npa = np.arange(math.prod(shape), dtype=np.int32).reshape(shape) + kwargs = {"urlpath": tmp_path / f"basic{ishape}.b2nd", "mode": "w"} if on_disk else {} + a = blosc2.asarray(npa, **kwargs) + n = shape[0] if shape else 1 + for _ in range(150): + items = tuple(_random_key_item(rng, n) for _ in range(int(rng.integers(0, len(shape) + 3)))) + key = items[0] if len(items) == 1 and rng.random() < 0.5 else items + outcomes = [] + for general in (False, True): + with monkeypatch.context() as m: + if general: + m.setattr(blosc2.NDArray, "get_basic_numpy", lambda self, key: None) + try: + outcomes.append(("ok", a[key])) + except Exception as e: + outcomes.append((type(e), str(e))) + (fast, got), (general, want) = outcomes + assert fast == general, key + if fast == "ok": + assert type(got) is type(want), key + assert got.shape == want.shape, key + assert got.dtype == want.dtype, key + np.testing.assert_array_equal(got, want) + if a.get_basic_numpy(key) is not None: # the fast path took it: compare with NumPy + np.testing.assert_array_equal(got, npa[key]) + assert got.shape == np.asarray(npa[key]).shape, key + + def _assert_getitem_matches_numpy(npa, a, key): try: want = npa[key] @@ -965,6 +1018,21 @@ def test_getitem_0d_integer_array(key): _assert_getitem_matches_numpy(npa, blosc2.asarray(npa), key) +# timedelta64 is a NumPy integer type, but no index (it has no __index__): IndexError, as in NumPy +@pytest.mark.parametrize( + "key", [np.timedelta64(1, "ns"), (np.timedelta64(1, "ps"),), (1, np.timedelta64(2, "ns"))] +) +def test_timedelta64_key(key): + npa = np.arange(20).reshape(4, 5) + a = blosc2.asarray(npa) + _assert_getitem_matches_numpy(npa, a, key) + with pytest.raises(IndexError): + npa[key] = 7 + with pytest.raises(IndexError): + a[key] = 7 + np.testing.assert_array_equal(a[...], npa) + + # NumPy places the dimension of scalar bools among the other advanced indices (integers count as # such) when they are adjacent, else in front; also after an explicit None. Shapes used to differ. @pytest.mark.parametrize( diff --git a/tests/ndarray/test_setitem.py b/tests/ndarray/test_setitem.py index 78b2cb8e1..79f7edfa4 100644 --- a/tests/ndarray/test_setitem.py +++ b/tests/ndarray/test_setitem.py @@ -149,3 +149,56 @@ def test_setitem_fancy_index(): out_numpy[idx] = value np.testing.assert_array_equal(out[:], out_numpy) + + +def _random_setitem_key(rng, n, ndim): + items = [] + for _ in range(int(rng.integers(0, ndim + 2))): + r = rng.random() + if r < 0.35: + i = int(rng.integers(-n - 2, n + 2)) + items.append(i if rng.random() < 0.5 else np.int64(i)) + elif r < 0.75: + bounds = [None, *rng.integers(-n - 3, n + 3, size=2).tolist()] + step = [None, None, 1, np.int64(1), 2, -1, True][int(rng.integers(7))] + items.append(slice(bounds[int(rng.integers(3))], bounds[int(rng.integers(3))], step)) + elif r < 0.85: + items.append(Ellipsis) + elif r < 0.9: + items.append(None) + else: + items.append([True, [0], 1.5][int(rng.integers(3))]) + return items[0] if len(items) == 1 and rng.random() < 0.5 else tuple(items) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_setitem_basic_key_fast_path_matches_general_path(tmp_path, monkeypatch, on_disk): + # NDArray.get_basic_start_stop() resolves basic keys for __setitem__; every key must write the + # same data (or raise the same exception) with and without it + rng = np.random.default_rng(2) + for ishape, shape in enumerate([(), (7,), (6, 5), (4, 3, 5)]): + npa = np.arange(int(np.prod(shape)), dtype=np.int32).reshape(shape) + n = shape[0] if shape else 1 + for _ in range(60): + key = _random_setitem_key(rng, n, len(shape)) + try: + values = [-1, np.full(np.shape(npa[key]), -2, dtype=np.int64)] + except Exception: + values = [-1] + for value in values: + outcomes = [] + for general in (False, True): + kwargs = {"urlpath": tmp_path / f"set{ishape}.b2nd", "mode": "w"} if on_disk else {} + a = blosc2.asarray(npa, **kwargs) + with monkeypatch.context() as m: + if general: + m.setattr(blosc2.NDArray, "get_basic_start_stop", lambda self, key: None) + try: + a[key] = value + outcomes.append(a[()]) + except Exception as e: + outcomes.append((type(e), str(e))) + if isinstance(outcomes[1], tuple): + assert outcomes[0] == outcomes[1], key + else: + np.testing.assert_array_equal(outcomes[0], outcomes[1]) diff --git a/tests/test_swmr.py b/tests/test_swmr.py index 135e746c1..f2cb340af 100644 --- a/tests/test_swmr.py +++ b/tests/test_swmr.py @@ -189,3 +189,29 @@ def test_cross_process_growth_hammer(tmp_path): assert writer.returncode == 0, f"writer process failed with exit code {writer.returncode}" assert seen > 0 + + +@pytest.mark.parametrize( + "key", [15, -1, slice(None), slice(-12, None), (Ellipsis, 3), (None, slice(8, 13)), (np.int64(19), 2)] +) +def test_reader_basic_keys_follow_resize(tmp_path, key): + # Keys are resolved against the shape another handle just set (getitem's fast path) + urlpath = tmp_path / "keys.b2nd" + w = create_array(urlpath) + w[:] = np.arange(100, dtype=DTYPE).reshape(10, 10) + r = blosc2.open(str(urlpath)) + _ = r[0] + + w.resize((20, 10)) + w[10:20] = np.arange(100, 200, dtype=DTYPE).reshape(10, 10) + np.testing.assert_array_equal(r[key], np.arange(200, dtype=DTYPE).reshape(20, 10)[key]) + + w.resize((5, 10)) + expected = np.arange(50, dtype=DTYPE).reshape(5, 10) + try: + want = expected[key] + except IndexError: + with pytest.raises(IndexError): + r[key] + else: + np.testing.assert_array_equal(r[key], want) From c7a562768e56b99e80d935a225ccb817a6700247 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:42:38 +0100 Subject: [PATCH 09/12] Lower the per-call overhead of creating arrays and of small writes Creating a small array cost ~30 us, mostly Python bookkeeping: - compute_chunks_blocks() and _check_ndarray_kwargs() converted CParams and Storage with dataclasses.asdict(), which deep-copies; a shallow dict (vars() for Storage, whose fields are plain values) is enough; - every SChunk wrapper, so every NDArray, built its cparams and dparams eagerly under the handle lock; they are now built on first use; - create_cparams_from_kwargs() looked up Filter.BYTEDELTA.value (a slow enum descriptor) twice per filter slot; now once per call; - is_inside_new_expr() and is_inside_ne_evaluate(), which arange, linspace, eye, reshape, fromiter and lazy functions such as clip run on every call, used inspect.stack(), which reads the source of every frame; walking the frames gives the same answer in under 1 us. __setitem__ also wrapped plain NumPy values in a SimpleProxy, whose constructor runs compute_chunks_blocks(); other array types keep the wrapping, NumPy arrays give the same outcome without it (checked for contiguous, strided, Fortran-order, 0-d, broadcast and converted values). Measured (Ryzen 9 7950X): zeros((100, 100)) 31 -> 20 us, eye() 311 -> 179 us, reshape() 437 -> 276 us, clip(a, 0, 1) 84 -> 12 us; a 3x3 write into a 100x100 array ~12 us less (asdict() ~5 us, SimpleProxy ~7 us). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/blosc2_ext.pyx | 6 +++-- src/blosc2/core.py | 8 ++++--- src/blosc2/ndarray.py | 18 +++++++-------- src/blosc2/schunk.py | 9 ++++++-- src/blosc2/utils.py | 17 ++++++++++---- tests/ndarray/test_setitem.py | 43 +++++++++++++++++++++++++++++++++++ 6 files changed, 81 insertions(+), 20 deletions(-) diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 4e7fcbd7c..fc127fc1b 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -1387,10 +1387,12 @@ cdef create_cparams_from_kwargs(blosc2_cparams *cparams, kwargs): filters = kwargs.get('filters', blosc2.cparams_dflts['filters']) if len(filters) > BLOSC2_MAX_FILTERS: raise ValueError(f"filters list cannot exceed {BLOSC2_MAX_FILTERS}") + # Looked up once: an enum's value is a slow descriptor + cdef int bytedelta = blosc2.Filter.BYTEDELTA.value for i, filter in enumerate(filters): cparams.filters[i] = filter.value if isinstance(filter, Enum) else filter # Bytedelta does not work on typesize 1 - if cparams.filters[i] == blosc2.Filter.BYTEDELTA.value and typesize == 1: + if cparams.filters[i] == bytedelta and typesize == 1: cparams.filters[i] = 0 if "filters_meta" not in kwargs: @@ -1404,7 +1406,7 @@ cdef create_cparams_from_kwargs(blosc2_cparams *cparams, kwargs): for i, meta in enumerate(filters_meta): # We still may want to encode negative values meta_value = meta if meta < 0 else meta - if meta_value == 0 and cparams.filters[i] == blosc2.Filter.BYTEDELTA.value: + if meta_value == 0 and cparams.filters[i] == bytedelta: # bytedelta typesize cannot be zero when using compress2 cparams.filters_meta[i] = typesize else: diff --git a/src/blosc2/core.py b/src/blosc2/core.py index 1e3dea7bd..b29842d42 100644 --- a/src/blosc2/core.py +++ b/src/blosc2/core.py @@ -24,7 +24,7 @@ import sys import urllib.parse import urllib.request -from dataclasses import asdict +from dataclasses import asdict, fields from functools import cache, lru_cache from typing import TYPE_CHECKING, ClassVar @@ -1984,7 +1984,9 @@ def compute_chunks_blocks( # noqa: C901 cparams = kwargs.get("cparams") or blosc2.CParams() # just get defaults if isinstance(cparams, blosc2.CParams): - cparams = asdict(cparams) + # Shallow is enough (only keys are set below, filters only on a deep copy), and asdict() + # costs more than the rest of this function + cparams = {f.name: getattr(cparams, f.name) for f in fields(cparams)} # Typesize in dtype always has preference over typesize in cparams itemsize = cparams["typesize"] = np.dtype(dtype).itemsize @@ -1992,7 +1994,7 @@ def compute_chunks_blocks( # noqa: C901 # Get the default blocksize for the compression params # Check if we need STUNE for lossy codecs/filters that have specific blocksize requirements codec = cparams.get("codec") - filters = cparams.get("filters", None) + filters = cparams.get("filters") needs_stune = codec in ( blosc2.Codec.ZFP_RATE, blosc2.Codec.ZFP_PREC, diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 87ad6e9a4..985faf1b6 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -8,7 +8,6 @@ from __future__ import annotations import builtins -import inspect import itertools import math import operator @@ -44,6 +43,7 @@ get_local_slice, get_selection, incomplete_lazyfunc, + is_called_from, npbinvert, nplshift, nprshift, @@ -242,9 +242,7 @@ def is_inside_new_expr() -> bool: """ Whether the current code is being executed during the creation of new expression. """ - # Get the current call stack - stack = inspect.stack() - return builtins.any(frame_info.function in {"_new_expr", "open_lazyarray"} for frame_info in stack) + return is_called_from(("_new_expr", "open_lazyarray")) def _normalize_slice(s): @@ -4872,9 +4870,10 @@ def __setitem__( 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() - value = ( - value if np.isscalar(value) else blosc2.as_simpleproxy(value) - ) # convert to SimpleProxy for e.g. JAX, Tensorflow, PyTorch + # 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): + value = blosc2.as_simpleproxy(value) if basic is not None: start, stop, mask = basic @@ -7177,8 +7176,9 @@ def _check_ndarray_kwargs(**kwargs): # noqa: C901 else: kwargs = {**kwargs, **storage} else: - # Add the default storage values as long as they are not already passed - storage_dflts = asdict(blosc2.Storage(urlpath=kwargs.get("urlpath"))) # urlpath can affect defaults + # Add the default storage values as long as they are not already passed. Its fields + # are plain values, so vars() gives what asdict() does, only much faster. + storage_dflts = vars(blosc2.Storage(urlpath=kwargs.get("urlpath"))) # urlpath can affect defaults # If a key appears in both operands, the one from the right-hand operand wins kwargs = storage_dflts | kwargs diff --git a/src/blosc2/schunk.py b/src/blosc2/schunk.py index ff69826a5..2907f1447 100644 --- a/src/blosc2/schunk.py +++ b/src/blosc2/schunk.py @@ -450,8 +450,9 @@ def __init__( # noqa: C901 default = DeserializeMode.SAFE if opens_existing else DeserializeMode.FULL set_deserialize(self, default if deserialize is None else deserialize) self._vlmeta = vlmeta(self, super().c_schunk) - self._cparams = super().get_cparams() - self._dparams = super().get_dparams() + # Built on first use: most handles (e.g. those of NDArrays) never need them + self._cparams = None + self._dparams = None def __enter__(self) -> SChunk: """Enter a context manager and return this super-chunk.""" @@ -470,6 +471,8 @@ def cparams(self) -> blosc2.CParams: """ :class:`blosc2.CParams` instance with the compression parameters. """ + if self._cparams is None: + self._cparams = super().get_cparams() return self._cparams @cparams.setter @@ -484,6 +487,8 @@ def dparams(self) -> blosc2.DParams: """ :class:`blosc2.DParams` instance with the decompression parameters. """ + if self._dparams is None: + self._dparams = super().get_dparams() return self._dparams @dparams.setter diff --git a/src/blosc2/utils.py b/src/blosc2/utils.py index 77f2d3390..5b1fcaa97 100644 --- a/src/blosc2/utils.py +++ b/src/blosc2/utils.py @@ -8,7 +8,6 @@ import ast import builtins import contextlib -import inspect import math import sys import warnings @@ -1050,13 +1049,23 @@ def process_key(key, shape): return key, mask +def is_called_from(names) -> bool: + """Whether a function with one of these names is on the call stack.""" + # Walk the call stack by hand: inspect.stack() also reads the source of every frame, + # which costs hundreds of microseconds per call + frame = sys._getframe() + while frame is not None: + if frame.f_code.co_name in names: + return True + frame = frame.f_back + return False + + def is_inside_ne_evaluate() -> bool: """ Whether the current code is being executed from an ne_evaluate call """ - # Get the current call stack - stack = inspect.stack() - return builtins.any(frame_info.function in {"ne_evaluate"} for frame_info in stack) + return is_called_from(("ne_evaluate",)) def incomplete_lazyfunc(func) -> None: diff --git a/tests/ndarray/test_setitem.py b/tests/ndarray/test_setitem.py index 79f7edfa4..3f2a7cdb3 100644 --- a/tests/ndarray/test_setitem.py +++ b/tests/ndarray/test_setitem.py @@ -151,6 +151,49 @@ def test_setitem_fancy_index(): np.testing.assert_array_equal(out[:], out_numpy) +# NumPy values are written directly, without the SimpleProxy wrapping of other array types; the +# outcome (data or exception) must be the same as with the wrapping +@pytest.mark.parametrize( + "key", + [ + (slice(1, 4), slice(2, 5)), + (2, slice(None)), + (slice(0, 6, 2), slice(1, 7, 3)), + ([0, 2], slice(None)), + (slice(None), [1, 5]), + Ellipsis, + (1, 2), + (slice(5, 2), 1), + (None, slice(1, 3), 0), + ], +) +def test_setitem_numpy_value_same_as_proxied(key): + base = np.arange(48, dtype=np.float32).reshape(6, 8) + shape = np.shape(base[key]) + values = [ + np.full(shape, 7.5, np.float32), + np.full(shape, 7, np.int64), # converted + np.full((*shape, 2), 9, np.float32)[..., 0], # not contiguous + np.asfortranarray(np.full(shape, 3, np.float32)), + np.array(2.5, np.float32), # 0-d + np.full(shape[-1:], 4, np.float32), # broadcast + np.zeros((1, *shape), np.float64), + ] + for value in values: + outcomes = [] + for v in (value, blosc2.SimpleProxy(value)): + a = blosc2.asarray(base, chunks=(3, 4), blocks=(2, 2)) + try: + a[key] = v + outcomes.append(a[:]) + except Exception as e: + outcomes.append((type(e), str(e))) + if isinstance(outcomes[0], tuple): + assert outcomes[0] == outcomes[1] + else: + np.testing.assert_array_equal(outcomes[0], outcomes[1]) + + def _random_setitem_key(rng, n, ndim): items = [] for _ in range(int(rng.integers(0, ndim + 2))): From 5c48fa878e66251f80c044a5b7a79d756ca64848 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:45:01 +0100 Subject: [PATCH 10/12] Write scalars and NDArray values chunk by chunk in NDArray.__setitem__ a[region] = scalar built np.full() of the whole region, and a[region] = ndarray decompressed all of the source first: an uncompressed copy of the region in memory (a fill of a large on-disk array could exceed RAM), then one set_slice. Scalars: every chunk the region fully covers gets one shared special (repeated-value) chunk from SChunk._special_chunk(), and the rest is written in boxes of bounded size (32 MB, or one chunk) from one buffer sized to the largest box. Edge chunks get a special chunk only for zero, as writing them leaves zeros in their padding (which a resize brings to light). Chunks go last to first: a frame compacts its payload behind an updated chunk, and that tail then holds the tiny special chunks already written. Lossy codecs and filters, and dictionaries, keep compressing the values. NDArray values: read box by box along the destination's chunk boundaries into reused buffers and written with one set_slice per box; when both arrays encode chunks the same way (dtype, chunks, blocks, lossless cparams) and the region starts at a chunk boundary, covered chunks without padding are copied compressed. A value that may share storage with the destination (same super-chunk or same file) is still read in full first. Both hold the handle locks (HandleLocks with the source) for the whole write, so that other threads see it as atomic as one set_slice. Regions smaller than a chunk, and arrays whose frame takes the cross-process lock (locking=True or BLOSC_LOCKING, destination or source), keep the single set_slice, which holds that lock from the shape refresh to the last chunk. A shared flat chunk index helper now also serves detect_aligned_chunks(), and __setitem__ marks indexes stale through one helper. Measured (300x512x512 float32, default chunks, Ryzen 9 7950X): a[:] = 0.0 41 -> 0.2 ms with no extra memory (was 303 MB); the first a[:] = 1 into an existing on-disk array 588 -> 29 ms; a[:] = b with the same encoding 142 -> 18 ms and 306 -> 28 MB peak. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/blosc2_ext.pyx | 12 + src/blosc2/ndarray.py | 378 +++++++++++++++++++++--- src/blosc2/utils.py | 9 + tests/ndarray/test_setitem_chunkwise.py | 344 +++++++++++++++++++++ tests/test_locking.py | 218 +++++++++++++- 5 files changed, 910 insertions(+), 51 deletions(-) create mode 100644 tests/ndarray/test_setitem_chunkwise.py diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index fc127fc1b..71ee9b7eb 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -2600,6 +2600,18 @@ cdef class SChunk: if rc < 0: raise RuntimeError("Could not unlock the super-chunk frame") + @property + def _locks_frame(self): + """Whether operations take the cross-process frame lock: on disk with the default I/O, and + `locking` or the BLOSC_LOCKING environment variable set (what C-Blosc2 checks when it + opens the frame).""" + cdef blosc2_storage *storage = self.schunk.storage + if storage.urlpath == NULL or storage.io == NULL or storage.io.id != BLOSC2_IO_FILESYSTEM: + return False + if storage.io.params != NULL and ( storage.io.params).locking: + return True + return os.environ.get("BLOSC_LOCKING", "") not in ("", "0") + @_handle_locked def refresh(self): cdef int rc = blosc2_schunk_refresh(self.schunk) diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 985faf1b6..e835bbe95 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -11,6 +11,7 @@ import itertools import math import operator +import os import sys import weakref from abc import abstractmethod @@ -39,6 +40,7 @@ from .linalg import matmul from .utils import ( + chunk_flat_index, get_chunks_idx, get_local_slice, get_selection, @@ -3778,38 +3780,14 @@ def detect_aligned_chunks( if s.stop is not None and s.stop % chunks[i] != 0: return [] - # Parse the slice boundaries - start_indices = [] - end_indices = [] - n_chunks = [] - - for i, s in enumerate(key): - start = s.start if s.start is not None else 0 - stop = s.stop if s.stop is not None else shape[i] - chunk_size = chunks[i] - start_idx = start // chunk_size - end_idx = stop // chunk_size - start_indices.append(start_idx) - end_indices.append(end_idx) - # Total chunk count along this dim: ceil, not floor -- a trailing - # partial chunk (shape[i] not a multiple of chunk_size) still counts - # as one chunk. Floor division here undercounts it, which corrupts - # the flat chunk-index math below for any aligned slice with a - # nonzero start in an earlier dimension (silently returns a - # different chunk's data instead of the requested one). - n_chunks.append(math.ceil(shape[i] / chunk_size)) - # Get all chunk combinations in the slice - indices = [range(start, end) for start, end in zip(start_indices, end_indices, strict=False)] - result = [] - - for combination in product(*indices): - flat_index = 0 - multiplier = 1 - for idx, n in zip(reversed(range(len(n_chunks))), reversed(n_chunks), strict=False): - flat_index += combination[idx] * multiplier - multiplier *= n - result.append(flat_index) + indices = [] + for s, n, c in zip(key, shape, chunks, strict=True): + start = s.start if s.start is not None else 0 + stop = s.stop if s.stop is not None else n + indices.append(range(start // c, stop // c)) + grid, _ = get_chunks_idx(shape, chunks) + result = [chunk_flat_index(combination, grid) for combination in product(*indices)] # Check if chunks are consecutive if requested if consecutive and result: @@ -3825,6 +3803,173 @@ def detect_aligned_chunks( return sorted(result) +_LOSSLESS_CODECS = frozenset( + (blosc2.Codec.BLOSCLZ, blosc2.Codec.LZ4, blosc2.Codec.LZ4HC, blosc2.Codec.ZLIB, blosc2.Codec.ZSTD) +) +_LOSSLESS_FILTERS = frozenset( + ( + blosc2.Filter.NOFILTER, + blosc2.Filter.SHUFFLE, + blosc2.Filter.BITSHUFFLE, + blosc2.Filter.DELTA, + blosc2.Filter.BYTEDELTA, + ) +) + + +def _is_lossless(cparams: blosc2.CParams) -> bool: + """Whether compressing with *cparams* gives back exactly the input (no lossy codec or filter).""" + return ( + cparams.codec in _LOSSLESS_CODECS + and builtins.all(f in _LOSSLESS_FILTERS for f in cparams.filters) + and not cparams.use_dict + ) + + +_CHUNK_ENCODING = ( + "codec", + "codec_meta", + "clevel", + "use_dict", + "typesize", + "blocksize", + "splitmode", + "filters", + "filters_meta", +) + + +def _chunks_interchangeable(dst: NDArray, src: NDArray) -> bool: + """Whether the compressed chunks of `src` can stand for those of `dst` (and vice versa): they + lay out the same items the same way (dtype, chunks, blocks) and encode them the same lossless + way.""" + if src.dtype != dst.dtype or src.chunks != dst.chunks or src.blocks != dst.blocks: + return False + a, b = dst.cparams, src.cparams + return _is_lossless(a) and builtins.all(getattr(a, f) == getattr(b, f) for f in _CHUNK_ENCODING) + + +def _resolved_urlpath(a: NDArray, path: str) -> str: + """`path`, the urlpath of `a`, with symlinks resolved. Cached per handle: resolving costs one + lstat per path component, and on NFS one of the file itself also flushes its dirty pages.""" + cached = a.__dict__.get("_resolved_urlpath") + if cached is None or cached[0] != path: + cached = a._resolved_urlpath = (path, os.path.realpath(path)) + return cached[1] + + +def _same_storage(a: NDArray, b: NDArray) -> bool: + """Whether two arrays may share their data (the same super-chunk, or the same file).""" + if a._handle_lock is b._handle_lock: # one lock per C super-chunk, shared by views + return True + pa, pb = a.schunk.urlpath, b.schunk.urlpath + if not (isinstance(pa, str) and isinstance(pb, str)): + return False + return pa == pb or _resolved_urlpath(a, pa) == _resolved_urlpath(b, pb) + + +# Upper bound for the temporary buffers of chunk-wise writes (unless a single chunk is larger) +_CHUNKWISE_MAX_NBYTES = 32 * 2**20 + + +def _region_boxes( + shape, chunks, start, stop, max_items, split_covered=False, reverse=False, split_edges=True +): + """Split the non-empty region ``[start, stop)`` along chunk boundaries into boxes. + + Yields ``(lo, hi, nchunk)``. Boxes hold at most `max_items` items (or one chunk, when that is + larger), so that writing each with one ``set_slice`` keeps memory bounded while staying with + few calls (every call costs some frame I/O on disk). With `split_covered`, a chunk the region + fully covers is always yielded alone, with its index in `nchunk` (None otherwise); an edge + chunk (one with padding beyond the shape) only with `split_edges`. Boxes come in C order, or + the reverse one. + """ + ndim = len(shape) + grid, _ = get_chunks_idx(shape, chunks) + strides = [math.prod(grid[d + 1 :]) for d in range(ndim)] # of the chunk indices + + def covers_a_chunk(lo, hi): + for s, c, l_, h in zip(shape, chunks, lo, hi, strict=True): + i = -(-l_ // c) # first chunk starting inside the box + end = builtins.min((i + 1) * c, s) if split_edges else (i + 1) * c + if i * c >= h or end > h: + return False + return True + + def whole_chunks(lo, hi): + """The chunks of a box made of whole chunks, in the order in which visit() splits it.""" + ranges = [range(l_ // c, -(-h // c)) for l_, h, c in zip(lo, hi, chunks, strict=True)] + if reverse: + ranges = [r[::-1] for r in ranges] + los = [[i * c for i in r] for r, c in zip(ranges, chunks, strict=True)] + his = [ + [builtins.min((i + 1) * c, s) for i in r] for r, c, s in zip(ranges, chunks, shape, strict=True) + ] + offsets = [[i * st for i in r] for r, st in zip(ranges, strides, strict=True)] + return zip(product(*los), product(*his), map(builtins.sum, product(*offsets)), strict=True) + + def visit(lo, hi): + covered = split_covered and covers_a_chunk(lo, hi) + nitems = math.prod(h - l_ for l_, h in zip(lo, hi, strict=True)) + if not covered and nitems <= max_items: + yield lo, hi, None + return + d = next((d for d in range(ndim) if lo[d] // chunks[d] != (hi[d] - 1) // chunks[d]), None) + if d is None: # a single chunk + nchunk = None + if covered: + nchunk = chunk_flat_index([l_ // c for l_, c in zip(lo, chunks, strict=True)], grid) + yield lo, hi, nchunk + return + if covered and builtins.all( + l_ % c == 0 and (h % c == 0 or (split_edges and h == s)) + for l_, h, c, s in zip(lo, hi, chunks, shape, strict=True) + ): + yield from whole_chunks(lo, hi) + return + c = chunks[d] + # Slabs as thick as the budget allows; covered chunks have to be isolated one by one + step = 1 if covered else builtins.max(1, max_items // (nitems // (hi[d] - lo[d]) * c)) + idx = range(lo[d] // c, (hi[d] - 1) // c + 1, step) + for i in reversed(idx) if reverse else idx: + sub_lo = lo[:d] + (builtins.max(lo[d], i * c),) + lo[d + 1 :] + sub_hi = hi[:d] + (builtins.min(hi[d], (i + step) * c),) + hi[d + 1 :] + yield from visit(sub_lo, sub_hi) + + yield from visit(tuple(start), tuple(stop)) + + +def _chunkwise_plan( + shape, + chunks, + start, + stop, + itemsize, + split_covered, + reverse=False, + buffer_covered=False, + split_edges=True, +): + """Plan writing the region ``[start, stop)`` box by box, with boxes of at most + _CHUNKWISE_MAX_NBYTES of `itemsize`-byte items (or one chunk, when that is larger). + + Returns the boxes from _region_boxes as ``(lo, hi, nchunk, nitems)``, and the largest + `nitems` of the boxes that are written from a buffer: all of them with `buffer_covered`, else + all but the covered chunks (`nchunk` not None), whose `nitems` is left out (None). + """ + max_items = builtins.max(math.prod(chunks), _CHUNKWISE_MAX_NBYTES // itemsize) + boxes, buf_items = [], 0 + for lo, hi, nchunk in _region_boxes( + shape, chunks, start, stop, max_items, split_covered, reverse, split_edges + ): + n = None + if nchunk is None or buffer_covered: + n = math.prod(map(operator.sub, hi, lo)) + buf_items = builtins.max(buf_items, n) + boxes.append((lo, hi, nchunk, n)) + return boxes, buf_items + + class NDOuterIterator: def __init__(self, ndarray: NDArray | NDField, cache_size=1): self.ndarray = ndarray @@ -4819,7 +4964,7 @@ def _remember_read(self, key, nparr: np.ndarray) -> None: self._last_read.clear() self._last_read[make_key_hashable(key)] = nparr - def __setitem__( + def __setitem__( # noqa: C901 self, key: int | slice | Sequence[slice | int | np.bool_ | np.ndarray[int | np.bool_] | None] | None, value: object, @@ -4890,26 +5035,28 @@ def __setitem__( _slice = ndindex.ndindex(()).expand(self.shape) # just get whole array else: # do nothing return self - result = self._get_set_findex_default(_slice, value=value) - from . import indexing - - indexing.mark_indexes_stale(self) - return result + 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) 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__.") - result = self._get_set_nonunit_steps((start, stop, step, mask), value=value) - from . import indexing - - indexing.mark_indexes_stale(self) - return result + 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)] + if ( + isinstance(value, NDArray) + 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 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: @@ -4919,12 +5066,157 @@ def __setitem__( # when using complex functions (e.g. conj) with real arrays value = value.real.astype(self.dtype) - result = super().set_slice((start, stop), value) + return self._mark_indexes_stale(super().set_slice((start, stop), value)) + + def _mark_indexes_stale(self, result): + """Mark the indexes of this array stale (its data changed), and return `result`.""" from . import indexing indexing.mark_indexes_stale(self) return result + def _fill_chunkwise(self, start, stop, value) -> bool: + """Fill the region ``[start, stop)`` with a scalar in bounded pieces. + + Every chunk the region fully covers gets one shared special (repeated value) chunk instead + of being compressed (edge chunks only for a zero value: writing them leaves zeros in their + padding, which a resize brings to light), and the rest is written in boxes of bounded + size instead of one buffer as large as the region. Returns False without writing + anything when the region is smaller than a chunk (one buffer and one ``set_slice`` call is + cheaper there), or when the frame takes the cross-process lock: one ``set_slice`` holds + that from refreshing the shape to the last chunk, so the write is atomic for other + processes and lands where the key says even if another handle resizes the array + meanwhile. (Holding it across a Python loop would make other handles of the file in this + process wait for it holding the GIL.) + """ + if self.ndim == 0: + return False + region = math.prod(sp - st for st, sp in zip(start, stop, strict=True)) + chunk_nitems = math.prod(self.chunks) + if region == 0 or region < chunk_nitems or self.schunk._locks_frame: + return False + fill = np.full(1, value, dtype=self.dtype) # same conversion as filling the whole region + const_chunk = None + if self.schunk.typesize == self.dtype.itemsize and _is_lossless(self.cparams): + # A lossy codec would not give back `value` exactly; keep compressing for those + const_chunk = self._constant_chunk(fill) + # Last chunks first: a frame compacts its payload behind an updated chunk, and that + # tail then holds the (tiny) special chunks already written instead of the old ones. + boxes, nbuf = _chunkwise_plan( + self.shape, + self.chunks, + start, + stop, + self.dtype.itemsize, + const_chunk is not None, + reverse=True, + split_edges=fill.tobytes() == bytes(fill.nbytes), + ) + # One buffer, as large as the largest box written from it (often just a few edge chunks) + buf = np.full(nbuf, fill[0], dtype=self.dtype) + # Holding the handle lock for the whole fill keeps it atomic for other threads, as a single + # set_slice was + with self._handle_lock: + for lo, hi, nchunk, n in boxes: + if nchunk is not None: + self.schunk.update_chunk(nchunk, const_chunk) + else: + super().set_slice((lo, hi), buf[:n].reshape(tuple(map(operator.sub, hi, lo)))) + return True + + def _set_array_chunkwise(self, start, stop, mask, value: NDArray) -> bool: + """Copy the NDArray `value` into the region ``[start, stop)`` in bounded pieces. + + Instead of decompressing all of `value` first, read it box by box (boxes of bounded size, + along this array's chunk boundaries) and write each box with one ``set_slice``. When both + arrays encode chunks the same way (dtype, chunks, blocks, cparams) and the region starts at + a chunk boundary, the covered chunks are copied compressed. Returns False without writing + anything when `value` does not have the shape of the region (minus integer-indexed + dimensions), when the region is smaller than a chunk, when `value` may share storage + with this array or takes the cross-process frame lock (it must be read in full before + writing then), or when this array takes that lock (see _fill_chunkwise). + """ + if self.ndim == 0: + return False + region = tuple(sp - st for st, sp in zip(start, stop, strict=True)) + chunk_nitems = math.prod(self.chunks) + nitems = math.prod(region) + if nitems == 0 or nitems < chunk_nitems or self.schunk._locks_frame or _same_storage(self, value): + return False + if value.schunk._locks_frame: + # One read keeps `value` consistent + return False + value.refresh() # as reading it with value[...] would + kept = (True,) * self.ndim if value.shape == region else tuple(not m for m in mask) + if value.shape != tuple(s for s, k in zip(region, kept, strict=True) if k): + return False + copy_compressed = ( + value.shape == region + and builtins.all(st % c == 0 for st, c in zip(start, self.chunks, strict=True)) + and _chunks_interchangeable(self, value) + ) + # Covered chunks are best copied one by one (get_slice and set_slice take their fast paths + # for single chunks), unless that would decompress blocks of `value` more than once. + it = iter(value.blocks) + vblocks = [next(it) if k else 1 for k in kept] # per dimension of the region + one_by_one = copy_compressed or builtins.all( + st // c == (sp - 1) // c or (c % b == 0 and st % b == 0) + for st, sp, c, b in zip(start, stop, self.chunks, vblocks, strict=True) + ) + if copy_compressed: + vgrid, _ = get_chunks_idx(value.shape, self.chunks) + vsize, size = value.dtype.itemsize, self.dtype.itemsize + # Reused raw byte buffers, as large as the largest box read into them (fresh ones would + # page-fault in on every box) + itemsize = builtins.max(vsize, size) + # The padding of an edge chunk of `value` may hold anything (full() repeats its value + # there), where writing leaves zeros: only chunks without padding are copied compressed + boxes, buf_items = _chunkwise_plan( + self.shape, + self.chunks, + start, + stop, + itemsize, + one_by_one, + buffer_covered=not copy_compressed, + split_edges=not copy_compressed, + ) + src_raw = dst_raw = None + # Holding the handle locks (in a fixed order) for the whole copy keeps it atomic for other + # threads, as a single set_slice was + with blosc2_ext.HandleLocks((self, value)): + for lo, hi, nchunk, n in boxes: + if nchunk is not None and copy_compressed: + # A covered chunk starting at a chunk boundary of `value` holds the same items + vcoords = [(l_ - st) // c for l_, st, c in zip(lo, start, self.chunks, strict=True)] + vchunk = chunk_flat_index(vcoords, vgrid) + self.schunk.update_chunk(nchunk, value.schunk.get_chunk(vchunk)) + continue + if src_raw is None: + src_raw = np.empty(buf_items * vsize, dtype=np.uint8) + data = src_raw[: n * vsize] + vkey = [(l_ - st, h - st) for l_, h, st, k in zip(lo, hi, start, kept, strict=True) if k] + blosc2_ext.NDArray.get_slice_numpy(value, data, tuple(zip(*vkey, strict=True))) + if value.dtype != self.dtype: + if dst_raw is None: + dst_raw = np.empty(buf_items * size, dtype=np.uint8) + src, data = data.view(value.dtype), dst_raw[: n * size] + try: + np.copyto(data.view(self.dtype), src, casting="unsafe") + except ComplexWarning: + np.copyto(data.view(self.dtype), src.real, casting="unsafe") + super().set_slice((lo, hi), data) + return True + + def _constant_chunk(self, fill: np.ndarray) -> bytes: + """Return a special chunk for this array where every item equals ``fill[0]``.""" + schunk = self.schunk + if fill.tobytes() == bytes(fill.nbytes): + return schunk._special_chunk(schunk.chunksize, SpecialValue.ZERO.value) + # Pass a void view: datetime dtypes do not export a buffer + void = fill.view(np.dtype((np.void, fill.itemsize))) + return schunk._special_chunk(schunk.chunksize, SpecialValue.VALUE.value, void) + def __iter__(self): """Iterate over the (outer) elements of the array. diff --git a/src/blosc2/utils.py b/src/blosc2/utils.py index 5b1fcaa97..b44572685 100644 --- a/src/blosc2/utils.py +++ b/src/blosc2/utils.py @@ -1006,6 +1006,15 @@ def get_chunks_idx(shape, chunks): return chunks_idx, nchunks +def chunk_flat_index(coords, grid) -> int: + """The index, in C order, of the chunk at `coords` in a grid of `grid` chunks per dimension + (as from get_chunks_idx).""" + nchunk = 0 + for i, n in zip(coords, grid, strict=True): + nchunk = nchunk * n + i + return nchunk + + def process_key(key, shape): # Fast path: a plain tuple (or scalar) of slices/ints, with no ellipsis, # newaxis, arrays, or strided/negative steps, needs none of ndindex's diff --git a/tests/ndarray/test_setitem_chunkwise.py b/tests/ndarray/test_setitem_chunkwise.py new file mode 100644 index 000000000..86e4c72c5 --- /dev/null +++ b/tests/ndarray/test_setitem_chunkwise.py @@ -0,0 +1,344 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +# Assignments of scalars and NDArrays that are written chunk by chunk (bounded memory, special +# chunks, compressed chunk copies). + +import importlib +import os + +import numpy as np +import pytest + +import blosc2 + +ndarray_mod = importlib.import_module("blosc2.ndarray") + +SHAPE, CHUNKS, BLOCKS = (23, 17, 9), (6, 5, 4), (3, 2, 4) +REGIONS = [ + (slice(None),), + (Ellipsis,), + (slice(6, 18),), # whole chunks + (slice(1, 22), slice(2, 16), slice(1, 8)), # full and partial chunks + (slice(18, 23),), # the edge chunks only + (slice(None), slice(None), slice(8, 9)), # partial chunks only + (slice(-5, None), slice(15, None)), + (3, slice(None)), # integer index +] + + +@pytest.fixture(params=[None, 3000], ids=["budget-dflt", "budget-small"]) +def budget(request, monkeypatch): + """Size bound of the buffers of chunk-wise writes: forces boxes of one or several chunks.""" + if request.param is not None: + monkeypatch.setattr(ndarray_mod, "_CHUNKWISE_MAX_NBYTES", request.param) + return request.param + + +def _random(shape, dtype, rng): + dtype = np.dtype(dtype) + if dtype.kind in "fc": + return (rng.normal(size=shape) * 100).astype(dtype) + if dtype.kind in "iumM": + return rng.integers(0, 100, size=shape).astype(dtype) + if dtype.kind == "b": + return rng.integers(0, 2, size=shape).astype(dtype) + return np.zeros(shape, dtype=dtype) + + +def _make(x, tmp_path, name, **kwargs): + if tmp_path is None: + return blosc2.asarray(x, **kwargs) + return blosc2.asarray(x, urlpath=str(tmp_path / name), mode="w", **kwargs) + + +@pytest.mark.parametrize( + ("dtype", "value"), + [ + (np.float32, 3.5), + (np.float32, 0.0), + (np.float32, -0.0), + (np.float64, np.nan), + (np.int64, -7), + (np.uint8, 255), + (np.bool_, True), + (np.bool_, False), + (np.complex128, 1 - 2j), + ("U7", "héllo"), + ("U100", "a longer string value"), + ("S5", b"ab"), + ("m8[s]", 5), + ("M8[ms]", 1577836800000), + ([("a", "f8"), ("b", "i4")], np.array((1.5, 2), dtype=[("a", "f8"), ("b", "i4")])[()]), + ], +) +@pytest.mark.parametrize("key", REGIONS) +@pytest.mark.parametrize("disk", [False, True]) +def test_scalar_fill(dtype, value, key, disk, budget, tmp_path): + x = _random(SHAPE, dtype, np.random.default_rng(0)) + a = _make(x, tmp_path if disk else None, "a.b2nd", chunks=CHUNKS, blocks=BLOCKS) + a[key] = value + x[key] = value + if disk: + a = blosc2.open(a.urlpath) + np.testing.assert_array_equal(a[...], x) + # Partial reads only decompress some blocks of the (special) chunks + sub = (slice(2, 13), slice(3, 9), slice(1, 6)) + np.testing.assert_array_equal(a[sub], x[sub]) + + +@pytest.mark.parametrize(("start", "stop"), [(5, 10), (8, 10), (0, 10), (3, 9), (4, 8), (1, 6)]) +def test_scalar_fill_1d_edges(start, stop, budget): + x = np.arange(10.0) + a = blosc2.asarray(x, chunks=(4,), blocks=(2,)) + a[start:stop] = 7.0 + x[start:stop] = 7.0 + np.testing.assert_array_equal(a[:], x) + np.testing.assert_array_equal(a[2:9], x[2:9]) + + +@pytest.mark.parametrize( + ("value", "special"), [(0.0, blosc2.SpecialValue.ZERO), (2.5, blosc2.SpecialValue.VALUE)] +) +@pytest.mark.parametrize("disk", [False, True]) +def test_scalar_fill_special_chunks(value, special, disk, tmp_path): + x = _random(SHAPE, np.float32, np.random.default_rng(1)) + a = _make(x, tmp_path if disk else None, "a.b2nd", chunks=CHUNKS, blocks=BLOCKS) + a[6:18] = value # chunk rows 1 and 2 are fully covered + info = list(a.schunk.iterchunks_info()) + grid = [-(-s // c) for s, c in zip(SHAPE, CHUNKS, strict=True)] + covered = [i for i in info if 1 <= i.nchunk // (grid[1] * grid[2]) <= 2] + assert len(covered) == 2 * grid[1] * grid[2] + # Edge chunks only get special chunks of zeros: writing them leaves zeros in their padding + for i in covered: + edge = (i.nchunk // grid[2]) % grid[1] == grid[1] - 1 or i.nchunk % grid[2] == grid[2] - 1 + assert i.special == (blosc2.SpecialValue.NOT_SPECIAL if edge and value != 0 else special) + assert all(i.special == blosc2.SpecialValue.NOT_SPECIAL for i in info if i.nchunk < grid[1] * grid[2]) + x[6:18] = value + np.testing.assert_array_equal(a[...], x) + + +@pytest.mark.parametrize("value", ["scalar", "zero", "ndarray"]) +@pytest.mark.parametrize("disk", [False, True]) +def test_edge_chunk_padding(value, disk, tmp_path): + # Writing whole edge chunks leaves zeros in their padding, which a resize brings to light; full() + # repeats its value there instead, as do special chunks of other values than zero + shape, chunks, blocks = (10, 7), (4, 4), (2, 2) + kwargs = {"urlpath": str(tmp_path / "a.b2nd"), "mode": "w"} if disk else {} + a = blosc2.full(shape, 3, dtype=np.int32, chunks=chunks, blocks=blocks, **kwargs) + if value == "ndarray": + a[:] = blosc2.full(shape, 5, dtype=np.int32, chunks=chunks, blocks=blocks) + else: + a[:] = 5 if value == "scalar" else 0 + a.resize((12, 9)) + expected = np.zeros((12, 9), dtype=np.int32) + expected[:10, :7] = 0 if value == "zero" else 5 + np.testing.assert_array_equal(a[...], expected) + + +@pytest.mark.parametrize( + "cparams", + [ + {"filters": [blosc2.Filter.TRUNC_PREC], "filters_meta": [10]}, + {"codec": blosc2.Codec.ZSTD, "use_dict": True}, + ], +) +def test_scalar_fill_lossy_or_dict(cparams): + # Lossy filters must not get exact special chunks: the result has to be what compressing the + # filled region gives, as with an explicit buffer + x = _random((40, 30), np.float64, np.random.default_rng(2)) + a = blosc2.asarray(x, chunks=(10, 10), blocks=(5, 5), cparams=cparams) + b = blosc2.asarray(x, chunks=(10, 10), blocks=(5, 5), cparams=cparams) + a[5:35] = np.pi + b[5:35] = np.full((30, 30), np.pi) + np.testing.assert_array_equal(a[...], b[...]) + assert all(i.special == blosc2.SpecialValue.NOT_SPECIAL for i in a.schunk.iterchunks_info()) + + +def test_scalar_fill_conversion_errors(): + a = blosc2.zeros((40, 30), dtype=np.uint8, chunks=(10, 10), blocks=(5, 5)) + with pytest.raises(OverflowError): + a[...] = 300 + with pytest.raises(ValueError): + a[...] = "abc" + np.testing.assert_array_equal(a[...], 0) + + +SOURCES = [ + ("same", {"chunks": CHUNKS, "blocks": BLOCKS}, None), + ("clevel1", {"chunks": CHUNKS, "blocks": BLOCKS, "cparams": {"clevel": 1}}, None), + ("lz4", {"chunks": CHUNKS, "blocks": BLOCKS, "cparams": {"codec": blosc2.Codec.LZ4}}, None), + ("other-chunks", {"chunks": (4, 7, 9), "blocks": (2, 7, 3)}, None), + ("float64", {"chunks": CHUNKS, "blocks": BLOCKS}, np.float64), + ("int16", {}, np.int16), +] + + +@pytest.mark.parametrize(("name", "kwargs", "dtype"), SOURCES, ids=[s[0] for s in SOURCES]) +@pytest.mark.parametrize("key", REGIONS) +@pytest.mark.parametrize(("dst_disk", "src_disk"), [(False, False), (True, False), (False, True)]) +def test_ndarray_value(name, kwargs, dtype, key, dst_disk, src_disk, budget, tmp_path): + rng = np.random.default_rng(3) + x = _random(SHAPE, np.float32, rng) + a = _make(x, tmp_path if dst_disk else None, "a.b2nd", chunks=CHUNKS, blocks=BLOCKS) + v = _random(x[key].shape, dtype or np.float32, rng) + if "chunks" in kwargs and v.ndim == len(kwargs["chunks"]): + kwargs = dict( + kwargs, chunks=tuple(min(c, s) for c, s in zip(kwargs["chunks"], v.shape, strict=True)) + ) + kwargs["blocks"] = tuple(min(b, c) for b, c in zip(kwargs["blocks"], kwargs["chunks"], strict=True)) + else: + kwargs = {k: val for k, val in kwargs.items() if k not in ("chunks", "blocks")} + b = _make(v, tmp_path if src_disk else None, "b.b2nd", **kwargs) + a[key] = b + x[key] = v + if dst_disk: + a = blosc2.open(a.urlpath) + np.testing.assert_array_equal(a[...], x) + sub = (slice(2, 13), slice(3, 9), slice(1, 6)) + np.testing.assert_array_equal(a[sub], x[sub]) + + +@pytest.mark.parametrize("disk", [False, True]) +def test_ndarray_value_compressed_copy(disk, tmp_path): + rng = np.random.default_rng(4) + x, y = _random(SHAPE, np.float32, rng), _random(SHAPE, np.float32, rng) + a = _make(x, tmp_path if disk else None, "a.b2nd", chunks=CHUNKS, blocks=BLOCKS) + b = blosc2.asarray(y, chunks=CHUNKS, blocks=BLOCKS) + a[...] = b + # Same chunks, blocks, dtype and cparams: the chunks are copied as they are, but for the edge + # ones (their padding is written anew) + grid = [-(-s // c) for s, c in zip(SHAPE, CHUNKS, strict=True)] + inner = [n for n, idx in enumerate(np.ndindex(*grid)) if all(np.array(idx) < np.array(grid) - 1)] + assert inner + assert all(a.schunk.get_chunk(n) == b.schunk.get_chunk(n) for n in inner) + np.testing.assert_array_equal(a[...], y) + # A region starting at a chunk boundary, with a source of the same chunking + c = blosc2.asarray(y[6:18, 5:15], chunks=CHUNKS, blocks=BLOCKS) + a[6:18, 5:15] = c + x = y.copy() + np.testing.assert_array_equal(a[...], x) + a[...] = 0 + a[6:, 5:] = blosc2.asarray(y[6:, 5:], chunks=CHUNKS, blocks=BLOCKS) + x[:] = 0 + x[6:, 5:] = y[6:, 5:] + np.testing.assert_array_equal(a[...], x) + + +def test_ndarray_value_shared_storage(tmp_path): + # Values sharing storage with the destination are read in full before writing + x = _random((60, 50), np.float32, np.random.default_rng(5)) + a = blosc2.asarray(x, chunks=(7, 50), blocks=(7, 10)) + a[...] = a + np.testing.assert_array_equal(a[...], x) + view = blosc2.expand_dims(a, axis=0) # shares the super-chunk of `a` + view[0] = a + a[...] = blosc2.squeeze(view, axis=0) + np.testing.assert_array_equal(a[...], x) + # Two handles on the same file + p = str(tmp_path / "alias.b2nd") + a = blosc2.asarray(x, urlpath=p, mode="w", chunks=(7, 50), blocks=(7, 10)) + a[...] = blosc2.open(p, mode="r") + np.testing.assert_array_equal(blosc2.open(p)[...], x) + # also through a symlink (resolved paths are cached per handle) + link = tmp_path / "link.b2nd" + try: + os.symlink(p, link) + except OSError: + return + b = blosc2.open(str(link), mode="r") + other = blosc2.asarray(x, urlpath=str(tmp_path / "other.b2nd"), mode="w") + for _ in range(2): + assert ndarray_mod._same_storage(a, b) + assert not ndarray_mod._same_storage(a, other) + a[...] = b + np.testing.assert_array_equal(blosc2.open(p)[...], x) + + +def test_ndarray_value_squeezed_and_kinds(): + rng = np.random.default_rng(6) + x = _random((5, 40, 30), np.float32, rng) + a = blosc2.asarray(x, chunks=(1, 10, 10), blocks=(1, 5, 5)) + v = _random((40, 30), np.float32, rng) + a[3] = blosc2.asarray(v, chunks=(10, 10), blocks=(5, 5)) + a[:, 7] = blosc2.asarray(v[:5]) + x[3] = v + x[:, 7] = v[:5] + np.testing.assert_array_equal(a[...], x) + + dt = np.dtype([("a", "f8"), ("b", "i4")]) + xs = np.zeros((30, 20), dt) + xs["a"] = rng.normal(size=(30, 20)) + xs["b"] = 3 + s = blosc2.zeros((30, 20), dtype=dt, chunks=(4, 20), blocks=(2, 10)) + s[...] = blosc2.asarray(xs, chunks=(8, 5), blocks=(4, 5)) + np.testing.assert_array_equal(s[...], xs) + + xu = np.array([["abc", "défgh"] * 10] * 30, dtype="U5") + u = blosc2.zeros((30, 20), dtype="U7", chunks=(4, 20), blocks=(2, 10)) + u[...] = blosc2.asarray(xu, chunks=(8, 5), blocks=(4, 5)) + np.testing.assert_array_equal(u[...], xu.astype("U7")) + + # complex -> float keeps the real part, as the assignment of a NumPy array does + xc = (rng.normal(size=(30, 20)) + 1j).astype(np.complex64) + f = blosc2.zeros((30, 20), dtype=np.float32, chunks=(4, 20), blocks=(2, 10)) + f[...] = blosc2.asarray(xc) + np.testing.assert_array_equal(f[...], xc.real) + + +def test_chunkwise_buffers_fit_the_boxes(): + # The temporary buffers are as large as the largest box written from them, not as the budget + import tracemalloc + + x = np.arange(100 * 100 * 100, dtype=np.float64).reshape(100, 100, 100) + a = blosc2.asarray(x, chunks=(10, 100, 100), blocks=(5, 50, 100)) # 800 KB chunks + b = blosc2.asarray(x[::-1].copy(), chunks=(10, 100, 100), blocks=(5, 50, 100), cparams={"clevel": 1}) + tracemalloc.start() + try: + a[9:91] = 1.0 # two one-row edge boxes (80 KB) and special chunks + _, fill_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + a[...] = b # chunk by chunk (one-chunk boxes) + _, copy_peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert fill_peak < 200_000 + assert copy_peak < 1_200_000 + np.testing.assert_array_equal(a[...], x[::-1]) + + +@pytest.mark.parametrize( + ("dtype", "value", "kwargs"), + [ + (np.float64, 3.5, {}), + (np.float64, 0.0, {}), + (np.int16, -3, {"chunks": (7, 5, 3), "blocks": (3, 2, 2)}), + (np.uint8, 9, {"cparams": blosc2.CParams(codec=blosc2.Codec.LZ4, clevel=9)}), + ( + np.float32, + 1.25, + { + "cparams": blosc2.CParams( + codec=blosc2.Codec.ZSTD, clevel=0, splitmode=blosc2.SplitMode.ALWAYS_SPLIT + ) + }, + ), + ("M8[ms]", 1577836800000, {}), + ([("a", "f8"), ("b", "i4")], np.array((1.5, 2), dtype=[("a", "f8"), ("b", "i4")])[()], {}), + ], +) +def test_constant_chunk_bytes(dtype, value, kwargs): + # The special chunks of scalar fills are those of a temporary SChunk with the array's cparams + a = blosc2.zeros(SHAPE, dtype=dtype, **({"chunks": CHUNKS, "blocks": BLOCKS} | kwargs)) + fill = np.full(1, value, dtype=a.dtype) + tmp = blosc2.SChunk(chunksize=a.schunk.chunksize, cparams=a.cparams) + nitems = a.schunk.chunksize // a.schunk.typesize + if fill.tobytes() == bytes(fill.nbytes): + tmp.fill_special(nitems, blosc2.SpecialValue.ZERO) + else: + tmp.fill_special(nitems, blosc2.SpecialValue.VALUE, fill.view(np.dtype((np.void, fill.itemsize)))) + assert a._constant_chunk(fill) == tmp.get_chunk(0) diff --git a/tests/test_locking.py b/tests/test_locking.py index 6f2757acc..41015218d 100644 --- a/tests/test_locking.py +++ b/tests/test_locking.py @@ -8,6 +8,7 @@ # Tests for the opt-in cross-process file locking for disk-based containers # (the `locking` storage parameter, backed by a `.b2lock` sidecar file). +import math import os import subprocess import sys @@ -1077,23 +1078,29 @@ def test_cross_process_vlmeta_and_growth_interleaved(tmp_path): import numpy as np import blosc2 -urlpath, value, iters = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]) +urlpath, value, iters, kind = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), sys.argv[4] w = blosc2.open(urlpath, mode="a", locking=True) buf = np.full(w.shape, value, dtype=np.int64) +if kind == "scalar": + buf = value +elif kind == "ndarray": + buf = blosc2.asarray(buf, chunks=w.chunks, blocks=w.blocks) for i in range(iters): w[:, :] = buf sys.exit(0) """ -def test_cross_process_overlapping_slice_atomic(tmp_path): +# Scalars and NDArrays, which are written chunk by chunk without locking, take one +# b2nd_set_slice_cbuffer call with it too +@pytest.mark.parametrize("kind", ["numpy", "scalar", "ndarray"]) +def test_cross_process_overlapping_slice_atomic(tmp_path, kind): # Two writer processes repeatedly overwrite the *same* (multi-chunk) # NDArray region with their own distinguishable constant value via - # __setitem__ (b2nd_set_slice_cbuffer at the C level); a locked reader - # sampling the whole region under holding_lock() must only ever see one - # writer's complete pass, never a chunk-wise mix of both (the atomicity - # b2nd_set_slice_cbuffer's exclusive-lock bracket provides -- python - # inherits it through __setitem__ with no code changes of its own). + # __setitem__; a locked reader sampling the whole region under + # holding_lock() must only ever see one writer's complete pass, never a + # chunk-wise mix of both (the atomicity b2nd_set_slice_cbuffer's + # exclusive-lock bracket provides). urlpath = str(tmp_path / "array-overlap.b2nd") nrows, ncols = 200, 50 iters = 150 @@ -1111,7 +1118,7 @@ def test_cross_process_overlapping_slice_atomic(tmp_path): writers = [ subprocess.Popen( - [sys.executable, "-c", OVERLAPPING_SLICE_WRITER_SCRIPT, urlpath, str(v), str(iters)] + [sys.executable, "-c", OVERLAPPING_SLICE_WRITER_SCRIPT, urlpath, str(v), str(iters), kind] ) for v in (1, 2) ] @@ -1138,6 +1145,201 @@ def test_cross_process_overlapping_slice_atomic(tmp_path): blosc2.remove_urlpath(urlpath) +@pytest.mark.parametrize("kind", ["scalar", "ndarray", "ndarray-cast"]) +def test_chunkwise_write_holds_the_handle_lock(tmp_path, monkeypatch, kind): + # Scalars and NDArrays are written chunk by chunk (without locking); for the whole write, as + # with one set_slice, other threads must not get the handle lock + import threading + + urlpath = tmp_path / "array-chunkwise.b2nd" + shape, chunks, blocks = (40, 30), (10, 10), (5, 5) + a = blosc2.zeros(shape, dtype=np.int64, chunks=chunks, blocks=blocks, urlpath=str(urlpath)) + handle_lock = a._handle_lock + probes = [] + + def probe(): + got = [] + + def try_handle_lock(): + got.append(handle_lock.acquire(blocking=False)) + if got[0]: + handle_lock.release() + + t = threading.Thread(target=try_handle_lock) + t.start() + t.join() + probes.append(got[0]) + + def probing(f): + def wrapper(*args, **kwargs): + probe() + return f(*args, **kwargs) + + return wrapper + + keys = [(slice(None), slice(None)), (slice(1, 40), slice(0, 25))] + x = np.arange(math.prod(shape)).reshape(shape) + if kind == "scalar": + values = [7, 8] + else: + dtype = np.int64 if kind == "ndarray" else np.int32 + values = [blosc2.asarray(x[key].astype(dtype), chunks=chunks, blocks=blocks) for key in keys] + # Hooks in the middle of the writes: covered chunks are stored with update_chunk, casts copy + monkeypatch.setattr(blosc2.SChunk, "update_chunk", probing(blosc2.SChunk.update_chunk)) + if kind == "ndarray-cast": + monkeypatch.setattr(np, "copyto", probing(np.copyto)) + expected = np.zeros(shape, dtype=np.int64) + for key, value in zip(keys, values, strict=True): + a[key] = value + expected[key] = value if kind == "scalar" else x[key] + monkeypatch.undo() + np.testing.assert_array_equal(a[:], expected) + assert len(probes) > 1 + assert probes == [False] * len(probes) + del a + blosc2.remove_urlpath(str(urlpath)) + + +def test_chunkwise_write_of_a_locked_ndarray(tmp_path, monkeypatch): + # An NDArray value that takes its own frame lock is read in one piece (consistently), and a + # destination that takes one is written in one piece: no chunk-wise copy + monkeypatch.delenv("BLOSC_LOCKING", raising=False) + shape, chunks = (40, 30), (10, 10) + x = np.arange(math.prod(shape)).reshape(shape) + bpath, apath = str(tmp_path / "b.b2nd"), str(tmp_path / "a.b2nd") + b = blosc2.asarray(x, chunks=chunks, urlpath=bpath, locking=True) + c = blosc2.open(bpath) + assert b.schunk._locks_frame + assert not c.schunk._locks_frame + for value, locking, ncalls in [(b, True, 0), (b, False, 0), (c, True, 0), (c, False, 12)]: + a = blosc2.zeros(shape, dtype=x.dtype, chunks=chunks, urlpath=apath, mode="w", locking=locking) + calls = [] + with monkeypatch.context() as m: + m.setattr(blosc2.SChunk, "update_chunk", lambda *args, calls=calls: calls.append(args)) + a[:] = value + # Without the lock, the chunks are copied compressed + assert len(calls) == ncalls + if ncalls == 0: + np.testing.assert_array_equal(a[:], x) + del a + blosc2.remove_urlpath(apath) + blosc2.remove_urlpath(bpath) + + +SAME_FILE_THREADS_SCRIPT = """ +import sys +import threading +import time + +import blosc2 + +h1 = blosc2.open(sys.argv[1], mode="a", locking=True) +h2 = blosc2.open(sys.argv[1], mode="a", locking=True) +src = None +if sys.argv[2] == "ndarray": + src = blosc2.full(h1.shape, 3.0, dtype=h1.dtype, chunks=h1.chunks, blocks=h1.blocks) +stop = time.monotonic() + 1 + + +def fill(): + i = 0 + while time.monotonic() < stop: + h1[:] = float(i) if src is None else src + i += 1 + + +def write_one(): + while time.monotonic() < stop: + h2[0, 0] = 1.0 + + +threads = [threading.Thread(target=fill), threading.Thread(target=write_one)] +for t in threads: + t.start() +for t in threads: + t.join() +""" + + +@pytest.mark.parametrize("kind", ["scalar", "ndarray"]) +def test_same_file_handles_write_in_threads(tmp_path, kind): + # Two locked handles of one file, written from two threads: a writer must not hold the frame + # lock while the other thread waits for it holding the GIL (both would hang) + urlpath = str(tmp_path / "threads.b2nd") + blosc2.zeros((400, 400), chunks=(50, 50), blocks=(10, 10), urlpath=urlpath, mode="w", locking=True) + try: + subprocess.run( + [sys.executable, "-c", SAME_FILE_THREADS_SCRIPT, urlpath, kind], check=True, timeout=60 + ) + except subprocess.TimeoutExpired: + pytest.fail("writing two handles of a locked file from two threads deadlocked") + blosc2.remove_urlpath(urlpath) + + +@pytest.mark.parametrize("kind", ["scalar", "ndarray"]) +def test_write_after_another_handle_resized(tmp_path, monkeypatch, kind): + # Another handle resizes the array after this one has followed the shape for its key, before it + # writes: the write must land where the key says in the new shape, as set_slice does it under + # the frame lock, refreshing the shape there + urlpath = str(tmp_path / "resized.b2nd") + blosc2.zeros((4, 4), dtype=np.int32, chunks=(2, 2), blocks=(2, 2), urlpath=urlpath, locking=True) + h1 = blosc2.open(urlpath, mode="a", locking=True) + h2 = blosc2.open(urlpath, mode="a", locking=True) + get_basic_start_stop = blosc2.NDArray.get_basic_start_stop + + def resized_meanwhile(self, key): + basic = get_basic_start_stop(self, key) + if h2.shape == (4, 4): + h2.resize((4, 8)) + return basic + + monkeypatch.setattr(blosc2.NDArray, "get_basic_start_stop", resized_meanwhile) + value = 7 + if kind == "ndarray": + value = blosc2.full((2, 4), 7, dtype=np.int32, chunks=(2, 2), blocks=(2, 2)) + h1[2:4, 0:4] = value + monkeypatch.undo() + expected = np.zeros((4, 8), dtype=np.int32) + expected[2:4, 0:4] = 7 + np.testing.assert_array_equal(h2[...], expected) + + +CROSS_COPY_SCRIPT = """ +import sys +import blosc2 + +src, dst, iters = sys.argv[1], sys.argv[2], int(sys.argv[3]) +s = blosc2.open(src, mode="r", locking=True) +d = blosc2.open(dst, mode="a", locking=True) +for i in range(iters): + d[:] = s +sys.exit(0) +""" + + +def test_cross_process_copies_both_ways(tmp_path): + # a[:] = b in one process and b[:] = a in another: neither may wait for the lock of its source + # while holding the lock of its destination + paths = [str(tmp_path / f"{name}.b2nd") for name in "ab"] + for i, path in enumerate(paths): + blosc2.full((64, 64), i, dtype=np.int64, chunks=(16, 16), urlpath=path, mode="w", locking=True) + procs = [ + subprocess.Popen([sys.executable, "-c", CROSS_COPY_SCRIPT, src, dst, "100"]) + for src, dst in (paths, paths[::-1]) + ] + try: + for p in procs: + p.wait(timeout=120) + finally: + for p in procs: + if p.poll() is None: + p.kill() + p.wait() + assert all(p.returncode == 0 for p in procs), "a copying process failed or deadlocked" + for path in paths: + blosc2.remove_urlpath(path) + + # --------------------------------------------------------------------------- # EmbedStore under locking: transactional writes + key-map re-sync # --------------------------------------------------------------------------- From 97e1c82214f802838bfd85fef1012a44df390302 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:46:09 +0100 Subject: [PATCH 11/12] Cast whole chunks with reused buffers in asarray/astype asarray() of an NDArray, and so astype(), sliced every chunk with __getitem__ and cast it into fresh arrays: chunk-sized allocations per chunk, which glibc trims and page-faults back in each time. When source and destination have the same shape, chunks and blocks (the default for astype), chunks without padding are now decompressed into one reused buffer, cast item by item into another and compressed from there; edge chunks keep the per-chunk path (their padding may hold anything, and casting it could warn). Chunks encoded the same way in both arrays are copied compressed (e.g. asarray(a, copy=True, urlpath=...)). Measured (300x512x512 float32, Ryzen 9 7950X, this change toggled): astype(np.uint8) 110 -> 56 ms, with nnU-Net-like chunks 286 -> 267 ms; astype(np.float64) unchanged (compression-bound). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/ndarray.py | 46 ++++++++++++++++++++++++- tests/ndarray/test_setitem_chunkwise.py | 41 +++++++++++++++++++++- 2 files changed, 85 insertions(+), 2 deletions(-) diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index e835bbe95..4c028c956 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -7250,7 +7250,7 @@ def _ndarray_asarray_requires_copy( return builtins.any(key in user_kwargs for key in copy_keys) -def asarray( +def asarray( # noqa: C901 array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs: Any ) -> NDArray | blosc2.UTF8Array: """Convert the `array` to an `NDArray`. @@ -7353,6 +7353,8 @@ def asarray( # Create the empty array ndarr = empty(shape, dtype, chunks=chunks, blocks=blocks, **kwargs) + if isinstance(array, NDArray) and _copy_chunkwise(array, ndarr): + return ndarr behaved = are_partitions_behaved(shape, chunks, blocks) # Get the coordinates of the chunks @@ -7382,6 +7384,48 @@ def asarray( return ndarr +def _copy_chunkwise(src: NDArray, dst: NDArray) -> bool: + """Fill `dst` (just created) with the items of `src`, casting them to the dtype of `dst`. + + Both must have the same shape, chunks and blocks, so that their chunks lay out the same items + the same way; returns False, writing nothing, otherwise. Chunks without padding are + decompressed into one reused buffer, cast into another and compressed from there: no slicing + and no allocations per chunk. Chunks encoded the same way in both arrays are copied as they are. + """ + if src.ndim == 0 or src.shape != dst.shape or src.chunks != dst.chunks or src.blocks != dst.blocks: + return False + extshape = src.ext_chunks + nitems = math.prod(extshape) + if ( + nitems * src.dtype.itemsize != src.schunk.chunksize + or nitems * dst.dtype.itemsize != dst.schunk.chunksize + ): + return False + if _chunks_interchangeable(dst, src): + for nchunk in range(src.schunk.nchunks): + dst.schunk.update_chunk(nchunk, src.schunk.get_chunk(nchunk)) + return True + # Raw byte buffers: datetime dtypes do not export a buffer themselves + src_raw = np.empty(src.schunk.chunksize, dtype=np.uint8) + dst_raw = np.empty(dst.schunk.chunksize, dtype=np.uint8) + src_buf, dst_buf = src_raw.view(src.dtype), dst_raw.view(dst.dtype) + grid, _ = get_chunks_idx(src.shape, src.chunks) + for nchunk, coords in enumerate(np.ndindex(grid)): + valid = tuple( + builtins.min(c, s - i * c) for i, c, s in zip(coords, src.chunks, src.shape, strict=True) + ) + if valid == extshape: + # Casting goes item by item, so the block layout inside the chunk does not matter + src.schunk.decompress_chunk(nchunk, src_raw) + np.copyto(dst_buf, src_buf, casting="unsafe") + dst.schunk.update_data(nchunk, dst_raw, copy=False) + else: + # Padding may hold anything (casting it could warn): go through the items only + box = tuple(slice(i * c, i * c + v) for i, c, v in zip(coords, src.chunks, valid, strict=True)) + dst[box] = np.require(src[box], dtype=dst.dtype, requirements="C") + return True + + def array(array: Sequence | blosc2.Array, copy: bool | None = True, **kwargs: Any) -> NDArray: """Create an :class:`NDArray` from an array-like object. diff --git a/tests/ndarray/test_setitem_chunkwise.py b/tests/ndarray/test_setitem_chunkwise.py index 86e4c72c5..48f446a91 100644 --- a/tests/ndarray/test_setitem_chunkwise.py +++ b/tests/ndarray/test_setitem_chunkwise.py @@ -6,7 +6,7 @@ ####################################################################### # Assignments of scalars and NDArrays that are written chunk by chunk (bounded memory, special -# chunks, compressed chunk copies). +# chunks, compressed chunk copies), and the chunk-wise casts of astype(). import importlib import os @@ -290,6 +290,45 @@ def test_ndarray_value_squeezed_and_kinds(): np.testing.assert_array_equal(f[...], xc.real) +# asarray() casts arrays from 16 MB on chunk by chunk +BIG = (65, 256, 256) + + +@pytest.mark.parametrize( + ("chunks", "blocks"), + [(None, None), ((32, 100, 64), (8, 50, 32)), ((13, 256, 256), (13, 16, 256))], +) +@pytest.mark.parametrize("dtype", [np.float64, np.uint8, np.int16, np.float32]) +def test_astype_chunkwise(chunks, blocks, dtype): + x = _random(BIG, np.float32, np.random.default_rng(7)) + if np.dtype(dtype).kind in "iu": + # Casting floats out of an integer type's range is undefined: on aarch64 NumPy warns, and its + # contiguous and strided casts give different values + info = np.iinfo(dtype) + x = np.clip(x, info.min, info.max) + a = blosc2.asarray(x, chunks=chunks, blocks=blocks) + b = blosc2.astype(a, dtype) + assert b.dtype == np.dtype(dtype) + assert b.chunks == a.chunks + assert b.blocks == a.blocks + np.testing.assert_array_equal(b[...], x.astype(dtype)) + sub = (slice(3, 50), slice(10, 200), slice(7, 250)) + np.testing.assert_array_equal(b[sub], x[sub].astype(dtype)) + + +def test_astype_chunkwise_kinds(tmp_path): + rng = np.random.default_rng(8) + x = rng.integers(-100, 100, size=BIG).astype(np.int64) + a = blosc2.asarray(x, chunks=(16, 128, 256), blocks=(4, 32, 256)) + np.testing.assert_array_equal(blosc2.astype(a, "m8[s]")[...], x.astype("m8[s]")) + c = blosc2.astype(a, np.complex64, urlpath=str(tmp_path / "c.b2nd"), mode="w") + np.testing.assert_array_equal(blosc2.open(c.urlpath)[...], x.astype(np.complex64)) + # A copy with the same chunk encoding copies the compressed chunks + d = blosc2.asarray(a, copy=True, urlpath=str(tmp_path / "d.b2nd"), mode="w") + assert all(a.schunk.get_chunk(n) == d.schunk.get_chunk(n) for n in range(a.schunk.nchunks)) + np.testing.assert_array_equal(blosc2.open(d.urlpath)[...], x) + + def test_chunkwise_buffers_fit_the_boxes(): # The temporary buffers are as large as the largest box written from them, not as the budget import tracemalloc From 10fa50ae635ef2af8fbfe9c73a0eee75e92d8cb2 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Wed, 4 Nov 2026 08:53:55 +0100 Subject: [PATCH 12/12] Preserve per-handle callbacks in chunk-wise copies and fills Co-Authored-By: Codex --- src/blosc2/blosc2_ext.pyx | 12 ++ src/blosc2/ndarray.py | 10 +- tests/ndarray/test_chunkwise_callbacks.py | 140 ++++++++++++++++++++++ 3 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 tests/ndarray/test_chunkwise_callbacks.py diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 71ee9b7eb..ac7716d8a 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -1993,6 +1993,18 @@ cdef class SChunk: def c_schunk(self): return self.schunk + @property + def _has_prefilter(self): + """Whether writes through this handle run a prefilter.""" + with self._handle_lock: + return self.schunk.storage.cparams.prefilter != NULL + + @property + def _has_postfilter(self): + """Whether reads through this handle run a postfilter.""" + with self._handle_lock: + return self.schunk.storage.dparams.postfilter != NULL + @property def chunksize(self): return self.schunk.chunksize diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 4c028c956..2b517307b 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -3843,6 +3843,9 @@ def _chunks_interchangeable(dst: NDArray, src: NDArray) -> bool: """Whether the compressed chunks of `src` can stand for those of `dst` (and vice versa): they lay out the same items the same way (dtype, chunks, blocks) and encode them the same lossless way.""" + # Copying compressed bytes skips both the source postfilter and the destination prefilter. + if src.schunk._has_postfilter or dst.schunk._has_prefilter: + return False if src.dtype != dst.dtype or src.chunks != dst.chunks or src.blocks != dst.blocks: return False a, b = dst.cparams, src.cparams @@ -5089,7 +5092,9 @@ def _fill_chunkwise(self, start, stop, value) -> bool: meanwhile. (Holding it across a Python loop would make other handles of the file in this process wait for it holding the GIL.) """ - if self.ndim == 0: + # Special chunks skip compression, including its prefilter. Keep a single ordinary + # set_slice call with callbacks so partial chunks are filtered exactly as before. + if self.ndim == 0 or self.schunk._has_prefilter: return False region = math.prod(sp - st for st, sp in zip(start, stop, strict=True)) chunk_nitems = math.prod(self.chunks) @@ -5136,7 +5141,8 @@ def _set_array_chunkwise(self, start, stop, mask, value: NDArray) -> bool: with this array or takes the cross-process frame lock (it must be read in full before writing then), or when this array takes that lock (see _fill_chunkwise). """ - if self.ndim == 0: + # One ordinary write preserves prefilter calls on the containing partial chunks too. + if self.ndim == 0 or self.schunk._has_prefilter: return False region = tuple(sp - st for st, sp in zip(start, stop, strict=True)) chunk_nitems = math.prod(self.chunks) diff --git a/tests/ndarray/test_chunkwise_callbacks.py b/tests/ndarray/test_chunkwise_callbacks.py new file mode 100644 index 000000000..4bc4e17ec --- /dev/null +++ b/tests/ndarray/test_chunkwise_callbacks.py @@ -0,0 +1,140 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +import numpy as np +import pytest + +import blosc2 + +# Whole arrays, aligned whole chunks, mixed whole/partial chunks, and one partial chunk. +KEYS = [slice(None), slice(8, 24), slice(4, 28), slice(4, 8)] + + +def make_array(values, chunks=(8,), blocks=(4,)): + return blosc2.asarray( + values, chunks=chunks, blocks=blocks, cparams={"nthreads": 1}, dparams={"nthreads": 1} + ) + + +@pytest.mark.parametrize("key", KEYS) +def test_assignment_applies_source_postfilter(key): + expected = np.zeros(32, dtype=np.int32) + values = np.arange(expected[key].size, dtype=np.int32) + src = make_array(values) + dst = make_array(expected) + + @src.schunk.postfilter(np.int32, np.int32) + def copy_source_postfilter(inp, out, offset): + out[:] = inp + 100 + + try: + expected[key] = values + 100 + dst[key] = src + np.testing.assert_array_equal(dst[:], expected) + finally: + src.schunk.remove_postfilter("copy_source_postfilter") + + +@pytest.mark.parametrize("scalar", [True, False], ids=["scalar", "array"]) +@pytest.mark.parametrize("key", KEYS) +def test_assignment_applies_destination_prefilter(key, scalar): + expected = np.zeros(32, dtype=np.int32) + dst = make_array(expected) + reference = make_array(expected) + values = 7 if scalar else np.arange(expected[key].size, dtype=np.int32) + src = values if scalar else make_array(values) + + @dst.schunk.prefilter(np.int32, np.int32) + def copy_destination_prefilter(inp, out, offset): + out[:] = inp + 100 + + @reference.schunk.prefilter(np.int32, np.int32) + def reference_destination_prefilter(inp, out, offset): + out[:] = inp + 100 + + try: + # NumPy input takes the ordinary set_slice path. A partial write recompresses its + # containing chunks, so their existing items also pass through the prefilter. + reference[key] = np.full(expected[key].shape, values, dtype=np.int32) if scalar else values + dst[key] = src + np.testing.assert_array_equal(dst[:], reference[:]) + np.testing.assert_array_equal(dst[key], np.broadcast_to(values + 100, expected[key].shape)) + finally: + dst.schunk.remove_prefilter("copy_destination_prefilter") + reference.schunk.remove_prefilter("reference_destination_prefilter") + + +def test_large_asarray_copy_applies_source_postfilter(): + # Exactly 16 MiB reaches asarray's chunk-wise path, rather than its small-array path. + values = np.zeros(2**22, dtype=np.int32) + src = make_array(values, chunks=(2**18,), blocks=(2**16,)) + + @src.schunk.postfilter(np.int32, np.int32) + def large_copy_source_postfilter(inp, out, offset): + out[:] = inp + 100 + + try: + dst = blosc2.asarray(src, copy=True, cparams={"nthreads": 1}, dparams={"nthreads": 1}) + np.testing.assert_array_equal(dst[:], values + 100) + finally: + src.schunk.remove_postfilter("large_copy_source_postfilter") + + +@pytest.mark.parametrize("callback", ["prefilter", "postfilter"]) +def test_unrelated_callbacks_keep_chunkwise_fast_paths(callback, monkeypatch): + owner = make_array(np.zeros(32, dtype=np.int32)) + src = make_array(np.arange(32, dtype=np.int32)) + dst = make_array(np.zeros(32, dtype=np.int32)) + + @getattr(owner.schunk, callback)(np.int32, np.int32) + def unrelated_copy_callback(inp, out, offset): + out[:] = inp + 100 + + special_chunks = [] + copied_chunks = [] + original_constant_chunk = blosc2.NDArray._constant_chunk + original_get_chunk = blosc2.SChunk.get_chunk + + def constant_chunk(self, fill): + special_chunks.append(self) + return original_constant_chunk(self, fill) + + def get_chunk(self, nchunk): + if self is src.schunk: + copied_chunks.append(nchunk) + return original_get_chunk(self, nchunk) + + monkeypatch.setattr(blosc2.NDArray, "_constant_chunk", constant_chunk) + monkeypatch.setattr(blosc2.SChunk, "get_chunk", get_chunk) + try: + dst[:] = 7 + np.testing.assert_array_equal(dst[:], np.full(32, 7, dtype=np.int32)) + assert special_chunks == [dst] + dst[:] = src + np.testing.assert_array_equal(dst[:], np.arange(32, dtype=np.int32)) + assert copied_chunks == [0, 1, 2, 3] + finally: + getattr(owner.schunk, "remove_" + callback)("unrelated_copy_callback") + + +@pytest.mark.parametrize("callback", ["prefilter", "postfilter"]) +def test_handle_callback_flags_follow_registration(callback): + owner = make_array(np.zeros(32, dtype=np.int32)) + other = make_array(np.zeros(32, dtype=np.int32)) + flag = "_has_" + callback + assert not getattr(owner.schunk, flag) + + @getattr(owner.schunk, callback)(np.int32, np.int32) + def handle_copy_callback(inp, out, offset): + out[:] = inp + + try: + assert getattr(owner.schunk, flag) + assert not getattr(other.schunk, flag) + finally: + getattr(owner.schunk, "remove_" + callback)("handle_copy_callback") + assert not getattr(owner.schunk, flag)