From f133d846fe1f0dd44a7e73b159fd14a2c76b95cd Mon Sep 17 00:00:00 2001
From: Hugo Linsenmaier <hlinsenmaier@gmail.com>
Date: Tue, 9 Jan 2024 19:54:47 -0800
Subject: [PATCH] Use numpy

---
 python/cugraph/cugraph/internals/__init__.py  |  1 +
 .../cugraph/internals/callbacks_implems.hpp   | 18 +++++++++--
 python/cugraph/cugraph/internals/holder.py    |  2 ++
 .../cugraph/cugraph/internals/internals.pyx   | 23 +++++++++++++-
 test.py                                       | 30 +++++++++++++++++--
 5 files changed, 68 insertions(+), 6 deletions(-)
 create mode 100644 python/cugraph/cugraph/internals/holder.py

diff --git a/python/cugraph/cugraph/internals/__init__.py b/python/cugraph/cugraph/internals/__init__.py
index 232afde7165..cb1ccd1cff9 100644
--- a/python/cugraph/cugraph/internals/__init__.py
+++ b/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
diff --git a/python/cugraph/cugraph/internals/callbacks_implems.hpp b/python/cugraph/cugraph/internals/callbacks_implems.hpp
index e7cf44dfbcc..bc5d99f9127 100644
--- a/python/cugraph/cugraph/internals/callbacks_implems.hpp
+++ b/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);
   }
 
diff --git a/python/cugraph/cugraph/internals/holder.py b/python/cugraph/cugraph/internals/holder.py
new file mode 100644
index 00000000000..780e0e5be36
--- /dev/null
+++ b/python/cugraph/cugraph/internals/holder.py
@@ -0,0 +1,2 @@
+class ArrayInterfaceHolder:
+    pass
diff --git a/python/cugraph/cugraph/internals/internals.pyx b/python/cugraph/cugraph/internals/internals.pyx
index 388fb628038..564e49c9711 100644
--- a/python/cugraph/cugraph/internals/internals.pyx
+++ b/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
diff --git a/test.py b/test.py
index 4e136e83f8b..6e60ea2a3a9 100644
--- a/test.py
+++ b/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())