diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index de90e254c..dd0e3af8f 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -6861,6 +6861,30 @@ cdef class NDArray: _check_rc(rc, "Error while getting the sparse selection") return arr + cdef _check_oindex_buffer(self, arr, key, str role): + """Raise ValueError if the buffer of `arr` (the `role` "destination" or "source" in the + message) cannot hold the orthogonal selection `key`. + + The C copy moves typesize bytes per item, but an array's dtype can have smaller items: + arrays that c-blosc2 2.7-2.16 wrote without a dtype store "|u1" whatever the typesize, + or the dtype string is damaged. A buffer allocated for the dtype is too small then. + """ + cdef int64_t typesize = self.array.sc.typesize + nbytes = typesize # a Python int: it cannot wrap around for huge selections + for k in key: + nbytes *= len(k) + cdef Py_buffer view + PyObject_GetBuffer(arr, &view, PyBUF_SIMPLE) + cdef Py_ssize_t buflen = view.len + PyBuffer_Release(&view) + if buflen < nbytes: + msg = (f"{role} buffer is smaller than the requested orthogonal selection " + f"({buflen} < {nbytes} bytes)") + itemsize = getattr(arr, "itemsize", typesize) + if itemsize != typesize: + msg += f": its itemsize is {itemsize}, but the array's typesize is {typesize}" + raise ValueError(msg) + @_handle_locked def get_oindex_numpy(self, arr, key): """ @@ -6868,6 +6892,7 @@ cdef class NDArray: """ if len(key) != self.array.ndim: raise ValueError(f"Key must have {self.array.ndim} dimensions, got {len(key)}.") + self._check_oindex_buffer(arr, key, "destination") cdef int64_t[B2ND_MAX_DIM] buffershape_ cdef int64_t** key_ cdef int64_t buffersize_ = self.array.sc.typesize @@ -6901,6 +6926,7 @@ cdef class NDArray: """ if len(key) != self.array.ndim: raise ValueError(f"Key must have {self.array.ndim} dimensions, got {len(key)}.") + self._check_oindex_buffer(arr, key, "source") cdef int64_t[B2ND_MAX_DIM] buffershape_ cdef int64_t** key_ cdef int64_t buffersize_ = self.array.sc.typesize diff --git a/tests/data/legacy_u1_dtype/example_day_month_temp.b2nd b/tests/data/legacy_u1_dtype/example_day_month_temp.b2nd new file mode 100644 index 000000000..3f2a09c24 Binary files /dev/null and b/tests/data/legacy_u1_dtype/example_day_month_temp.b2nd differ diff --git a/tests/ndarray/test_oindex_typesize.py b/tests/ndarray/test_oindex_typesize.py new file mode 100644 index 000000000..010c37816 --- /dev/null +++ b/tests/ndarray/test_oindex_typesize.py @@ -0,0 +1,120 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +# Orthogonal selections copy typesize bytes per item in C into a NumPy buffer allocated for +# the array's dtype. Arrays whose dtype has smaller items than the typesize must raise +# instead of overflowing that buffer. + +import shutil +from pathlib import Path + +import numpy as np +import pytest + +import blosc2 + +# Written by c-blosc2's examples/b2nd/example_frame_generator.c (c-blosc2 2.7-2.16 store +# "|u1" in the b2nd metalayer of arrays created without a dtype): float32 items, typesize 4 +LEGACY_U1_FILE = Path(__file__).parent.parent / "data" / "legacy_u1_dtype" / "example_day_month_temp.b2nd" + + +def with_dtype_string(arr, old, new): + """`arr` with dtype string `old` replaced by `new` (same length) in its b2nd metalayer""" + assert len(old) == len(new) + cframe = arr.to_cframe() + # The metalayers come before the chunks, so the first match is the dtype + return blosc2.ndarray_from_cframe(cframe.replace(old.encode(), new.encode(), 1), copy=True) + + +def assert_oindex_raises(arr): + """The paths to the C orthogonal selection raise ValueError and leave the array unchanged""" + before = arr.tobytes() + full = [list(range(n)) for n in arr.shape] + match = f"typesize is {arr.schunk.typesize}" + with pytest.raises(ValueError, match=match): + arr.get_oselection_numpy(full) + with pytest.raises(ValueError, match=match): + arr.get_oselection_numpy([[0]] * arr.ndim) # one item: a few bytes too many + with pytest.raises(ValueError, match=match): + arr.oindex[tuple(full)] = np.zeros(arr.shape, dtype=arr.dtype) + if arr.ndim > 1: + with pytest.raises(ValueError, match=match): + arr[[0, 1]] # through take(axis=0) + assert arr.tobytes() == before + + +def test_legacy_u1_file(tmp_path): + path = tmp_path / LEGACY_U1_FILE.name + shutil.copy(LEGACY_U1_FILE, path) + arr = blosc2.open(path, mode="a") + assert arr.dtype == np.uint8 + assert arr.schunk.typesize == 4 + assert arr.shape == (400, 3) + assert_oindex_raises(arr) + assert path.read_bytes() == LEGACY_U1_FILE.read_bytes() + + +@pytest.mark.parametrize( + ("dtype", "shape", "chunks", "blocks", "stored"), + [ + # As c-blosc2's plugins/test_data/example_item_prices.b2nd + (np.float32, (12, 25, 250), (8, 10, 50), (4, 5, 10), "|u1"), + (np.float64, (100,), (30,), (7,), "|u1"), + # A damaged dtype string + (np.float64, (6, 10), (4, 4), (2, 2), "i4", "S7", "