Use numpy

python/cugraph/cugraph/internals/__init__.py

@@ -13,3 +13,4 @@
 # limitations under the License.
 
 from cugraph.internals.internals import GraphBasedDimRedCallback
+from cugraph.internals.holder import ArrayInterfaceHolder

python/cugraph/cugraph/internals/callbacks_implems.hpp

@@ -50,11 +50,23 @@ class DefaultGraphBasedDimRedCallback : public GraphBasedDimRedCallback {
     }
   }
 
+  PyObject* get_numpy_array(void* array) {
+    PyObject* pycl = (PyObject*)this->pyCallbackClass;
+
+    if (isFloat) {
+      return PyObject_CallMethod(pycl, "get_numpy_array", "(l(ll)s)", array, 10,
+                                 1, "float32");
+    } else {
+      return PyObject_CallMethod(pycl, "get_numpy_array", "(l(ll)s)", array, 10,
+                                 1, "float64");
+    }
+  }
+
   void push(void* array) override {
-    // PyObject* numba_array = get_numba_array(array);
+    PyObject* numpy_array = get_numpy_array(array);
     PyObject* res =
-        PyObject_CallMethod(this->pyCallbackClass, "push", "(O)", array);
-    Py_DECREF(numba_array);
+        PyObject_CallMethod(this->pyCallbackClass, "push", "(O)", numpy_array);
+    Py_DECREF(numpy_array);
     Py_DECREF(res);
   }
 

python/cugraph/cugraph/internals/holder.py

@@ -0,0 +1,2 @@
+class ArrayInterfaceHolder:
+    pass

python/cugraph/cugraph/internals/internals.pyx

@@ -21,7 +21,7 @@
 from libc.stdint cimport uintptr_t
 from numba.cuda.api import from_cuda_array_interface
 import numpy as np
-
+from cugraph.internals.holder import ArrayInterfaceHolder
 
 cdef extern from "Python.h":
     cdef cppclass PyObject
@@ -41,6 +41,27 @@ cdef extern from "callbacks_implems.hpp" namespace "cugraph::internals":
 
 
 cdef class PyCallback:
+    def get_numpy_array(self, positions, shape, typestr):
+
+        sizeofType = 4 if typestr == "float32" else 8
+        desc = {
+            'shape': shape,
+            'strides': (sizeofType, shape[0]*sizeofType),
+            'typestr': typestr,
+            'data': (positions, False),
+            'order': 'C',
+            'version': 3
+        }
+        print(desc)
+        tmp = np.array([0])
+        print(tmp.__array_interface__)
+        # tmp.___array_interface__ = desc
+        # print(positions)
+        holder = ArrayInterfaceHolder()
+        holder.__array_interface__ = desc
+        view = np.array(holder, copy=False)
+        return view
+
     def get_numba_array(self, positions, shape, typestr):
 
         sizeofType = 4 if typestr == "float32" else 8

test.py

@@ -1,6 +1,6 @@
 import cugraph
 from numba import cuda
-
+import numpy as np
 
 # Define the parameters
 ITERATIONS = 500
@@ -16,11 +16,37 @@
 dummy = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
 
 
+class ArrayInterfaceHolder:
+    pass
+
+
+def get_numpy_array(positions, shape, typestr):
+    sizeofType = 4 if typestr == "float32" else 8
+    desc = {
+        "shape": shape,
+        "strides": (sizeofType, shape[0] * sizeofType),
+        "typestr": typestr,
+        "data": [positions],
+        "order": "C",
+        "version": 1,
+    }
+
+    holder = ArrayInterfaceHolder()
+    holder.__array_interface__ = desc
+    view = np.array(holder, copy=False)
+    return view
+
+
+# numpy_array = get_numpy_array(172383, (3, 1), "float32")
+# print(numpy_array)
+
+
 class CustomCallback(GraphBasedDimRedCallback):
     def push(self, array):
-        print(array.copy_to_host())
+        print(array)
         for i in range(10):
             array[i] = dummy[i]
+        print(array)
 
     def on_preprocess_end(self, positions):
         print(positions.copy_to_host())