Sampling primitive updates to support last-N sampling (#5622)

This PR

* refactors sampling primitive functions to minimize code duplication between random sampling and last-N (top-K) sampling
* added `per_v_top_k_select_transform_outgoing_e` (in addition to the existing `per_v_random_select_transform_outgoing_e`) to support last-N sampling.

`per_v_top_k_select_transform_outgoing_e` selects neighbors with the K highest bias values (bias value should be non-negative and bias 0 neighbors cannot be selected similar to biased random sampling).

Authors:
  - Seunghwa Kang (https://github.com/seunghwak)

Approvers:
  - Chuck Hastings (https://github.com/ChuckHastings)

URL: #5622

Author: seunghwak

cpp/include/cugraph/prims/detail/per_v_select_transform_e.cuh

@@ -0,0 +1,84 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <cugraph/export.hpp>
+
+#include <raft/core/device_span.hpp>
+#include <raft/core/handle.hpp>
+
+#include <rmm/device_uvector.hpp>
+
+#include <cub/cub.cuh>
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+
+namespace CUGRAPH_EXPORT cugraph {
+namespace detail {
+
+template <typename key_t, typename value_t>
+void device_segmented_sort_pairs(raft::handle_t const& handle,
+                                 raft::device_span<key_t const> keys_in,
+                                 raft::device_span<key_t> keys_out,
+                                 raft::device_span<value_t const> values_in,
+                                 raft::device_span<value_t> values_out,
+                                 raft::device_span<size_t const> begin_offsets,
+                                 raft::device_span<size_t const> end_offsets)
+{
+  assert(keys_in.size() == keys_out.size());
+  assert(values_in.size() == values_out.size());
+  assert(keys_in.size() == values_in.size());
+  assert(begin_offsets.size() == end_offsets.size());
+
+  size_t tmp_storage_bytes{0};
+  cub::DeviceSegmentedSort::SortPairs(static_cast<void*>(nullptr),
+                                      tmp_storage_bytes,
+                                      keys_in.data(),
+                                      keys_out.data(),
+                                      values_in.data(),
+                                      values_out.data(),
+                                      keys_in.size(),
+                                      begin_offsets.size(),
+                                      begin_offsets.data(),
+                                      end_offsets.data(),
+                                      handle.get_stream());
+  rmm::device_uvector<std::byte> d_tmp_storage(tmp_storage_bytes, handle.get_stream());
+  cub::DeviceSegmentedSort::SortPairs(d_tmp_storage.data(),
+                                      tmp_storage_bytes,
+                                      keys_in.data(),
+                                      keys_out.data(),
+                                      values_in.data(),
+                                      values_out.data(),
+                                      keys_in.size(),
+                                      begin_offsets.size(),
+                                      begin_offsets.data(),
+                                      end_offsets.data(),
+                                      handle.get_stream());
+}
+
+// offsets is a CSR offset array of size num_segments + 1
+template <typename key_t, typename value_t>
+void device_segmented_sort_pairs(raft::handle_t const& handle,
+                                 raft::device_span<key_t const> keys_in,
+                                 raft::device_span<key_t> keys_out,
+                                 raft::device_span<value_t const> values_in,
+                                 raft::device_span<value_t> values_out,
+                                 raft::device_span<size_t const> offsets)
+{
+  assert(offsets.size() >= size_t{1});
+  device_segmented_sort_pairs(
+    handle,
+    keys_in,
+    keys_out,
+    values_in,
+    values_out,
+    raft::device_span<size_t const>{offsets.data(), offsets.size() - 1},
+    raft::device_span<size_t const>{offsets.data() + 1, offsets.size() - 1});
+}
+
+}  // namespace detail
+}  // namespace CUGRAPH_EXPORT cugraph

cpp/include/cugraph/prims/detail/sample_and_compute_local_nbr_indices.cuh

@@ -819,6 +819,11 @@ if(BUILD_CUGRAPH_MG_TESTS)
     ConfigureTestMG(MG_PER_V_RANDOM_SELECT_TRANSFORM_OUTGOING_E_TEST
                     prims/mg_per_v_random_select_transform_outgoing_e.cu)
 
+    ###############################################################################################
+    # - MG PRIMS PER_V_TOP_K_SELECT_TRANSFORM_OUTGOING_E tests ------------------------------------
+    ConfigureTestMG(MG_PER_V_TOP_K_SELECT_TRANSFORM_OUTGOING_E_TEST
+                    prims/mg_per_v_top_k_select_transform_outgoing_e.cu)
+
     ###############################################################################################
     # - MG PRIMS PER_V_PAIR_TRANSFORM_SRC_DST_NBR_INTERSECTION tests ------------------------------
     ConfigureTestMG(MG_PER_V_PAIR_TRANSFORM_SRC_DST_NBR_INTERSECTION_TEST

cpp/include/cugraph/prims/per_v_random_select_transform_outgoing_e.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -610,6 +610,7 @@ class Tests_MGPerVRandomSelectTransformOutgoingE
 
               if (sg_nbr_bias_first) {
                 auto lower_it = thrust::lower_bound(thrust::seq, sg_nbr_first, sg_nbr_last, sg_dst);
+                if ((lower_it == sg_nbr_last) || (*lower_it != sg_dst)) { return true; }
                 auto upper_it = thrust::upper_bound(thrust::seq, sg_nbr_first, sg_nbr_last, sg_dst);
                 bool found    = false;
                 for (auto it = (*sg_nbr_bias_first + cuda::std::distance(sg_nbr_first, lower_it));