WIP. Add get_slice_nchunks

blosc2/__init__.py

@@ -106,6 +106,7 @@ class Tuner(Enum):
     get_cbuffer_sizes,
     get_clib,
     get_compressor,
+    get_slice_nchunks,
     load_array,
     load_tensor,
     ndarray_from_cframe,
@@ -227,4 +228,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,37 @@ def ndarray_from_cframe(cframe, copy=False):
     if not copy:
         ndarray._schunk._avoid_cframe_free(True)
     return ndarray
+
+
+def b2nd_get_slice_nchunks(array: NDArray, key):
+    start, stop = key
+    ndim = array.ndim
+    cdef int64_t[B2ND_MAX_DIM] start_, stop_
+    for i in range(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)
+    # Try using py buffer (?)
+    _check_rc(rc, "Error while getting slice nchunks")
+    nchunks_array = np.zeros([rc], dtype=np.int64)
+    for i in range(rc):
+        nchunks_array[i] = chunks_idx[i]
+    free(chunks_idx)
+    return nchunks_array
+
+
+def schunk_get_slice_nchunks(schunk: SChunk, key):
+    start, stop = key
+    cdef int64_t start_ = start
+    cdef int64_t stop_ = stop
+
+    cdef int64_t *chunks_idx
+    rc = blosc2_get_slice_nchunks(schunk.schunk, &start_, &stop_, &chunks_idx)
+    # Try using py buffer (?)
+    _check_rc(rc, "Error while getting slice nchunks")
+    nchunks_array = np.zeros([rc], dtype=np.int64)
+    for i in range(rc):
+        nchunks_array[i] = chunks_idx[i]
+    free(chunks_idx)
+    return nchunks_array

blosc2/core.py

@@ -16,6 +16,7 @@
 
 import blosc2
 from blosc2 import blosc2_ext
+from ndarray import process_key, get_ndarray_start_stop
 
 
 def _check_typesize(typesize):
@@ -1561,3 +1562,14 @@ def backward(input, output, meta, schunk):
     if id in blosc2.ufilters_registry:
         raise ValueError("Id already in use")
     blosc2_ext.register_filter(id, forward, backward, name)
+
+
+def get_slice_nchunks(schunk, key):
+    if isinstance(schunk, blosc2.NDArray):
+        array = schunk
+        key, mask = process_key(key, array.shape)
+        start, stop, _ = get_ndarray_start_stop(array.ndim, key, array.shape)
+        key = (start, stop)
+        return blosc2_ext.b2nd_get_slice_nchunks(array, key)
+    else:
+        return blosc2_ext.schunk_get_slice_nchunks(schunk, key)

examples/ndarray/getitem.py

@@ -25,3 +25,4 @@
 # Set a slice
 a[slices] = np.ones_like(buffer) - buffer
 print(a[...])
+print(blosc2.get_slice_nchunks(a, slices))