diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index de90e254c..475072943 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -3681,6 +3681,7 @@ def compress2(src, **kwargs): _dest = dest cctx = blosc2_create_cctx(cparams) if cctx == NULL: + PyBuffer_Release(&buf) raise RuntimeError("Could not create the compression context") size = _compress_ctx(cctx, buf.buf, buf.len, _dest, len_dest) blosc2_free_ctx(cctx) @@ -4426,6 +4427,7 @@ cdef class SChunk: if dst is not None: PyObject_GetBuffer(dst, &buf, PyBUF_SIMPLE) if buf.len == 0: + PyBuffer_Release(&buf) raise ValueError("The dst length must be greater than 0") size = blosc2_schunk_decompress_chunk(self.schunk, nchunk, buf.buf, buf.len) PyBuffer_Release(&buf) @@ -4612,6 +4614,7 @@ cdef class SChunk: if out is not None: PyObject_GetBuffer(out, &buf, PyBUF_SIMPLE) if buf.len < nbytes: + PyBuffer_Release(&buf) raise ValueError("Not enough space for writing the slice in out") rc = blosc2_schunk_get_slice_buffer(self.schunk, start, stop, buf.buf) PyBuffer_Release(&buf) @@ -5871,15 +5874,26 @@ def open(urlpath, mode, offset, **kwargs): is_ndarray = schunk_is_ndarray(schunk) - cdef b2nd_array_t *array - if is_ndarray: - _check_rc(b2nd_from_schunk(schunk, &array), - "Could not create array from schunk") - - kwargs["urlpath"] = urlpath - kwargs["contiguous"] = schunk.storage.contiguous - if mode != "w" and kwargs is not None: - check_schunk_params(schunk, kwargs) + cdef b2nd_array_t *array = NULL + cdef int rc + try: + if is_ndarray: + rc = b2nd_from_schunk(schunk, &array) + if rc < 0: + array = NULL # the schunk is still ours + _check_rc(rc, "Could not create array from schunk") + + kwargs["urlpath"] = urlpath + kwargs["contiguous"] = schunk.storage.contiguous + if mode != "w" and kwargs is not None: + check_schunk_params(schunk, kwargs) + except BaseException: + # Nothing owns the schunk (or the array made from it) yet: free it, which also closes the file + if array != NULL: + b2nd_free(array) + else: + blosc2_schunk_free(schunk) + raise cparams = kwargs.get("cparams") dparams = kwargs.get("dparams", blosc2.DParams()) @@ -6015,6 +6029,7 @@ def schunk_from_cframe(cframe, copy=False): PyObject_GetBuffer(cframe, &buf, PyBUF_SIMPLE) cdef blosc2_schunk *schunk_ = blosc2_schunk_from_buffer(buf.buf, buf.len, copy) if schunk_ == NULL: + PyBuffer_Release(&buf) raise RuntimeError("Could not get the schunk from the cframe") schunk = blosc2.SChunk(_schunk=PyCapsule_New(schunk_, "blosc2_schunk*", NULL)) PyBuffer_Release(&buf) @@ -6321,6 +6336,23 @@ cdef enum: cdef object _np_integer = np.integer +cdef list _drop_masked(seq, mask): + """The items of `seq` where `mask` is False.""" + return [x for x, m in zip(seq, mask) if not m] + + +cdef _full_rank_chunkshape(chunkshape, mask): + """Chunks or blocks given for the result of a slice, which has no dims for those in `mask`, with + 1 inserted in those dims. Other lengths go to compute_chunks_blocks() unchanged.""" + if chunkshape is None: + return None + chunkshape = list(chunkshape) if isinstance(chunkshape, (tuple, list)) else [chunkshape] + if len(chunkshape) != len(mask) - sum(mask): + return chunkshape + kept = iter(chunkshape) + return [1 if m else next(kept) for m in mask] + + cdef class NDArray: cdef b2nd_array_t* array cdef readonly object _handle_lock @@ -6873,25 +6905,28 @@ cdef class NDArray: cdef int64_t buffersize_ = self.array.sc.typesize cdef int64_t[B2ND_MAX_DIM] sel_size - key_ = malloc(len(key) * sizeof(int64_t *)) - - for i in range(self.array.ndim): - buffershape_[i] = len(key[i]) - buffersize_ *= buffershape_[i] - sel_size[i] = len(key[i]) - key_[i] = malloc(sel_size[i] * sizeof(int64_t)) - for j in range(len(key[i])): - key_[i][j] = key[i][j] - cdef Py_buffer buf - PyObject_GetBuffer(arr, &buf, PyBUF_SIMPLE) - - _check_rc(b2nd_get_orthogonal_selection(self.array, key_, sel_size, buf.buf, - buffershape_, buffersize_), "Error while getting orthogonal selection") - PyBuffer_Release(&buf) - for i in range(len(key)): - free(key_[i]) # Free the allocated memory for each key - free(key_) + cdef int rc + # calloc: the finally frees every entry, also the ones not allocated yet + key_ = calloc(len(key), sizeof(int64_t *)) + try: + for i in range(self.array.ndim): + buffershape_[i] = len(key[i]) + buffersize_ *= buffershape_[i] + sel_size[i] = len(key[i]) + key_[i] = malloc(sel_size[i] * sizeof(int64_t)) + for j in range(len(key[i])): + key_[i][j] = key[i][j] + + PyObject_GetBuffer(arr, &buf, PyBUF_SIMPLE) + rc = b2nd_get_orthogonal_selection(self.array, key_, sel_size, buf.buf, + buffershape_, buffersize_) + PyBuffer_Release(&buf) + _check_rc(rc, "Error while getting orthogonal selection") + finally: + for i in range(len(key)): + free(key_[i]) # Free the allocated memory for each key + free(key_) return arr @_handle_locked @@ -6906,86 +6941,96 @@ cdef class NDArray: cdef int64_t buffersize_ = self.array.sc.typesize cdef int64_t[B2ND_MAX_DIM] sel_size - key_ = malloc(len(key) * sizeof(int64_t *)) - - for i in range(self.array.ndim): - buffershape_[i] = len(key[i]) - buffersize_ *= buffershape_[i] - sel_size[i] = len(key[i]) - key_[i] = malloc(sel_size[i] * sizeof(int64_t)) - for j in range(len(key[i])): - key_[i][j] = key[i][j] - cdef Py_buffer buf - PyObject_GetBuffer(arr, &buf, PyBUF_SIMPLE) - - _check_rc(b2nd_set_orthogonal_selection(self.array, key_, sel_size, buf.buf, - buffershape_, buffersize_), "Error while getting orthogonal selection") - PyBuffer_Release(&buf) - for i in range(len(key)): - free(key_[i]) # Free the allocated memory for each key - free(key_) + cdef int rc + # calloc: the finally frees every entry, also the ones not allocated yet + key_ = calloc(len(key), sizeof(int64_t *)) + try: + for i in range(self.array.ndim): + buffershape_[i] = len(key[i]) + buffersize_ *= buffershape_[i] + sel_size[i] = len(key[i]) + key_[i] = malloc(sel_size[i] * sizeof(int64_t)) + for j in range(len(key[i])): + key_[i][j] = key[i][j] + + PyObject_GetBuffer(arr, &buf, PyBUF_SIMPLE) + rc = b2nd_set_orthogonal_selection(self.array, key_, sel_size, buf.buf, + buffershape_, buffersize_) + PyBuffer_Release(&buf) + _check_rc(rc, "Error while getting orthogonal selection") + finally: + for i in range(len(key)): + free(key_[i]) # Free the allocated memory for each key + free(key_) return arr @_handle_locked def get_slice(self, key, mask, **kwargs): + """The slice from `key` (start, stop) as a new array without the dims in `mask` (those + indexed by an int, which have length 1).""" start, stop = key shape = tuple(sp - st for sp, st in zip(stop, start)) - chunks = kwargs.pop("chunks", None) - blocks = kwargs.pop("blocks", None) - if blocks and len(shape) != len(blocks): - for i in range(len(shape)): - if shape[i] == 1: - blocks.insert(i, 1) - if chunks and len(shape) != len(chunks): - for i in range(len(shape)): - if shape[i] == 1: - chunks.insert(i, 1) - chunks, blocks = blosc2.compute_chunks_blocks(shape, chunks, blocks, self.dtype) - - # shape will be overwritten by get_slice - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, - self.dtype, kwargs) - if ctx == NULL: - raise RuntimeError("Error while creating the context") - ndim = self.ndim + chunks, blocks = blosc2.compute_chunks_blocks( + shape, _full_rank_chunkshape(kwargs.pop("chunks", None), mask), + _full_rank_chunkshape(kwargs.pop("blocks", None), mask), self.dtype) + drop_dims = any(mask) + if drop_dims: + # Chunk and block 1 in the dropped dims: then the chunks of the slice are chunks of the result + chunks = [1 if m else c for c, m in zip(chunks, mask)] + blocks = [1 if m else b for b, m in zip(blocks, mask)] cdef int64_t[B2ND_MAX_DIM] start_, stop_ - for i in range(ndim): + for i in range(self.ndim): start_[i] = start[i] stop_[i] = stop[i] - cdef b2nd_array_t *array - _check_rc(b2nd_get_slice(ctx, &array, self.array, start_, stop_), - "Error while getting the slice") - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") - - cdef c_bool mask_[B2ND_MAX_DIM] - for i in range(ndim): - mask_[i] = mask[i] - _check_rc(b2nd_squeeze_index(array, &array, mask_), "Error while squeezing sliced array") - ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), - _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) - - - return ndarray + # Without dims to drop, the slice is the result. Else the slice is built in memory and its + # chunks are moved into the result (b2nd_squeeze_index() would give a view that needs the + # slice to be kept alive). + slice_kwargs = kwargs + if drop_dims: + slice_kwargs = {k: kwargs[k] for k in ("cparams", "dparams") if k in kwargs} + slice_kwargs["contiguous"] = False # not a frame: the chunks are in sc.data + # shape will be overwritten by get_slice + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, self.dtype, slice_kwargs) + if ctx.ptr == NULL: + raise RuntimeError("Error while creating the context") + cdef b2nd_array_t *sliced = NULL + _check_rc(b2nd_get_slice(ctx.ptr, &sliced, self.array, start_, stop_), "Error while getting the slice") + if not drop_dims: + return blosc2.NDArray(_schunk=PyCapsule_New(sliced.sc, "blosc2_schunk*", NULL), + _array=PyCapsule_New(sliced, "b2nd_array_t*", NULL)) + + cdef NDArray result + cdef int64_t nchunk + try: + result = empty(_drop_masked(shape, mask), _drop_masked(chunks, mask), + _drop_masked(blocks, mask), self.dtype, **kwargs) + # The result owns each chunk after a successful update (in memory: no copy, no allocation) + for nchunk in range(sliced.sc.nchunks): + _check_rc(blosc2_schunk_update_chunk(result.array.sc, nchunk, sliced.sc.data[nchunk], False), + "Error while moving the sliced array") + sliced.sc.data[nchunk] = NULL + finally: + b2nd_free(sliced) # cannot fail; must not raise here + return result @_handle_locked def set_slice(self, key, ndarray): ndim = self.ndim start, stop = key - cdef Py_buffer buf - PyObject_GetBuffer(ndarray, &buf, PyBUF_SIMPLE) - cdef int64_t[B2ND_MAX_DIM] buffershape_, start_, stop_ for i in range(ndim): start_[i] = start[i] stop_[i] = stop[i] buffershape_[i] = stop[i] - start[i] - _check_rc(b2nd_set_slice_cbuffer(buf.buf, buffershape_, buf.len, start_, stop_, self.array), - "Error while setting the slice") + cdef Py_buffer buf + PyObject_GetBuffer(ndarray, &buf, PyBUF_SIMPLE) + cdef int rc = b2nd_set_slice_cbuffer(buf.buf, buffershape_, buf.len, start_, stop_, self.array) PyBuffer_Release(&buf) + _check_rc(rc, "Error while setting the slice") return self @@ -7009,17 +7054,16 @@ cdef class NDArray: kwargs["contiguous"] = kwargs.get("contiguous", self.array.sc.storage.contiguous) chunks, blocks = blosc2.compute_chunks_blocks(self.shape, chunks, blocks, dtype, **kwargs) - cdef b2nd_context_t *ctx = create_b2nd_context(self.shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(self.shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") cdef b2nd_array_t *array - _check_rc(b2nd_copy(ctx, self.array, &array), + _check_rc(b2nd_copy(ctx.ptr, self.array, &array), "Error while copying the array") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") return ndarray @@ -7704,7 +7748,24 @@ cdef class NDArray: _check_rc(b2nd_free(self.array), "Error while freeing the array") -cdef b2nd_context_t* create_b2nd_context(shape, chunks, blocks, dtype, kwargs): +cdef class _B2ndContext: + """A b2nd context, freed with this object, and what it points to: the dparams and the metalayer + names and contents (b2nd_create_ctx() copies the storage struct and the cparams only).""" + cdef b2nd_context_t *ptr + cdef blosc2_dparams dparams + cdef blosc2_metalayer metalayers[B2ND_MAX_METALAYERS] + cdef int nmetalayers + + def __dealloc__(self): + cdef int i + if self.ptr != NULL: + b2nd_free_ctx(self.ptr) + for i in range(self.nmetalayers): + free(self.metalayers[i].name) + free(self.metalayers[i].content) + + +cdef _B2ndContext create_b2nd_context(shape, chunks, blocks, dtype, kwargs): if isinstance(dtype, list) and len(dtype) > 0 and isinstance(dtype[0], tuple): # Extract just the field names and basic dtype info fields = [] @@ -7768,11 +7829,11 @@ cdef b2nd_context_t* create_b2nd_context(shape, chunks, blocks, dtype, kwargs): raise ValueError("NDArray must already exist") # Create storage + cdef _B2ndContext ctx = _B2ndContext.__new__(_B2ndContext) cdef blosc2_storage storage - cdef blosc2_cparams *cparams = malloc(sizeof(blosc2_cparams)) - cdef blosc2_dparams *dparams = malloc(sizeof(blosc2_dparams)) - storage.cparams = cparams - storage.dparams = dparams + cdef blosc2_cparams cparams = BLOSC2_CPARAMS_DEFAULTS # create_cparams_from_kwargs() leaves some fields alone + storage.cparams = &cparams + storage.dparams = &ctx.dparams create_storage(&storage, kwargs) # Shapes @@ -7787,33 +7848,36 @@ cdef b2nd_context_t* create_b2nd_context(shape, chunks, blocks, dtype, kwargs): # Metalayers meta = kwargs.get('meta', None) - cdef blosc2_metalayer[B2ND_MAX_METALAYERS] metalayers if meta is None: - return b2nd_create_ctx(&storage, len(shape), shape_, chunkshape, blockshape, str_dtype, - B2ND_DEFAULT_DTYPE_FORMAT, NULL, 0) + ctx.ptr = b2nd_create_ctx(&storage, len(shape), shape_, chunkshape, blockshape, str_dtype, + B2ND_DEFAULT_DTYPE_FORMAT, NULL, 0) else: nmetalayers = len(meta) + if nmetalayers > B2ND_MAX_METALAYERS: + raise ValueError(f"An NDArray can have at most {B2ND_MAX_METALAYERS} metalayers, got {nmetalayers}") for i, (name, content) in enumerate(meta.items()): name2 = name.encode("utf-8") if isinstance(name, str) else name # do a copy - metalayers[i].name = strdup(name2) + ctx.metalayers[i].name = strdup(name2) + ctx.metalayers[i].content = NULL + ctx.nmetalayers = i + 1 # freed with ctx from here on content = packb(content, default=encode_tuple, strict_types=True, use_bin_type=True) - metalayers[i].content = malloc(len(content)) - memcpy(metalayers[i].content, content, len(content)) - metalayers[i].content_len = len(content) + ctx.metalayers[i].content = malloc(len(content)) + memcpy(ctx.metalayers[i].content, content, len(content)) + ctx.metalayers[i].content_len = len(content) - return b2nd_create_ctx(&storage, len(shape), shape_, chunkshape, blockshape, str_dtype, - B2ND_DEFAULT_DTYPE_FORMAT, metalayers, nmetalayers) + ctx.ptr = b2nd_create_ctx(&storage, len(shape), shape_, chunkshape, blockshape, str_dtype, + B2ND_DEFAULT_DTYPE_FORMAT, ctx.metalayers, nmetalayers) + return ctx def uninit(shape, chunks, blocks, dtype, **kwargs): - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") cdef b2nd_array_t *array - _check_rc(b2nd_uninit(ctx, &array), "Could not build uninit array") - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") + _check_rc(b2nd_uninit(ctx.ptr, &array), "Could not build uninit array") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) ndarray._schunk.mode = kwargs.get("mode", "a") @@ -7822,13 +7886,12 @@ def uninit(shape, chunks, blocks, dtype, **kwargs): def nans(shape, chunks, blocks, dtype, **kwargs): - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") cdef b2nd_array_t *array - _check_rc(b2nd_nans(ctx, &array), "Could not build nans array") - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") + _check_rc(b2nd_nans(ctx.ptr, &array), "Could not build nans array") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) ndarray._schunk.mode = kwargs.get("mode", "a") @@ -7837,13 +7900,12 @@ def nans(shape, chunks, blocks, dtype, **kwargs): def empty(shape, chunks, blocks, dtype, **kwargs): - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") cdef b2nd_array_t *array - _check_rc(b2nd_empty(ctx, &array), "Could not build empty array") - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") + _check_rc(b2nd_empty(ctx.ptr, &array), "Could not build empty array") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) ndarray._schunk.mode = kwargs.get("mode", "a") @@ -7852,23 +7914,22 @@ def empty(shape, chunks, blocks, dtype, **kwargs): def zeros(shape, chunks, blocks, dtype, **kwargs): - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") cdef b2nd_array_t *array - _check_rc(b2nd_zeros(ctx, &array), "Could not build zeros array") + _check_rc(b2nd_zeros(ctx.ptr, &array), "Could not build zeros array") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") ndarray._schunk.mode = kwargs.get("mode", "a") return ndarray def full(shape, chunks, blocks, fill_value, dtype, **kwargs): - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") dtype = np.dtype(dtype) @@ -7877,28 +7938,27 @@ def full(shape, chunks, blocks, fill_value, dtype, **kwargs): PyObject_GetBuffer(nparr, &val, PyBUF_SIMPLE) cdef b2nd_array_t *array - _check_rc(b2nd_full(ctx, &array, val.buf), "Could not create full array") + cdef int rc = b2nd_full(ctx.ptr, &array, val.buf) PyBuffer_Release(&val) + _check_rc(rc, "Could not create full array") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") ndarray._schunk.mode = kwargs.get("mode", "a") return ndarray def from_buffer(buf, shape, chunks, blocks, dtype, **kwargs): - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") cdef b2nd_array_t *array - _check_rc(b2nd_from_cbuffer(ctx, &array, buf, len(buf)), + _check_rc(b2nd_from_cbuffer(ctx.ptr, &array, buf, len(buf)), "Error while creating the NDArray") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") ndarray._schunk.mode = kwargs.get("mode", "a") return ndarray @@ -7906,25 +7966,23 @@ def from_buffer(buf, shape, chunks, blocks, dtype, **kwargs): def asarray(ndarray, chunks, blocks, **kwargs): interface = ndarray.__array_interface__ - cdef Py_buffer buf - PyObject_GetBuffer(ndarray, &buf, PyBUF_SIMPLE) - shape = interface["shape"] dtype = interface["typestr"] if dtype.startswith("|V") and "descr" in interface: # Structured dtype dtype = interface["descr"] - cdef b2nd_context_t *ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(shape, chunks, blocks, dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context") + cdef Py_buffer buf + PyObject_GetBuffer(ndarray, &buf, PyBUF_SIMPLE) cdef b2nd_array_t *array - _check_rc(b2nd_from_cbuffer(ctx, &array, buf.buf, buf.len), - "Error while creating the NDArray") + cdef int rc = b2nd_from_cbuffer(ctx.ptr, &array, buf.buf, buf.len) PyBuffer_Release(&buf) + _check_rc(rc, "Error while creating the NDArray") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") ndarray._schunk.mode = kwargs.get("mode", "a") return ndarray @@ -7941,6 +7999,7 @@ def ndarray_from_cframe(cframe, copy=False): cdef int rc rc = b2nd_from_cframe(buf.buf, buf.len, copy, &array) if rc < 0: + PyBuffer_Release(&buf) raise RuntimeError("Could not get the NDArray from the cframe") ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, "blosc2_schunk*", NULL), _array=PyCapsule_New(array, "b2nd_array_t*", NULL)) @@ -7993,16 +8052,15 @@ def concat(arr1: NDArray, arr2: NDArray, axis: int, **kwargs): Concatenate two NDArray objects along a specified axis. """ cdef c_bool copy = kwargs.pop("copy", True) - cdef b2nd_context_t *ctx = create_b2nd_context(arr1.shape, arr1.chunks, arr1.blocks, arr1.dtype, kwargs) - if ctx == NULL: + cdef _B2ndContext ctx = create_b2nd_context(arr1.shape, arr1.chunks, arr1.blocks, arr1.dtype, kwargs) + if ctx.ptr == NULL: raise RuntimeError("Error while creating the context for concatenation") cdef b2nd_array_t *array cdef int rc = 0 with HandleLocks((arr1, arr2)): - rc = b2nd_concatenate(ctx, arr1.array, arr2.array, axis, copy, &array) + rc = b2nd_concatenate(ctx.ptr, arr1.array, arr2.array, axis, copy, &array) _check_rc(rc, "Error while concatenating the arrays") # before the locks' release() runs - _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context") if copy: # We have copied the concatenated data into a new array diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index bf739cb7d..9c146df39 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -5951,7 +5951,9 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray: Notes ----- - There is a fast path for slices that are aligned with underlying chunks. + The result is a new array (not a view), with the storage given in `kwargs`. + There is a fast path for slices that are aligned with underlying chunks + (unless `chunks` or `blocks` are given). Aligned means that the slices are made entirely with complete chunks. """ if "cparams" not in kwargs: @@ -5964,28 +5966,23 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray: key, mask = process_key(key, self.shape) start, stop, step, _ = get_ndarray_start_stop(self.ndim, key, self.shape) - # Fast path for slices made with aligned chunks - if step == (1,) * self.ndim: + def kept(seq): # without the dims indexed by an int (mask), which the result drops + return [x for x, m in zip(seq, mask, strict=True) if not m] + + # Fast path for slices made with aligned chunks (unless other chunks or blocks are asked for) + if step == (1,) * self.ndim and kwargs.get("chunks") is None and kwargs.get("blocks") is None: aligned_chunks = detect_aligned_chunks(key, self.shape, self.chunks, consecutive=False) if aligned_chunks: - # print("Aligned chunks detected", aligned_chunks) - # Create a new ndarray for the key slice - new_shape = [ - sp - st for sp, st in zip([k.stop for k in key], [k.start for k in key], strict=False) - ] + # Those dims have chunk and block 1 (aligned), so the chunks of self are chunks of + # the result too: no need to decompress and compress them again newarr = blosc2.empty( - shape=new_shape, + shape=kept([k.stop - k.start for k in key]), dtype=self.dtype, - chunks=self.chunks, - blocks=self.blocks, - **kwargs, + **(kwargs | {"chunks": kept(self.chunks), "blocks": kept(self.blocks)}), ) - # Get the chunks from the original array and update the new array - # No need for chunks to decompress and compress again for order, nchunk in enumerate(aligned_chunks): - chunk = self.schunk.get_chunk(nchunk) - newarr.schunk.update_chunk(order, chunk) - return newarr.squeeze(axis=np.where(mask)[0]) # remove any dummy dims introduced + newarr.schunk.update_chunk(order, self.schunk.get_chunk(nchunk)) + return newarr key = (start, stop) ndslice = super().get_slice(key, mask, **kwargs) @@ -5994,11 +5991,7 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray: # TODO: perhaps add a step param in the get_slice method in the future? if step != (1,) * self.ndim: nparr = ndslice[...] - if len(step) == 1: - nparr = nparr[:: step[0]] - else: - slice_ = tuple(slice(None, None, st) for st in step) - nparr = nparr[slice_] + nparr = nparr[tuple(slice(None, None, st) for st in kept(step))] return asarray(nparr, **kwargs) return ndslice diff --git a/tests/ndarray/test_slice.py b/tests/ndarray/test_slice.py index 3234df718..fd76ac9f4 100644 --- a/tests/ndarray/test_slice.py +++ b/tests/ndarray/test_slice.py @@ -5,6 +5,10 @@ # SPDX-License-Identifier: BSD-3-Clause ####################################################################### +import ctypes +import gc +import sys + import numpy as np import pytest @@ -124,6 +128,16 @@ def test_aligned_chunks_consecutive_non_exact(): ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(10, 40), slice(25, 75 - 1), slice(100, 200)), np.int32), ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(20, 35), slice(50, 75), slice(100, 300 - 1)), np.int32), ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(20 + 1, 25), slice(25, 50), slice(50, 100)), np.int32), + # int indices drop their dims + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (2, slice(3, 17), slice(5, 29)), np.int32), + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (slice(1, 3), slice(3, 17), 4), np.int32), + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (2, 5, slice(5, 29)), np.int32), + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (2, 5, 7), np.int32), + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (2, slice(5, 6), slice(5, 29)), np.int32), + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (slice(2, 3), slice(3, 17), slice(5, 29)), np.int32), + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (2, slice(8, 16), slice(8, 16)), np.int32), # aligned + ((2, 2, 2), (1, 1, 1), (1, 1, 1), (1, 0, 1), np.int32), # aligned + ((7, 20, 30), (1, 8, 8), (1, 4, 4), (2, slice(3, 17, 2), slice(5, 29)), np.int32), # step ] @@ -136,6 +150,7 @@ def test_slice(shape, chunks, blocks, slices, dtype): np_slice = a[slices] assert b.shape == np_slice.shape np.testing.assert_almost_equal(b[...], np_slice) + assert b.base is None # a new array, not a view @pytest.mark.parametrize(argnames, argvalues) @@ -179,3 +194,101 @@ def test_slice_chunks_blocks(shape, chunks, blocks, chunks2, blocks2, slices, dt b = a.slice(slices, chunks=chunks2, blocks=blocks2) np_slice = a[slices] np.testing.assert_almost_equal(b[...], np_slice) + + +# --- int indices and storage arguments -------------------------------------- +# The result of slice() owns its data and has the storage asked for. + +_INT_KEY = (2, slice(3, 17), slice(5, 29)) # not chunk-aligned +_SLICES_KEY = (slice(1, 3), slice(3, 17), slice(5, 29)) # not chunk-aligned +_ALIGNED_INT_KEY = (2, slice(8, 16), slice(8, 16)) +_ALIGNED_SLICES_KEY = (slice(2, 4), slice(8, 16), slice(8, 16)) + + +def _array_3d(): + nparray = np.arange(7 * 20 * 30, dtype=np.int32).reshape(7, 20, 30) + return nparray, blosc2.asarray(nparray, chunks=(1, 8, 8), blocks=(1, 4, 4)) + + +_ALL_KEYS = [_INT_KEY, _SLICES_KEY, _ALIGNED_INT_KEY, _ALIGNED_SLICES_KEY] + + +@pytest.mark.parametrize("key", _ALL_KEYS) +def test_slice_storage(tmp_path, key): + nparray, a = _array_3d() + urlpath = str(tmp_path / "slice.b2nd") + b = a.slice(key, urlpath=urlpath, mode="w", meta={"m": [1, 2]}) + assert b.schunk.urlpath == urlpath + assert b.schunk.contiguous + for c in (b, blosc2.open(urlpath)): + assert c.shape == nparray[key].shape + assert c.schunk.meta["m"] == [1, 2] + np.testing.assert_array_equal(c[...], nparray[key]) + assert a.slice(key, contiguous=True).schunk.contiguous + + +@pytest.mark.parametrize("key", _ALL_KEYS) +def test_slice_cparams_dparams(key): + nparray, a = _array_3d() + cparams = {"codec": blosc2.Codec.LZ4, "clevel": 1, "filters": [blosc2.Filter.NOFILTER], "nthreads": 1} + b = a.slice(key, cparams=cparams, dparams={"nthreads": 1}) + assert (b.cparams.codec, b.cparams.clevel, b.cparams.filters[-1]) == ( + blosc2.Codec.LZ4, + 1, + blosc2.Filter.NOFILTER, + ) + assert b.cparams.nthreads == 1 + assert b.dparams.nthreads == 1 + np.testing.assert_array_equal(b[...], nparray[key]) + + +@pytest.mark.parametrize( + ("key", "chunks", "blocks", "chunks2", "blocks2"), + [ + (_INT_KEY, (7, 12), (7, 4), (7, 12), (7, 4)), + (_INT_KEY, [2, 7, 12], [2, 7, 4], (7, 12), (7, 4)), # with the dropped dim + ((2, slice(5, 6), slice(5, 29)), [1, 12], [1, 4], (1, 12), (1, 4)), + (_ALIGNED_INT_KEY, (4, 8), (2, 4), (4, 8), (2, 4)), + ], +) +def test_slice_int_indices_chunks_blocks(key, chunks, blocks, chunks2, blocks2): + nparray, a = _array_3d() + b = a.slice(key, chunks=chunks, blocks=blocks) + assert b.chunks == chunks2 + assert b.blocks == blocks2 + np.testing.assert_array_equal(b[...], nparray[key]) + + +class _Mallinfo2(ctypes.Structure): + # all fields: mallinfo2() returns the struct by value + _fields_ = [ + (name, ctypes.c_size_t) + for name in "arena ordblks smblks hblks hblkhd usmblks fsmblks uordblks fordblks keepcost".split() + ] + + +_mallinfo2 = getattr(ctypes.CDLL(None), "mallinfo2", None) if sys.platform.startswith("linux") else None +if _mallinfo2 is not None: # glibc >= 2.33 + _mallinfo2.restype = _Mallinfo2 + + +def _malloc_in_use(): + gc.collect() + info = _mallinfo2() + return info.uordblks + info.hblkhd + + +@pytest.mark.skipif(_mallinfo2 is None, reason="needs glibc's mallinfo2()") +@pytest.mark.parametrize("key", [_INT_KEY, _SLICES_KEY]) +def test_slice_frees_sliced_array(key): + # The sliced array (about 75 KiB here) was kept for every slice that was not chunk-aligned + _, a = _array_3d() + for _ in range(20): # warm-up + a.slice(key) + growth = [] + for _ in range(3): + before = _malloc_in_use() + for _ in range(100): + a.slice(key) + growth.append((_malloc_in_use() - before) / 100) + assert min(growth) < 256 diff --git a/tests/test_error_path_leaks.py b/tests/test_error_path_leaks.py new file mode 100644 index 000000000..274faaa7f --- /dev/null +++ b/tests/test_error_path_leaks.py @@ -0,0 +1,280 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +"""Calls that fail must not keep C memory, Python buffers or files. + +The C memory checks count the bytes that glibc's malloc has in use (mallinfo2), so they only run +on glibc. Python objects come from pymalloc arenas, which malloc does not count: after a warm-up, +repeating a call that frees what it allocates leaves the count where it was.""" + +import ctypes +import gc +import os +import sys +import weakref + +import numpy as np +import pytest + +import blosc2 + + +class _Mallinfo2(ctypes.Structure): + _fields_ = [ + (name, ctypes.c_size_t) + for name in ( + "arena", + "ordblks", + "smblks", + "hblks", + "hblkhd", + "usmblks", + "fsmblks", + "uordblks", + "fordblks", + "keepcost", + ) + ] + + +def _get_mallinfo2(): + if not sys.platform.startswith("linux"): + return None + try: + mallinfo2 = ctypes.CDLL(None).mallinfo2 # glibc >= 2.33 + except (OSError, AttributeError): + return None + mallinfo2.restype = _Mallinfo2 + return mallinfo2 + + +_mallinfo2 = _get_mallinfo2() + + +def _c_bytes_left_per_call(fn, ncalls=300, nrounds=3): + """Smallest growth of malloc's in-use bytes per call of `fn` over `nrounds` rounds.""" + for _ in range(50): # warm-up: caches and other one-time allocations + fn() + growth = [] + for _ in range(nrounds): + gc.collect() + info = _mallinfo2() + before = info.uordblks + info.hblkhd + for _ in range(ncalls): + fn() + gc.collect() + info = _mallinfo2() + growth.append((info.uordblks + info.hblkhd - before) / ncalls) + return min(growth) + + +@pytest.fixture +def missing_dir_path(tmp_path): + # A file cannot be created there: the C call that makes the array fails + return str(tmp_path / "no_such_dir" / "array.b2nd") + + +@pytest.fixture +def arr(): + return blosc2.asarray( + np.arange(2000, dtype=np.int32).reshape(100, 20), chunks=(50, 20), blocks=(10, 20) + ) + + +@pytest.mark.skipif(_mallinfo2 is None, reason="needs glibc's mallinfo2()") +@pytest.mark.parametrize( + ("call", "exc", "match", "limit"), + [ + # The b2nd context (about 1 kB) was leaked when the call raised + pytest.param( + lambda arr, path: blosc2.concat([arr, arr], urlpath=path, mode="w"), + RuntimeError, + "Error while concatenating the arrays", + 256, + id="concat", + ), + pytest.param( + lambda arr, path: blosc2.empty((100,), dtype=np.int32, urlpath=path, mode="w"), + RuntimeError, + "Could not build empty array", + 256, + id="empty", + ), + pytest.param( + lambda arr, path: blosc2.full((100,), "abc", dtype=np.int32), + ValueError, + "invalid literal", + 256, + id="full", + ), + pytest.param( + lambda arr, path: arr.copy(urlpath=path, mode="w"), + RuntimeError, + "Error while copying the array", + 256, + id="copy", + ), + # The index arrays and the output buffer (about 180 bytes) were leaked + pytest.param( + lambda arr, path: arr.oindex[[0, 1000], [0]], + RuntimeError, + "Error while getting orthogonal selection", + 64, + id="oindex", + ), + ], +) +def test_failed_call_frees_c_memory(call, exc, match, limit, arr, missing_dir_path): + def fail(): + with pytest.raises(exc, match=match): + call(arr, missing_dir_path) + + assert _c_bytes_left_per_call(fail) < limit + + +@pytest.mark.skipif(_mallinfo2 is None, reason="needs glibc's mallinfo2()") +@pytest.mark.parametrize("meta", [None, {"m": 1, "n": [1, 2, 3]}]) +def test_new_array_frees_its_parameters(meta): + # Every array creation leaked the cparams and dparams it was made with (about 200 bytes), and + # the names and contents of its metalayers + assert _c_bytes_left_per_call(lambda: blosc2.empty((100,), dtype=np.int32, meta=meta)) < 64 + + +def test_too_many_metalayers(): + # One more than fits overflowed a fixed-size table (the process aborted) + with pytest.raises(ValueError, match="at most 15 metalayers"): + blosc2.empty((10,), dtype=np.int32, meta={f"m{i}": i for i in range(16)}) + assert len(blosc2.empty((10,), dtype=np.int32, meta={f"m{i}": i for i in range(15)}).schunk.meta) == 16 + + +def _assert_released(make_input, call): + x = make_input() + ref = weakref.ref(x) + call(x) + del x + gc.collect() + assert ref() is None, "the failed call still holds a buffer of its input" + + +def test_asarray_error_releases_input(missing_dir_path): + def call(x): + with pytest.raises(RuntimeError, match="Error while creating the NDArray"): + blosc2.asarray(x, urlpath=missing_dir_path, mode="w") + + _assert_released(lambda: np.arange(1000, dtype=np.int32), call) + + +def test_setitem_error_releases_value(arr): + def call(x): + with pytest.raises(RuntimeError, match="Error while setting the slice"): + arr[0:10, 0:20] = x + + _assert_released(lambda: np.zeros(4, dtype=np.int32), call) + + +def test_oindex_setitem_error_releases_value(arr): + def call(x): + with pytest.raises(RuntimeError, match="orthogonal selection"): + arr.oindex[[0, 1000], [0, 1]] = x + + _assert_released(lambda: np.zeros((2, 2), dtype=np.int32), call) + + +@pytest.mark.parametrize("func", [blosc2.schunk_from_cframe, blosc2.ndarray_from_cframe]) +def test_from_cframe_error_releases_input(func): + ba = bytearray(2000) + with pytest.raises(RuntimeError, match="from the cframe"): + func(ba) + ba.extend(b"x") # BufferError while a buffer of it is still exported + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="c-blosc2's blosc2_create_cctx() frees its _aligned_malloc'ed context with free() on this error path", +) +def test_compress2_error_releases_input(): + ba = bytearray(1000) + with pytest.raises(RuntimeError, match="Could not create the compression context"): + blosc2.compress2(ba, filters=[10]) # no such filter + ba.extend(b"x") # BufferError while a buffer of it is still exported + + +def test_schunk_buffer_errors_release_buffer(): + # A bytearray cannot be resized (BufferError) while a buffer of it is still exported + schunk = blosc2.SChunk(chunksize=400, data=np.arange(1000, dtype=np.int32)) + dst = bytearray() + with pytest.raises(ValueError, match="The dst length must be greater than 0"): + schunk.decompress_chunk(0, dst) + dst.extend(b"x") + out = bytearray(4) + with pytest.raises(ValueError, match="Not enough space"): + schunk.get_slice(0, 100, out=out) + out.extend(b"x") + + +def _open_fds(): + for d in ("/proc/self/fd", "/dev/fd"): + if os.path.isdir(d): + return len(os.listdir(d)) + pytest.skip("cannot count open file descriptors here") + + +def _bad_b2nd_meta(urlpath): + # A frame with a b2nd metalayer that cannot be decoded + blosc2.SChunk(chunksize=400, urlpath=urlpath, mode="w", meta={"b2nd": [1, 2, 3]}) + + +def _ndarray(urlpath): + blosc2.asarray(np.arange(100), urlpath=urlpath, mode="w") + + +def _schunk(urlpath): + blosc2.SChunk(chunksize=400, urlpath=urlpath, mode="w", data=np.arange(100)) + + +@pytest.mark.skipif(blosc2.IS_WASM, reason="no file descriptors to count") +@pytest.mark.parametrize( + ("make", "kwargs", "exc", "match"), + [ + pytest.param(_bad_b2nd_meta, {}, RuntimeError, "Could not create array from schunk", id="b2nd-meta"), + pytest.param( + _bad_b2nd_meta, + {"mmap_mode": "r"}, + RuntimeError, + "Could not create array from schunk", + id="b2nd-meta-mmap", + ), + # check_schunk_params() raises after the array (or super-chunk) was made + pytest.param( + _ndarray, + {"cparams": {"typesize": 2}, "typesize": 2}, + ValueError, + "Cannot change typesize", + id="params-ndarray", + ), + pytest.param( + _schunk, + {"cparams": {"typesize": 2}, "typesize": 2}, + ValueError, + "Cannot change typesize", + id="params-schunk", + ), + ], +) +def test_open_error_closes_file(tmp_path, make, kwargs, exc, match): + # open() raised after it opened the file and kept it open + urlpath = str(tmp_path / "frame.b2frame") + make(urlpath) + with pytest.raises(exc, match=match): + blosc2.open(urlpath, **kwargs) + gc.collect() + nfds = _open_fds() + for _ in range(5): + with pytest.raises(exc, match=match): + blosc2.open(urlpath, **kwargs) + gc.collect() + assert _open_fds() <= nfds