[FEA] Support Edge ID Lookup in PyLibcuGraph (#4687)

Support Edge ID lookup in `pylibcugraph`.  Also fixes some bugs in the C API (i.e. lookup table not being cleaned up correctly, container being incorrectly dereferenced as graph).

Verified in rapidsai/cugraph-gnn#50

Authors:
  - Alex Barghi (https://github.com/alexbarghi-nv)

Approvers:
  - Chuck Hastings (https://github.com/ChuckHastings)
  - Seunghwa Kang (https://github.com/seunghwak)
  - Rick Ratzel (https://github.com/rlratzel)

URL: #4687

cpp/include/cugraph_c/lookup_src_dst.h

@@ -136,6 +136,14 @@ cugraph_type_erased_device_array_view_t* cugraph_lookup_result_get_dsts(
  */
 void cugraph_lookup_result_free(cugraph_lookup_result_t* result);
 
+/**
+ * @ingroup samplingC
+ * @brief     Free a sampling lookup map
+ *
+ * @param [in] container     The sampling lookup map (a.k.a. container).
+ */
+void cugraph_lookup_container_free(cugraph_lookup_container_t* container);
+
 #ifdef __cplusplus
 }
 #endif

cpp/src/c_api/lookup_src_dst.cpp

@@ -307,23 +307,26 @@ extern "C" cugraph_error_code_t cugraph_lookup_endpoints_from_edge_ids_and_types
 {
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
-      reinterpret_cast<cugraph::c_api::cugraph_graph_t const*>(lookup_container)->vertex_type_,
+      reinterpret_cast<cugraph::c_api::cugraph_lookup_container_t const*>(lookup_container)
+        ->vertex_type_,
     CUGRAPH_INVALID_INPUT,
     "vertex type of graph and lookup_container must match",
     *error);
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_type_ ==
-      reinterpret_cast<cugraph::c_api::cugraph_graph_t const*>(lookup_container)->edge_type_,
+      reinterpret_cast<cugraph::c_api::cugraph_lookup_container_t const*>(lookup_container)
+        ->edge_type_,
     CUGRAPH_INVALID_INPUT,
     "edge type of graph and lookup_container must match",
     *error);
 
-  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_type_id_type_ ==
-                 reinterpret_cast<cugraph::c_api::cugraph_graph_t const*>(lookup_container)
-                   ->edge_type_id_type_,
-               CUGRAPH_INVALID_INPUT,
-               "edge type id type of graph and lookup_container must match",
-               *error);
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_type_id_type_ ==
+      reinterpret_cast<cugraph::c_api::cugraph_lookup_container_t const*>(lookup_container)
+        ->edge_type_id_type_,
+    CUGRAPH_INVALID_INPUT,
+    "edge type id type of graph and lookup_container must match",
+    *error);
 
   lookup_using_edge_ids_and_types_functor functor(
     handle, graph, lookup_container, edge_ids_to_lookup, edge_types_to_lookup);
@@ -341,23 +344,26 @@ extern "C" cugraph_error_code_t cugraph_lookup_endpoints_from_edge_ids_and_singl
 {
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
-      reinterpret_cast<cugraph::c_api::cugraph_graph_t const*>(lookup_container)->vertex_type_,
+      reinterpret_cast<cugraph::c_api::cugraph_lookup_container_t const*>(lookup_container)
+        ->vertex_type_,
     CUGRAPH_INVALID_INPUT,
     "vertex type of graph and lookup_container must match",
     *error);
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_type_ ==
-      reinterpret_cast<cugraph::c_api::cugraph_graph_t const*>(lookup_container)->edge_type_,
+      reinterpret_cast<cugraph::c_api::cugraph_lookup_container_t const*>(lookup_container)
+        ->edge_type_,
     CUGRAPH_INVALID_INPUT,
     "edge type of graph and lookup_container must match",
     *error);
 
-  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_type_id_type_ ==
-                 reinterpret_cast<cugraph::c_api::cugraph_graph_t const*>(lookup_container)
-                   ->edge_type_id_type_,
-               CUGRAPH_INVALID_INPUT,
-               "edge type id type of graph and lookup_container must match",
-               *error);
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_type_id_type_ ==
+      reinterpret_cast<cugraph::c_api::cugraph_lookup_container_t const*>(lookup_container)
+        ->edge_type_id_type_,
+    CUGRAPH_INVALID_INPUT,
+    "edge type id type of graph and lookup_container must match",
+    *error);
 
   lookup_using_edge_ids_of_single_type_functor functor(
     handle, graph, lookup_container, edge_ids_to_lookup, edge_type_to_lookup);
