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..ac7716d8a 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 @@ -27,8 +30,13 @@ from cpython cimport ( PyBytes_FromStringAndSize, PyObject_GetBuffer, ) +from cpython.bool cimport PyBool_Check +from cpython.long cimport PyLong_Check +from cpython.number cimport PyIndex_Check +from cpython.pyport cimport PY_SSIZE_T_MAX, PY_SSIZE_T_MIN from cpython.ref cimport Py_INCREF, Py_DECREF from cpython.pycapsule cimport PyCapsule_GetPointer, PyCapsule_New +from cpython.slice cimport PySlice_AdjustIndices, PySlice_Check, PySlice_Unpack from cython.operator cimport dereference from libc.stdint cimport uintptr_t from libc.stdlib cimport free, malloc, realloc, calloc @@ -186,6 +194,7 @@ cdef extern from "blosc2.h": ctypedef enum: BLOSC2_SPECIAL_ZERO BLOSC2_SPECIAL_NAN + BLOSC2_SPECIAL_VALUE BLOSC2_SPECIAL_UNINIT ctypedef enum: @@ -365,6 +374,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) @@ -511,6 +524,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 +619,7 @@ cdef extern from "b2nd.h": int64_t chunk_array_strides[B2ND_MAX_DIM] char *dtype int8_t dtype_format + int64_t last_tick ctypedef struct b2nd_context_t: pass @@ -622,7 +639,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, @@ -872,7 +889,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 +1278,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 +1387,12 @@ cdef create_cparams_from_kwargs(blosc2_cparams *cparams, kwargs): filters = kwargs.get('filters', blosc2.cparams_dflts['filters']) if len(filters) > BLOSC2_MAX_FILTERS: raise ValueError(f"filters list cannot exceed {BLOSC2_MAX_FILTERS}") + # Looked up once: an enum's value is a slow descriptor + cdef int bytedelta = blosc2.Filter.BYTEDELTA.value for i, filter in enumerate(filters): cparams.filters[i] = filter.value if isinstance(filter, Enum) else filter # Bytedelta does not work on typesize 1 - if cparams.filters[i] == blosc2.Filter.BYTEDELTA.value and typesize == 1: + if cparams.filters[i] == bytedelta and typesize == 1: cparams.filters[i] = 0 if "filters_meta" not in kwargs: @@ -1369,7 +1406,7 @@ cdef create_cparams_from_kwargs(blosc2_cparams *cparams, kwargs): for i, meta in enumerate(filters_meta): # We still may want to encode negative values meta_value = meta if meta < 0 else meta - if meta_value == 0 and cparams.filters[i] == blosc2.Filter.BYTEDELTA.value: + if meta_value == 0 and cparams.filters[i] == bytedelta: # bytedelta typesize cannot be zero when using compress2 cparams.filters_meta[i] = typesize else: @@ -1699,6 +1736,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 +1758,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) @@ -1848,6 +1993,18 @@ cdef class SChunk: def c_schunk(self): return self.schunk + @property + def _has_prefilter(self): + """Whether writes through this handle run a prefilter.""" + with self._handle_lock: + return self.schunk.storage.cparams.prefilter != NULL + + @property + def _has_postfilter(self): + """Whether reads through this handle run a postfilter.""" + with self._handle_lock: + return self.schunk.storage.dparams.postfilter != NULL + @property def chunksize(self): return self.schunk.chunksize @@ -1881,6 +2038,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 +2072,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 +2116,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 +2132,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 +2160,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 +2172,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 +2201,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 +2247,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 +2264,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 +2276,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 +2300,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 +2311,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 +2329,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 +2340,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 +2356,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 +2383,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 +2391,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 +2406,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 +2439,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 +2464,7 @@ cdef class SChunk: if rc < 0: raise RuntimeError("Error while getting the slice") + @_handle_locked def set_slice(self, value, start=0, stop=None): cdef int64_t nitems = self.schunk.nbytes // self.schunk.typesize start, stop = self._massage_key(start, stop, nitems) @@ -2370,6 +2570,7 @@ cdef class SChunk: if rc < 0: raise RuntimeError("Error while setting the slice") + @_handle_locked def to_cframe(self): cdef c_bool needs_free cdef uint8_t *cframe @@ -2387,9 +2588,11 @@ cdef class SChunk: return out + @_handle_locked def _avoid_cframe_free(self, avoid_cframe_free): blosc2_schunk_avoid_cframe_free(self.schunk, avoid_cframe_free) + @_handle_locked def lock(self): """Take the exclusive frame lock and hold it across several operations. @@ -2402,12 +2605,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 +2650,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 +2685,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 +2709,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 +2739,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 +2775,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 +2858,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 +3048,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 +3550,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 +3572,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 +3586,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 +3606,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 +3618,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 +3637,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 +3646,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 = [] @@ -3631,12 +3878,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 +3954,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 +4077,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 +4183,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 +4496,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 +4514,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 +4547,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 +4581,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 +4623,7 @@ cdef class NDArray: return ndarray + @_handle_locked def set_slice(self, key, ndarray): ndim = self.ndim start, stop = key @@ -4138,6 +4642,7 @@ cdef class NDArray: return self + @_handle_locked def tobytes(self): buffersize = self.size * self.array.sc.typesize buffer = bytes(buffersize) @@ -4146,6 +4651,7 @@ cdef class NDArray: return buffer + @_handle_locked def to_cframe(self): cdef c_bool needs_free cdef uint8_t *cframe @@ -4160,6 +4666,7 @@ cdef class NDArray: return out + @_handle_locked def copy(self, dtype, **kwargs): chunks = kwargs.pop("chunks", self.chunks) blocks = kwargs.pop("blocks", self.blocks) @@ -4180,6 +4687,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 +4695,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 +4865,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 +5048,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 +5076,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 +5099,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 +5128,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 +5377,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 +5394,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 +5424,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 +5438,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 +5456,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/lazyexpr.py b/src/blosc2/lazyexpr.py index 6cf988169..188d4c612 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -1983,10 +1983,13 @@ def _miniexpr_eligible_operand(op): nct = res_eval.schunk.nchunks if bpc > 0 and _cb_anchor.shape[0] >= nct * bpc: chunk_has_candidate = _cb_anchor[: nct * bpc].reshape(nct, bpc).any(axis=1) + # The prefilter reads the operands' super-chunks from C, outside their methods + operand_locks = blosc2.blosc2_ext.HandleLocks(_pref_ops.values()) for nchunk in range(res_eval.schunk.nchunks): if chunk_has_candidate is not None and not chunk_has_candidate[nchunk]: continue - res_eval.schunk.update_data(nchunk, data, copy=False) + with operand_locks: + res_eval.schunk.update_data(nchunk, data, copy=False) # _cb_anchor (if any) stays referenced until fast_eval returns, which # is after the finally below removes the prefilter — so the buffer # safely outlives every prefilter dereference. @@ -2924,8 +2927,11 @@ def reduce_slices( # noqa: C901 data = np.empty(res_eval.schunk.chunksize, dtype=np.uint8) chunk_data = np.empty(res_eval.schunk.chunksize + blosc2.MAX_OVERHEAD, dtype=np.uint8) # Exercise prefilter for each chunk + # The prefilter reads the operands' super-chunks from C, outside their methods + operand_locks = blosc2.blosc2_ext.HandleLocks(operands.values()) for nchunk in range(res_eval.schunk.nchunks): - res_eval.schunk._prefilter_data(nchunk, data, chunk_data) + with operand_locks: + res_eval.schunk._prefilter_data(nchunk, data, chunk_data) except Exception as e: use_miniexpr = False if callable(expression) and _is_dsl_kernel_expression(expression): @@ -4483,29 +4489,29 @@ def _collect_flat_indices_from_bool_ndarray(bool_ndarray): Flat indices of True positions (int64). """ chunk_len = bool_ndarray.chunks[0] + nitems = bool_ndarray.shape[0] all_indices = [] for info in bool_ndarray.iterchunks_info(): + offset = info.nchunk * chunk_len + # Items of this chunk inside the array: the last chunk's padding is not data, and it + # can hold anything (e.g. the fill value of a blosc2.full() array) + nvalid = builtins.min(chunk_len, nitems - offset) # Skip special-value chunks that are entirely False if info.special == blosc2.SpecialValue.ZERO: continue if info.special == blosc2.SpecialValue.VALUE: - if not info.repeated_value: # repeated_value is False/0 + if not bool(info.repeated_value): # repeated_value is False/0 continue # repeated_value is True: all elements in this chunk are True - offset = info.nchunk * chunk_len - all_indices.append(np.arange(offset, offset + chunk_len, dtype=np.int64)) + all_indices.append(np.arange(offset, offset + nvalid, dtype=np.int64)) continue # Normal chunk: decompress and scan for True positions raw = bool_ndarray.schunk.decompress_chunk(info.nchunk) - arr = np.frombuffer(raw, dtype=np.bool_) - # Truncate to the logical chunk size (buffer may include padding) - if len(arr) > chunk_len: - arr = arr[:chunk_len] + arr = np.frombuffer(raw, dtype=np.bool_)[:nvalid] idx = np.flatnonzero(arr) if len(idx) > 0: - offset = info.nchunk * chunk_len all_indices.append(idx + offset) if not all_indices: diff --git a/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/ndarray.py b/src/blosc2/ndarray.py index 98ab50b0a..2b517307b 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) @@ -3725,38 +3780,14 @@ def detect_aligned_chunks( if s.stop is not None and s.stop % chunks[i] != 0: return [] - # Parse the slice boundaries - start_indices = [] - end_indices = [] - n_chunks = [] - - for i, s in enumerate(key): - start = s.start if s.start is not None else 0 - stop = s.stop if s.stop is not None else shape[i] - chunk_size = chunks[i] - start_idx = start // chunk_size - end_idx = stop // chunk_size - start_indices.append(start_idx) - end_indices.append(end_idx) - # Total chunk count along this dim: ceil, not floor -- a trailing - # partial chunk (shape[i] not a multiple of chunk_size) still counts - # as one chunk. Floor division here undercounts it, which corrupts - # the flat chunk-index math below for any aligned slice with a - # nonzero start in an earlier dimension (silently returns a - # different chunk's data instead of the requested one). - n_chunks.append(math.ceil(shape[i] / chunk_size)) - # Get all chunk combinations in the slice - indices = [range(start, end) for start, end in zip(start_indices, end_indices, strict=False)] - result = [] - - for combination in product(*indices): - flat_index = 0 - multiplier = 1 - for idx, n in zip(reversed(range(len(n_chunks))), reversed(n_chunks), strict=False): - flat_index += combination[idx] * multiplier - multiplier *= n - result.append(flat_index) + indices = [] + for s, n, c in zip(key, shape, chunks, strict=True): + start = s.start if s.start is not None else 0 + stop = s.stop if s.stop is not None else n + indices.append(range(start // c, stop // c)) + grid, _ = get_chunks_idx(shape, chunks) + result = [chunk_flat_index(combination, grid) for combination in product(*indices)] # Check if chunks are consecutive if requested if consecutive and result: @@ -3772,6 +3803,176 @@ 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.""" + # Copying compressed bytes skips both the source postfilter and the destination prefilter. + if src.schunk._has_postfilter or dst.schunk._has_prefilter: + return False + if src.dtype != dst.dtype or src.chunks != dst.chunks or src.blocks != dst.blocks: + return False + a, b = dst.cparams, src.cparams + 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,160 @@ 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.) + """ + # Special chunks skip compression, including its prefilter. Keep a single ordinary + # set_slice call with callbacks so partial chunks are filtered exactly as before. + if self.ndim == 0 or self.schunk._has_prefilter: + return False + region = math.prod(sp - st for st, sp in zip(start, stop, strict=True)) + chunk_nitems = math.prod(self.chunks) + 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). + """ + # One ordinary write preserves prefilter calls on the containing partial chunks too. + if self.ndim == 0 or self.schunk._has_prefilter: + return False + region = tuple(sp - st for st, sp in zip(start, stop, strict=True)) + chunk_nitems = math.prod(self.chunks) + 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 +7256,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 +7359,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 +7390,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 +7518,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_chunkwise_callbacks.py b/tests/ndarray/test_chunkwise_callbacks.py new file mode 100644 index 000000000..4bc4e17ec --- /dev/null +++ b/tests/ndarray/test_chunkwise_callbacks.py @@ -0,0 +1,140 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +import numpy as np +import pytest + +import blosc2 + +# Whole arrays, aligned whole chunks, mixed whole/partial chunks, and one partial chunk. +KEYS = [slice(None), slice(8, 24), slice(4, 28), slice(4, 8)] + + +def make_array(values, chunks=(8,), blocks=(4,)): + return blosc2.asarray( + values, chunks=chunks, blocks=blocks, cparams={"nthreads": 1}, dparams={"nthreads": 1} + ) + + +@pytest.mark.parametrize("key", KEYS) +def test_assignment_applies_source_postfilter(key): + expected = np.zeros(32, dtype=np.int32) + values = np.arange(expected[key].size, dtype=np.int32) + src = make_array(values) + dst = make_array(expected) + + @src.schunk.postfilter(np.int32, np.int32) + def copy_source_postfilter(inp, out, offset): + out[:] = inp + 100 + + try: + expected[key] = values + 100 + dst[key] = src + np.testing.assert_array_equal(dst[:], expected) + finally: + src.schunk.remove_postfilter("copy_source_postfilter") + + +@pytest.mark.parametrize("scalar", [True, False], ids=["scalar", "array"]) +@pytest.mark.parametrize("key", KEYS) +def test_assignment_applies_destination_prefilter(key, scalar): + expected = np.zeros(32, dtype=np.int32) + dst = make_array(expected) + reference = make_array(expected) + values = 7 if scalar else np.arange(expected[key].size, dtype=np.int32) + src = values if scalar else make_array(values) + + @dst.schunk.prefilter(np.int32, np.int32) + def copy_destination_prefilter(inp, out, offset): + out[:] = inp + 100 + + @reference.schunk.prefilter(np.int32, np.int32) + def reference_destination_prefilter(inp, out, offset): + out[:] = inp + 100 + + try: + # NumPy input takes the ordinary set_slice path. A partial write recompresses its + # containing chunks, so their existing items also pass through the prefilter. + reference[key] = np.full(expected[key].shape, values, dtype=np.int32) if scalar else values + dst[key] = src + np.testing.assert_array_equal(dst[:], reference[:]) + np.testing.assert_array_equal(dst[key], np.broadcast_to(values + 100, expected[key].shape)) + finally: + dst.schunk.remove_prefilter("copy_destination_prefilter") + reference.schunk.remove_prefilter("reference_destination_prefilter") + + +def test_large_asarray_copy_applies_source_postfilter(): + # Exactly 16 MiB reaches asarray's chunk-wise path, rather than its small-array path. + values = np.zeros(2**22, dtype=np.int32) + src = make_array(values, chunks=(2**18,), blocks=(2**16,)) + + @src.schunk.postfilter(np.int32, np.int32) + def large_copy_source_postfilter(inp, out, offset): + out[:] = inp + 100 + + try: + dst = blosc2.asarray(src, copy=True, cparams={"nthreads": 1}, dparams={"nthreads": 1}) + np.testing.assert_array_equal(dst[:], values + 100) + finally: + src.schunk.remove_postfilter("large_copy_source_postfilter") + + +@pytest.mark.parametrize("callback", ["prefilter", "postfilter"]) +def test_unrelated_callbacks_keep_chunkwise_fast_paths(callback, monkeypatch): + owner = make_array(np.zeros(32, dtype=np.int32)) + src = make_array(np.arange(32, dtype=np.int32)) + dst = make_array(np.zeros(32, dtype=np.int32)) + + @getattr(owner.schunk, callback)(np.int32, np.int32) + def unrelated_copy_callback(inp, out, offset): + out[:] = inp + 100 + + special_chunks = [] + copied_chunks = [] + original_constant_chunk = blosc2.NDArray._constant_chunk + original_get_chunk = blosc2.SChunk.get_chunk + + def constant_chunk(self, fill): + special_chunks.append(self) + return original_constant_chunk(self, fill) + + def get_chunk(self, nchunk): + if self is src.schunk: + copied_chunks.append(nchunk) + return original_get_chunk(self, nchunk) + + monkeypatch.setattr(blosc2.NDArray, "_constant_chunk", constant_chunk) + monkeypatch.setattr(blosc2.SChunk, "get_chunk", get_chunk) + try: + dst[:] = 7 + np.testing.assert_array_equal(dst[:], np.full(32, 7, dtype=np.int32)) + assert special_chunks == [dst] + dst[:] = src + np.testing.assert_array_equal(dst[:], np.arange(32, dtype=np.int32)) + assert copied_chunks == [0, 1, 2, 3] + finally: + getattr(owner.schunk, "remove_" + callback)("unrelated_copy_callback") + + +@pytest.mark.parametrize("callback", ["prefilter", "postfilter"]) +def test_handle_callback_flags_follow_registration(callback): + owner = make_array(np.zeros(32, dtype=np.int32)) + other = make_array(np.zeros(32, dtype=np.int32)) + flag = "_has_" + callback + assert not getattr(owner.schunk, flag) + + @getattr(owner.schunk, callback)(np.int32, np.int32) + def handle_copy_callback(inp, out, offset): + out[:] = inp + + try: + assert getattr(owner.schunk, flag) + assert not getattr(other.schunk, flag) + finally: + getattr(owner.schunk, "remove_" + callback)("handle_copy_callback") + assert not getattr(owner.schunk, flag) 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..472f483fe 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 + + +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) + + +@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[:])) diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 5a099cbac..763d4de15 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -2406,3 +2406,19 @@ def test_cumulative_block_layout(monkeypatch, blocks): expected = np.cumulative_sum(data, axis=0) np.testing.assert_array_equal(result[:], expected) np.testing.assert_array_equal(out[:], expected) + + +@pytest.mark.parametrize("source", ["contiguous", "cframe"]) +def test_frame_backed_inmemory_operands(source): + # In-memory frames (contiguous=True, or views of a cframe) keep no chunk table, so the + # miniexpr prefilter must fetch their chunks like it does for on-disk operands + na = np.arange(270, dtype=np.int16).reshape(10, 27) + if source == "contiguous": + a = blosc2.asarray(na, chunks=(3, 4), blocks=(1, 2), contiguous=True) + else: + a = blosc2.ndarray_from_cframe(blosc2.asarray(na, chunks=(3, 4), blocks=(1, 2)).to_cframe()) + assert a.schunk.contiguous + assert a.schunk.urlpath is None + b = blosc2.asarray(na * 2, chunks=(3, 4), blocks=(1, 2)) + np.testing.assert_array_equal((a + 1)[:], na + 1) + np.testing.assert_array_equal((a + b).compute()[:], na * 3) diff --git a/tests/ndarray/test_linalg.py b/tests/ndarray/test_linalg.py index 2827f8a23..a93b0b1c2 100644 --- a/tests/ndarray/test_linalg.py +++ b/tests/ndarray/test_linalg.py @@ -1283,3 +1283,31 @@ def test_matmul_broadcast_batch_chunks(): blosc2.asarray(b, chunks=(1, 2, 2)), blosc2.asarray(a, chunks=(1, 2, 2)), chunks=(1, 2, 2) ) np.testing.assert_allclose(res[:], np.matmul(b, a)) + + +@pytest.mark.parametrize("storage", ["contiguous", "disk"]) +def test_matmul_frame_backed_operands(storage, tmp_path): + # The fast matmul prefilter reads in-memory chunk tables; frames must take the general path + x = np.arange(64 * 64, dtype=np.float32).reshape(64, 64) / 100 + y = x.T.copy() + if storage == "contiguous": + a, b = blosc2.asarray(x, contiguous=True), blosc2.asarray(y, contiguous=True) + else: + a = blosc2.asarray(x, urlpath=str(tmp_path / "x.b2nd"), mode="w") + b = blosc2.asarray(y, urlpath=str(tmp_path / "y.b2nd"), mode="w") + np.testing.assert_allclose(blosc2.matmul(a, b)[:], x @ y, rtol=1e-5) + + +@pytest.mark.parametrize("storage", ["contiguous", "disk"]) +def test_set_pref_matmul_rejects_frame_backed_operands(storage, tmp_path): + # The prefilter itself refuses frames too, rather than leaving that to matmul's Python checks + x = np.ones((64, 64), dtype=np.float32) + if storage == "contiguous": + a = blosc2.asarray(x, contiguous=True) + else: + a = blosc2.asarray(x, urlpath=str(tmp_path / "x.b2nd"), mode="w") + b = blosc2.asarray(x) + out = blosc2.zeros((64, 64), dtype=np.float32) + for inputs in ({"x1": a, "x2": b}, {"x1": b, "x2": a}): + with pytest.raises(ValueError, match="frame-backed"): + out._set_pref_matmul(inputs, fp_accuracy=blosc2.FPAccuracy.DEFAULT) 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..48f446a91 --- /dev/null +++ b/tests/ndarray/test_setitem_chunkwise.py @@ -0,0 +1,383 @@ +####################################################################### +# 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)) + if np.dtype(dtype).kind in "iu": + # Casting floats out of an integer type's range is undefined: on aarch64 NumPy warns, and its + # contiguous and strided casts give different values + info = np.iinfo(dtype) + x = np.clip(x, info.min, info.max) + a = blosc2.asarray(x, chunks=chunks, blocks=blocks) + b = blosc2.astype(a, dtype) + assert b.dtype == np.dtype(dtype) + assert b.chunks == a.chunks + assert b.blocks == a.blocks + np.testing.assert_array_equal(b[...], x.astype(dtype)) + sub = (slice(3, 50), slice(10, 200), slice(7, 250)) + np.testing.assert_array_equal(b[sub], x[sub].astype(dtype)) + + +def test_astype_chunkwise_kinds(tmp_path): + rng = np.random.default_rng(8) + x = rng.integers(-100, 100, size=BIG).astype(np.int64) + a = blosc2.asarray(x, chunks=(16, 128, 256), blocks=(4, 32, 256)) + np.testing.assert_array_equal(blosc2.astype(a, "m8[s]")[...], x.astype("m8[s]")) + c = blosc2.astype(a, np.complex64, urlpath=str(tmp_path / "c.b2nd"), mode="w") + np.testing.assert_array_equal(blosc2.open(c.urlpath)[...], x.astype(np.complex64)) + # A copy with the same chunk encoding copies the compressed chunks + d = blosc2.asarray(a, copy=True, urlpath=str(tmp_path / "d.b2nd"), mode="w") + assert all(a.schunk.get_chunk(n) == d.schunk.get_chunk(n) for n in range(a.schunk.nchunks)) + np.testing.assert_array_equal(blosc2.open(d.urlpath)[...], x) + + +def test_chunkwise_buffers_fit_the_boxes(): + # The temporary buffers are as large as the largest box written from them, not as the budget + import tracemalloc + + 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)