CSR Data transfer from Daphne to Python in Scipy csr form

Author: LuloDuarte

src/api/python/daphne/operator/operation_node.py

@@ -29,6 +29,10 @@
 
 import numpy as np
 import pandas as pd
+try:
+    import scipy.sparse as sp
+except ImportError as e:
+    sp = e
 try:
     import torch as torch
 except ImportError as e:
@@ -188,11 +192,37 @@ def compute(self, type="shared memory", verbose=False, asTensorFlow=False, asPyT
                 self.clear_tmp()
             elif self._output_type == OutputType.MATRIX and type=="shared memory":
                 daphneLibResult = DaphneLib.getResult()
-                result = np.ctypeslib.as_array(
-                    ctypes.cast(daphneLibResult.address, ctypes.POINTER(self.getType(daphneLibResult.vtc))),
-                    shape=[daphneLibResult.rows, daphneLibResult.cols]
-                )
-                self.clear_tmp()
+                if not daphneLibResult.isSparse:
+                    # Dense Matrix
+                    daphneLibResult = DaphneLib.getResult()
+                    result = np.ctypeslib.as_array(
+                        ctypes.cast(daphneLibResult.address, ctypes.POINTER(self.getType(daphneLibResult.vtc))),
+                        shape=[daphneLibResult.rows, daphneLibResult.cols]
+                    )
+                else:
+                    # CSR Matrix
+                    VT = self.getType(daphneLibResult.vtc)
+                    
+                    # wrap each pointer into a numpy array
+                    indptr = np.ctypeslib.as_array(
+                        ctypes.cast(daphneLibResult.row_related, ctypes.POINTER(ctypes.c_size_t)),
+                        shape=(daphneLibResult.rows + 1,)
+                    )
+                    nnz = int(indptr[-1] - indptr[0])
+
+                    
+                    data = np.ctypeslib.as_array(
+                        ctypes.cast(daphneLibResult.data, ctypes.POINTER(VT)),
+                        shape=(nnz,)
+                    )
+
+                    indices = np.ctypeslib.as_array(
+                        ctypes.cast(daphneLibResult.col_related, ctypes.POINTER(ctypes.c_size_t)),
+                        shape=(nnz,)
+                    )
+                    # build scipy CSR
+                    result = sp.csr_matrix((data, indices, indptr), shape=(daphneLibResult.rows, daphneLibResult.cols))
+                self.clear_tmp() 
             elif self._output_type == OutputType.MATRIX and type=="files":
                 arr = np.genfromtxt(result, delimiter=',')
                 self.clear_tmp()

src/api/python/daphne/utils/daphnelib.py

@@ -20,18 +20,28 @@
 # Python representation of the struct DaphneLibResult.
 class DaphneLibResult(ctypes.Structure):
     _fields_ = [
-        # For matrices.
-        ("address", ctypes.c_void_p),
+        # For matrices
         ("rows", ctypes.c_int64),
         ("cols", ctypes.c_int64),
         ("vtc", ctypes.c_int64),
+        ("isSparse", ctypes.c_bool),
+
+        # For dense matrices
+        ("address", ctypes.c_void_p),
+
+        # For sparse matrices
+        ("data", ctypes.c_void_p),
+        ("row_related",ctypes.c_void_p),
+        ("col_related",ctypes.c_void_p),
+
         # For frames.
         ("vtcs", ctypes.POINTER(ctypes.c_int64)),
         ("labels", ctypes.POINTER(ctypes.c_char_p)),
         ("columns", ctypes.POINTER(ctypes.c_void_p)),
+
         # To pass error messages to Python code.
         ("error_message", ctypes.c_char_p)
     ]
 
 DaphneLib = ctypes.CDLL(os.path.join(PROTOTYPE_PATH, DAPHNELIB_FILENAME))
-DaphneLib.getResult.restype = DaphneLibResult
+DaphneLib.getResult.restype = DaphneLibResult

src/runtime/local/kernels/SaveDaphneLibResult.h

@@ -18,6 +18,7 @@
 #define SRC_RUNTIME_LOCAL_KERNELS_SAVEDAPHNELIBRESULT_H
 
 #include <runtime/local/context/DaphneContext.h>
+#include <runtime/local/datastructures/CSRMatrix.h>
 #include <runtime/local/datastructures/DenseMatrix.h>
 #include <runtime/local/datastructures/Frame.h>
 
@@ -64,6 +65,72 @@ template <typename VT> struct SaveDaphneLibResult<DenseMatrix<VT>> {
     }
 };
 
+// ----------------------------------------------------------------------------
+// CSRMatrix
+// ----------------------------------------------------------------------------
+
+template <typename VT> struct SaveDaphneLibResult<CSRMatrix<VT>> {
+    static void apply(const CSRMatrix<VT> *arg, DCTX(ctx)) {
+
+        const_cast<CSRMatrix<VT> *>(arg)->increaseRefCounter();
+
+        DaphneLibResult *daphneLibRes = ctx->getUserConfig().result_struct;
+        if (!daphneLibRes)
+            throw std::runtime_error("saveDaphneLibRes(): daphneLibRes is nullptr");
+
+        // Result is a sparse matrix
+        daphneLibRes->isSparse = true;
+
+        const size_t rows = arg->getNumRows();
+        const size_t cols = arg->getNumCols();
+        daphneLibRes->rows = rows;
+        daphneLibRes->cols = cols;
+        daphneLibRes->vtc = (int64_t)ValueTypeUtils::codeFor<VT>;
+
+        // original raw pointers
+        auto origVals = arg->getValuesSharedPtr().get();
+        auto origCols = arg->getColIdxsSharedPtr().get();
+        auto origOffsets = arg->getRowOffsetsSharedPtr().get();
+
+        // first pass: count *actual* non-zeros
+        size_t actualNNZ = 0;
+        for (size_t r = 0; r < rows; r++) {
+            auto start = origOffsets[r];
+            auto end = origOffsets[r + 1];
+            for (size_t i = start; i < end; i++)
+                if (origVals[i] != VT(0))
+                    actualNNZ++;
+        }
+
+        // allocate new tight buffers
+        VT *cleanVals = new VT[actualNNZ];
+        size_t *cleanCols = new size_t[actualNNZ];
+        size_t *cleanOffs = new size_t[rows + 1];
+
+        // second pass: copy only non-zeros
+        size_t p = 0;
+        cleanOffs[0] = 0;
+        for (size_t r = 0; r < rows; r++) {
+            auto start = origOffsets[r];
+            auto end = origOffsets[r + 1];
+            for (size_t i = start; i < end; i++) {
+                VT v = origVals[i];
+                if (v != VT(0)) {
+                    cleanVals[p] = v;
+                    cleanCols[p] = origCols[i];
+                    p++;
+                }
+            }
+            cleanOffs[r + 1] = p;
+        }
+
+        // hand them off in the result struct
+        daphneLibRes->data = static_cast<void *>(cleanVals);
+        daphneLibRes->col_related = static_cast<void *>(cleanCols);
+        daphneLibRes->row_related = static_cast<void *>(cleanOffs);
+    }
+};
+
 // ----------------------------------------------------------------------------
 // Frame
 // ----------------------------------------------------------------------------