Merge pull request #4800 from rapidsai/branch-24.12

Forward-merge branch-24.12 into branch-25.02

cpp/include/cugraph/sampling_functions.hpp

@@ -306,7 +306,7 @@ struct sampling_flags_t {
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_label_offsets Optional device span of labels associated with each starting
+ * @param starting_vertex_labels Optional device span of labels associated with each starting
  * vertex for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get sampling
  * outputs of each vertex label.  This should be the same on each rank.
@@ -340,7 +340,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_label_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -385,7 +385,7 @@ homogeneous_uniform_neighbor_sample(
  * corresponding edge can never be selected.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_label_offsets Optional device span of labels associated with each starting
+ * @param starting_vertex_labels Optional device span of labels associated with each starting
  * vertex for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get sampling
  * outputs of each vertex label.  This should be the same on each rank.
@@ -421,7 +421,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_label_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -462,7 +462,7 @@ homogeneous_biased_neighbor_sample(
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_label_offsets Optional device span of labels associated with each starting
+ * @param starting_vertex_labels Optional device span of labels associated with each starting
  * vertex for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get sampling
  * outputs of each vertex label.  This should be the same on each rank.
@@ -498,7 +498,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_label_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   edge_type_t num_edge_types,
@@ -545,7 +545,7 @@ heterogeneous_uniform_neighbor_sample(
  * corresponding edge can never be selected.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_label_offsets Optional device span of labels associated with each starting
+ * @param starting_vertex_labels Optional device span of labels associated with each starting
  * vertex for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get sampling
  * outputs of each vertex label.  This should be the same on each rank.
@@ -583,7 +583,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_label_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   edge_type_t num_edge_types,

cpp/include/cugraph_c/sampling_algorithms.h

@@ -551,6 +551,8 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
  * @param [in]  start_vertices Device array of start vertices for the sampling
  * @param [in]  starting_vertex_label_offsets Device array of the offsets for each label in
  * the seed list. This parameter is only used with the retain_seeds option.
+ * @param [in]  vertex_type_offsets Device array of the offsets for each vertex type in the
+ * graph.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
  * algorithm. We only support fan_out values of type INT32
  * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
@@ -570,6 +572,7 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* starting_vertex_label_offsets,
+  const cugraph_type_erased_device_array_view_t* vertex_type_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -598,6 +601,8 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
  * @param [in]  start_vertices Device array of start vertices for the sampling
  * @param [in]  starting_vertex_label_offsets Device array of the offsets for each label in
  * the seed list. This parameter is only used with the retain_seeds option.
+ * @param [in]  vertex_type_offsets Device array of the offsets for each vertex type in the
+ * graph.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
  * algorithm. We only support fan_out values of type INT32
  * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
@@ -618,6 +623,7 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* starting_vertex_label_offsets,
+  const cugraph_type_erased_device_array_view_t* vertex_type_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -735,6 +741,16 @@ cugraph_type_erased_device_array_view_t* cugraph_sample_result_get_hop(
 cugraph_type_erased_device_array_view_t* cugraph_sample_result_get_label_hop_offsets(
   const cugraph_sample_result_t* result);
 
+/**
+ * @ingroup samplingC
+ * @brief     Get the label-type-hop offsets from the sampling algorithm result
+ *
+ * @param [in]   result   The result from a sampling algorithm
+ * @return type erased array pointing to the label-type-hop offsets
+ */
+cugraph_type_erased_device_array_view_t* cugraph_sample_result_get_label_type_hop_offsets(
+  const cugraph_sample_result_t* result);
+
 /**
  * @ingroup samplingC
  * @brief     Get the index from the sampling algorithm result