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 ) 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 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/__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 64affa623..055c3f24b 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -14,8 +14,11 @@ import ast import atexit import pathlib import sys +import functools import time import warnings +import weakref +from threading import RLock import _ctypes @@ -24,11 +27,17 @@ from cpython cimport ( Py_buffer, PyBUF_SIMPLE, PyBuffer_Release, + PyBytes_AsString, PyBytes_FromStringAndSize, PyObject_GetBuffer, ) -from cpython.ref cimport Py_INCREF, Py_DECREF +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 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 from libc.stdint cimport uintptr_t from libc.stdlib cimport free, malloc, realloc, calloc @@ -186,6 +195,7 @@ cdef extern from "blosc2.h": ctypedef enum: BLOSC2_SPECIAL_ZERO BLOSC2_SPECIAL_NAN + BLOSC2_SPECIAL_VALUE BLOSC2_SPECIAL_UNINIT ctypedef enum: @@ -365,6 +375,10 @@ 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 int blosc2_set_maskout(blosc2_context *ctx, c_bool *maskout, int nblocks) @@ -481,7 +495,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 @@ -511,6 +527,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) @@ -603,6 +622,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 @@ -622,7 +642,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, @@ -632,8 +652,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) @@ -872,7 +890,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): @@ -1255,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) @@ -1352,10 +1388,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: @@ -1369,7 +1407,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: @@ -1699,6 +1737,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: @@ -1707,11 +1759,105 @@ 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 + + +# 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 + 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 = _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) @@ -1881,6 +2027,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 +2061,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 +2105,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 +2121,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 +2149,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) @@ -2009,15 +2161,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 @@ -2046,21 +2190,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) @@ -2069,6 +2236,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") @@ -2079,6 +2253,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 +2265,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 +2289,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 +2300,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 +2318,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 +2329,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 @@ -2164,21 +2345,23 @@ 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") return rc + @_handle_locked def update_chunk(self, nchunk, chunk): cdef const uint8_t[:] typed_view_chunk mem_view_chunk = memoryview(chunk) @@ -2189,6 +2372,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 +2380,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) @@ -2210,22 +2395,24 @@ 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") 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 +2428,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 +2453,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,26 +2559,15 @@ 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 - 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): 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 +2580,26 @@ 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") + @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) _check_rc(rc, "Error while refreshing the schunk") @@ -2433,6 +2625,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 +2660,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 +2684,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 +2714,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 +2750,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 +2833,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. @@ -2820,7 +3023,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] @@ -3320,6 +3525,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 +3547,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,15 +3561,18 @@ 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) if rc < 0: raise RuntimeError("Could not delete the vlmeta") + @_handle_locked 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 +3581,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 +3593,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 +3612,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 @@ -3401,10 +3621,12 @@ 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 +@_handle_locked def meta_keys(self): cdef blosc2_schunk *schunk = self.c_schunk keys = [] @@ -3549,6 +3771,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) @@ -3631,12 +3891,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, @@ -3698,11 +3967,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: @@ -3814,24 +4090,52 @@ 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), 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: + return 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 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 = _new_handle_lock() + if schunk is not None: + (schunk)._handle_lock = self._handle_lock @property def c_array(self): @@ -3892,29 +4196,237 @@ cdef class NDArray: 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 + 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) + # 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: - PyThread_release_lock(self.read_lock) - PyBuffer_Release(&view) - _check_rc(rc, "Error while getting the buffer") - - return arr + PyBuffer_Release(&view) + return rc + @_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 +4509,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 +4527,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 +4560,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 +4594,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 +4636,7 @@ cdef class NDArray: return ndarray + @_handle_locked def set_slice(self, key, ndarray): ndim = self.ndim start, stop = key @@ -4138,6 +4655,7 @@ cdef class NDArray: return self + @_handle_locked def tobytes(self): buffersize = self.size * self.array.sc.typesize buffer = bytes(buffersize) @@ -4146,20 +4664,11 @@ cdef class NDArray: return buffer + @_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): chunks = kwargs.pop("chunks", self.chunks) blocks = kwargs.pop("blocks", self.blocks) @@ -4180,6 +4689,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 +4697,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 +4867,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,7 +5050,15 @@ cdef class NDArray: "status": _me_compile_status_name(rc), } + @_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 @@ -4558,6 +5078,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 +5101,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 +5130,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 +5379,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 +5396,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 +5426,10 @@ 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 + 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") if copy: @@ -4916,6 +5440,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 +5458,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..b29842d42 100644 --- a/src/blosc2/core.py +++ b/src/blosc2/core.py @@ -24,7 +24,7 @@ import sys import urllib.parse import urllib.request -from dataclasses import asdict +from dataclasses import asdict, fields from functools import cache, lru_cache from typing import TYPE_CHECKING, ClassVar @@ -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 @@ -1449,6 +1452,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) @@ -1977,7 +1984,9 @@ def compute_chunks_blocks( # noqa: C901 cparams = kwargs.get("cparams") or blosc2.CParams() # just get defaults if isinstance(cparams, blosc2.CParams): - cparams = asdict(cparams) + # Shallow is enough (only keys are set below, filters only on a deep copy), and asdict() + # costs more than the rest of this function + cparams = {f.name: getattr(cparams, f.name) for f in fields(cparams)} # Typesize in dtype always has preference over typesize in cparams itemsize = cparams["typesize"] = np.dtype(dtype).itemsize @@ -1985,7 +1994,7 @@ def compute_chunks_blocks( # noqa: C901 # Get the default blocksize for the compression params # Check if we need STUNE for lossy codecs/filters that have specific blocksize requirements codec = cparams.get("codec") - filters = cparams.get("filters", None) + filters = cparams.get("filters") needs_stune = codec in ( blosc2.Codec.ZFP_RATE, blosc2.Codec.ZFP_PREC, diff --git a/src/blosc2/indexing.py b/src/blosc2/indexing.py index 892ff36ea..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", {}) @@ -5098,6 +5114,8 @@ def get_index(array: blosc2.NDArray, field: str | None = None, name: str | None 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/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 732c1141e..6e7a92075 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -54,6 +54,8 @@ 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 ( check_smaller_shape, @@ -294,6 +296,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 +325,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) @@ -334,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 = ( @@ -1328,7 +1354,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 +1378,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 +1407,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 +1471,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 +1519,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 +1542,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 +1562,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: @@ -1698,13 +1734,63 @@ 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. + + 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 @@ -1806,11 +1892,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 +1910,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: @@ -1862,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 @@ -1931,6 +2023,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 +2055,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 @@ -1979,10 +2080,16 @@ 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) - 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) + # The prefilter reads the operands' super-chunks from C, outside their methods + operand_locks = blosc2.blosc2_ext.HandleLocks(_pref_ops.values()) + 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. @@ -2015,10 +2122,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): @@ -2028,7 +2136,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, @@ -2814,7 +2930,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 @@ -2862,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") @@ -2920,8 +3036,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): @@ -2965,16 +3084,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 ) @@ -2984,10 +3106,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 @@ -2996,12 +3118,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) @@ -3497,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}") @@ -3515,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 @@ -3625,6 +3759,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) @@ -3909,7 +4046,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 @@ -4141,50 +4278,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) @@ -4193,9 +4349,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 @@ -4204,6 +4360,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, @@ -4214,56 +4382,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, @@ -4479,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: @@ -4613,7 +4732,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() @@ -5691,7 +5810,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/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/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/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 98ab50b0a..e952000cc 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -8,9 +8,11 @@ from __future__ import annotations import builtins -import inspect import itertools import math +import operator +import os +import sys import weakref from abc import abstractmethod from collections import OrderedDict, namedtuple @@ -38,10 +40,12 @@ from .linalg import matmul from .utils import ( + chunk_flat_index, get_chunks_idx, get_local_slice, get_selection, incomplete_lazyfunc, + is_called_from, npbinvert, nplshift, nprshift, @@ -240,9 +244,43 @@ 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): + """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): @@ -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) @@ -305,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. @@ -322,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) @@ -3725,38 +3783,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: @@ -3772,6 +3806,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 @@ -4545,33 +4746,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 +4801,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: @@ -4645,6 +4867,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) @@ -4663,17 +4893,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 +4933,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) @@ -4727,13 +4959,15 @@ 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 __setitem__( + 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__( # noqa: C901 self, key: int | slice | Sequence[slice | int | np.bool_ | np.ndarray[int | np.bool_] | None] | None, value: object, @@ -4768,21 +5002,31 @@ 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 + # 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 + 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 @@ -4794,27 +5038,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 - - 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_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__.") + 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: @@ -4824,12 +5069,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. @@ -6863,7 +7253,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`. @@ -6966,6 +7356,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 @@ -6995,6 +7387,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. @@ -7081,8 +7515,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..2907f1447 100644 --- a/src/blosc2/schunk.py +++ b/src/blosc2/schunk.py @@ -450,8 +450,9 @@ def __init__( # noqa: C901 default = DeserializeMode.SAFE if opens_existing else DeserializeMode.FULL set_deserialize(self, default if deserialize is None else deserialize) self._vlmeta = vlmeta(self, super().c_schunk) - self._cparams = super().get_cparams() - self._dparams = super().get_dparams() + # Built on first use: most handles (e.g. those of NDArrays) never need them + self._cparams = None + self._dparams = None def __enter__(self) -> SChunk: """Enter a context manager and return this super-chunk.""" @@ -470,6 +471,8 @@ def cparams(self) -> blosc2.CParams: """ :class:`blosc2.CParams` instance with the compression parameters. """ + if self._cparams is None: + self._cparams = super().get_cparams() return self._cparams @cparams.setter @@ -484,6 +487,8 @@ def dparams(self) -> blosc2.DParams: """ :class:`blosc2.DParams` instance with the decompression parameters. """ + if self._dparams is None: + self._dparams = super().get_dparams() return self._dparams @dparams.setter @@ -815,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/src/blosc2/utils.py b/src/blosc2/utils.py index 77f2d3390..b44572685 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 @@ -1007,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 @@ -1050,13 +1058,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_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 diff --git a/tests/ndarray/test_concurrent_access.py b/tests/ndarray/test_concurrent_access.py new file mode 100644 index 000000000..186870e50 --- /dev/null +++ b/tests/ndarray/test_concurrent_access.py @@ -0,0 +1,227 @@ +####################################################################### +# 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 itertools +import os +import signal +import threading +import time + +import numpy as np +import pytest + +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): + """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 + + +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 write_region(array, k): + array[8:24, 0:100] = VALUES[8:24, 0:100] + (k % 2) + + +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}, + ) + + def longest_gap(): + """Longest time a ticker thread did not run during one read, as a fraction of the read.""" + 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] + 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))) == [] + + +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_vlmeta(a, k): + a.schunk.vlmeta[f"k{k % 5}"] = list(range(k % 50)) + + +def swap_dparams(a, k): + a.schunk.dparams = blosc2.DParams(nthreads=1 + k % 2) + + +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(): + 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) + + assert run_for(SECONDS, 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( + 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() diff --git a/tests/ndarray/test_getitem.py b/tests/ndarray/test_getitem.py index 437cab4f5..bad79ff34 100644 --- a/tests/ndarray/test_getitem.py +++ b/tests/ndarray/test_getitem.py @@ -912,3 +912,286 @@ 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) + + +# 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 + + +@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] + 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) + + +# 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( + "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) diff --git a/tests/ndarray/test_indexing.py b/tests/ndarray/test_indexing.py index 23aab884e..3495b50d4 100644 --- a/tests/ndarray/test_indexing.py +++ b/tests/ndarray/test_indexing.py @@ -1886,6 +1886,37 @@ 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") + + 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)) + 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 + arr, urlpath = _make_scalar_persistent_array(tmp_path) + 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() diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 5a099cbac..b4a916c59 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -2406,3 +2406,412 @@ 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( + ("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 + # 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) + + +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) 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) 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)) diff --git a/tests/ndarray/test_setitem.py b/tests/ndarray/test_setitem.py index 78b2cb8e1..3f2a7cdb3 100644 --- a/tests/ndarray/test_setitem.py +++ b/tests/ndarray/test_setitem.py @@ -149,3 +149,99 @@ def test_setitem_fancy_index(): out_numpy[idx] = value 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))): + 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/ndarray/test_setitem_chunkwise.py b/tests/ndarray/test_setitem_chunkwise.py new file mode 100644 index 000000000..0a23cb0e7 --- /dev/null +++ b/tests/ndarray/test_setitem_chunkwise.py @@ -0,0 +1,378 @@ +####################################################################### +# 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), and the chunk-wise casts of astype(). + +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) + + +# 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 + + 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 # --------------------------------------------------------------------------- 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) 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 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)