@@ -387,3 +393,10 @@ extern "C" void cugraph_lookup_result_free(cugraph_lookup_result_t* result)
   delete internal_pointer->dsts_;
   delete internal_pointer;
 }
+
+extern "C" void cugraph_lookup_container_free(cugraph_lookup_container_t* container)
+{
+  auto internal_ptr = reinterpret_cast<cugraph::c_api::cugraph_lookup_container_t*>(container);
+  // The graph should presumably own the other structures.
+  delete internal_ptr;
+}

python/pylibcugraph/pylibcugraph/CMakeLists.txt

@@ -65,6 +65,7 @@ set(cython_sources
     all_pairs_sorensen_coefficients.pyx
     all_pairs_overlap_coefficients.pyx
     all_pairs_cosine_coefficients.pyx
+    edge_id_lookup_table.pyx
 )
 set(linked_libraries cugraph::cugraph;cugraph::cugraph_c)
 

python/pylibcugraph/pylibcugraph/__init__.py

@@ -21,6 +21,8 @@
 
 from pylibcugraph.graph_properties import GraphProperties
 
+from pylibcugraph.edge_id_lookup_table import EdgeIdLookupTable
+
 from pylibcugraph.eigenvector_centrality import eigenvector_centrality
 
 from pylibcugraph.katz_centrality import katz_centrality

python/pylibcugraph/pylibcugraph/_cugraph_c/lookup_src_dst.pxd

@@ -70,3 +70,5 @@ cdef extern from "cugraph_c/lookup_src_dst.h":
         const cugraph_lookup_result_t* result)
 
     cdef void cugraph_lookup_result_free(cugraph_lookup_result_t* result)
+
+    cdef void cugraph_lookup_container_free(cugraph_lookup_container_t* container)

python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pxd

@@ -0,0 +1,34 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from pylibcugraph._cugraph_c.error cimport (
+    cugraph_error_code_t,
+    cugraph_error_t,
+)
+from pylibcugraph._cugraph_c.lookup_src_dst cimport (
+    cugraph_lookup_container_t,
+)
+from pylibcugraph.resource_handle cimport (
+    ResourceHandle,
+)
+from pylibcugraph.graphs cimport (
+    _GPUGraph,
+)
+
+cdef class EdgeIdLookupTable:
+    cdef ResourceHandle handle,
+    cdef _GPUGraph graph,
+    cdef cugraph_lookup_container_t* lookup_container_c_ptr

python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx

