add support for vertex type at the plc layer

cpp/src/c_api/neighbor_sampling.cpp

@@ -1230,7 +1230,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
             rmm::device_uvector<vertex_t> vertex_type_offsets(2, handle_.get_stream());
 
-            if (vertex_type_offsets_ != nullptr) {
+            if (vertex_type_offsets_ == nullptr) {
               // If no 'vertex_type_offsets' is provided, all vertices are assumed to have
               // a vertex type of value 1.
               cugraph::detail::stride_fill(handle_.get_stream(),

python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx

@@ -87,6 +87,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
                                          _GPUGraph input_graph,
                                          start_vertex_list,
                                          starting_vertex_label_offsets,
+                                         vertex_type_offsets,
                                          h_fan_out,
                                          num_edge_types,
                                          bool_t with_replacement,
@@ -123,6 +124,9 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
         Offsets of each label within the start vertex list. Expanding
         'starting_vertex_label_offsets' must lead to an array of
         len(start_vertex_list)
+    
+    vertex_type_offsets: device array type (Optional)
+        Offsets for each vertex type in the graph.
 
     h_fan_out: numpy array type
         Device array containing the branching out (fan-out) degrees per
@@ -247,6 +251,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
 
     assert_CAI_type(start_vertex_list, "start_vertex_list")
     assert_CAI_type(starting_vertex_label_offsets, "starting_vertex_label_offsets", True)
+    assert_CAI_type(vertex_type_offsets, "vertex_type_offsets", True)
 
     assert_AI_type(h_fan_out, "h_fan_out")
 
@@ -276,6 +281,11 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
     if starting_vertex_label_offsets is not None:
         cai_starting_vertex_label_offsets_ptr = \
             starting_vertex_label_offsets.__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 cugraph_type_erased_device_array_view_t* start_vertex_list_ptr = \
@@ -293,6 +303,16 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
                 len(starting_vertex_label_offsets),
                 SIZE_T
             )
+
+    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),
+                SIZE_T
+            )
+
 
     cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
     if retain_seeds:
@@ -354,6 +374,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
         <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
         start_vertex_list_ptr,
         starting_vertex_label_offsets_ptr,
+        vertex_type_offsets_ptr,
         fan_out_ptr,
         num_edge_types,
         sampling_options,

python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx

@@ -84,6 +84,7 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
                                           _GPUGraph input_graph,
                                           start_vertex_list,
                                           starting_vertex_label_offsets,
+                                          vertex_type_offsets,
                                           h_fan_out,
                                           num_edge_types,
                                           bool_t with_replacement,
@@ -118,6 +119,9 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
         Offsets of each label within the start vertex list. Expanding
         'starting_vertex_label_offsets' must lead to an array of
         len(start_vertex_list)
+    
+    vertex_type_offsets: device array type (Optional)
+        Offsets for each vertex type in the graph.
 
     h_fan_out: numpy array type
         Device array containing the branching out (fan-out) degrees per
@@ -242,6 +246,7 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
 
     assert_CAI_type(start_vertex_list, "start_vertex_list")
     assert_CAI_type(starting_vertex_label_offsets, "starting_vertex_label_offsets", True)
+    assert_CAI_type(vertex_type_offsets, "vertex_type_offsets", True)
 
     assert_AI_type(h_fan_out, "h_fan_out")
 
@@ -270,6 +275,11 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
     if starting_vertex_label_offsets is not None:
         cai_starting_vertex_label_offsets_ptr = \
             starting_vertex_label_offsets.__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 cugraph_type_erased_device_array_view_t* start_vertex_list_ptr = \
@@ -288,6 +298,15 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
                 SIZE_T
             )
 
+    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),
+                SIZE_T
+            )
+
     cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
     if retain_seeds:
         if starting_vertex_label_offsets is None:
@@ -347,6 +366,7 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
         c_graph_ptr,
         start_vertex_list_ptr,
         starting_vertex_label_offsets_ptr,
+        vertex_type_offsets_ptr,
         fan_out_ptr,
         num_edge_types,
         sampling_options,