Support Heterogeneous Renumbering in the C/PLC APIs

cpp/include/cugraph_c/sampling_algorithms.h

@@ -312,6 +312,24 @@ void cugraph_sampling_set_prior_sources_behavior(cugraph_sampling_options_t* opt
  */
 void cugraph_sampling_set_dedupe_sources(cugraph_sampling_options_t* options, bool_t value);
 
+/**
+ * @ingroup samplingC
+ * @brief   Set the number of vertex types in the graph.
+ *
+ * @param options - opaque pointer to the sampling options struct
+ * @param value - the number of vertex types to set
+ */
+void cugraph_sampling_set_num_vertex_types(cugraph_sampling_options_t* options, size_t value);
+
+/**
+ * @ingroup samplingC
+ * @brief   Set the number of edge types in the graph.
+ *
+ * @param options - opaque pointer to the sampling options struct
+ * @param value - the number of edge types to set
+ */
+void cugraph_sampling_set_num_edge_types(cugraph_sampling_options_t* options, size_t value);
+
 /**
  * @ingroup samplingC
  * @brief     Free sampling options object
@@ -368,6 +386,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
+  const cugraph_type_erased_device_array_view_t* vertex_type_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
@@ -427,6 +446,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
+  const cugraph_type_erased_device_array_view_t* vertex_type_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
@@ -584,6 +604,26 @@ cugraph_type_erased_device_array_view_t* cugraph_sample_result_get_renumber_map(
 cugraph_type_erased_device_array_view_t* cugraph_sample_result_get_renumber_map_offsets(
   const cugraph_sample_result_t* result);
 
+/**
+ * @ingroup samplingC
+ * @brief     Get the edge renumber map
+ *
+ * @param [in]   result   The result from a sampling algorithm
+ * @return type erased array pointing to the edge renumber map
+ */
+cugraph_type_erased_device_array_view_t* cugraph_sample_result_get_edge_renumber_map(
+  const cugraph_sample_result_t* result);
+
+/**
+ * @ingroup samplingC
+ * @brief     Get the edge renumber map offsets
+ *
+ * @param [in]   result   The result from a sampling algorithm
+ * @return type erased array pointing to the edge renumber map offsets
+ */
+cugraph_type_erased_device_array_view_t* cugraph_sample_result_get_edge_renumber_map_offsets(
+  const cugraph_sample_result_t* result);
+
 /**
  * @ingroup samplingC
  * @brief     Free a sampling result

cpp/tests/c_api/biased_neighbor_sample_test.c

@@ -142,6 +142,7 @@ int generic_biased_neighbor_sample_test(const cugraph_resource_handle_t* handle,
                                             NULL,
                                             NULL,
                                             NULL,
+                                            NULL,
                                             h_fan_out_view,
                                             rng_state,
                                             sampling_options,
@@ -572,6 +573,7 @@ int test_biased_neighbor_sample_with_labels(const cugraph_resource_handle_t* han
                                             NULL,
                                             NULL,
                                             NULL,
+                                            NULL,
                                             h_fan_out_view,
                                             rng_state,
                                             sampling_options,

cpp/tests/c_api/create_graph_test.c

@@ -271,6 +271,7 @@ int test_create_sg_graph_csr()
                                              NULL,
                                              NULL,
                                              NULL,
+                                             NULL,
                                              h_fan_out_view,
                                              rng_state,
                                              sampling_options,

cpp/tests/c_api/uniform_neighbor_sample_test.c

@@ -141,6 +141,7 @@ int generic_uniform_neighbor_sample_test(const cugraph_resource_handle_t* handle
                                              NULL,
                                              NULL,
                                              NULL,
+                                             NULL,
                                              h_fan_out_view,
                                              rng_state,
                                              sampling_options,
@@ -662,6 +663,7 @@ int test_uniform_neighbor_sample_with_labels(const cugraph_resource_handle_t* ha
                                              NULL,
                                              NULL,
                                              NULL,
+                                             NULL,
                                              h_fan_out_view,
                                              rng_state,
                                              sampling_options,

python/pylibcugraph/pylibcugraph/_cugraph_c/algorithms.pxd

@@ -246,6 +246,16 @@ cdef extern from "cugraph_c/algorithms.h":
             const cugraph_sample_result_t* result
         )
 
+    cdef cugraph_type_erased_device_array_view_t* \
+        cugraph_sample_result_get_edge_renumber_map(
+            const cugraph_sample_result_t* result
+        )
+
+    cdef cugraph_type_erased_device_array_view_t* \
+        cugraph_sample_result_get_edge_renumber_map_offsets(
+            const cugraph_sample_result_t* result
+        )
+
     cdef void \
         cugraph_sample_result_free(
             const cugraph_sample_result_t* result
@@ -288,6 +298,16 @@ cdef extern from "cugraph_c/algorithms.h":
             cugraph_error_t** error,
         )
 
+    cdef void cugraph_sampling_set_num_vertex_types(
+        cugraph_sampling_options_t* options,
+        size_t value
+    );
+
+    cdef void cugraph_sampling_set_num_edge_types(
+        cugraph_sampling_options_t* options,
+        size_t value
+    );
+
     cdef void \
         cugraph_sampling_set_renumber_results(
             cugraph_sampling_options_t* options,
@@ -336,6 +356,16 @@ cdef extern from "cugraph_c/algorithms.h":
             cugraph_compression_type_t value,
         )
 
+    cdef void cugraph_sampling_set_num_edge_types(
+        cugraph_sampling_options_t* options,
+        size_t value,
+    )
+
+    cdef void cugraph_sampling_set_num_vertex_types(
+        cugraph_sampling_options_t* options,
+        size_t value,
+    )
+
     cdef void \
         cugraph_sampling_options_free(
             cugraph_sampling_options_t* options,

python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd

@@ -59,6 +59,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_device_array_view_t* label_list,
         const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
         const cugraph_type_erased_device_array_view_t* label_offsets,
+        const cugraph_type_erased_device_array_view_t* vertex_type_offsets,
         const cugraph_type_erased_host_array_view_t* fan_out,
         cugraph_rng_state_t* rng_state,
         const cugraph_sampling_options_t* options,
@@ -76,6 +77,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_device_array_view_t* label_list,
         const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
         const cugraph_type_erased_device_array_view_t* label_offsets,
+        const cugraph_type_erased_device_array_view_t* vertex_type_offsets,
         const cugraph_type_erased_host_array_view_t* fan_out,
         cugraph_rng_state_t* rng_state,
         const cugraph_sampling_options_t* options,

python/pylibcugraph/pylibcugraph/biased_neighbor_sample.pyx

@@ -52,6 +52,8 @@ from pylibcugraph._cugraph_c.algorithms cimport (
     cugraph_sampling_set_compress_per_hop,
     cugraph_sampling_set_compression_type,
     cugraph_sampling_set_retain_seeds,
+    cugraph_sampling_set_num_vertex_types,
+    cugraph_sampling_set_num_edge_types,
 )
 from pylibcugraph._cugraph_c.sampling_algorithms cimport (
     cugraph_biased_neighbor_sample,
@@ -93,6 +95,7 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
                             label_list=None,
                             label_to_output_comm_rank=None,
                             label_offsets=None,
+                            vertex_type_offsets=None,
                             biases=None,
                             prior_sources_behavior=None,
                             deduplicate_sources=False,
@@ -101,6 +104,8 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
                             retain_seeds=False,
                             compression='COO',
                             compress_per_hop=False,
+                            num_vertex_types=None,
+                            num_edge_types=None,
                             random_state=None,
                             return_dict=False,):
     """
@@ -152,6 +157,15 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
     label_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
+        vertex_type_offsets: device array type (Optional)
+        Offsets of each vertex type within the graph.
+        Vertices must be numbered in ascending order
+        by vertex type in order to properly renumber
+        heterogeneous graphs.  i.e. if there are two
+        vertex types, the 0th with 10 vertices and 1st
+        with 32 vertices, the offsets array should be
+        [0, 10, 42]
+
     biases: list[float32/64] (Optional)
         Edge biases.  If not provided, uses the weight property.
         Currently unsupported.
@@ -188,6 +202,15 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
         If True, will create a separate compressed edgelist per hop within
         a batch.
 
+    num_vertex_types: int (Optional)
+        If provided, sets the number of vertex types in the graph.
+        Otherwise, it is assumed that there is only one vertex
+        type.
+
+    num_edge_types: int (Optional)
+        If provided, sets the number of edge types in the graph.
+        Otherwise, it is assumed that there is only one edge type.
+
     random_state: int (Optional)
         Random state to use when generating samples.  Optional argument,
         defaults to a hash of process id, time, and hostname.
@@ -259,6 +282,11 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
         cai_label_to_output_comm_rank_ptr = \
             label_to_output_comm_rank.__cuda_array_interface__['data'][0]
 
+    cdef uintptr_t cai_vertex_type_offsets_ptr
+    if vertex_type_offsets is not None:
+        cai_vertex_type_offsets_ptr = \
+            vertex_type_offsets.__cuda_array_interface__['data'][0]
+
     cdef uintptr_t cai_label_offsets_ptr
     if label_offsets is not None:
         cai_label_offsets_ptr = \
@@ -300,6 +328,15 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
                 get_c_type_from_numpy_type(label_to_output_comm_rank.dtype)
             )
 
+    cdef cugraph_type_erased_device_array_view_t* vertex_type_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if vertex_type_offsets is not None:
+        vertex_type_offsets_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_vertex_type_offsets_ptr,
+                len(vertex_type_offsets),
+                get_c_type_from_numpy_type(vertex_type_offsets.dtype)
+            )
+
     cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
     if retain_seeds:
         if label_offsets is None:
@@ -366,6 +403,11 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
     cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
     cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
 
+    if num_vertex_types:
+        cugraph_sampling_set_num_vertex_types(sampling_options, num_vertex_types)
+    if num_edge_types:
+        cugraph_sampling_set_num_edge_types(sampling_options, num_edge_types)
+
     error_code = cugraph_biased_neighbor_sample(
         c_resource_handle_ptr,
         c_graph_ptr,
@@ -375,6 +417,7 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
         label_list_ptr,
         label_to_output_comm_rank_ptr,
         label_offsets_ptr,
+        vertex_type_offsets_ptr,
         fan_out_ptr,
         rng_state_ptr,
         sampling_options,
@@ -413,6 +456,8 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
     if renumber:
         cupy_renumber_map = result.get_renumber_map()
         cupy_renumber_map_offsets = result.get_renumber_map_offsets()
+        cupy_edge_renumber_map = result.get_edge_renumber_map()
+        cupy_edge_renumber_map_offsets = result.get_edge_renumber_map_offsets()
 
         if return_dict:
             return {
@@ -426,11 +471,13 @@ def biased_neighbor_sample(ResourceHandle resource_handle,
                 'label_hop_offsets': cupy_label_hop_offsets,
                 'hop_id': None,
                 'renumber_map': cupy_renumber_map,
-                'renumber_map_offsets': cupy_renumber_map_offsets
+                'renumber_map_offsets': cupy_renumber_map_offsets,
+                'edge_renumber_map': cupy_edge_renumber_map,
+                'edge_renumber_map_offsets': cupy_edge_renumber_map_offsets,
             }
         else:
             cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
-            return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets)
+            return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets, cupy_edge_renumber_map, cupy_edge_renumber_map_offsets)
     else:
         cupy_hop_ids = result.get_hop_ids()
         if return_dict:

python/pylibcugraph/pylibcugraph/internal_types/sampling_result.pyx

@@ -34,6 +34,8 @@ from pylibcugraph._cugraph_c.algorithms cimport (
     cugraph_sample_result_get_offsets, # deprecated
     cugraph_sample_result_get_renumber_map,
     cugraph_sample_result_get_renumber_map_offsets,
+    cugraph_sample_result_get_edge_renumber_map,
+    cugraph_sample_result_get_edge_renumber_map_offsets,
     cugraph_sample_result_free,
 )
 from pylibcugraph.utils cimport (
@@ -257,3 +259,28 @@ cdef class SamplingResult:
 
         return create_cupy_array_view_for_device_ptr(device_array_view_ptr,
                                                      self)
+    def get_edge_renumber_map(self):
+        if self.c_sample_result_ptr is NULL:
+            raise ValueError("pointer not set, must call set_ptr() with a "
+                             "non-NULL value first.")
+        cdef cugraph_type_erased_device_array_view_t* device_array_view_ptr = (
+            cugraph_sample_result_get_edge_renumber_map(self.c_sample_result_ptr)
+        )
+        if device_array_view_ptr is NULL:
+            return None
+
+        return create_cupy_array_view_for_device_ptr(device_array_view_ptr,
+                                                     self)
+
+    def get_edge_renumber_map_offsets(self):
+        if self.c_sample_result_ptr is NULL:
+            raise ValueError("pointer not set, must call set_ptr() with a "
+                             "non-NULL value first.")
+        cdef cugraph_type_erased_device_array_view_t* device_array_view_ptr = (
+            cugraph_sample_result_get_edge_renumber_map_offsets(self.c_sample_result_ptr)
+        )
+        if device_array_view_ptr is NULL:
+            return None
+
+        return create_cupy_array_view_for_device_ptr(device_array_view_ptr,
+                                                     self)