diff --git a/cpp/include/cugraph_c/lookup_src_dst.h b/cpp/include/cugraph_c/lookup_src_dst.h index f4d63572e82..64051743981 100644 --- a/cpp/include/cugraph_c/lookup_src_dst.h +++ b/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 diff --git a/cpp/src/c_api/lookup_src_dst.cpp b/cpp/src/c_api/lookup_src_dst.cpp index 1be2137ef2f..3b87791ac50 100644 --- a/cpp/src/c_api/lookup_src_dst.cpp +++ b/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(graph)->vertex_type_ == - reinterpret_cast(lookup_container)->vertex_type_, + reinterpret_cast(lookup_container) + ->vertex_type_, CUGRAPH_INVALID_INPUT, "vertex type of graph and lookup_container must match", *error); CAPI_EXPECTS( reinterpret_cast(graph)->edge_type_ == - reinterpret_cast(lookup_container)->edge_type_, + reinterpret_cast(lookup_container) + ->edge_type_, CUGRAPH_INVALID_INPUT, "edge type of graph and lookup_container must match", *error); - CAPI_EXPECTS(reinterpret_cast(graph)->edge_type_id_type_ == - reinterpret_cast(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(graph)->edge_type_id_type_ == + reinterpret_cast(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(graph)->vertex_type_ == - reinterpret_cast(lookup_container)->vertex_type_, + reinterpret_cast(lookup_container) + ->vertex_type_, CUGRAPH_INVALID_INPUT, "vertex type of graph and lookup_container must match", *error); CAPI_EXPECTS( reinterpret_cast(graph)->edge_type_ == - reinterpret_cast(lookup_container)->edge_type_, + reinterpret_cast(lookup_container) + ->edge_type_, CUGRAPH_INVALID_INPUT, "edge type of graph and lookup_container must match", *error); - CAPI_EXPECTS(reinterpret_cast(graph)->edge_type_id_type_ == - reinterpret_cast(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(graph)->edge_type_id_type_ == + reinterpret_cast(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(container); + // The graph should presumably own the other structures. + delete internal_ptr; +} diff --git a/python/pylibcugraph/pylibcugraph/CMakeLists.txt b/python/pylibcugraph/pylibcugraph/CMakeLists.txt index 9f1b9924336..3a53c7d16c3 100644 --- a/python/pylibcugraph/pylibcugraph/CMakeLists.txt +++ b/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) diff --git a/python/pylibcugraph/pylibcugraph/__init__.py b/python/pylibcugraph/pylibcugraph/__init__.py index 26fa3f64ddd..9c04a528fd8 100644 --- a/python/pylibcugraph/pylibcugraph/__init__.py +++ b/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 diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/lookup_src_dst.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/lookup_src_dst.pxd index 710ca7d113b..e8a2bbf47ae 100644 --- a/python/pylibcugraph/pylibcugraph/_cugraph_c/lookup_src_dst.pxd +++ b/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) diff --git a/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pxd b/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pxd new file mode 100644 index 00000000000..9bbd19963a7 --- /dev/null +++ b/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 diff --git a/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx b/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx new file mode 100644 index 00000000000..49ccdbdd168 --- /dev/null +++ b/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( + self.handle.c_resource_handle_ptr, + 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( + self.handle.c_resource_handle_ptr, + 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((result_ptr)) + return { + 'sources': lr.get_sources(), + 'destinations': lr.get_destinations(), + } diff --git a/python/pylibcugraph/pylibcugraph/internal_types/CMakeLists.txt b/python/pylibcugraph/pylibcugraph/internal_types/CMakeLists.txt index 22f07939db0..1b0d6ec71a4 100644 --- a/python/pylibcugraph/pylibcugraph/internal_types/CMakeLists.txt +++ b/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) diff --git a/python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pxd b/python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pxd new file mode 100644 index 00000000000..68dd2362a00 --- /dev/null +++ b/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) diff --git a/python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pyx b/python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pyx new file mode 100644 index 00000000000..5f7165ce988 --- /dev/null +++ b/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) diff --git a/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py b/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py new file mode 100644 index 00000000000..2910a5f8d4d --- /dev/null +++ b/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py @@ -0,0 +1,80 @@ +# 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. + +import cupy + +from pylibcugraph import ( + SGGraph, + ResourceHandle, + GraphProperties, + EdgeIdLookupTable, +) + + +# ============================================================================= +# Pytest fixtures +# ============================================================================= +# fixtures used in this test module are defined in conftest.py + + +# ============================================================================= +# Tests +# ============================================================================= + + +def test_lookup_table(): + # Vertex id array + vtcs = cupy.arange(6, dtype="int64") + + # Edge ids are unique per edge type and start from 0 + # Each edge type has the same src/dst vertex type here, + # just as it would in a GNN application. + srcs = cupy.array([0, 1, 5, 4, 3, 2, 2, 0, 5, 4, 4, 5]) + dsts = cupy.array([1, 5, 0, 3, 2, 1, 3, 3, 2, 3, 1, 4]) + etps = cupy.array([0, 2, 6, 7, 4, 3, 4, 1, 7, 7, 6, 8], dtype="int32") + eids = cupy.array([0, 0, 0, 0, 0, 0, 1, 0, 1, 2, 1, 0]) + + wgts = cupy.ones((len(srcs),), dtype="float32") + + graph = SGGraph( + resource_handle=ResourceHandle(), + graph_properties=GraphProperties(is_symmetric=False, is_multigraph=True), + src_or_offset_array=srcs, + dst_or_index_array=dsts, + vertices_array=vtcs, + weight_array=wgts, + edge_id_array=eids, + edge_type_array=etps, + store_transposed=False, + renumber=False, + do_expensive_check=True, + ) + + table = EdgeIdLookupTable(ResourceHandle(), graph) + + assert table is not None + + found_edges = table.lookup_vertex_ids(cupy.array([0, 1, 2, 3, 4]), 7) + assert (found_edges["sources"] == cupy.array([4, 5, 4, -1, -1])).all() + assert (found_edges["destinations"] == cupy.array([3, 2, 3, -1, -1])).all() + + found_edges = table.lookup_vertex_ids(cupy.array([0]), 5) + assert (found_edges["sources"] == cupy.array([-1])).all() + assert (found_edges["destinations"] == cupy.array([-1])).all() + + found_edges = table.lookup_vertex_ids(cupy.array([3, 1, 0, 5]), 6) + assert (found_edges["sources"] == cupy.array([-1, 4, 5, -1])).all() + assert (found_edges["destinations"] == cupy.array([-1, 1, 0, -1])).all() + + # call __dealloc__() + del table