Merge pull request #460 from barakugav/ffi-ptr

Add `as_ffi_ptr` to NDArray

Author: FrancescAlted

src/blosc2/__init__.py

@@ -328,6 +328,7 @@ class Tuner(Enum):
     arctan,
     arctan2,
     arctanh,
+    array_from_ffi_ptr,
     conj,
     contains,
     cos,

src/blosc2/blosc2_ext.pyx

@@ -2615,6 +2615,9 @@ cdef class NDArray:
         if self.array.shape[0] == 1 and self.ndim == 1:
             self.array.ndim = 0
 
+    def as_ffi_ptr(self):
+        return PyCapsule_New(self.array, <char *> "b2nd_array_t*", NULL)
+
     cdef udf_udata *_fill_udf_udata(self, func_id, inputs_id):
         cdef udf_udata *udata = <udf_udata *> malloc(sizeof(udf_udata))
         udata.py_func = <char *> malloc(strlen(func_id) + 1)
@@ -2901,6 +2904,10 @@ def asarray(ndarray, chunks, blocks, **kwargs):
 
     return ndarray
 
+def array_from_ffi_ptr(array_ptr):
+    array = <b2nd_array_t *> PyCapsule_GetPointer(array_ptr, <char *> "b2nd_array_t*")
+    return blosc2.NDArray(_schunk=PyCapsule_New(array.sc, <char *> "blosc2_schunk*", NULL),
+                             _array=array_ptr)
 
 def ndarray_from_cframe(cframe, copy=False):
     cdef Py_buffer buf

src/blosc2/ndarray.py

@@ -2169,6 +2169,24 @@ def sort(self, order: str | list[str] | None = None, **kwargs: Any) -> NDArray:
         """
         return sort(self, order, **kwargs)
 
+    def as_ffi_ptr(self):
+        """Returns the pointer to the raw FFI blosc2::b2nd_array_t object.
+
+        This function is useful for passing the array to C functions.
+        """
+        return super().as_ffi_ptr()
+
+
+def array_from_ffi_ptr(array_ptr) -> NDArray:
+    """
+    Create an NDArray from a raw FFI pointer.
+
+    This function is useful for passing arrays across FFI boundaries.
+    This function move the ownership of the underlying `b2nd_array_t*` object to the new NDArray, and it will be freed
+    when the object is destroyed.
+    """
+    return blosc2_ext.array_from_ffi_ptr(array_ptr)
+
 
 def sin(ndarr: NDArray | NDField | blosc2.C2Array | blosc2.LazyExpr, /) -> blosc2.LazyExpr:
     """