From 659abb9a66e29b766c7af957a70ab1c3626282fc Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Tue, 1 Oct 2024 14:21:23 -0700 Subject: [PATCH 1/4] [FEA] Support Edge ID Lookup in PyLibcuGraph --- cpp/include/cugraph_c/lookup_src_dst.h | 8 ++ cpp/src/c_api/lookup_src_dst.cpp | 45 ++++--- .../pylibcugraph/pylibcugraph/CMakeLists.txt | 1 + python/pylibcugraph/pylibcugraph/__init__.py | 2 + .../_cugraph_c/lookup_src_dst.pxd | 2 + .../pylibcugraph/edge_id_lookup_table.pxd | 34 ++++++ .../pylibcugraph/edge_id_lookup_table.pyx | 114 ++++++++++++++++++ .../internal_types/CMakeLists.txt | 1 + .../internal_types/edge_id_lookup_result.pxd | 30 +++++ .../internal_types/edge_id_lookup_result.pyx | 63 ++++++++++ 10 files changed, 284 insertions(+), 16 deletions(-) create mode 100644 python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pxd create mode 100644 python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx create mode 100644 python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pxd create mode 100644 python/pylibcugraph/pylibcugraph/internal_types/edge_id_lookup_result.pyx 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..077e1f3a30b --- /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 find( + 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) From 6082e0d2849f31cdb0f168896a9e5befc0421986 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Thu, 17 Oct 2024 12:48:17 -0700 Subject: [PATCH 2/4] add plc test for lookup table --- .../pylibcugraph/tests/test_lookup_table.py | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py 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..18faa54c098 --- /dev/null +++ b/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py @@ -0,0 +1,98 @@ +# 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 +import numpy as np + +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(): + + vtcs = cupy.arange(6, dtype="int64") + vtps = np.array([0, 0, 1, 1, 2, 2]) + + e_lookup = { + (0, 0): [0, 0], + (0, 1): [1, 0], + (0, 2): [2, 0], + (1, 0): [3, 0], + (1, 1): [4, 0], + (1, 2): [5, 0], + (2, 0): [6, 0], + (2, 1): [7, 0], + (2, 2): [8, 0], + } + + srcs = np.array([0, 1, 5, 4, 3, 2, 2, 0, 5, 4, 4, 5]) + dsts = np.array([1, 5, 0, 3, 2, 1, 3, 3, 2, 3, 1, 4]) + wgts = cupy.ones((len(srcs),), dtype="float32") + + eids = [] + etps = [] + for i in range(len(srcs)): + key = (int(vtps[srcs[i]]), int(vtps[dsts[i]])) + etps.append(e_lookup[key][0]) + eids.append(e_lookup[key][1]) + + e_lookup[key][1] += 1 + + eids = cupy.array(eids) + etps = cupy.array(etps, dtype="int32") + + graph = SGGraph( + resource_handle=ResourceHandle(), + graph_properties=GraphProperties(is_symmetric=False, is_multigraph=True), + src_or_offset_array=cupy.array(srcs), + dst_or_index_array=cupy.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, + ) + + # call __dealloc__() + table = EdgeIdLookupTable(ResourceHandle(), graph) + + assert table is not None + + found_edges = table.find(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.find(cupy.array([0]), 5) + assert (found_edges["sources"] == cupy.array([-1])).all() + assert (found_edges["destinations"] == cupy.array([-1])).all() + + found_edges = table.find(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() From 9c4c15b733cc0f13af2f417b0490b33e84d6c914 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Thu, 17 Oct 2024 12:49:21 -0700 Subject: [PATCH 3/4] check dealloc --- python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py b/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py index 18faa54c098..c2e49a4fe6d 100644 --- a/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py +++ b/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py @@ -80,7 +80,6 @@ def test_lookup_table(): do_expensive_check=True, ) - # call __dealloc__() table = EdgeIdLookupTable(ResourceHandle(), graph) assert table is not None @@ -96,3 +95,6 @@ def test_lookup_table(): found_edges = table.find(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 From 340a4881e2b35f4a6851dac5242e1a66eb34f96d Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Fri, 18 Oct 2024 08:54:24 -0700 Subject: [PATCH 4/4] simplify test, add comment, change function name --- .../pylibcugraph/edge_id_lookup_table.pyx | 2 +- .../pylibcugraph/tests/test_lookup_table.py | 48 ++++++------------- 2 files changed, 15 insertions(+), 35 deletions(-) diff --git a/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx b/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx index 077e1f3a30b..49ccdbdd168 100644 --- a/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx +++ b/python/pylibcugraph/pylibcugraph/edge_id_lookup_table.pyx @@ -77,7 +77,7 @@ cdef class EdgeIdLookupTable: if self.lookup_container_c_ptr is not NULL: cugraph_lookup_container_free(self.lookup_container_c_ptr) - def find( + def lookup_vertex_ids( self, edge_ids, int edge_type diff --git a/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py b/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py index c2e49a4fe6d..2910a5f8d4d 100644 --- a/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py +++ b/python/pylibcugraph/pylibcugraph/tests/test_lookup_table.py @@ -12,7 +12,6 @@ # limitations under the License. import cupy -import numpy as np from pylibcugraph import ( SGGraph, @@ -34,43 +33,24 @@ def test_lookup_table(): - + # Vertex id array vtcs = cupy.arange(6, dtype="int64") - vtps = np.array([0, 0, 1, 1, 2, 2]) - - e_lookup = { - (0, 0): [0, 0], - (0, 1): [1, 0], - (0, 2): [2, 0], - (1, 0): [3, 0], - (1, 1): [4, 0], - (1, 2): [5, 0], - (2, 0): [6, 0], - (2, 1): [7, 0], - (2, 2): [8, 0], - } - - srcs = np.array([0, 1, 5, 4, 3, 2, 2, 0, 5, 4, 4, 5]) - dsts = np.array([1, 5, 0, 3, 2, 1, 3, 3, 2, 3, 1, 4]) - wgts = cupy.ones((len(srcs),), dtype="float32") - eids = [] - etps = [] - for i in range(len(srcs)): - key = (int(vtps[srcs[i]]), int(vtps[dsts[i]])) - etps.append(e_lookup[key][0]) - eids.append(e_lookup[key][1]) + # 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]) - e_lookup[key][1] += 1 - - eids = cupy.array(eids) - etps = cupy.array(etps, dtype="int32") + 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=cupy.array(srcs), - dst_or_index_array=cupy.array(dsts), + src_or_offset_array=srcs, + dst_or_index_array=dsts, vertices_array=vtcs, weight_array=wgts, edge_id_array=eids, @@ -84,15 +64,15 @@ def test_lookup_table(): assert table is not None - found_edges = table.find(cupy.array([0, 1, 2, 3, 4]), 7) + 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.find(cupy.array([0]), 5) + 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.find(cupy.array([3, 1, 0, 5]), 6) + 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()