@@ -0,0 +1,114 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from pylibcugraph._cugraph_c.resource_handle cimport (
+    cugraph_resource_handle_t,
+)
+from pylibcugraph._cugraph_c.error cimport (
+    cugraph_error_code_t,
+    cugraph_error_t,
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+    cugraph_type_erased_device_array_view_create,
+    cugraph_type_erased_device_array_view_free,
+    cugraph_type_erased_host_array_view_t,
+    cugraph_type_erased_host_array_view_create,
+    cugraph_type_erased_host_array_view_free,
+)
+from pylibcugraph._cugraph_c.graph cimport (
+    cugraph_graph_t,
+)
+from pylibcugraph._cugraph_c.lookup_src_dst cimport (
+    cugraph_lookup_container_t,
+    cugraph_build_edge_id_and_type_to_src_dst_lookup_map,
+    cugraph_lookup_container_free,
+    cugraph_lookup_endpoints_from_edge_ids_and_single_type,
+    cugraph_lookup_result_t,
+)
+from pylibcugraph.utils cimport (
+    assert_success,
+    assert_CAI_type,
+    assert_AI_type,
+    get_c_type_from_numpy_type,
+    create_cugraph_type_erased_device_array_view_from_py_obj
+)
+from pylibcugraph.resource_handle cimport (
+    ResourceHandle,
+)
+from pylibcugraph.graphs cimport (
+    _GPUGraph,
+)
+from pylibcugraph.internal_types.edge_id_lookup_result cimport (
+    EdgeIdLookupResult,
+)
+
+cdef class EdgeIdLookupTable:
+    def __cinit__(self, ResourceHandle resource_handle, _GPUGraph graph):
+        self.handle = resource_handle
+        self.graph = graph
+
+        cdef cugraph_error_code_t error_code
+        cdef cugraph_error_t* error_ptr
+
+        error_code = cugraph_build_edge_id_and_type_to_src_dst_lookup_map(
+            <cugraph_resource_handle_t*>self.handle.c_resource_handle_ptr,
+            <cugraph_graph_t*>self.graph.c_graph_ptr,
+            &self.lookup_container_c_ptr,
+            &error_ptr,
+        )
+
+        assert_success(error_code, error_ptr, "cugraph_build_edge_id_and_type_to_src_dst_lookup_map")
+
+    def __dealloc__(self):
+        if self.lookup_container_c_ptr is not NULL:
+            cugraph_lookup_container_free(self.lookup_container_c_ptr)
+
+    def lookup_vertex_ids(
+        self,
+        edge_ids,
+        int edge_type
+    ):
+        """
+        For a single edge type, finds the source and destination vertex ids corresponding
+        to the provided edge ids.
+        """
+
+        cdef cugraph_error_code_t error_code
+        cdef cugraph_error_t* error_ptr
+        cdef cugraph_lookup_result_t* result_ptr
+
+        cdef cugraph_type_erased_device_array_view_t* edge_ids_c_ptr
+        edge_ids_c_ptr = create_cugraph_type_erased_device_array_view_from_py_obj(edge_ids)
+
+        error_code = cugraph_lookup_endpoints_from_edge_ids_and_single_type(
+            <cugraph_resource_handle_t*>self.handle.c_resource_handle_ptr,
+            <cugraph_graph_t*>self.graph.c_graph_ptr,
+            self.lookup_container_c_ptr,
+            edge_ids_c_ptr,
+            edge_type,
+            &result_ptr,
+            &error_ptr,
+        )
+
+        assert_success(error_code, error_ptr, "cugraph_lookup_endpoints_from_edge_ids_and_single_type")
+
+        lr = EdgeIdLookupResult()
+        lr.set_ptr(<cugraph_lookup_result_t*>(result_ptr))
+        return {
+            'sources': lr.get_sources(),
+            'destinations': lr.get_destinations(),
+        }

python/pylibcugraph/pylibcugraph/internal_types/CMakeLists.txt

@@ -15,6 +15,7 @@
 set(cython_sources
     sampling_result.pyx
     coo.pyx
+    edge_id_lookup_result.pyx
 )
 set(linked_libraries cugraph::cugraph;cugraph::cugraph_c)
 

python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pxd

@@ -0,0 +1,30 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+
+from pylibcugraph._cugraph_c.lookup_src_dst cimport (
+    cugraph_lookup_result_t
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+)
+
+cdef class EdgeIdLookupResult:
+    cdef cugraph_lookup_result_t* result_c_ptr
+
+    cdef get_array(self, cugraph_type_erased_device_array_view_t* ptr)
+
+    cdef set_ptr(self, cugraph_lookup_result_t* ptr)

python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pyx

@@ -0,0 +1,63 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from pylibcugraph._cugraph_c.lookup_src_dst cimport (
+    cugraph_lookup_result_t,
+    cugraph_lookup_result_free,
+    cugraph_lookup_result_get_dsts,
+    cugraph_lookup_result_get_srcs,
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+)
+from pylibcugraph.utils cimport (
+    create_cupy_array_view_for_device_ptr,
+)
+
+cdef class EdgeIdLookupResult:
+    def __cinit__(self):
+        """
+        Sets this object as the owner of the given pointer.
+        """
+        self.result_c_ptr = NULL
+
+    cdef set_ptr(self, cugraph_lookup_result_t* ptr):
+        self.result_c_ptr = ptr
+
+    def __dealloc__(self):
+        if self.result_c_ptr is not NULL:
+            cugraph_lookup_result_free(self.result_c_ptr)
+
+    cdef get_array(self, cugraph_type_erased_device_array_view_t* ptr):
+        if ptr is NULL:
+            return None
+
+        return create_cupy_array_view_for_device_ptr(
+            ptr,
+            self,
+        )
+
+    def get_sources(self):
+        if self.result_c_ptr is NULL:
+            return None
+        cdef cugraph_type_erased_device_array_view_t* ptr = cugraph_lookup_result_get_srcs(self.result_c_ptr)
+        return self.get_array(ptr)
+
+    def get_destinations(self):
+        if self.result_c_ptr is NULL:
+            return None
+        cdef cugraph_type_erased_device_array_view_t* ptr = cugraph_lookup_result_get_dsts(self.result_c_ptr)
+        return self.get_array(ptr)