From 7590301f8393e49cecde06d15faaecc557993b7e Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Wed, 28 Oct 2026 16:18:31 +0100 Subject: [PATCH 01/18] Skip the index store in mark_indexes_stale for unindexed arrays NDArray.__setitem__ calls mark_indexes_stale() after every write. Loading the store resolves the urlpath and stats every cached persistent path, and on NFS each stat() of a file with buffered writes first pushes them to the server. Return early when no cached store has indexes and the array's vlmeta has no index store. Co-Authored-By: Claude Opus 5.5 --- src/blosc2/indexing.py | 18 +++++++++++++++++ tests/ndarray/test_indexing.py | 35 ++++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/blosc2/indexing.py b/src/blosc2/indexing.py index 892ff36ea..da6ee128f 100644 --- a/src/blosc2/indexing.py +++ b/src/blosc2/indexing.py @@ -5097,7 +5097,25 @@ def get_index(array: blosc2.NDArray, field: str | None = None, name: str | None return Index(array, token) +def _may_have_indexes(array: blosc2.NDArray) -> bool: + """Cheap pre-check for :func:`mark_indexes_stale`, which runs after every write. + + Resolving a persistent array's store costs a ``realpath()`` and ``stat()`` calls on + its file, and on NFS each ``stat()`` of a file with buffered writes first pushes them + to the server: a synchronous round trip per ``__setitem__``. A file whose vlmeta has + no index store, in a process that caches no store with indexes either, has nothing + to mark. + """ + if not _is_persistent_array(array): + return True + if any(store.get("indexes") for store in _PERSISTENT_INDEXES.values()): + return True + return INDEXES_VLMETA_KEY in array.schunk.vlmeta + + def mark_indexes_stale(array: blosc2.NDArray) -> None: + if not _may_have_indexes(array): + return store = _load_store(array) if not store["indexes"]: return diff --git a/tests/ndarray/test_indexing.py b/tests/ndarray/test_indexing.py index 23aab884e..89beaf6de 100644 --- a/tests/ndarray/test_indexing.py +++ b/tests/ndarray/test_indexing.py @@ -1886,6 +1886,41 @@ def test_invalidation_on_mark_indexes_stale(tmp_path): assert indexing._HOT_CACHE_BYTES == 0 +def test_setitem_without_indexes_skips_index_store(tmp_path, monkeypatch): + # Loading the store costs realpath() and stat() calls on the file for every write, and on + # NFS each stat() first pushes the writes still buffered for that file to the server + _clear_caches() + urlpath = str(tmp_path / "plain.b2nd") + arr = blosc2.zeros(1_000, dtype=np.int64, chunks=(250,), blocks=(50,), urlpath=urlpath, mode="w") + monkeypatch.setattr( + indexing, + "_load_store", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("an unindexed array has no store")), + ) + arr[:3] = -1 + arr[10:20] = 7 + np.testing.assert_array_equal(arr[:3], np.full(3, -1, dtype=np.int64)) + np.testing.assert_array_equal(arr[10:20], np.full(10, 7, dtype=np.int64)) + + +def test_setitem_marks_on_disk_index_stale_with_cold_caches(tmp_path): + # The index is only known from the file's vlmeta here, not from any in-process cache + urlpath = str(tmp_path / "indexed.b2nd") + arr = blosc2.asarray( + np.arange(1_000, dtype=np.int64), chunks=(250,), blocks=(50,), urlpath=urlpath, mode="w" + ) + arr.create_index(kind=blosc2.IndexKind.FULL) + del arr + _clear_caches() + + arr = blosc2.open(urlpath, mode="a") + arr[:3] = -1 + assert arr.indexes[0]["stale"] is True + del arr + _clear_caches() + assert blosc2.open(urlpath, mode="r").indexes[0]["stale"] is True + + def test_invalidation_on_append(tmp_path): arr, _ = _make_persistent_array(tmp_path) _clear_caches() From 499933b039ceca6af94a8e3329325f7598eea3f0 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Wed, 28 Oct 2026 21:32:42 +0100 Subject: [PATCH 02/18] 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 078cd6227b8bf2249f845d5ac78ef374e2ab5e95 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Wed, 28 Oct 2026 22:26:22 +0100 Subject: [PATCH 03/18] 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 732c1141e..af0d421d8 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -1979,10 +1979,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. @@ -2920,8 +2923,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 03db12aa546eb3de14c9b2ff22b1d1dad0219e6a Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Thu, 29 Oct 2026 07:38:38 +0100 Subject: [PATCH 04/18] Tidy the mark_indexes_stale pre-check and its tests Move _may_have_indexes next to the store helpers, iterate over a snapshot of the persistent store cache, and use plain helpers in the two tests. Co-Authored-By: Claude Opus 5.5 --- src/blosc2/indexing.py | 32 ++++++++++++++++---------------- tests/ndarray/test_indexing.py | 16 ++++++---------- 2 files changed, 22 insertions(+), 26 deletions(-) diff --git a/src/blosc2/indexing.py b/src/blosc2/indexing.py index da6ee128f..cc344d3cd 100644 --- a/src/blosc2/indexing.py +++ b/src/blosc2/indexing.py @@ -555,6 +555,22 @@ def _load_store(array: blosc2.NDArray) -> dict: return store +def _may_have_indexes(array: blosc2.NDArray) -> bool: + """Cheap pre-check for :func:`mark_indexes_stale`, which runs after every write. + + Resolving a persistent array's store costs ``realpath()`` and ``stat()`` calls on its + file, and on NFS each ``stat()`` of a file with buffered writes first pushes them to + the server: a synchronous round trip per ``__setitem__``. A file whose vlmeta has no + index store, in a process that caches no store with indexes either, has nothing to + mark. + """ + if not _is_persistent_array(array): + return True + if any(store.get("indexes") for store in tuple(_PERSISTENT_INDEXES.values())): + return True + return INDEXES_VLMETA_KEY in array.schunk.vlmeta + + def _save_store(array: blosc2.NDArray, store: dict) -> None: store.setdefault("version", INDEX_FORMAT_VERSION) store.setdefault("indexes", {}) @@ -5097,22 +5113,6 @@ def get_index(array: blosc2.NDArray, field: str | None = None, name: str | None return Index(array, token) -def _may_have_indexes(array: blosc2.NDArray) -> bool: - """Cheap pre-check for :func:`mark_indexes_stale`, which runs after every write. - - Resolving a persistent array's store costs a ``realpath()`` and ``stat()`` calls on - its file, and on NFS each ``stat()`` of a file with buffered writes first pushes them - to the server: a synchronous round trip per ``__setitem__``. A file whose vlmeta has - no index store, in a process that caches no store with indexes either, has nothing - to mark. - """ - if not _is_persistent_array(array): - return True - if any(store.get("indexes") for store in _PERSISTENT_INDEXES.values()): - return True - return INDEXES_VLMETA_KEY in array.schunk.vlmeta - - def mark_indexes_stale(array: blosc2.NDArray) -> None: if not _may_have_indexes(array): return diff --git a/tests/ndarray/test_indexing.py b/tests/ndarray/test_indexing.py index 89beaf6de..3495b50d4 100644 --- a/tests/ndarray/test_indexing.py +++ b/tests/ndarray/test_indexing.py @@ -1892,11 +1892,11 @@ def test_setitem_without_indexes_skips_index_store(tmp_path, monkeypatch): _clear_caches() urlpath = str(tmp_path / "plain.b2nd") arr = blosc2.zeros(1_000, dtype=np.int64, chunks=(250,), blocks=(50,), urlpath=urlpath, mode="w") - monkeypatch.setattr( - indexing, - "_load_store", - lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("an unindexed array has no store")), - ) + + def fail_load_store(array): + raise AssertionError("an unindexed array has no store to load") + + monkeypatch.setattr(indexing, "_load_store", fail_load_store) arr[:3] = -1 arr[10:20] = 7 np.testing.assert_array_equal(arr[:3], np.full(3, -1, dtype=np.int64)) @@ -1905,11 +1905,7 @@ def test_setitem_without_indexes_skips_index_store(tmp_path, monkeypatch): def test_setitem_marks_on_disk_index_stale_with_cold_caches(tmp_path): # The index is only known from the file's vlmeta here, not from any in-process cache - urlpath = str(tmp_path / "indexed.b2nd") - arr = blosc2.asarray( - np.arange(1_000, dtype=np.int64), chunks=(250,), blocks=(50,), urlpath=urlpath, mode="w" - ) - arr.create_index(kind=blosc2.IndexKind.FULL) + arr, urlpath = _make_scalar_persistent_array(tmp_path) del arr _clear_caches() From 441dc57569510bfcec9818f0c705ea6180a3b586 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:39:51 +0100 Subject: [PATCH 05/18] 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 359144a924ca26b411bbfad70db4d7ae4b99b68a Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:40:36 +0100 Subject: [PATCH 06/18] 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 e8f72af22433e47248dc3d33b04a9f5663cc98a1 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:41:47 +0100 Subject: [PATCH 07/18] 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 d43eab991..d376d802c 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 @@ -608,6 +613,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 @@ -4004,6 +4010,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 @@ -4081,40 +4099,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 567d56ed9..5d78923cb 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 6c47400cbea1cbd7b895aef96b66bc0cd6b83220 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:42:38 +0100 Subject: [PATCH 08/18] 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 d376d802c..4a81291e3 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -1369,10 +1369,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: @@ -1386,7 +1388,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 21a4eafbc..bbc1c6a7a 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 @@ -1981,7 +1981,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 @@ -1989,7 +1991,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 aa058e044..76649266e 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 b30eefa848684e2f76de9aeead249d0fc95121c0 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:43:56 +0100 Subject: [PATCH 09/18] 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 4a81291e3..9cfc4729d 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -194,6 +194,7 @@ cdef extern from "blosc2.h": ctypedef enum: BLOSC2_SPECIAL_ZERO BLOSC2_SPECIAL_NAN + BLOSC2_SPECIAL_VALUE BLOSC2_SPECIAL_UNINIT ctypedef enum: @@ -521,6 +522,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) @@ -1718,6 +1722,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: @@ -2128,15 +2146,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 @@ -2165,22 +2175,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) @@ -2189,6 +2221,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") @@ -2291,16 +2330,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") @@ -2340,16 +2380,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 76649266e..2907f1447 100644 --- a/src/blosc2/schunk.py +++ b/src/blosc2/schunk.py @@ -820,13 +820,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 88d3ffc6ca3e488711ebcac600afc1e8fec22615 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:45:01 +0100 Subject: [PATCH 10/18] 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 9cfc4729d..408cae03b 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -2586,6 +2586,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 35078a62cdf80fe83600902181568ef197b94118 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:46:09 +0100 Subject: [PATCH 11/18] 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 | 36 ++++++++++++++++++- 2 files changed, 80 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..0a23cb0e7 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,40 @@ 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)) + 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 554a307aabb890f0ac0e4328bdc93d5d118d87dc Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:47:07 +0100 Subject: [PATCH 12/18] Write cframes straight into the returned bytes NDArray.to_cframe() and SChunk.to_cframe() got the frame from b2nd_to_cframe() / blosc2_schunk_to_buffer() in a C buffer and copied it into a new bytes object. They now call blosc2_schunk_to_buffer_alloc() with a callback that creates the bytes object of the final size, and the frame is written into it directly: no second copy of the whole frame. SChunk.to_cframe() still releases the GIL only under set_releasegil(True); the callback takes it back to create the bytes. Needs blosc2_schunk_to_buffer_alloc() from the next c-blosc2 release (or the companion c-blosc2 PR). Measured with it (300x512x512 float32 in memory, default chunks): to_cframe() 54 -> 15 ms and peak memory 456 -> 230 MB. Adds round-trip tests of to_cframe() / from_cframe() for every storage kind and for special, updated and resized chunks. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/blosc2_ext.pyx | 76 +++++++++------- tests/ndarray/test_cframe.py | 172 +++++++++++++++++++++++++++++++++++ 2 files changed, 217 insertions(+), 31 deletions(-) create mode 100644 tests/ndarray/test_cframe.py diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 408cae03b..65c10cfec 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -27,6 +27,7 @@ from cpython cimport ( Py_buffer, PyBUF_SIMPLE, PyBuffer_Release, + PyBytes_AsString, PyBytes_FromStringAndSize, PyObject_GetBuffer, ) @@ -34,7 +35,7 @@ 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.ref cimport PyObject, Py_INCREF, Py_DECREF, Py_XDECREF from cpython.pycapsule cimport PyCapsule_GetPointer, PyCapsule_New from cpython.slice cimport PySlice_AdjustIndices, PySlice_Check, PySlice_Unpack from cython.operator cimport dereference @@ -492,7 +493,9 @@ cdef extern from "blosc2.h": blosc2_schunk *blosc2_schunk_open_offset(const char* urlpath, int64_t offset) nogil blosc2_schunk* blosc2_schunk_open_offset_udio(const char* urlpath, int64_t offset, const blosc2_io *udio) nogil - int64_t blosc2_schunk_to_buffer(blosc2_schunk* schunk, uint8_t** cframe, c_bool* needs_free) nogil + int64_t blosc2_schunk_to_buffer_alloc(blosc2_schunk* schunk, + uint8_t* (*alloc)(int64_t length, void* user_data) noexcept nogil, + void* user_data) nogil void blosc2_schunk_avoid_cframe_free(blosc2_schunk *schunk, c_bool avoid_cframe_free) int blosc2_schunk_lock(blosc2_schunk *schunk) nogil int blosc2_schunk_unlock(blosc2_schunk *schunk) nogil @@ -647,8 +650,6 @@ cdef extern from "b2nd.h": int b2nd_get_sparse_cbuffer(b2nd_array_t *array, int64_t ncoords, const int64_t *coords, void *buffer, int64_t buffersize) int b2nd_from_cframe(uint8_t *cframe, int64_t cframe_len, c_bool copy, b2nd_array_t ** array); - int b2nd_to_cframe(const b2nd_array_t *array, uint8_t ** cframe, int64_t *cframe_len, - c_bool *needs_free); int b2nd_squeeze(b2nd_array_t *array, b2nd_array_t **view) int b2nd_squeeze_index(b2nd_array_t *array, b2nd_array_t **view, const c_bool *index) @@ -2546,21 +2547,7 @@ cdef class SChunk: @_handle_locked def to_cframe(self): - cdef c_bool needs_free - cdef uint8_t *cframe - cdef int64_t cframe_len - if RELEASEGIL: - with nogil: - cframe_len = blosc2_schunk_to_buffer(self.schunk, &cframe, &needs_free) - else: - cframe_len = blosc2_schunk_to_buffer(self.schunk, &cframe, &needs_free) - if cframe_len < 0: - raise RuntimeError("Error while getting the cframe") - out = PyBytes_FromStringAndSize(cframe, cframe_len) - if needs_free: - free(cframe) - - return out + return schunk_to_cframe(self.schunk, RELEASEGIL) @_handle_locked def _avoid_cframe_free(self, avoid_cframe_free): @@ -3770,6 +3757,44 @@ cdef schunk_is_ndarray(blosc2_schunk* schunk): return blosc2_meta_exists(schunk, meta) >= 0 +cdef struct cframe_dest: + PyObject* out + c_bool alloc_failed + + +cdef uint8_t* cframe_alloc(int64_t length, void* user_data) noexcept with gil: + # The bytes object that to_cframe() returns, for the frame to be written into + cdef cframe_dest* dest = user_data + try: + out = PyBytes_FromStringAndSize(NULL, length) + except MemoryError: + dest.alloc_failed = True + return NULL + Py_INCREF(out) + dest.out = out + return PyBytes_AsString(out) + + +cdef schunk_to_cframe(blosc2_schunk* schunk, c_bool release_gil): + """Serialize `schunk` straight into a new bytes object (a cframe).""" + cdef cframe_dest dest + dest.out = NULL + dest.alloc_failed = False + cdef int64_t cframe_len + if release_gil: + with nogil: + cframe_len = blosc2_schunk_to_buffer_alloc(schunk, cframe_alloc, &dest) + else: + cframe_len = blosc2_schunk_to_buffer_alloc(schunk, cframe_alloc, &dest) + out = dest.out if dest.out != NULL else None + Py_XDECREF(dest.out) + if dest.alloc_failed: + raise MemoryError() + if cframe_len < 0: + raise RuntimeError("Error while getting the cframe") + return out + + def schunk_from_cframe(cframe, copy=False): cdef Py_buffer buf PyObject_GetBuffer(cframe, &buf, PyBUF_SIMPLE) @@ -4627,18 +4652,7 @@ cdef class NDArray: @_handle_locked def to_cframe(self): - cdef c_bool needs_free - cdef uint8_t *cframe - cdef int64_t cframe_len; - cdef int rc; - rc = b2nd_to_cframe(self.array, &cframe, &cframe_len, &needs_free) - if rc < 0: - raise RuntimeError("Error while getting the cframe") - out = PyBytes_FromStringAndSize(cframe, cframe_len) - if needs_free: - free(cframe) - - return out + return schunk_to_cframe(self.array.sc, False) @_handle_locked def copy(self, dtype, **kwargs): diff --git a/tests/ndarray/test_cframe.py b/tests/ndarray/test_cframe.py new file mode 100644 index 000000000..4bbedbc1f --- /dev/null +++ b/tests/ndarray/test_cframe.py @@ -0,0 +1,172 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +# Serialization of arrays to cframes (to_cframe) and back (ndarray_from_cframe / schunk_from_cframe), for +# every storage kind and for chunks that are special (zeros, NaNs, uninitialized), updated or resized. + +import numpy as np +import pytest + +import blosc2 + +STORAGES = ["memory", "memory_contiguous", "file", "directory"] + + +def storage_kwargs(storage, tmp_path): + if storage == "memory": + return {} + if storage == "memory_contiguous": + return {"contiguous": True} + if storage == "file": + return {"urlpath": str(tmp_path / "src.b2nd"), "contiguous": True, "mode": "w"} + return {"urlpath": str(tmp_path / "src_dir.b2nd"), "contiguous": False, "mode": "w"} + + +def data2d(): + return (np.sin(np.arange(200 * 300, dtype=np.float32) / 50) * 100).reshape(200, 300) + + +def make_plain(**kw): + return blosc2.asarray(data2d(), chunks=(50, 300), blocks=(10, 300), **kw) + + +def make_updated(**kw): + a = blosc2.asarray(data2d(), chunks=(50, 300), blocks=(10, 300), **kw) + a[10:20] = 0 + a[60:70, :100] = np.arange(1000, dtype=np.float32).reshape(10, 100) + a.resize((260, 300)) # zero chunks at the end + a[255:260] = 7 + return a + + +def make_shrunk(**kw): + a = blosc2.asarray(data2d(), chunks=(30, 300), blocks=(10, 300), **kw) + a.resize((95, 300)) + return a + + +def make_meta(**kw): + a = blosc2.asarray(data2d(), meta={"m1": [1, 2, 3], "m2": "hello"}, **kw) + a.vlmeta["v1"] = {"a": 1, "b": [1.5, 2]} + a.vlmeta["v2"] = b"x" * 1000 + return a + + +def make_mixed_special(**kw): + a = blosc2.zeros((1000,), dtype=np.int16, chunks=(100,), blocks=(20,), **kw) + a[150:450] = np.arange(300, dtype=np.int16) + a[900:] = 3 + return a + + +def make_zeros(**kw): + return blosc2.zeros((500, 500), dtype=np.float64, chunks=(100, 500), **kw) + + +def make_nans(**kw): + return blosc2.full((500, 500), np.nan, dtype=np.float32, chunks=(100, 500), **kw) + + +def make_empty(**kw): + return blosc2.zeros((0, 10), dtype=np.float32, **kw) + + +def make_scalar(**kw): + return blosc2.asarray(np.array(3.5), **kw) + + +def make_many_chunks(**kw): + # more than 2048 chunks: the frame offsets take several blocks + return blosc2.asarray(np.arange(300_000, dtype=np.int32), chunks=(100,), blocks=(100,), **kw) + + +MAKERS = [ + make_plain, + make_updated, + make_shrunk, + make_meta, + make_mixed_special, + make_zeros, + make_nans, + make_empty, + make_scalar, + make_many_chunks, +] + + +def check_same(a, b): + assert b.shape == a.shape + assert b.dtype == a.dtype + assert b.chunks == a.chunks + assert b.blocks == a.blocks + if a.ndim == 0: + np.testing.assert_array_equal(b[()], a[()]) + else: + np.testing.assert_array_equal(b[...], a[...]) + assert b.schunk.meta.getall() == a.schunk.meta.getall() + assert b.vlmeta.getall() == a.vlmeta.getall() + + +@pytest.mark.parametrize("storage", STORAGES) +@pytest.mark.parametrize("make", MAKERS) +def test_roundtrip(make, storage, tmp_path): + a = make(**storage_kwargs(storage, tmp_path)) + cframe = a.to_cframe() + assert isinstance(cframe, bytes) + for copy in (True, False): + check_same(a, blosc2.ndarray_from_cframe(cframe, copy=copy)) + check_same(a, blosc2.from_cframe(cframe)) + # Saving goes through the same frame building + path = str(tmp_path / "saved.b2nd") + a.save(path, mode="w") + check_same(a, blosc2.open(path)) + copied = a.copy(urlpath=str(tmp_path / "copied.b2nd"), mode="w") + check_same(a, copied) + + +@pytest.mark.parametrize("storage", ["memory", "file", "directory"]) +@pytest.mark.parametrize("make", MAKERS) +def test_reserialize(make, storage, tmp_path): + # Arrays that are not a frame in memory get a new frame built; it must not depend on where the + # chunks come from (special chunks are plain chunks in a sparse super-chunk, but offsets in a frame) + a = make(**storage_kwargs(storage, tmp_path)) + if a.schunk.nchunks > 2048: + pytest.skip("the offsets take several blocks, compressed by several threads in a varying order") + cframe = a.to_cframe() + assert blosc2.ndarray_from_cframe(cframe, copy=True).to_cframe() == cframe + assert a.to_cframe() == cframe + if storage == "memory": + path = str(tmp_path / "saved.b2nd") + a.save(path, mode="w") + assert blosc2.open(path).to_cframe() == cframe + + +@pytest.mark.parametrize("storage", STORAGES) +def test_schunk_roundtrip(storage, tmp_path): + kw = storage_kwargs(storage, tmp_path) + kw.pop("mode", None) + if "urlpath" in kw: + blosc2.remove_urlpath(kw["urlpath"]) + schunk = blosc2.SChunk(chunksize=4000, cparams={"typesize": 4}, **kw) + schunk.append_data(np.arange(1000, dtype=np.int32)) + schunk.append_data(np.zeros(1000, dtype=np.int32)) + schunk.append_data(np.arange(500, dtype=np.int32)) # a shorter last chunk + schunk.vlmeta["k"] = "v" + cframe = schunk.to_cframe() + for copy in (True, False): + back = blosc2.schunk_from_cframe(cframe, copy=copy) + assert back.nchunks == 3 + assert back[:] == schunk[:] + assert back.vlmeta["k"] == "v" + + +def test_schunk_empty(): + schunk = blosc2.SChunk(chunksize=4000, cparams={"typesize": 4}) + cframe = schunk.to_cframe() + back = blosc2.schunk_from_cframe(cframe, copy=True) + assert back.nchunks == 0 + assert back.to_cframe() == cframe From b28fb0e21366aab0497ebca66c74714873cf92bb Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:48:20 +0100 Subject: [PATCH 13/18] Speed up lazy-expression getitem, std/var and chunked reductions - expr[:] and expr[()] with miniexpr ran the prefilter into a temporary uncompressed NDArray and read that back into NumPy. A clevel=0 chunk is a header followed by the raw prefilter output, so when the chunks tile the flat array (blocks C-contiguous in their chunk, chunks in the array, only the last one partial) each chunk is now produced in its final place in the NumPy result, last to first, its header landing in bytes not computed yet. - Full std()/var() of float32/float64 (no axis, no where, same dtype) are two single-pass reductions: the mean, then the mean of (x - m) * (x - m) with m inlined as a scalar, instead of a broadcast (1, ..., 1) operand that kept the second pass off the fast paths. Other cases, such as dtype=np.float64 for float32 data, keep the general path. std() and var() share one implementation. - Chunked reductions (axis=..., argmin, cumulative_sum, ...) and chunked evaluation decompress the chunks of operands whose blocks are C-contiguous runs of their chunk straight into reused chunk-sized buffers, compute flat chunk indices without NumPy and skip ndindex for whole arrays. The unused `aligned` plumbing is gone. - where(cond, x, y) with a scalar x or y is fused before the operands are specialized, so the scalar is inlined and miniexpr is used (it always fell back to numexpr). - No dtype inference for the bare-array LazyExpr that the NDArray reductions build, and one regex for all constructor checks. Measured (300x512x512 float32 in memory, 28 threads, Ryzen 9 7950X): (a + b)[:] 120 -> 62 ms, a.std() 68 -> 21 ms, (a.mean(), a.std()) 78 -> 31 ms, a.sum(axis=0) 28 -> 22 ms, where(a > 100, a, 0).compute() 184 -> 153 ms; 100x100: a.std() 0.61 -> 0.38 ms. std/var may differ from before in the last ulps (same accuracy as NumPy). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/lazyexpr.py | 399 +++++++++++++++++++------------ src/blosc2/ndarray.py | 27 ++- tests/ndarray/test_lazyexpr.py | 96 ++++++++ tests/ndarray/test_reductions.py | 75 ++++++ 4 files changed, 431 insertions(+), 166 deletions(-) diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index af0d421d8..6d378156a 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -54,6 +54,7 @@ write_b2object_user_vlmeta, ) from .dsl_kernel import DSLKernel, DSLSyntaxError, DSLValidator, specialize_miniexpr_inputs +from .ndarray import blocks_contiguous from .proxy_source import convert_dtype from .utils import ( check_smaller_shape, @@ -294,6 +295,9 @@ def _get_result(expression, chunk_operands, ne_args, where=None, indices=None, _ functions = blosc2_funcs _TRANSIENT_MASK_CPARAMS = blosc2.CParams(codec=blosc2.Codec.LZ4, clevel=5, filters=[blosc2.Filter.SHUFFLE]) _constructor_call_patterns = {name: re.compile(rf"\b{re.escape(name)}\s*\(") for name in constructors} +_any_constructor_call_pattern = re.compile( + rf"\b(?:{'|'.join(re.escape(name) for name in constructors)})\s*\(" +) def _restore_code_unit_shuffle(cparams: blosc2.CParams, dtype) -> None: @@ -320,6 +324,10 @@ def _has_constructor_call(expression: str, constructor: str) -> bool: return _constructor_call_patterns[constructor].search(expression) is not None +def _has_any_constructor_call(expression: str) -> bool: + return _any_constructor_call_pattern.search(expression) is not None + + def _find_constructor_call(expression: str, constructor: str) -> re.Match | None: return _constructor_call_patterns[constructor].search(expression) @@ -1328,7 +1336,7 @@ def do_slices_intersect(slice1: list | tuple, slice2: list | tuple) -> bool: def get_chunk(arr, info, nchunk): - reduc, aligned, low_mem, chunks_idx = info + reduc, low_mem, chunks_idx = info if low_mem: # We don't want to uncompress the chunk, so keep it compressed and @@ -1352,14 +1360,6 @@ def get_chunk(arr, info, nchunk): slice(c * s, min((c + 1) * s, shape[i])) for i, (c, s) in enumerate(zip(coords, chunks, strict=True)) ) - chunks_ = tuple(s.stop - s.start for s in slice_) - - if aligned: - # Decompress the whole chunk and return it - buff = arr.schunk.decompress_chunk(nchunk) - bsize = arr.dtype.itemsize * math.prod(chunks_) - return np.frombuffer(buff[:bsize], dtype=arr.dtype).reshape(chunks_) - return arr[slice_] @@ -1389,10 +1389,10 @@ def read_chunks_worker(arrs, info, queue, stop): max_workers = max(1, min(len(arrs), int(getattr(blosc2, "nthreads", 1) or 1))) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: my_chunk_iter = range(arrs[0].schunk.nchunks) - if len(info) == 5: + if len(info) == 4: if info[-1] is not None: my_chunk_iter = sliced_chunk_iter(chunks_, (), shape, axis=info[-1], nchunk=True) - info = info[:4] + info = info[:3] for i, nchunk in enumerate(my_chunk_iter): futures = [executor.submit(get_chunk, arr, info, nchunk) for arr in arrs] # result() keeps operand order and propagates the first exception raised @@ -1453,12 +1453,43 @@ def read_nchunk(arrs, info): iter_chunks = None +def _chunk_buffers(operands, shape, chunks): + """Reusable chunk-sized buffers for the NDArray operands of `shape` and `chunks` whose + decompressed chunks are C-order. + + That is the case when blocks are C-contiguous runs of their chunk (as with the default + partitions), so a plain decompress_chunk() gives the (padded) chunk, and is cheaper than a + slice read. + """ + return { + key: np.empty(value.chunks, dtype=value.dtype) + for key, value in operands.items() + if isinstance(value, blosc2.NDArray) + and value.ndim > 0 + and value.shape == shape + and value.chunks == chunks + and 0 not in value.blocks + and blocks_contiguous(value.chunks, value.blocks) + } + + def fill_chunk_operands( - operands, slice_, chunks_, full_chunk, aligned, nchunk, iter_disk, chunk_operands, reduc=False, axis=None + operands, + slice_, + chunks_, + full_chunk, + nchunk, + iter_disk, + chunk_operands, + reduc=False, + axis=None, + chunk_buffers=None, ): """Retrieve the chunk operands for evaluating an expression. This function provides an optimized path for full chunks and a slower path for partial chunks. + Operands with a buffer in `chunk_buffers` (see `_chunk_buffers`) get their chunk decompressed + into it, whether the chunk is full or not. """ global iter_chunks @@ -1470,28 +1501,16 @@ def fill_chunk_operands( if iter_chunks is None: # Initialize the iterator for reading the chunks # Take any operand (all should have the same shape and chunks) - key, arr = next(iter(operands.items())) + arr = next(iter(operands.values())) chunks_idx, _ = get_chunks_idx(arr.shape, arr.chunks) - info = (reduc, aligned[key], low_mem, chunks_idx, axis) + info = (reduc, low_mem, chunks_idx, axis) iter_chunks = read_nchunk(list(operands.values()), info) # Run the asynchronous file reading function from a synchronous context chunks = next(iter_chunks) for i, (key, value) in enumerate(operands.items()): - # Chunks are already decompressed, so we can use them directly - if not low_mem: - if full_chunk: - chunk_operands[key] = chunks[i] - else: - chunk_operands[key] = value[slice_] - continue - # Otherwise, we need to decompress them - if aligned[key]: - buff = blosc2.decompress2(chunks[i]) - bsize = value.dtype.itemsize * math.prod(chunks_) - chunk_operands[key] = np.frombuffer(buff[:bsize], dtype=value.dtype).reshape(chunks_) - else: - chunk_operands[key] = value[slice_] + # Chunks are already decompressed (unless low_mem), so full ones can be used directly + chunk_operands[key] = chunks[i] if full_chunk and not low_mem else value[slice_] return # Get the starts and stops for the slice @@ -1505,6 +1524,11 @@ def fill_chunk_operands( if value.shape == (): chunk_operands[key] = value[()] continue + buf = chunk_buffers.get(key) if chunk_buffers else None + if buf is not None: + value.schunk.decompress_chunk(nchunk, buf) + chunk_operands[key] = buf if full_chunk else buf[tuple(slice(0, n) for n in chunks_)] + continue if not full_chunk or not isinstance(value, blosc2.NDArray): # The chunk is not a full one, or has padding, or is not a blosc2.NDArray, # so we need to go the slow path @@ -1520,13 +1544,7 @@ def fill_chunk_operands( value.get_slice_numpy(chunk_operands[key], (starts, stops)) continue - if aligned[key]: - # Decompress the whole chunk and store it - buff = value.schunk.decompress_chunk(nchunk) - bsize = value.dtype.itemsize * math.prod(chunks_) - chunk_operands[key] = np.frombuffer(buff[:bsize], dtype=value.dtype).reshape(chunks_) - else: - chunk_operands[key] = value[slice_] + chunk_operands[key] = value[slice_] def _apply_jit_backend_pragma(expression: str, inputs: dict, jit_backend: str | None) -> str: @@ -1707,6 +1725,65 @@ def _miniexpr_integer_atan2(expression, operands): ) +def _c_run_axis(outer, inner) -> int: + """Axis along which `inner` tiles are C-contiguous runs of `outer`, or -1 if they are not. + + That is the case when the leading dims of the tile are 1, then one dim is free, and the rest + span the container. + """ + k = 0 + while k < len(inner) - 1 and inner[k] == 1: + k += 1 + return k if tuple(inner[k + 1 :]) == tuple(outer[k + 1 :]) else -1 + + +def _chunks_tile_flat(shape, chunks, blocks) -> bool: + """Whether the uncompressed chunks of an array, in order, are consecutive runs of its C-order buffer. + + Blocks must be C-contiguous runs of their chunk, and chunks C-contiguous runs of the array; only + the very last chunk may be partial, so that its padding spills past the end of the array instead + of into another chunk. + """ + if len(shape) == 0 or not blocks_contiguous(chunks, blocks): + return False + k = _c_run_axis(shape, chunks) + return k >= 0 and (shape[k] % chunks[k] == 0 or math.prod(shape[:k]) == 1) + + +# Byte offset and bit of the "memcpyed" flag in a Blosc2 chunk header +_BLOSC2_CHUNK_FLAGS = 2 +_BLOSC_MEMCPYED = 0x2 + + +def _prefilter_to_numpy(arr, data, operand_locks): + """Run the prefilter set on uncompressed `arr` for every chunk, straight into a new NumPy array. + + A clevel=0 chunk is a header followed by the raw prefilter output, so when the chunks tile the + flat array (see `_chunks_tile_flat`) each one can be produced in its final place, with its header + landing in the bytes just in front of it. Going from the last chunk to the first, those bytes + belong to a chunk not computed yet (or to the spare header room before the first one), so the + result needs neither a temporary array nor a copy. Only a partial last chunk, whose padding + would spill past the end of the array, is produced aside and its data copied in. Returns None + if a chunk does not have the expected layout (nothing usable has been produced then). + """ + schunk = arr.schunk + csize = schunk.chunksize + hdr = blosc2.MAX_OVERHEAD + nbytes = math.prod(arr.shape) * arr.dtype.itemsize + raw = np.empty(hdr + nbytes, dtype=np.uint8) + for nchunk in range(schunk.nchunks - 1, -1, -1): + start = nchunk * csize + partial = start + csize > nbytes + dest = np.empty(hdr + csize, dtype=np.uint8) if partial else raw[start : start + hdr + csize] + with operand_locks: + size = schunk._prefilter_data(nchunk, data, dest) + if size != hdr + csize or not dest[_BLOSC2_CHUNK_FLAGS] & _BLOSC_MEMCPYED: + return None + if partial: + raw[hdr + start :] = dest[hdr : hdr + nbytes - start] + return raw[hdr:].view(arr.dtype).reshape(arr.shape) + + @blosc2.jit_config.jit_execution def fast_eval( # noqa: C901 expression: str | Callable[[tuple, np.ndarray, tuple[int]], None], @@ -1806,11 +1883,7 @@ def fast_eval( # noqa: C901 blocks = kwargs.pop("blocks", None) if blocks is None: blocks = basearr.blocks - # Check whether the partitions are aligned and behaved - aligned = { - k: False if not hasattr(k, "chunks") else blosc2.are_partitions_aligned(k.shape, k.chunks, k.blocks) - for k in operands - } + # Check whether the partitions are behaved behaved = blosc2.are_partitions_behaved(shape, chunks, blocks) # Check that all operands are NDArray for fast path @@ -1828,13 +1901,23 @@ def fast_eval( # noqa: C901 expr_string_miniexpr = expr_string operands_miniexpr = operands + # Fuse where(cond, x, y) before specializing, so that scalar x/y get inlined as well; + # fusing afterwards would reference a dropped operand and make miniexpr fail to compile. + where_fused = ( + use_miniexpr + and isinstance(expr_string, str) + and not is_dsl + and where is not None + and len(where) == 2 + ) + if where_fused: + expr_string_miniexpr = f"where({expr_string}, _where_x, _where_y)" if use_miniexpr and isinstance(expr_string, str): - try: - expr_string_miniexpr, operands_miniexpr = specialize_miniexpr_inputs(expr_string, operands) - except Exception: - # If specialization fails, keep original expression/operands and let normal checks decide. - expr_string_miniexpr = expr_string - operands_miniexpr = operands + # If specialization fails, keep original expression/operands and let normal checks decide. + with contextlib.suppress(Exception): + expr_string_miniexpr, operands_miniexpr = specialize_miniexpr_inputs( + expr_string_miniexpr, operands + ) # Check whether we can use miniexpr if use_miniexpr: @@ -1931,6 +2014,15 @@ def _miniexpr_eligible_operand(op): if use_miniexpr: cparams = kwargs.pop("cparams", None) + # getitem results can be written straight into the NumPy output (see _prefilter_to_numpy) + getitem_direct = ( + getitem + and cparams is None + and candidate_blocks is None + and not prune_chunks + and _chunks_tile_flat(shape, chunks, blocks) + ) + getitem_out = None if cparams is None and getitem: # getitem output is throwaway scratch (returned as a NumPy array and # discarded), so compressing it buys nothing but a round trip. @@ -1954,7 +2046,7 @@ def _miniexpr_eligible_operand(op): k: (np.ascontiguousarray(v) if isinstance(v, np.ndarray) else v) for k, v in _pref_ops.items() } - if where is not None and len(where) == 2: + if where is not None and len(where) == 2 and not where_fused: _pref_expr = f"where({_pref_expr}, _where_x, _where_y)" # _cb_anchor keeps the contiguous bitmap alive for the whole # update_data loop (the prefilter dereferences it per block). Pass @@ -1981,11 +2073,14 @@ def _miniexpr_eligible_operand(op): 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 - with operand_locks: - res_eval.schunk.update_data(nchunk, data, copy=False) + if getitem_direct: + getitem_out = _prefilter_to_numpy(res_eval, data, operand_locks) + if getitem_out is None: + for nchunk in range(res_eval.schunk.nchunks): + if chunk_has_candidate is not None and not chunk_has_candidate[nchunk]: + continue + 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. @@ -2018,10 +2113,11 @@ def _miniexpr_eligible_operand(op): pass else: if getitem: - return res_eval[()] + return getitem_out if getitem_out is not None else res_eval[()] return res_eval chunk_operands = {} + chunk_buffers = None if iter_disk else _chunk_buffers(operands, shape, chunks) # Check which chunks intersect with _slice all_chunks = get_intersecting_chunks((), shape, chunks) # if _slice is (), returns all chunks for nchunk, chunk_slice in enumerate(all_chunks): @@ -2031,7 +2127,15 @@ def _miniexpr_eligible_operand(op): full_chunk = chunks_ == chunks # slice is same as chunk fill_chunk_operands( - operands, cslice, chunks_, full_chunk, aligned, nchunk, iter_disk, chunk_operands + operands, + cslice, + chunks_, + full_chunk, + nchunk, + iter_disk, + chunk_operands, + # Partial chunks keep their contiguous slice reads, which UDFs may rely on + chunk_buffers=chunk_buffers if full_chunk else None, ) # Since ne_evaluate() can return a dtype larger than the one in computed in the expression, @@ -2817,7 +2921,6 @@ def reduce_slices( # noqa: C901 same_chunks = all(operand.chunks == o.chunks for o in operands.values() if hasattr(o, "chunks")) same_blocks = all(operand.blocks == o.blocks for o in operands.values() if hasattr(o, "blocks")) fast_path = same_shape and same_chunks and same_blocks and (0 not in operand.chunks) - aligned = dict.fromkeys(operands.keys(), False) iter_disk = False if fast_path: chunks = operand.chunks @@ -2971,16 +3074,19 @@ def reduce_slices( # noqa: C901 intersecting_chunks = get_intersecting_chunks(_slice, shape, chunks) out_init = False res_out_init = False - ratio = ( - np.ceil(np.asarray(shape) / np.asarray(chunks)).astype(np.int64) - if 0 not in chunks - else np.asarray(shape) + # Flat chunk index = sum(chunk coordinate * stride) over the chunk grid + chunk_strides = [] + if 0 not in chunks: + nchunks_dim, _ = get_chunks_idx(shape, chunks) + chunk_strides = [math.prod(nchunks_dim[i + 1 :]) for i in range(len(nchunks_dim))] + chunk_buffers = ( + _chunk_buffers(operands, shape, chunks) if fast_path and not iter_disk and _slice == () else None ) for chunk_slice in intersecting_chunks: cslice = chunk_slice.raw nchunk = ( - builtins.sum(c.start // chunks[i] * np.prod(ratio[i + 1 :]) for i, c in enumerate(cslice)) + builtins.sum(c.start // chunks[i] * chunk_strides[i] for i, c in enumerate(cslice)) if 0 not in chunks else 0 ) @@ -2990,10 +3096,10 @@ def reduce_slices( # noqa: C901 cslice = step_handler(cslice, _slice) offset = tuple(s.start for s in cslice) # offset for the udf starts = [s.start if s.start is not None else 0 for s in cslice] - unit_steps = np.all([s.step == 1 for s in cslice]) + unit_steps = builtins.all(s.step == 1 for s in cslice) cslice_shape = tuple(s.stop - s.start for s in cslice) - # get local index of part of out that is to be updated - cslice_subidx = ndindex.ndindex(cslice).as_subindex(_slice).raw # if _slice is (), just gives cslice + # get local index of part of out that is to be updated (with no _slice, that is cslice) + cslice_subidx = cslice if _slice == () else ndindex.ndindex(cslice).as_subindex(_slice).raw if _slice == () and fast_path and unit_steps: # Fast path full_chunk = cslice_shape == chunks @@ -3002,12 +3108,12 @@ def reduce_slices( # noqa: C901 cslice, cslice_shape, full_chunk, - aligned, nchunk, iter_disk, chunk_operands, reduc=True, axis=reduce_args["axis"] if np.isscalar(reduce_args["axis"]) else None, + chunk_buffers=chunk_buffers, ) else: get_chunk_operands(operands, cslice, chunk_operands, shape) @@ -3631,6 +3737,9 @@ def __init__(self, new_op): # noqa: C901 op = None self.operands = {"o0": value1} self.expression = "o0" if op is None else f"{op}(o0)" + if op is None and isinstance(value1, blosc2.NDArray): + # A bare array (as the NDArray reductions build): no need to infer the dtype + self._dtype = value1.dtype return value2 = ( blosc2.SimpleProxy(value2) @@ -3915,7 +4024,7 @@ def shape(self): return None # Operands shape can change, so we always need to recompute this - if any(_has_constructor_call(self.expression, constructor) for constructor in constructors): + if _has_any_constructor_call(self.expression): # might have an expression with pure constructors opshapes = {k: v if not hasattr(v, "shape") else v.shape for k, v in self.operands.items()} _shape = infer_shape(self.expression, opshapes) # infer shape, includes constructors @@ -4147,50 +4256,69 @@ def mean( out = blosc2.asarray(out, **kwargs) return out - def std( - self, - axis=None, - dtype=None, - keepdims=False, - ddof=0, - where=None, - fp_accuracy: blosc2.FPAccuracy | None = None, - **kwargs, - ): + def _can_full_var(self, dtype): + """Whether _full_var computes the variance in `dtype`: for float32/float64 operands, where + miniexpr evaluates it without per-element conversions, into their own dtype (the inlined + mean would not make the deviations of float32 operands float64).""" + own = np.dtype(self.dtype) + return own.kind == "f" and own.itemsize in (4, 8) and (dtype is None or np.dtype(dtype) == own) + + def _full_var(self, dtype, keepdims, fp_accuracy, execution_kwargs): + """Variance (ddof=0) over all the elements, as two full reductions over the operands. + + The mean goes into the second pass as a Python scalar, which gets inlined into the + expression; a (1, ..., 1) array would be an operand to broadcast, and keep that reduction + off its fast paths. d * d rather than d ** 2 spares a pow() per element. Only for + float32/float64 into the same dtype (see _can_full_var). + """ + mean_value = np.asarray(self.mean(dtype=dtype, fp_accuracy=fp_accuracy, **execution_kwargs)).item() + # Fold the sign into the operator: miniexpr reductions mis-evaluate a negative literal on + # the right of a binary operator, as in "(o0 - -0.5)" + dev = self + (-mean_value) if math.copysign(1.0, mean_value) < 0 else self - mean_value + out = (dev * dev).mean(dtype=dtype, fp_accuracy=fp_accuracy, **execution_kwargs) + # Evaluators other than miniexpr may compute against the Python scalar in float64 + out = np.dtype(self.dtype if dtype is None else dtype).type(out) + return np.reshape(out, (1,) * self.ndim) if keepdims else out + + def _var(self, axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs, finish=None): + """The variance, as var() returns it; std() passes np.sqrt as `finish`.""" execution_kwargs = blosc2.jit_config.pop_execution_options(kwargs) where = self._normalize_where(where) item = kwargs.pop("item", ()) - if item == (): # fast path - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - expr = (self - mean_value) ** 2 + if item == () and axis is None and where is None and self._can_full_var(dtype): + out = self._full_var(dtype, keepdims, fp_accuracy, execution_kwargs) else: - mean_value = self.mean( + if item == (): # fast path + mean_value = self.mean( + axis=axis, + dtype=dtype, + keepdims=True, + where=where, + fp_accuracy=fp_accuracy, + **execution_kwargs, + ) + expr = (self - mean_value) ** 2 + else: + mean_value = self.mean( + axis=axis, + dtype=dtype, + keepdims=True, + where=where, + item=item, + fp_accuracy=fp_accuracy, + **execution_kwargs, + ) + # TODO: Not optimal because we load the whole slice in memory. Would have to write + # a bespoke var function that executed within slice_eval to avoid this probably. + expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 + out = expr.mean( axis=axis, dtype=dtype, - keepdims=True, + keepdims=keepdims, where=where, - item=item, fp_accuracy=fp_accuracy, **execution_kwargs, ) - # TODO: Not optimal because we load the whole slice in memory. Would have to write - # a bespoke std function that executed within slice_eval to avoid this probably. - expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 - out = expr.mean( - axis=axis, - dtype=dtype, - keepdims=keepdims, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) if ddof != 0: num_elements = ( self.get_num_elements(axis, item, fp_accuracy=fp_accuracy, **execution_kwargs) @@ -4199,9 +4327,9 @@ def std( axis=axis, dtype=np.int64, fp_accuracy=fp_accuracy, **execution_kwargs ) ) - out = np.sqrt(out * num_elements / (num_elements - ddof)) - else: - out = np.sqrt(out) + out = out * num_elements / (num_elements - ddof) + if finish is not None: + out = finish(out) out2 = kwargs.pop("out", None) if out2 is not None: out2[:] = out @@ -4210,6 +4338,18 @@ def std( out = blosc2.asarray(out, **kwargs) return out + def std( + self, + axis=None, + dtype=None, + keepdims=False, + ddof=0, + where=None, + fp_accuracy: blosc2.FPAccuracy | None = None, + **kwargs, + ): + return self._var(axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs, np.sqrt) + def var( self, axis=None, @@ -4220,56 +4360,7 @@ def var( fp_accuracy: blosc2.FPAccuracy | None = None, **kwargs, ): - execution_kwargs = blosc2.jit_config.pop_execution_options(kwargs) - where = self._normalize_where(where) - item = kwargs.pop("item", ()) - if item == (): # fast path - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - expr = (self - mean_value) ** 2 - else: - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - item=item, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - # TODO: Not optimal because we load the whole slice in memory. Would have to write - # a bespoke var function that executed within slice_eval to avoid this probably. - expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 - out = expr.mean( - axis=axis, - dtype=dtype, - keepdims=keepdims, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - if ddof != 0: - num_elements = ( - self.get_num_elements(axis, item, fp_accuracy=fp_accuracy, **execution_kwargs) - if where is None - else where.where(blosc2.ones(self.shape, dtype=np.int64), 0).sum( - axis=axis, dtype=np.int64, fp_accuracy=fp_accuracy, **execution_kwargs - ) - ) - out = out * num_elements / (num_elements - ddof) - out2 = kwargs.pop("out", None) - if out2 is not None: - out2[:] = out - return out2 - if kwargs != {} and not np.isscalar(out): - out = blosc2.asarray(out, **kwargs) - return out + return self._var(axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs) def min( self, @@ -4619,7 +4710,7 @@ def _compute_expr(self, item, kwargs): # noqa: C901 return chunked_eval(lazy_expr.expression, lazy_expr.operands, item, **kwargs) - if any(_has_constructor_call(self.expression, constructor) for constructor in constructors): + if _has_any_constructor_call(self.expression): expression = self.expression newexpr = expression newops = self.operands.copy() @@ -5697,7 +5788,7 @@ def lazyexpr( operands = seek_operands(operand_set, local_dict, global_dict, _frame_depth=_frame_depth) else: # No operands found in the expression. Maybe a constructor? - constructor = any(_has_constructor_call(expression, constructor) for constructor in constructors) + constructor = _has_any_constructor_call(expression) if not constructor: raise ValueError("No operands nor constructors found in the expression") # _new_expr will take care of the constructor, but needs an empty dict in operands diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 4c028c956..e952000cc 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -360,6 +360,21 @@ def are_partitions_aligned(shape, chunks, blocks): return builtins.all(c % b == 0 for c, b in zip(chunks, blocks, strict=True)) +def blocks_contiguous(chunks, blocks): + """Whether the blocks are C-contiguous runs of their chunk (no block may be 0 long). + + That is the case only if their leading dims are 1, then one dim divides the chunk's, and the + rest are equal to the chunk's; same rule as nchunk_fastpath() in C-Blosc2's b2nd.c. + """ + k = 0 + for i, (size, unit) in enumerate(zip(chunks, blocks, strict=True)): + if size % unit != 0 or (i > k and size != unit): + return False + if i == k and unit == 1: + k += 1 + return True + + def are_partitions_behaved(shape, chunks, blocks): """ Check if the partitions defined by chunks and blocks are well-behaved with respect to the shape. @@ -377,18 +392,6 @@ def check_contiguity(container, part): return False return builtins.all(size % unit == 0 for size, unit in zip(container[:-1], part[:-1], strict=True)) - def blocks_contiguous(chunks, blocks): - # A block is a C-contiguous run of its chunk only if its leading dims are 1, - # then one dim divides the chunk's, and the rest are equal to the chunk's. - # Same rule as nchunk_fastpath() in C-Blosc2's b2nd.c. - k = 0 - for i, (size, unit) in enumerate(zip(chunks, blocks, strict=True)): - if size % unit != 0 or (i > k and size != unit): - return False - if i == k and unit == 1: - k += 1 - return True - return blocks_contiguous(chunks, blocks) and check_contiguity(shape, chunks) diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 763d4de15..f8e35e158 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -2408,6 +2408,102 @@ def test_cumulative_block_layout(monkeypatch, blocks): np.testing.assert_array_equal(out[:], expected) +@pytest.mark.parametrize( + ("dtype", "x", "y"), + [ + (np.float32, "a", 0), + (np.float32, 1.5, "a"), + (np.int32, "a", 0.5), + (np.int16, "a", -1), + (np.float64, "a", "b"), + ], +) +@pytest.mark.parametrize("getitem", [False, True]) +def test_where_scalar_uses_miniexpr(monkeypatch, dtype, x, y, getitem): + import importlib + + lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") + monkeypatch.setattr(lazyexpr_mod, "try_miniexpr", True) + original_set_pref_expr = blosc2.NDArray._set_pref_expr + compiled = [] + + def wrapped_set_pref_expr(self, expression, inputs, *args, **kwargs): + res = original_set_pref_expr(self, expression, inputs, *args, **kwargs) + compiled.append(expression) + return res + + monkeypatch.setattr(blosc2.NDArray, "_set_pref_expr", wrapped_set_pref_expr) + na = np.arange(40 * 30, dtype=dtype).reshape(40, 30) + nb = (na // 3).astype(dtype) + a = blosc2.asarray(na, chunks=(16, 30), blocks=(4, 30)) + b = blosc2.asarray(nb, chunks=(16, 30), blocks=(4, 30)) + ops, nops = {"a": a, "b": b}, {"a": na, "b": nb} + expr = blosc2.where(a > 500, ops.get(x, x), ops.get(y, y)) + expected = np.where(na > 500, nops.get(x, x), nops.get(y, y)) + res = expr[:] if getitem else expr.compute()[:] + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) + # The scalar is inlined, so miniexpr (rather than the numexpr fallback) did the work + assert len(compiled) == 1 + assert "_where_x" in compiled[0] or "_where_y" in compiled[0] + + +@pytest.mark.parametrize( + ("shape", "chunks", "blocks", "direct"), + [ + ((100, 100), (100, 100), (25, 100), True), + ((10_000,), (3000,), (1000,), True), # partial last chunk + ((5,), (8,), (4,), True), # a single, partial chunk + ((37, 50, 60), (5, 50, 60), (1, 25, 60), True), + ((4, 12, 30), (1, 4, 30), (1, 2, 30), True), + ((1, 10, 30), (1, 4, 30), (1, 2, 30), True), # partial, but the very last chunk + ((4, 10, 30), (1, 4, 30), (1, 2, 30), False), # partial chunks inside the array + ((40, 30), (16, 30), (4, 10), False), # blocks not contiguous in their chunk + ((40, 30), (16, 15), (4, 15), False), # chunks not contiguous in the array + ((24, 20, 18), (8, 10, 9), (4, 5, 9), False), + ], +) +@pytest.mark.parametrize("dtype", [np.float32, np.int16, np.complex128]) +def test_getitem_into_numpy_layouts(monkeypatch, shape, chunks, blocks, direct, dtype): + import importlib + + lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") + monkeypatch.setattr(lazyexpr_mod, "try_miniexpr", True) + assert lazyexpr_mod._chunks_tile_flat(shape, chunks, blocks) is direct + original = lazyexpr_mod._prefilter_to_numpy + results = [] + + def wrapped(*args): + res = original(*args) + results.append(res is not None) + return res + + monkeypatch.setattr(lazyexpr_mod, "_prefilter_to_numpy", wrapped) + rng = np.random.default_rng(1) + na = (rng.random(shape) * 100).astype(dtype) + nb = (rng.random(shape) * 10).astype(dtype) + a = blosc2.asarray(na, chunks=chunks, blocks=blocks) + b = blosc2.asarray(nb, chunks=chunks, blocks=blocks) + if dtype is np.complex128: + second = (a * b, na * nb) + else: + second = ((a - b) > 40, (na - nb) > 40) + for expr, expected in [(a + b * 2, na + nb * 2), second]: + for res in (expr[:], expr[()]): + assert res.dtype == expected.dtype + assert res.flags.c_contiguous + assert res.flags.writeable + np.testing.assert_array_equal(res, expected) + base = res + while base.base is not None: + base = base.base + # The padding of a partial last chunk is not kept allocated + assert base.nbytes <= res.nbytes + blosc2.MAX_OVERHEAD + if dtype is np.complex128 and not results: + return # miniexpr is not used for complex numbers on some platforms + assert results == ([True] * 4 if direct else []) + + @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 diff --git a/tests/ndarray/test_reductions.py b/tests/ndarray/test_reductions.py index 89982a550..76da246b9 100644 --- a/tests/ndarray/test_reductions.py +++ b/tests/ndarray/test_reductions.py @@ -977,3 +977,78 @@ def test_reductions_where_ndarray_and_lazyexpr(): np.testing.assert_allclose(a.std(axis=0, where=mask), np.std(na, axis=0, where=nmask)) np.testing.assert_allclose(a.min(axis=0, where=mask), np.min(na, axis=0, where=nmask, initial=np.inf)) np.testing.assert_allclose(a.max(axis=0, where=mask), np.max(na, axis=0, where=nmask, initial=-np.inf)) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("op", ["std", "var"]) +@pytest.mark.parametrize("ddof", [0, 1]) +@pytest.mark.parametrize("keepdims", [False, True]) +@pytest.mark.parametrize("offset", [0, 1e4, -1e4]) +@pytest.mark.parametrize("miniexpr", [True, False]) +def test_std_var_full_reduction(monkeypatch, dtype, op, ddof, keepdims, offset, miniexpr): + # Full std/var reduce (x - mean) * (x - mean) with the mean inlined as a scalar + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", miniexpr) + rng = np.random.default_rng(42) + na = (offset + rng.normal(scale=3, size=(37, 40, 50))).astype(dtype) + a = blosc2.asarray(na, chunks=(10, 20, 50), blocks=(5, 10, 25)) + res = getattr(a, op)(ddof=ddof, keepdims=keepdims) + expected = getattr(na.astype(np.float64), op)(ddof=ddof, keepdims=keepdims) + assert np.asarray(res).dtype == getattr(na, op)().dtype + assert np.shape(res) == np.shape(expected) + np.testing.assert_allclose(res, expected, rtol=1e-5 if dtype == np.float32 else 1e-12) + # Same through a lazy expression and the module-level functions + res = getattr(blosc2, op)(a + 1, ddof=ddof, keepdims=keepdims) + assert np.asarray(res).dtype == getattr(na, op)().dtype + np.testing.assert_allclose(res, expected, rtol=1e-5 if dtype == np.float32 else 1e-12) + + +@pytest.mark.parametrize("op", ["std", "var"]) +def test_std_var_full_reduction_wider_dtype(op): + # A float64 result of float32 data takes float64 deviations from the mean, as NumPy does + rng = np.random.default_rng(1) + na = (rng.random(400_000) + 1e6).astype(np.float32) + a = blosc2.asarray(na) + res = getattr(a, op)(dtype=np.float64) + assert np.asarray(res).dtype == np.float64 + np.testing.assert_allclose(res, getattr(np, op)(na, dtype=np.float64), rtol=1e-6) + + +@pytest.mark.parametrize("value", [np.nan, np.inf]) +def test_std_var_full_reduction_nonfinite(value): + na = np.arange(1000, dtype=np.float64).reshape(10, 100) + na[3, 7] = value + a = blosc2.asarray(na, chunks=(5, 50), blocks=(5, 10)) + with np.errstate(invalid="ignore"): + np.testing.assert_array_equal(a.std(), np.std(na)) + np.testing.assert_array_equal(a.var(), np.var(na)) + + +@pytest.mark.parametrize( + ("chunks", "blocks"), + [ + ((10, 20, 50), (5, 20, 50)), # decompressed chunks are C-order; partial along two axes + ((10, 20, 50), (1, 4, 50)), + ((10, 20, 50), (5, 10, 25)), # blocks not contiguous in their chunk + ], +) +@pytest.mark.parametrize("axis", [0, 1, 2, (0, 2), None]) +@pytest.mark.parametrize("op", ["sum", "max", "argmin", "cumulative_sum"]) +def test_reduction_axis_chunk_layouts(chunks, blocks, axis, op): + if op in ("argmin", "cumulative_sum") and not (axis is None or isinstance(axis, int)): + pytest.skip("argmin and cumulative_sum take a single axis") + if op == "cumulative_sum" and axis is None: + pytest.skip("cumulative_sum needs an axis for N-D arrays") + rng = np.random.default_rng(3) + na = rng.integers(-1000, 1000, size=(37, 45, 50)).astype(np.float64) + nb = rng.integers(-1000, 1000, size=(37, 45, 50)).astype(np.float64) + a = blosc2.asarray(na, chunks=chunks, blocks=blocks) + b = blosc2.asarray(nb, chunks=chunks, blocks=blocks) + for expr, nexpr in [(a, na), (a + b, na + nb)]: + res = getattr(expr, op)(axis=axis) + expected = getattr(np, op)(nexpr, axis=axis) + np.testing.assert_array_equal(res[()] if isinstance(res, blosc2.NDArray) else res, expected) + if op in ("sum", "max"): + res = getattr(expr, op)(axis=axis, keepdims=True) + np.testing.assert_array_equal(res, getattr(np, op)(nexpr, axis=axis, keepdims=True)) From 55660b47a393b11b30191d9331dd7903b2f84cd9 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:49:21 +0100 Subject: [PATCH 14/18] Keep expressions that miniexpr gets wrong on the numexpr path The bundled miniexpr (3f4db93) returns wrong results, or crashes, for several classes of expressions, all reproducible on main: reductions with a constant subexpression ((o * -2).sum() is 0.0 for an array of ones); two-argument float functions and true division of integer or bool operands computed in integer arithmetic (bool / bool died of SIGFPE); constants typed by the result dtype ((a / 3) > 2 on integers, int32 < 2.5); % and remainder() with C sign rules; mixed-dtype typing (garbage values, and int16 ((i + i) * abs(i))[:] corrupted the heap); bool + bool added as numbers; grammar differences (-x ** 2, ** associativity, chained comparisons). The new miniexpr_guard module reproduces miniexpr's typing of an expression and flags these cases, replacing the check for integer arctan2 only; fast_eval() and reduce_slices() then evaluate them with numexpr, which matches NumPy (1.5-4x slower for those expressions, e.g. (a * -2).sum() 10 ms but wrong -> 39 ms; others unchanged, the check is a cached lookup). It agreed with NumPy on 10.5k random expressions, but is not a proof. check_dtype() also gave integer or bool result dtypes to hypot, copysign and nextafter of integer subexpressions, to true division involving bools ((a < b) / b, bool / bool) and to the bool //, % and ** that NumPy computes in int8, so their results were truncated on every evaluation path. It now asks NumPy (the operator or ufunc on empty operands, with a scalar next to an array kept as it is, so that it promotes as a Python scalar); int32 / int32 keeps its float32. Array operands' dtypes are read directly, which pays for that probe. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- src/blosc2/lazyexpr.py | 64 +++-- src/blosc2/miniexpr_guard.py | 443 +++++++++++++++++++++++++++++++++ tests/ndarray/test_lazyexpr.py | 297 ++++++++++++++++++++++ 3 files changed, 783 insertions(+), 21 deletions(-) create mode 100644 src/blosc2/miniexpr_guard.py diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 6d378156a..311c7e7b0 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -54,6 +54,7 @@ write_b2object_user_vlmeta, ) from .dsl_kernel import DSLKernel, DSLSyntaxError, DSLValidator, specialize_miniexpr_inputs +from .miniexpr_guard import miniexpr_known_bad from .ndarray import blocks_contiguous from .proxy_source import convert_dtype from .utils import ( @@ -342,6 +343,23 @@ def _find_constructor_call(expression: str, constructor: str) -> re.Match | None "//": operator.floordiv, "%": operator.mod, } +# NumPy operations of the numeric operators and functions whose result dtype is not just the +# promoted operand dtype: true division and the float functions turn integers (and bools) into +# floats, and //, %, ** and the shifts, which NumPy has no bool loops for, turn bools into int8 +_result_dtype_ufuncs = { + "/": operator.truediv, + "//": operator.floordiv, + "%": operator.mod, + "**": operator.pow, # not np.power: NumPy squares for x ** 2, which keeps bools int8 + "<<": operator.lshift, + ">>": operator.rshift, + "arctan2": np.arctan2, + "copysign": np.copysign, + "hypot": np.hypot, + "nextafter": np.nextafter, + "pow": np.power, + "power": np.power, +} logical_ops = ["&", "|", "^", "~"] not_complex_ops = ["maximum", "minimum", "<", "<=", ">", ">="] funcs_2args = ( @@ -1716,15 +1734,6 @@ def _raise_dsl_miniexpr_required(reason: str | None = None) -> None: raise RuntimeError(_dsl_miniexpr_required_message(reason)) -def _miniexpr_integer_atan2(expression, operands): - # ponytail: use the chunked fallback until miniexpr promotes integer atan2 inputs. - return ( - isinstance(expression, str) - and re.search(r"\b(?:arctan2|atan2)\s*\(", expression) - and any(getattr(op, "dtype", np.dtype(float)).kind in "biu" for op in operands.values()) - ) - - def _c_run_axis(outer, inner) -> int: """Axis along which `inner` tiles are C-contiguous runs of `outer`, or -1 if they are not. @@ -1945,7 +1954,7 @@ def fast_eval( # noqa: C901 if is_dsl and dsl_disable_reason is None: dsl_disable_reason = "cumulative scans are not supported by the DSL miniexpr path." if isinstance(expr_string_miniexpr, str): - if not is_dsl and _miniexpr_integer_atan2(expr_string_miniexpr, operands): + if not is_dsl and miniexpr_known_bad(expr_string_miniexpr, operands_miniexpr, dtype): use_miniexpr = False expr_string_miniexpr = _apply_jit_backend_pragma( expr_string_miniexpr, operands_miniexpr, jit_backend @@ -2968,7 +2977,8 @@ def reduce_slices( # noqa: C901 # Check whether we can use miniexpr if use_miniexpr and isinstance(expression, str): - if _miniexpr_integer_atan2(expression, operands): + reduced_expr = expression if where is None else f"where({expression}, _where_x, _where_y)" + if miniexpr_known_bad(reduced_expr, operands, dtype, reduction=True): use_miniexpr = False has_complex = any( isinstance(op, blosc2.NDArray) and blosc2.isdtype(op.dtype, "complex floating") @@ -3609,12 +3619,18 @@ def fuse_expressions(expr, new_base, dup_op): return new_expr +def _operand_dtype(value): + """blosc2.result_type(value), without building a NumPy array for an array operand.""" + dtype = None if np.isscalar(value) else getattr(value, "dtype", None) + return blosc2.result_type(value) if dtype is None else convert_dtype(dtype) + + def check_dtype(op, value1, value2): if op in ("contains", "startswith", "endswith"): return np.dtype(np.bool_) - v1_dtype = blosc2.result_type(value1) - v2_dtype = v1_dtype if value2 is None else blosc2.result_type(value2) + v1_dtype = _operand_dtype(value1) + v2_dtype = v1_dtype if value2 is None else _operand_dtype(value2) if op in not_complex_ops and (v1_dtype == np.complex128 or v2_dtype == np.complex128): # Ensure that throw exception for functions which don't support complex args raise ValueError(f"Invalid operand type for {op}: {v1_dtype, v2_dtype}") @@ -3627,14 +3643,20 @@ def check_dtype(op, value1, value2): if v2_dtype not in (np.bool_, np.int32, np.int64): raise ValueError(f"Invalid operand type for {op}: {v2_dtype}") - if op == "/": - if v1_dtype == np.int32 and v2_dtype == np.int32: - return blosc2.float32 - if np.issubdtype(v1_dtype, np.integer) and np.issubdtype(v2_dtype, np.integer): - return blosc2.float64 - - if op == "arctan2": - return np.arctan2(np.empty(0, v1_dtype), np.empty(0, v2_dtype)).dtype + if op == "/" and v1_dtype == np.int32 and v2_dtype == np.int32: + return blosc2.float32 + + ufunc = _result_dtype_ufuncs.get(op) + if ufunc is not None and value2 is not None and v1_dtype.kind in "biufc" and v2_dtype.kind in "biufc": + # Ask NumPy on empty operands, but with a scalar next to an array as it is, so that NumPy + # promotes it as usual; what it rejects is left to the evaluation + scalars = np.isscalar(value1) and np.isscalar(value2) + probes = [ + v if np.isscalar(v) and not scalars else np.empty(0, dt) + for v, dt in ((value1, v1_dtype), (value2, v2_dtype)) + ] + with contextlib.suppress(TypeError, ValueError, OverflowError): + return ufunc(*probes).dtype if (v1_dtype.kind in "Mm" or v2_dtype.kind in "Mm") and op in _datetime_binops: # result_type describes promotion, not operator semantics, and the two diff --git a/src/blosc2/miniexpr_guard.py b/src/blosc2/miniexpr_guard.py new file mode 100644 index 000000000..f04951b76 --- /dev/null +++ b/src/blosc2/miniexpr_guard.py @@ -0,0 +1,443 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +"""Which expressions miniexpr is known to get wrong, so that they can take another path. + +It models how miniexpr types and parses an expression, as far as needed to find the shapes it +mis-evaluates (see :func:`miniexpr_known_bad`). +""" + +import ast +import functools +import re + +import numpy as np + +# Two-argument functions that NumPy keeps integer for integer inputs. miniexpr computes every +# two-argument function in the promoted operand type, so all the others (arctan2, hypot, +# logaddexp, nextafter, ...) come out in integer arithmetic for two integer operands. +_ME_INT_BINARY_CALLS = frozenset(("fmax", "fmin", "fmod", "maximum", "minimum", "pow", "power")) +_ME_KINDS = "bifc" # promotion order of the operand kinds below +# Decimal int and float literals (group 1: float); others (complex, hex, 1_000) are left alone +_ME_LITERAL = re.compile(r"(? 3 divides integers too; + - % and remainder() follow C (sign of the dividend), not NumPy (sign of the divisor); + - in reductions, constant subexpressions other than plain literals (-2, (0 - 2), sin(1.0)) + are not folded and evaluate to garbage; + - the typing and parsing bugs that `_miniexpr_type_bug` models, which give wrong results or + crash: (a < c) + c for operands of other dtypes, bool + bool, int32 a < 2.5, ... + + Such expressions take the regular evaluation path instead. + """ + if not isinstance(expression, str): + return False + types = [] + for name, value in operands.items(): + try: + dt = np.dtype(getattr(value, "dtype", type(value))) + except TypeError: + types.append((name, "U" if isinstance(value, str | bytes) else "i", None)) + else: + types.append((name, *_me_dtype_info(dt)[:2])) + return _miniexpr_known_bad_cached( + expression, tuple(types), _me_dtype_info(np.dtype(dtype))[2], reduction + ) + + +@functools.lru_cache(maxsize=256) +def _me_dtype_info(dt): + # Kind for the checks below, miniexpr type code, and a key to rebuild the dtype from + return "i" if dt.kind in "uMm" else dt.kind, _ME_CODES.get(dt.name), dt.str + + +def _me_integral(const, kind, const_kind): + """Whether miniexpr computes with a subexpression as an integer (or bool); see `visit` below.""" + if const and kind in _ME_KINDS: + kind = const_kind or kind + return kind in "bi" + + +def _me_call_kind(func, args, const_kind): + """Result kind of a call, and whether it hits a miniexpr bug; `args` as from `visit` below.""" + bad = func in ("remainder", "mod") or ( + len(args) == 2 + and func not in _ME_INT_BINARY_CALLS + and all(_me_integral(c, k, const_kind) for c, k in args) + ) + kinds = [k for _, k in args] + if func == "where" and len(args) == 3: + return _promote_kind(kinds[1], kinds[2]), bad + # Float functions of integers are counted as integers: at worst a false positive, which only + # costs the faster path + return _promote_kind(*kinds), bad + + +def _me_literal_class(match): + # The checks only look at the kind of a literal, whether it is an integer, and whether it fits + # 0/1 or any integer dtype, so 2-127 and the non-integers below that stand for each other + text = match.group() + value = float(text) if match.group(1) is not None else int(text) + if 2 <= value <= 127 and float(value).is_integer(): + return "2.0" if match.group(1) is not None else "2" + if match.group(1) is not None and abs(value) < 127 and not value.is_integer(): + return "1.5" + return text + + +@functools.lru_cache(maxsize=1024) +def _miniexpr_known_bad_cached(expression, types, out, reduction): + # Expressions that differ in their inlined scalars only (a changing mean in _full_var, loop + # variables) share the work + return _miniexpr_known_bad_parsed(_ME_LITERAL.sub(_me_literal_class, expression), types, out, reduction) + + +@functools.lru_cache(maxsize=1024) +def _miniexpr_known_bad_parsed(expression, types, out, reduction): + source = expression.strip() + try: + tree = ast.parse(source, mode="eval") + except SyntaxError: + return False # not an expression miniexpr compiles either + kinds = {name: kind for name, kind, _ in types} + out_kind = np.dtype(out).kind + const_kind = "f" if out_kind in "fc" else "" if out_kind in "iu" else "b" + bad = False + + def visit(node): + # Returns (is_constant, kind); anything unknown counts as a non-constant integer, which is + # the conservative choice for the checks. + nonlocal bad + if isinstance(node, ast.Constant): + return True, {bool: "b", int: "i", float: "f", complex: "c"}.get(type(node.value), "U") + if isinstance(node, ast.Name): + if node.id in kinds: + return False, kinds[node.id] + return (True, "f") if node.id in ("nan", "inf", "pi", "e") else (False, "i") + if isinstance(node, ast.UnaryOp): + const, kind = visit(node.operand) + res = (const, "b" if isinstance(node.op, ast.Not) else kind) + elif isinstance(node, ast.BinOp): + (lconst, lkind), (rconst, rkind) = visit(node.left), visit(node.right) + div = isinstance(node.op, ast.Div) + if isinstance(node.op, ast.Mod) or ( + div and _me_integral(lconst, lkind, const_kind) and _me_integral(rconst, rkind, const_kind) + ): + bad = True + res = (lconst and rconst, _promote_kind(lkind, rkind, *("f" if div else ""))) + elif isinstance(node, ast.Compare | ast.BoolOp): + values = node.values if isinstance(node, ast.BoolOp) else [node.left, *node.comparators] + consts = [visit(v)[0] for v in values] # visit all of them, for their checks + res = (all(consts), "b") + elif isinstance(node, ast.Call): + func = node.func.id if isinstance(node.func, ast.Name) else getattr(node.func, "attr", "") + args = [visit(a) for a in node.args] + kind, call_bad = _me_call_kind(func, args, const_kind) + bad = bad or call_bad + res = (all(c for c, _ in args), kind) + else: + for child in ast.iter_child_nodes(node): + visit(child) + res = (False, "i") + if reduction and res[0]: + bad = True + return res + + visit(tree.body) + if bad or any(kind in "USV" for kind in kinds.values()): + return bad + codes = {name: code for name, _, code in types} + return _miniexpr_type_bug(source, tree.body, codes, _ME_CODES.get(np.dtype(out).name), reduction) + + +# What follows models how miniexpr types an expression, to find the shapes it gets wrong. Dtypes +# are NumPy type characters (with "q"/"Q" for 64-bit integers). +_ME_CODES = {np.dtype(code).name: code for code in "?bhiqBHIQfdFD"} +# miniexpr's own promotion table (type_promotion_table in miniexpr.c), one row of 13 per code; it is +# not NumPy's for mixed signed and unsigned integers (int64 with uint8 gives float64) +_ME_PROMOTION_ROWS = ( + "?bhiqBHIQfdFDbbhiqhiqdfdFDhhhiqiiqdfdFDiiiiqqqqdddDDqqqqqddddddDDBhiqdBHIQfdFDHiiqdHHIQfdFD" + "IqqqdIIIQddDDQddddQQQQddDDfffddffddfdFDdddddddddddDDFFFDDFFDDFDFDDDDDDDDDDDDDD" +) +_ME_FLOAT_MATH = frozenset( + "acos acosh arccos arccosh arcsin arcsinh arctan arctanh asin asinh atan atanh cbrt cos cosh " # noqa: SIM905 + "cospi erf erfc exp exp10 exp2 expm1 lgamma ln log log10 log1p log2 sin sinh sinpi sqrt tan " + "tanh tgamma".split() +) +_ME_CALLS = _ME_FLOAT_MATH | frozenset( + "abs arctan2 atan2 ceil conj copysign fdim floor fma fmax fmin fmod hypot imag ldexp logaddexp " # noqa: SIM905 + "ncr nextafter npr pow power real remainder rint round sign square trunc where".split() +) +_ME_BINOPS = { + ast.Add: "+", + ast.Sub: "-", + ast.Mult: "*", + ast.Div: "/", + ast.Mod: "%", + ast.Pow: "**", + ast.LShift: "<<", + ast.RShift: ">>", + ast.BitAnd: "&", + ast.BitOr: "|", + ast.BitXor: "^", + ast.FloorDiv: "//", +} + + +def _me_promote(*codes): + out = "?" + for code in codes: + out = _ME_PROMOTION_ROWS[13 * "?bhiqBHIQfdFD".index(out) + "?bhiqBHIQfdFD".index(code)] + return out + + +class _MeUnsupported(Exception): + pass + + +class _MeNode: + """A node of an expression as miniexpr builds it. + + pd: dtype it gets when parsed, which conversion nodes are inserted against; rt: computation + dtype that evaluation infers (infer_result_type()); ot: output dtype that evaluation infers + (infer_output_type()); func: whether it is a function node, the only ones that get conversion + nodes; conv_to: for binary nodes, the dtype those conversion nodes convert to. + """ + + __slots__ = ("args", "conv_to", "func", "kind", "ot", "pd", "rt", "value") + + def __init__(self, kind, args, pd, rt, ot, func=True, conv_to=None, value=None): + self.kind, self.args, self.pd, self.rt, self.ot = kind, args, pd, rt, ot + self.func, self.conv_to, self.value = func, conv_to, value + + +def _me_literal(value, out): + # A float literal is float32 for a float32 result and float64 otherwise, an integer one int32 + # (int64 beyond its range); then a non-integer result type replaces it, and an integer one does + # unless the literal is wider + if isinstance(value, bool) or not isinstance(value, int | float): + raise _MeUnsupported + if isinstance(value, float): + code = "f" if out == "f" else "d" + return code if out in "bhiqBHIQ" else out + code = "q" if abs(value) > 2**31 - 1 else (out if out in "bhiqBHIQ" else "i") + if out in "bhiqBHIQ": + return code if np.dtype(code).itemsize > np.dtype(out).itemsize else out + return out + + +def _me_node(node, codes, out): # noqa: C901 + if isinstance(node, ast.Constant): + code = _me_literal(node.value, out) + return _MeNode("const", [], code, code, code, func=False, value=node.value) + if isinstance(node, ast.Name): + if codes.get(node.id) is None: + raise _MeUnsupported + code = codes[node.id] + return _MeNode("var", [], code, code, code, func=False) + if isinstance(node, ast.UnaryOp): + if isinstance(node.op, ast.UAdd): + return _me_node(node.operand, codes, out) + x = _me_node(node.operand, codes, out) + if isinstance(node.op, ast.USub): + # negate() is a new node, which gets the default float64; optimize() folds it for a + # constant + if x.kind in ("const", "negconst"): + return _MeNode("negconst", [], "d", "d", "d", func=False, value=-x.value) + return _MeNode("neg", [x], "d", x.rt, x.ot) + if isinstance(node.op, ast.Invert): + return _MeNode("~", [x], x.pd, x.rt, "?" if x.pd == "?" else x.ot) + raise _MeUnsupported + if isinstance(node, ast.BinOp): + op = _ME_BINOPS.get(type(node.op)) + if op is None: + raise _MeUnsupported + left, right = _me_node(node.left, codes, out), _me_node(node.right, codes, out) + conv_to = pd = _me_promote(left.pd, right.pd) + for x, c in ((node.right, node.left), (node.left, node.right)): + negconst = isinstance(c, ast.UnaryOp) and isinstance(c.op, ast.USub) + if op == "+" and negconst and isinstance(c.operand, ast.Constant): + # x + -c and -c + x are rewritten to x - c, after the conversion of x to the + # float64 of -c has been inserted + left, right = _me_node(x, codes, out), _me_node(c.operand, codes, out) + op = "-" + pd = _me_promote(conv_to if left.func and left.pd != conv_to else left.pd, right.pd) + break + ot = "?" if pd == "?" else _me_promote(left.ot, right.ot) + res = _MeNode(op, [left, right], pd, _me_promote(left.rt, right.rt), ot, conv_to=conv_to) + if op == "**" and right.kind == "const" and right.value in (2, 3) and left.pd != pd: + res.kind = "**2" # optimize() rewrites it to x * x, of another dtype + return res + if isinstance(node, ast.Compare): + if len(node.ops) != 1: + raise _MeUnsupported + left, right = _me_node(node.left, codes, out), _me_node(node.comparators[0], codes, out) + rt, conv_to = _me_promote(left.rt, right.rt), _me_promote(left.pd, right.pd) + return _MeNode("cmp", [left, right], "?", rt, "?", conv_to=conv_to, value=type(node.ops[0])) + if isinstance(node, ast.Call): + func = node.func.id if isinstance(node.func, ast.Name) else None + if func not in _ME_CALLS or node.keywords: + raise _MeUnsupported + args = [_me_node(a, codes, out) for a in node.args] + if func in _ME_FLOAT_MATH and len(args) == 1: + rt, ot = (a if a in "fdFD" else "d" for a in (args[0].rt, args[0].ot)) + elif func in ("abs", "imag", "real") and len(args) == 1: + rt, ot = ({"F": "f", "D": "d"}.get(a, a) for a in (args[0].rt, args[0].ot)) + elif func == "where" and len(args) == 3: + rt, ot = _me_promote(*(a.rt for a in args)), _me_promote(args[1].ot, args[2].ot) + else: + rt, ot = _me_promote(*(a.rt for a in args)), _me_promote(*(a.ot for a in args)) + # Calls get no conversion nodes, and the default float64 + return _MeNode("call", args, "d", rt, ot) + raise _MeUnsupported + + +def _me_parsed_differently(source, tree): + """Whether miniexpr's grammar reads `source` differently from Python's. + + Its unary operators bind tighter than **, ** is left-associative, and comparisons do not chain. + """ + for node in ast.walk(tree): + if isinstance(node, ast.Compare) and len(node.ops) > 1: + return True + if isinstance(node, ast.UnaryOp) and isinstance(node.operand, ast.BinOp): + seg = ast.get_source_segment(source, node) + if seg is None or seg[1:].strip() == ast.get_source_segment(source, node.operand): + return True # -x ** 2, without parentheses around the operand + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Pow): + right = node.right + if isinstance(right, ast.BinOp) and isinstance(right.op, ast.Pow): + seg = ast.get_source_segment(source, node) + if seg is None or seg.endswith(ast.get_source_segment(source, right)): + return True # x ** y ** z, without parentheses around y ** z + return False + + +def _me_representable(value, code, cmp=None, right=True): + """Whether a constant survives the cast to the dtype `code` it gets computed in. + + For a comparison `cmp`, the truncation of a positive constant is harmless in x > 2.5 and + x <= 2.5 (and 2.5 < x, 2.5 >= x). + """ + if code in "fdFD": + return True + info = np.iinfo(np.dtype(code)) if code != "?" else None + if info is None or not info.min <= value <= info.max: + return info is None and value in (0, 1) + if float(value).is_integer(): + return True + return value > 0 and cmp in ((ast.Gt, ast.LtE) if right else (ast.Lt, ast.GtE)) + + +def _miniexpr_type_bug(source, tree, codes, out, reduction): # noqa: C901 + """Whether `tree` (parsed from `source`) is one of the shapes that miniexpr gets wrong. + + - A function node whose parse-time dtype differs from its binary parent's gets a conversion + node, which writes the parent's parse-time dtype into a buffer of the dtype the parent is + evaluated in: the computation dtype inferred for the enclosing evaluation. Where the two + differ, as for a comparison (bool when parsed, its operands' type when inferred), a call + (float64 when parsed) or a promotion through miniexpr's table, the result is garbage, or a + heap overflow for a narrower buffer: int16 (a + a) * abs(a) crashes, int8/uint8 (a < c) + c + is wrong. A comparison evaluated for a bool result gets its conversions wrong too, and a + few related shapes found by testing corrupt the heap (see the comments below). + - bool + bool is added as numbers (True + True gives 2), not or-ed. + - Constants are cast to the dtype a node is computed in, which truncates them: int32 a < 2.5 + compares with 2, and int8 a < 300 with 44. + - A reduction of a comparison of two bools never writes its result. + - Its grammar reads some of the Python that python-blosc2 renders differently (-x ** 2). + """ + if out is None: + return False + if _me_parsed_differently(source, tree): + return True + try: + root = _me_node(tree, codes, out) + except _MeUnsupported: + return False + bad = False + nconv = 0 + + def region(node, rtype): + # `node` is evaluated by the typed kernels, in `rtype` + nonlocal bad, nconv + if node.kind in ("+", "-") and node.args[0].pd == node.args[1].pd == "?": + bad = True + convs = [c for c in node.args if node.conv_to is not None and c.func and c.pd != node.conv_to] + if convs and node.kind == "cmp" and any(c.kind == "negconst" for c in node.args): + bad = True # crashes for a float32 operand: where((b + b) > -1, b, b) + for i, child in enumerate(node.args): + if child.kind in ("const", "negconst") and not _me_representable( + child.value, rtype, node.value if node.kind == "cmp" else None, i == 1 + ): + bad = True + if child in convs: + # an x ** 2 rewritten to another dtype crashes as a conversion source + bad = bad or node.conv_to != rtype or child.kind == "**2" + nconv += 1 + evaluate(child, child.pd) + else: + region(child, rtype) + + def bool_ok(node): + # Whether eval_bool_expr() handles `node` + if node.kind in ("const", "negconst", "var"): + return True + if node.kind == "cmp": + return node.rt not in "FD" + return node.kind in ("&", "|", "^", "~") and node.pd == "?" and all(bool_ok(c) for c in node.args) + + def bool_path(node): + # `node` is evaluated by eval_bool_expr() + nonlocal bad + for i, child in enumerate(node.args): + if node.kind != "cmp": + bool_path(child) + continue + if child.kind in ("const", "negconst") and not _me_representable( + child.value, node.rt, node.value, i == 1 + ): + bad = True + if child.func and child.pd != node.conv_to: + bad = True + evaluate(child, node.rt) + + def evaluate(node, dtype): + # private_eval() of `node` into a `dtype` result + nonlocal bad + if node.func: + if dtype == "?" and node.ot == "?" and bool_ok(node): + bool_path(node) + else: + before = nconv + region(node, node.rt) + # A float32 result computed in float64 with conversions inside corrupts the heap, + # as (b + b) - (b * -1.5) does + bad = bad or (dtype == "f" and node.rt == "d" and nconv > before) + + if reduction: + # eval_reduction() evaluates a comparison into its computation dtype, which it then fails + # to convert to bool for a bool one; anything else into its output dtype + if root.kind == "cmp": + bad = root.rt == "?" + region(root, root.rt) + else: + evaluate(root, root.ot) + else: + evaluate(root, out) + return bad diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index f8e35e158..b4a916c59 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -2518,3 +2518,300 @@ def test_frame_backed_inmemory_operands(source): 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) + + +def _numpy_eval(expression, operands): + namespace = {"where": np.where, "sin": np.sin, "hypot": np.hypot, "nextafter": np.nextafter} + namespace.update( + remainder=np.remainder, arctan2=np.arctan2, copysign=np.copysign, logaddexp=np.logaddexp + ) + return eval(expression, namespace, operands) + + +@pytest.mark.parametrize("dtype", [np.int64, np.float64, np.float32]) +@pytest.mark.parametrize( + "expression", ["o0 * -2", "-2 * o0", "o0 - -0.5", "o0 / -2", "o0 * (0 - 2)", "o0 + sin(1.0)", "o0 > -2"] +) +@pytest.mark.parametrize("op", ["sum", "max", "mean"]) +def test_miniexpr_reduction_constant_subexpressions(monkeypatch, dtype, expression, op): + # miniexpr does not fold constant subexpressions inside reductions (they evaluated to garbage) + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + na = np.arange(37 * 40 * 5, dtype=dtype).reshape(37, 40, 5) % 7 + a = blosc2.asarray(na, chunks=(10, 20, 5), blocks=(5, 10, 5)) + res = getattr(blosc2.lazyexpr(expression, {"o0": a}), op)() + expected = getattr(_numpy_eval(expression, {"o0": na}), op)() + np.testing.assert_allclose(res, expected, rtol=1e-5) + + +@pytest.mark.parametrize("dtype", [np.int16, np.int32, np.int64, np.uint32, np.bool_]) +@pytest.mark.parametrize( + "expression", ["a / b", "(a + b) / b", "where(f > 0, a, b) / b", "hypot(a, b)", "nextafter(a, b)"] +) +@pytest.mark.parametrize("how", ["getitem", "compute", "sum"]) +def test_miniexpr_integer_true_division(monkeypatch, dtype, expression, how): + # miniexpr computed these in integer arithmetic for two integer operands + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + if dtype is np.bool_ and expression != "a / b": + pytest.skip("only true division is defined for two bools here") + if dtype is np.int16 and expression == "nextafter(a, b)": + pytest.skip("the float32 result is computed in float64 on every path") + rng = np.random.default_rng(2) + na = rng.integers(0, 50, size=(30, 40)).astype(dtype) + nb = (rng.integers(1, 9, size=(30, 40)) if dtype is not np.bool_ else np.ones((30, 40))).astype(dtype) + nf = rng.normal(size=(30, 40)) + ops = { + name: blosc2.asarray(v, chunks=(10, 40), blocks=(5, 40)) + for name, v in zip("abf", (na, nb, nf), strict=True) + } + expr = blosc2.lazyexpr(expression, ops) + expected = _numpy_eval(expression, {"a": na, "b": nb, "f": nf}) + if how == "sum": + np.testing.assert_allclose(expr.sum(), expected.sum(), rtol=1e-5) + else: + res = expr[:] if how == "getitem" else expr.compute()[:] + np.testing.assert_allclose(res, expected.astype(res.dtype), rtol=1e-6) + + +@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.uint32, np.bool_]) +@pytest.mark.parametrize("func", ["arctan2", "copysign", "hypot", "logaddexp"]) +@pytest.mark.parametrize("how", ["getitem", "compute", "sum"]) +def test_miniexpr_integer_float_functions(monkeypatch, dtype, func, how): + # miniexpr computes every two-argument function in the promoted operand type: integer + # arithmetic for two integer operands, where NumPy gives a float (logaddexp gave 1, 2, 3, ...) + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + if dtype is np.bool_ and func == "logaddexp": + pytest.skip("logaddexp does not accept bools") + na = (np.arange(1, 1201) % 400 + 1).reshape(30, 40).astype(dtype) + nb = (np.arange(1, 1201) % 7).reshape(30, 40).astype(dtype) + ops = { + name: blosc2.asarray(v, chunks=(10, 40), blocks=(5, 40)) + for name, v in zip("ab", (na, nb), strict=True) + } + expression = f"{func}(a, b) + {func}(a, a)" + expr = blosc2.lazyexpr(expression, ops) + if how == "sum": + # bool inputs give float16 in NumPy, so sum up float64 values instead + expected = _numpy_eval(expression, {"a": na.astype(np.float64), "b": nb.astype(np.float64)}) + np.testing.assert_allclose(expr.sum(), expected.sum(), rtol=1e-5) + else: + res = expr[:] if how == "getitem" else expr.compute()[:] + expected = _numpy_eval(expression, {"a": na, "b": nb}) + np.testing.assert_allclose(res, expected.astype(res.dtype), rtol=1e-6) + + +@pytest.mark.parametrize( + "expression", + [ + "(a / 3) > 2", + "(a / 2.5) > 3", + "where(a / 3 > 2, a, b)", + "hypot(a, 2.5) > 3", + "arctan2(a, 2) > 1", + "logaddexp(a, 2) > 3", + ], +) +@pytest.mark.parametrize("how", ["getitem", "sum"]) +def test_miniexpr_integer_division_by_constant(monkeypatch, expression, how): + # miniexpr types numeric constants after the result: as integers or bools when that is not a + # float, so these divided integers (or took hypot of them) in integer arithmetic + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + na = np.arange(-600, 600, dtype=np.int64).reshape(30, 40) % 23 + nb = np.arange(1200, dtype=np.int64).reshape(30, 40) + ops = { + name: blosc2.asarray(v, chunks=(10, 40), blocks=(5, 40)) + for name, v in zip("ab", (na, nb), strict=True) + } + expr = blosc2.lazyexpr(expression, ops) + expected = _numpy_eval(expression, {"a": na, "b": nb}) + if how == "sum": + assert expr.sum() == expected.sum() + else: + np.testing.assert_array_equal(expr[:], expected) + + +_RESULT_DTYPE_EXPRESSIONS = [ + "hypot({n}, b)", + "copysign({n}, b)", + "nextafter({n}, b)", + "arctan2({n}, b)", + "hypot(a, b * 1)", + "hypot({n}, 2)", + "arctan2({n}, 2.5)", + "(a < b) / b", + "a / 2", + "a / b", + "{n} / b", + "{n} // b", + "{n} % b", + "a ** b", + "{n} ** 2", + "sin({n})", + "sqrt(a)", +] + + +@pytest.mark.parametrize("dtype", [np.bool_, np.int8, np.uint16, np.int32, np.int64]) +@pytest.mark.parametrize("expression", _RESULT_DTYPE_EXPRESSIONS) +@pytest.mark.parametrize("build", ["functions", "string"]) +def test_lazyexpr_result_dtype_matches_numpy(dtype, expression, build): + # Float functions and true division of integer (or bool) subexpressions were given an integer + # (or bool) result dtype, so the results were truncated on every evaluation path + if dtype is np.int32 and expression in ("a / b", "{n} / b"): + pytest.skip("int32 / int32 gives float32 on purpose") + if "nextafter" in expression and dtype not in (np.int32, np.int64): + pytest.skip("a float16 or float32 nextafter() is computed in float64 on every path") + # A subexpression of the operand dtype (bool + bool is an or in NumPy) + expression = expression.format(n="(a & b)" if dtype is np.bool_ else "(a + b)") + na = (np.arange(1200).reshape(30, 40) % 5).astype(dtype) + nb = (np.arange(1200).reshape(30, 40) % 3 + 1).astype(dtype) + a = blosc2.asarray(na, chunks=(10, 40), blocks=(5, 40)) + b = blosc2.asarray(nb, chunks=(10, 40), blocks=(5, 40)) + namespace = {"hypot": np.hypot, "copysign": np.copysign, "nextafter": np.nextafter} + namespace.update(arctan2=np.arctan2, sin=np.sin, sqrt=np.sqrt) + with np.errstate(all="ignore"): + expected = eval(expression, namespace, {"a": na, "b": nb}) + if build == "string": + expr = blosc2.lazyexpr(expression, {"a": a, "b": b}) + else: + expr = eval(expression, {name: getattr(blosc2, name) for name in namespace}, {"a": a, "b": b}) + assert expr.dtype == expected.dtype + rtol = 1e-3 if expected.dtype == np.float16 else 1e-6 + results = [expr[:], expr[...], expr.compute()[:], blosc2.where(b > 1, expr, 0)[:]] + for res, exp in zip(results, [expected] * 3 + [np.where(nb > 1, expected, 0)], strict=True): + assert res.dtype == exp.dtype + np.testing.assert_allclose(res, exp, rtol=rtol) + np.testing.assert_allclose(expr.sum(), expected.sum(dtype=np.float64), rtol=rtol) + + +def _typing_bug_data(dtype): + x = np.arange(1200).reshape(30, 40) + if dtype is np.bool_: + return x % 3 == 0 + if np.dtype(dtype).kind == "u": + return (x % 200).astype(dtype) + if np.dtype(dtype).kind == "i": + return (x % 251 - 120).astype(dtype) + return ((x % 17 - 8) * 0.5).astype(dtype) + + +_I8, _U8, _F32 = np.int8, np.uint8, np.float32 + + +@pytest.mark.parametrize( + ("expression", "dtypes"), + [ + # a comparison of operands of other dtypes inside arithmetic + ("(a < c) + c", (_I8, _I8, _U8)), + ("(a < c) * c", (_I8, _I8, _U8)), + ("(a < c) == b", (_I8, _I8, _U8)), + ("(a < c) + 1", (np.bool_, np.bool_, _U8)), + ("(a + b) * abs(c)", (np.int64,) * 3), + ("a + b > c", (_F32, _F32, np.float64)), + ("a * b < c", (np.int16, np.int16, np.int32)), + # bool + bool, which miniexpr adds as numbers + ("hypot(a + b, 2)", (np.bool_,) * 3), + ("sin(a + b)", (np.bool_,) * 3), + ("(a + b) * 1.5", (np.bool_,) * 3), + ("(a + b) == b", (np.bool_,) * 3), + ("a + b", (np.bool_,) * 3), + ("(a + b) + c", (np.bool_, np.bool_, np.int64)), + # constants truncated to the dtype they are compared in + ("a < 2.5", (np.int32,) * 3), + ("a >= 2.5", (np.int32,) * 3), + ("a == 2.5", (np.int32,) * 3), + ("a < 300", (_I8,) * 3), + # a comparison of two bools in a reduction, and miniexpr's grammar + ("a == b", (np.bool_,) * 3), + ("-a ** 2", (np.int64,) * 3), + ("a ** b ** c", (np.float64,) * 3), + ], +) +@pytest.mark.parametrize("how", ["getitem", "compute", "sum"]) +def test_miniexpr_typing_bugs(monkeypatch, expression, dtypes, how): + # Expressions that miniexpr types or parses differently from NumPy, which gave wrong results + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + nps = {name: _typing_bug_data(dt) for name, dt in zip("abc", dtypes, strict=True)} + if expression == "a ** b ** c": + nps = {k: v / 8 + 1 for k, v in nps.items()} + ops = {k: blosc2.asarray(v, chunks=(10, 40), blocks=(5, 40)) for k, v in nps.items()} + expr = blosc2.lazyexpr(expression, ops) + with np.errstate(all="ignore"): + expected = eval(expression, {"hypot": np.hypot, "sin": np.sin, "abs": np.abs}, nps) + rtol = 1e-3 if expected.dtype == np.float16 else 1e-6 + if how == "sum": + np.testing.assert_allclose(expr.sum(), expected.sum(dtype=np.float64), rtol=rtol) + return + res = expr[:] if how == "getitem" else expr.compute()[:] + np.testing.assert_allclose(res, expected.astype(res.dtype), rtol=rtol) + if res.dtype == np.bool_: + assert res.view(np.uint8).max() <= 1 # True + True stored as 2 + + +def test_miniexpr_typing_bugs_no_crash(): + # These overflowed miniexpr's buffers (a float64 conversion into a narrower one) and corrupted the heap + code = ( + "import numpy as np, blosc2\n" + "x = np.arange(30000)\n" + "nps = {'a': (x % 7 - 3).astype(np.int16), 'b': (x % 7).astype(np.uint32), 'f': (x % 7).astype(np.float32),\n" + " 'i': (x % 251 - 120).astype(np.int32), 'j': (x % 7).astype(np.int32)}\n" + "ops = {k: blosc2.asarray(v) for k, v in nps.items()}\n" + "for e in ['(a + a) * abs(a)', '(b * 3) * (-b)', '(f + f) - (f * -1.5)', 'where((f + f) > -1, f, f)',\n" + " 'abs(i) + (j < i)']:\n" + " res = blosc2.lazyexpr(e, ops)[:]\n" + " with np.errstate(all='ignore'):\n" + " assert np.allclose(res, eval(e, {'abs': np.abs, 'where': np.where}, nps)), e\n" + ) + import subprocess + import sys + + proc = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + + +def test_miniexpr_bool_true_division_no_crash(): + # bool / bool with a zero divisor died of SIGFPE inside miniexpr + code = ( + "import numpy as np, blosc2\n" + "a = blosc2.asarray(np.arange(1200).reshape(30, 40) % 3 > 0)\n" + "with np.errstate(all='ignore'):\n" + " r = blosc2.lazyexpr('a / a', {'a': a})[:]\n" + " blosc2.lazyexpr_mod = __import__('importlib').import_module('blosc2.lazyexpr')\n" + " blosc2.lazyexpr_mod.try_miniexpr = False\n" + " assert np.array_equal(r, blosc2.lazyexpr('a / a', {'a': a})[:], equal_nan=True)\n" + ) + import subprocess + import sys + + proc = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + + +@pytest.mark.parametrize("dtype", [np.int32, np.float64]) +@pytest.mark.parametrize("expression", ["a % b", "remainder(a, b)", "a % 3", "a % -3"]) +@pytest.mark.parametrize("how", ["getitem", "sum"]) +def test_miniexpr_remainder_sign(monkeypatch, dtype, expression, how): + # miniexpr's % follows C (sign of the dividend); NumPy's the sign of the divisor + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + rng = np.random.default_rng(4) + na = rng.integers(-50, 50, size=(30, 40)).astype(dtype) + nb = (rng.integers(1, 9, size=(30, 40)) * rng.choice([-1, 1], size=(30, 40))).astype(dtype) + ops = { + name: blosc2.asarray(v, chunks=(10, 40), blocks=(5, 40)) + for name, v in zip("ab", (na, nb), strict=True) + } + expr = blosc2.lazyexpr(expression, ops) + expected = _numpy_eval(expression, {"a": na, "b": nb}) + res = expr.sum() if how == "sum" else expr[:] + np.testing.assert_allclose(res, expected.sum() if how == "sum" else expected, rtol=1e-12) From c6609dbc6273864d667145f3bd070dacac1cfe39 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:50:41 +0100 Subject: [PATCH 15/18] 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 311c7e7b0..6e7a92075 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -4598,29 +4598,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 5d78923cb..bad79ff34 100644 --- a/tests/ndarray/test_getitem.py +++ b/tests/ndarray/test_getitem.py @@ -967,6 +967,18 @@ def test_basic_index_fast_path_matches_general_path(tmp_path, monkeypatch, on_di assert got.shape == np.asarray(npa[key]).shape, key +@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[:])) + + def _assert_getitem_matches_numpy(npa, a, key): try: want = npa[key] From cf1e115952c6682d7036ed4776746f2ae610cb64 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:51:45 +0100 Subject: [PATCH 16/18] 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 65c10cfec..055c3f24b 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -375,6 +375,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 @@ -1277,6 +1279,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 bbc1c6a7a..b29842d42 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 7510159345c1dc1ad395fb5291732da9a773e1e5 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 19:52:40 +0100 Subject: [PATCH 17/18] Add release notes for the NDArray performance work and fixes List the faster indexing, writes, casts, serialization and expressions, the bug fixes, and the compatibility notes (the next C-Blosc2 release is required; None keys return arrays owning their data; std/var rounding) in the 4.14.2 section. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- RELEASE_NOTES.md | 57 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 1d7bd308c..09f0104ba 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -127,6 +127,63 @@ XXX version-specific blurb XXX plotting requires an explicit bounded row window. Client caches are separate from a shared cat2lite server cache; repository allowances are per root. +### Faster NDArray indexing, writes and expressions + +- Basic keys (integers, unit-step slices, `...`, `None`) in + `NDArray.__getitem__` and `__setitem__` are resolved in Cython: `a[3, 5, 7]` + and row reads take ~30% less time, keys with NumPy integers up to 5x less. +- Creating small arrays costs a third less (`zeros((100, 100))` 31 -> 20 µs); + `arange`, `linspace`, `eye`, `reshape` and lazy functions such as `clip` no + longer call `inspect.stack()` (`clip()` 84 -> 12 µs). +- `a[region] = scalar` stores fully covered chunks as special chunks and writes + the rest in bounded pieces instead of building the region with `np.full()` + (`a[:] = 0.0` on 300 MB: 41 -> 0.2 ms, no extra memory). `a[region] = ndarray` + goes chunk by chunk too, copying compressed chunks when both arrays encode + them the same way (8x faster, a tenth of the memory). Arrays opened with + `locking=True` keep the single atomic write. +- `asarray()`/`astype()` of an NDArray cast whole chunks in reused buffers + (`astype(np.uint8)` 2x faster). +- `to_cframe()` writes the frame straight into the returned `bytes` (3.5x + faster, half the peak memory). +- Lazy expressions: `expr[:]` evaluates straight into the NumPy result + (`(a + b)[:]` 2x faster), full `std()`/`var()` are two single-pass reductions + (3x faster), chunked reductions decompress into reused buffers, and + `where(cond, x, y)` with a scalar `x` or `y` runs in miniexpr. +- `SChunk.decompress_chunk()` no longer reads the chunk from a frame twice. + +### Bug fixes + +- Fixed crashes of lazy expressions over in-memory frame-backed arrays + (`contiguous=True`, or `ndarray_from_cframe()` without a copy) and of the fast + `matmul` path with on-disk or frame-backed operands. +- Expressions that the bundled miniexpr evaluates wrongly or crashes on now take + the numexpr path: reductions with constant subexpressions, true division and + two-argument float functions of integers or bools, `%` with negative + operands, constants compared with integers, some mixed-dtype expressions, + `bool + bool`, and grammar differences such as `-x ** 2`. +- The result dtype of `hypot`, `copysign` and `nextafter` of integer + subexpressions and of true division involving bools follows NumPy (results + were truncated to integers). +- `NDArray.__getitem__` accepts the keys NumPy accepts that raised or gave a + wrong shape: huge slice bounds, bool steps, 0-d integer arrays, scalar bools + (placed and broadcast as in NumPy) and `None` next to array indices. +- `SChunk.update_special()` on an NDArray's super-chunk wrote chunks with the + wrong block size, so later reads of the array failed. +- `update_data()`/`insert_data()` with `copy=False` leaked their chunk when the + write failed. +- True positions of 1-D bool masks no longer run past the array length. +- `blosc2.nthreads` and the `CParams`/`DParams` defaults follow the + `BLOSC_NTHREADS` environment variable, which C-Blosc2 always applied; importing + blosc2 no longer overrides `NUMEXPR_NUM_THREADS`. + +### Compatibility notes + +- Requires the next C-Blosc2 release, for `blosc2_schunk_to_buffer_alloc()` and + for its rule that a failed chunk write without a copy leaves the chunk with + the caller. +- Keys with `None` return arrays that own their data (formerly views of a + temporary). `std()`/`var()` may differ from before in the last bits. + ## Changes from 4.14.0 to 4.14.1 Python-Blosc2 4.14.1 is a security and feature release introducing safe From f8e910f2d8b41ba7f5956b22429d9997cb7a53fd Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 30 Oct 2026 21:08:34 +0100 Subject: [PATCH 18/18] TEMP: build against the companion c-blosc2 PR branch Drop this commit before merging: once the c-blosc2 changes are merged and released, point BLOSC2_BUNDLED_VERSION back to Blosc/c-blosc2 at a release tag that contains them and bump BLOSC2_MIN_VERSION, since this branch uses blosc2_schunk_to_buffer_alloc() and the documented ownership of chunks after failed writes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SwMZj6jNaqsesMKnBde8Lr --- CMakeLists.txt | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index e9c5a078b..78a719366 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -17,6 +17,9 @@ project(python-blosc2) # Update to the latest minimum C-Blosc2 version, and to the actual bundled one set(BLOSC2_MIN_VERSION 3.3.0) set(BLOSC2_BUNDLED_VERSION v3.3.5) +# TEMPORARY: build against the companion c-blosc2 PR branch until it is merged and released +set(BLOSC2_BUNDLED_REPOSITORY https://github.com/FabianIsensee/c-blosc2) +set(BLOSC2_BUNDLED_VERSION 903bf73f002fd2b9298cd606a725dd31a88c5367) #set(BLOSC2_BUNDLED_VERSION a54e259aeb1c0e62d6bbea28934369d4831e4046) if(WIN32 AND NOT CMAKE_C_COMPILER_ID STREQUAL "Clang") @@ -219,7 +222,7 @@ else() set(BLOSC_INSTALL ON) include(FetchContent) FetchContent_Declare(blosc2 - GIT_REPOSITORY https://github.com/Blosc/c-blosc2 + GIT_REPOSITORY ${BLOSC2_BUNDLED_REPOSITORY} GIT_TAG ${BLOSC2_BUNDLED_VERSION} # SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../c-blosc2 )