Add get_slice_nchunks

Author: martaiborra

blosc2/__init__.py

@@ -129,7 +129,7 @@ class Tuner(Enum):
     unpack_array2,
     unpack_tensor,
 )
-from .ndarray import NDArray, asarray, copy, empty, frombuffer, full, uninit, zeros
+from .ndarray import NDArray, asarray, copy, empty, frombuffer, full, get_slice_nchunks, uninit, zeros
 from .schunk import SChunk, open
 from .version import __version__
 
@@ -227,4 +227,5 @@ class Tuner(Enum):
     "nthreads",
     "compute_chunks_blocks",
     "cpu_info",
+    "get_slice_nchunks",
 ]

blosc2/blosc2_ext.pyx

@@ -416,6 +416,8 @@ cdef extern from "blosc2.h":
 
     int blosc2_register_filter(blosc2_filter *filter)
 
+    int blosc2_get_slice_nchunks(blosc2_schunk * schunk, int64_t *start, int64_t *stop, int64_t ** chunks_idx)
+
 
 cdef extern from "b2nd.h":
     ctypedef enum:
@@ -2362,3 +2364,38 @@ def ndarray_from_cframe(cframe, copy=False):
     if not copy:
         ndarray._schunk._avoid_cframe_free(True)
     return ndarray
+
+
+def array_get_slice_nchunks(array: NDArray, key):
+    start, stop = key
+    cdef int64_t[B2ND_MAX_DIM] start_, stop_
+    for i in range(array.ndim):
+        start_[i] = start[i]
+        stop_[i] = stop[i]
+    cdef int64_t *chunks_idx
+    rc = blosc2_get_slice_nchunks(array.array.sc, start_, stop_, &chunks_idx)
+    _check_rc(rc, "Error while getting the chunk indexes")
+    res = np.empty(rc, dtype=np.int64)
+    for i in range(rc):
+        res[i] = chunks_idx[i]
+    free(chunks_idx)
+    return res
+
+
+def schunk_get_slice_nchunks(schunk: SChunk, key):
+    start, stop = key
+    nitems = schunk.nbytes // schunk.typesize
+    start, stop, _ = slice(start, stop, 1).indices(nitems)
+
+    cdef int64_t start_, stop_
+    start_ = start
+    stop_ = stop
+    cdef int64_t *chunks_idx
+    rc = blosc2_get_slice_nchunks(schunk.schunk, &start_, &stop_, &chunks_idx)
+    _check_rc(rc, "Error while getting the chunk indexes")
+
+    res = np.empty(rc, dtype=np.int64)
+    for i in range(rc):
+        res[i] = chunks_idx[i]
+    free(chunks_idx)
+    return res

blosc2/ndarray.py

@@ -667,3 +667,27 @@ def _check_ndarray_kwargs(**kwargs):
         raise ValueError("You cannot pass chunks in cparams, use `chunks` argument instead")
     if "cparams" in kwargs and "blocks" in kwargs["cparams"]:
         raise ValueError("You cannot pass chunks in cparams, use `blocks` argument instead")
+
+
+def get_slice_nchunks(schunk, key):
+    """
+    Get the unidimensional chunk indexes needed to get a
+    slice of a SChunk or a NDArray.
+
+    Parameters
+    ----------
+    schunk:
+    key
+
+    Returns
+    -------
+
+    """
+    if isinstance(schunk, NDArray):
+        array = schunk
+        key, _ = process_key(key, array.shape)
+        start, stop, _ = get_ndarray_start_stop(array.ndim, key, array.shape)
+        key = (start, stop)
+        return blosc2_ext.array_get_slice_nchunks(array, key)
+    else:
+        return blosc2_ext.schunk_get_slice_nchunks(schunk, key)

tests/ndarray/test_get_slice_nchunks.py

@@ -0,0 +1,31 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <[email protected]>
+# All rights reserved.
+#
+# This source code is licensed under a BSD-style license (found in the
+# LICENSE file in the root directory of this source tree)
+#######################################################################
+
+import numpy as np
+import pytest
+
+import blosc2
+
+argnames = "shape, chunks, blocks, slices, dtype"
+argvalues = [
+    ([456], [258], [73], slice(0, 1), np.int32),
+    ([456, 200], [258, 100], [73, 25], (slice(0), slice(0)), np.int64),
+    ([77, 134, 13], [31, 13, 5], [7, 8, 3], (slice(3, 7), slice(50, 100), 7), np.float64),
+    ([12, 13, 14, 15, 16], [5, 5, 5, 5, 5], [2, 2, 2, 2, 2], (slice(1, 3), ..., slice(3, 6)), np.float32),
+]
+
+
+@pytest.mark.parametrize(argnames, argvalues)
+def test_getitem(shape, chunks, blocks, slices, dtype):
+    a = blosc2.zeros(shape, dtype, chunks=chunks, blocks=blocks)
+    schunk = a.schunk
+    for i in range(schunk.nchunks):
+        chunk = np.full(schunk.chunksize // schunk.typesize, i, dtype=dtype)
+        schunk.update_data(i, chunk, True)
+
+    np.array_equal(np.unique(a[slices]), blosc2.get_slice_nchunks(a, slices))

tests/test_schunk_get_slice_nchunks.py

@@ -0,0 +1,46 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <[email protected]>
+# All rights reserved.
+#
+# This source code is licensed under a BSD-style license (found in the
+# LICENSE file in the root directory of this source tree)
+#######################################################################
+
+import numpy as np
+import pytest
+
+import blosc2
+
+
+@pytest.mark.parametrize(
+    "contiguous, urlpath, cparams, nchunks, start, stop",
+    [
+        (True, None, {"typesize": 4}, 10, 0, 100),
+        (True, "b2frame", {"typesize": 4}, 1, 7, 23),
+        (
+            False,
+            None,
+            {"splitmode": blosc2.SplitMode.ALWAYS_SPLIT, "nthreads": 5, "typesize": 4},
+            5,
+            21,
+            200 * 2 * 100,
+        ),
+        (False, "b2frame", {"codec": blosc2.Codec.LZ4HC, "typesize": 4}, 7, None, None),
+        (True, None, {"blocksize": 200 * 100, "typesize": 4}, 5, -2456, -234),
+        (True, "b2frame", {"blocksize": 200 * 100, "typesize": 4}, 4, 2456, -234),
+        (False, None, {"blocksize": 100 * 100, "typesize": 4}, 2, -200 * 100 + 234, 40000),
+    ],
+)
+def test_schunk_get_slice(contiguous, urlpath, cparams, nchunks, start, stop):
+    storage = {"contiguous": contiguous, "urlpath": urlpath, "cparams": cparams}
+    schunk = blosc2.SChunk(chunksize=200 * 100 * 4, mode="w", **storage)
+    for i in range(nchunks):
+        chunk = np.full(schunk.chunksize // schunk.typesize, i, dtype=np.int32)
+        schunk.append_data(chunk)
+
+    aux = np.empty(200 * 100 * nchunks, dtype=np.int32)
+    res = aux[start:stop]
+    sl = schunk.get_slice(start, stop, res)
+    np.array_equal(np.unique(res), blosc2.get_slice_nchunks(schunk, (start, stop)))
+
+    blosc2.remove_urlpath(urlpath)