cugraph/sampling/
mod.rs

1//! Random-walk and sampling algorithms.
2
3use std::{ffi::CString, num::NonZeroUsize, ptr};
4
5use cugraph_sys::{
6    cugraph_biased_random_walks, cugraph_compression_type_t, cugraph_compression_type_t_COO,
7    cugraph_compression_type_t_CSC, cugraph_compression_type_t_CSR, cugraph_compression_type_t_DCSC,
8    cugraph_compression_type_t_DCSR, cugraph_error_t, cugraph_homogeneous_biased_neighbor_sample,
9    cugraph_homogeneous_biased_temporal_neighbor_sample, cugraph_homogeneous_uniform_neighbor_sample,
10    cugraph_homogeneous_uniform_temporal_neighbor_sample, cugraph_prior_sources_behavior_t,
11    cugraph_prior_sources_behavior_t_CARRY_OVER, cugraph_prior_sources_behavior_t_DEFAULT,
12    cugraph_prior_sources_behavior_t_EXCLUDE, cugraph_random_walk_result_get_max_path_length,
13    cugraph_random_walk_result_get_paths, cugraph_random_walk_result_get_weights, cugraph_random_walk_result_t,
14    cugraph_sample_result_get_edge_end_time, cugraph_sample_result_get_edge_id,
15    cugraph_sample_result_get_edge_renumber_map, cugraph_sample_result_get_edge_renumber_map_offsets,
16    cugraph_sample_result_get_edge_start_time, cugraph_sample_result_get_edge_type,
17    cugraph_sample_result_get_edge_weight, cugraph_sample_result_get_hop, cugraph_sample_result_get_label_hop_offsets,
18    cugraph_sample_result_get_label_type_hop_offsets, cugraph_sample_result_get_major_offsets,
19    cugraph_sample_result_get_majors, cugraph_sample_result_get_minors, cugraph_sample_result_get_renumber_map,
20    cugraph_sample_result_get_renumber_map_offsets, cugraph_sample_result_get_start_labels, cugraph_sample_result_t,
21    cugraph_sampling_options_create, cugraph_sampling_options_free, cugraph_sampling_options_t,
22    cugraph_sampling_set_compress_per_hop, cugraph_sampling_set_compression_type, cugraph_sampling_set_dedupe_sources,
23    cugraph_sampling_set_disjoint_sampling, cugraph_sampling_set_prior_sources_behavior,
24    cugraph_sampling_set_renumber_results, cugraph_sampling_set_retain_seeds, cugraph_sampling_set_return_hops,
25    cugraph_sampling_set_temporal_sampling_comparison, cugraph_sampling_set_with_replacement,
26    cugraph_select_random_vertices, cugraph_temporal_sampling_comparison_t,
27    cugraph_temporal_sampling_comparison_t_LAST, cugraph_temporal_sampling_comparison_t_MONOTONICALLY_DECREASING,
28    cugraph_temporal_sampling_comparison_t_MONOTONICALLY_INCREASING,
29    cugraph_temporal_sampling_comparison_t_STRICTLY_DECREASING,
30    cugraph_temporal_sampling_comparison_t_STRICTLY_INCREASING, cugraph_type_erased_device_array_t,
31    cugraph_uniform_random_walks,
32};
33
34use crate::{
35    array::{DeviceArray, DeviceArrayView, HostArrayView, RawDeviceArrayView},
36    cuda::{DeviceResourceKind, drop_device_resource},
37    error::{Error, ErrorKind, Result, check_error, check_host_error, require_non_null},
38    graph::{Forward, Graph},
39    handle::{EnteredContext, ResourceHandleLease},
40    params::PositiveF64,
41    random::RngState,
42    structure::{IntegerValues, ScoreValues, SizeValues, VertexIds},
43    types::{CugraphElement, DataType, to_bool_t},
44};
45
46/// Parameters shared by uniform and biased random walks.
47#[derive(Debug, Clone, Copy)]
48#[must_use]
49pub struct RandomWalkParams {
50    /// Maximum number of edges traversed per walk.
51    pub max_length: NonZeroUsize,
52}
53
54impl Default for RandomWalkParams {
55    fn default() -> Self {
56        Self { max_length: NonZeroUsize::new(1).expect("literal is non-zero") }
57    }
58}
59
60/// Parameters for node2vec random walks.
61#[derive(Debug, Clone, Copy)]
62#[must_use]
63pub struct Node2VecRandomWalkParams {
64    /// Maximum number of edges traversed per walk.
65    pub max_length: NonZeroUsize,
66    /// Return parameter.
67    pub p: PositiveF64,
68    /// In/out parameter.
69    pub q: PositiveF64,
70}
71
72impl Default for Node2VecRandomWalkParams {
73    fn default() -> Self {
74        Self {
75            max_length: NonZeroUsize::new(1).expect("literal is non-zero"),
76            p: PositiveF64::new(1.0).expect("literal is positive"),
77            q: PositiveF64::new(1.0).expect("literal is positive"),
78        }
79    }
80}
81
82/// Sampling output compression mode.
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
84#[must_use]
85pub enum SamplingCompression {
86    #[default]
87    Coo,
88    Csr,
89    Csc,
90    Dcsr,
91    Dcsc,
92}
93
94impl SamplingCompression {
95    fn to_c(self) -> cugraph_compression_type_t {
96        match self {
97            Self::Coo => cugraph_compression_type_t_COO,
98            Self::Csr => cugraph_compression_type_t_CSR,
99            Self::Csc => cugraph_compression_type_t_CSC,
100            Self::Dcsr => cugraph_compression_type_t_DCSR,
101            Self::Dcsc => cugraph_compression_type_t_DCSC,
102        }
103    }
104}
105
106/// How prior-hop sources contribute to the next sampling frontier.
107#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
108#[must_use]
109pub enum PriorSourcesBehavior {
110    #[default]
111    Default,
112    CarryOver,
113    Exclude,
114}
115
116impl PriorSourcesBehavior {
117    fn to_c(self) -> cugraph_prior_sources_behavior_t {
118        match self {
119            Self::Default => cugraph_prior_sources_behavior_t_DEFAULT,
120            Self::CarryOver => cugraph_prior_sources_behavior_t_CARRY_OVER,
121            Self::Exclude => cugraph_prior_sources_behavior_t_EXCLUDE,
122        }
123    }
124}
125
126/// Temporal edge-time comparison used by temporal neighbor sampling.
127#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
128#[must_use]
129pub enum TemporalSamplingComparison {
130    #[default]
131    StrictlyIncreasing,
132    MonotonicallyIncreasing,
133    StrictlyDecreasing,
134    MonotonicallyDecreasing,
135    Last,
136}
137
138impl TemporalSamplingComparison {
139    fn to_c(self) -> cugraph_temporal_sampling_comparison_t {
140        match self {
141            Self::StrictlyIncreasing => cugraph_temporal_sampling_comparison_t_STRICTLY_INCREASING,
142            Self::MonotonicallyIncreasing => cugraph_temporal_sampling_comparison_t_MONOTONICALLY_INCREASING,
143            Self::StrictlyDecreasing => cugraph_temporal_sampling_comparison_t_STRICTLY_DECREASING,
144            Self::MonotonicallyDecreasing => cugraph_temporal_sampling_comparison_t_MONOTONICALLY_DECREASING,
145            Self::Last => cugraph_temporal_sampling_comparison_t_LAST,
146        }
147    }
148}
149
150/// Options shared by homogeneous neighbor-sampling APIs.
151#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
152#[must_use]
153pub struct SamplingOptions {
154    pub retain_seeds: bool,
155    pub renumber_results: bool,
156    pub compress_per_hop: bool,
157    pub with_replacement: bool,
158    pub return_hops: bool,
159    pub compression: SamplingCompression,
160    pub prior_sources_behavior: PriorSourcesBehavior,
161    pub dedupe_sources: bool,
162    pub temporal_sampling_comparison: TemporalSamplingComparison,
163    pub disjoint_sampling: bool,
164}
165
166/// Parameters for homogeneous neighbor sampling.
167#[derive(Debug, Clone, PartialEq, Eq)]
168#[must_use]
169pub struct NeighborSamplingParams {
170    /// Per-hop fanout values. cuGraph requires `i32` fanout entries.
171    pub fan_out: Vec<i32>,
172    pub options: SamplingOptions,
173    pub do_expensive_check: bool,
174}
175
176impl NeighborSamplingParams {
177    pub fn new(fan_out: Vec<i32>) -> Self {
178        Self { fan_out, options: SamplingOptions::default(), do_expensive_check: false }
179    }
180}
181
182impl Default for NeighborSamplingParams {
183    fn default() -> Self {
184        Self::new(vec![1])
185    }
186}
187
188struct NativeSamplingOptions {
189    raw: *mut cugraph_sampling_options_t,
190    _nosend_nosync: std::marker::PhantomData<*mut ()>,
191}
192
193impl NativeSamplingOptions {
194    fn new(options: &SamplingOptions) -> Result<Self> {
195        let mut raw: *mut cugraph_sampling_options_t = ptr::null_mut();
196        let mut error: *mut cugraph_error_t = ptr::null_mut();
197
198        // SAFETY: `raw` and `error` are valid out-pointers.
199        unsafe {
200            let code = cugraph_sampling_options_create(&mut raw, &mut error);
201            check_host_error(code, error, "NativeSamplingOptions::new")?;
202        }
203        let raw = require_non_null(raw, "NativeSamplingOptions::new", "sampling_options")?;
204
205        // SAFETY: `raw` is a live sampling-options object owned by this
206        // wrapper, and setters only mutate that object.
207        unsafe {
208            cugraph_sampling_set_retain_seeds(raw, to_bool_t(options.retain_seeds));
209            cugraph_sampling_set_renumber_results(raw, to_bool_t(options.renumber_results));
210            cugraph_sampling_set_compress_per_hop(raw, to_bool_t(options.compress_per_hop));
211            cugraph_sampling_set_with_replacement(raw, to_bool_t(options.with_replacement));
212            cugraph_sampling_set_return_hops(raw, to_bool_t(options.return_hops));
213            cugraph_sampling_set_compression_type(raw, options.compression.to_c());
214            cugraph_sampling_set_prior_sources_behavior(raw, options.prior_sources_behavior.to_c());
215            cugraph_sampling_set_dedupe_sources(raw, to_bool_t(options.dedupe_sources));
216            cugraph_sampling_set_temporal_sampling_comparison(raw, options.temporal_sampling_comparison.to_c());
217            cugraph_sampling_set_disjoint_sampling(raw, to_bool_t(options.disjoint_sampling));
218        }
219
220        Ok(Self { raw, _nosend_nosync: std::marker::PhantomData })
221    }
222
223    fn as_ptr(&self) -> *const cugraph_sampling_options_t {
224        self.raw
225    }
226}
227
228impl Drop for NativeSamplingOptions {
229    fn drop(&mut self) {
230        if !self.raw.is_null() {
231            // SAFETY: `self.raw` is owned by this wrapper and freed exactly
232            // once here. Upstream's sampling-options object contains host
233            // scalar configuration only and does not own a CUDA resource.
234            unsafe { cugraph_sampling_options_free(self.raw) };
235        }
236    }
237}
238
239/// Result of a random-walk computation.
240pub struct RandomWalkResult {
241    raw: *mut cugraph_random_walk_result_t,
242    device: i32,
243    has_weights: bool,
244    // Keeps a runtime-bound cuGraph handle alive until after Drop frees `raw`.
245    _resource_handle_lease: Option<ResourceHandleLease>,
246    _nosend_nosync: std::marker::PhantomData<*mut ()>,
247}
248
249impl std::fmt::Debug for RandomWalkResult {
250    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
251        f.debug_struct("RandomWalkResult")
252            .field("max_path_length", &self.max_path_length())
253            .field("has_weights", &self.has_weights)
254            .finish_non_exhaustive()
255    }
256}
257
258impl RandomWalkResult {
259    /// # Safety
260    ///
261    /// `raw` must be a valid, non-null pointer returned by a cuGraph
262    /// random-walk C API. Ownership is transferred to this struct.
263    pub(crate) unsafe fn from_raw(
264        raw: *mut cugraph_random_walk_result_t,
265        device: i32,
266        has_weights: bool,
267        resource_handle_lease: Option<ResourceHandleLease>,
268    ) -> Result<Self> {
269        let raw = require_non_null(raw, "RandomWalkResult::from_raw", "random_walk_result")?;
270        Ok(Self {
271            raw,
272            device,
273            has_weights,
274            _resource_handle_lease: resource_handle_lease,
275            _nosend_nosync: std::marker::PhantomData,
276        })
277    }
278
279    /// Maximum number of edges requested per walk.
280    pub fn max_path_length(&self) -> usize {
281        // SAFETY: `self.raw` is a live random-walk result handle.
282        unsafe { cugraph_random_walk_result_get_max_path_length(self.raw) }
283    }
284
285    /// Dense row-major vertex path matrix.
286    ///
287    /// The row count is the number of start vertices and each row contains
288    /// `max_path_length() + 1` vertex ids. Native cuGraph uses a negative
289    /// sentinel in later columns when a walk terminates early.
290    pub fn paths(&self) -> Result<VertexIds<'_>> {
291        VertexIds::from_raw(self.raw_view(cugraph_random_walk_result_get_paths, "RandomWalkResult::paths", "paths")?)
292    }
293
294    /// Dense row-major edge-weight matrix, if the input graph had weights.
295    ///
296    /// Each row contains `max_path_length()` edge weights. Unweighted graphs
297    /// produce no native weight column and return `None`.
298    pub fn weights(&self) -> Result<Option<ScoreValues<'_>>> {
299        if !self.has_weights {
300            return Ok(None);
301        }
302
303        let raw = self.raw_view(cugraph_random_walk_result_get_weights, "RandomWalkResult::weights", "weights")?;
304        ScoreValues::from_raw(raw).map(Some)
305    }
306
307    fn raw_view(
308        &self,
309        getter: unsafe extern "C" fn(
310            *mut cugraph_random_walk_result_t,
311        ) -> *mut cugraph_sys::cugraph_type_erased_device_array_view_t,
312        operation: &'static str,
313        pointer_name: &'static str,
314    ) -> Result<RawDeviceArrayView<'_>> {
315        let view_ptr = require_non_null(
316            // SAFETY: `self.raw` is owned by this wrapper and remains live for
317            // the duration of this call.
318            unsafe { getter(self.raw) },
319            operation,
320            pointer_name,
321        )?;
322        // SAFETY: random-walk getters allocate view metadata for the caller.
323        // The backing device arrays are owned by `self.raw`.
324        unsafe { RawDeviceArrayView::from_owned_view(view_ptr, Some(self.device), operation, pointer_name) }
325    }
326}
327
328impl Drop for RandomWalkResult {
329    fn drop(&mut self) {
330        drop_device_resource(
331            &mut self.raw,
332            self.device,
333            "RandomWalkResult::drop",
334            DeviceResourceKind::RandomWalkResult,
335        );
336    }
337}
338
339/// Result of a homogeneous neighbor-sampling computation.
340pub struct SampleResult {
341    raw: *mut cugraph_sample_result_t,
342    device: i32,
343    // Keeps a runtime-bound cuGraph handle alive until after Drop frees `raw`.
344    _resource_handle_lease: Option<ResourceHandleLease>,
345    _nosend_nosync: std::marker::PhantomData<*mut ()>,
346}
347
348impl std::fmt::Debug for SampleResult {
349    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
350        f.debug_struct("SampleResult").finish_non_exhaustive()
351    }
352}
353
354impl SampleResult {
355    /// # Safety
356    ///
357    /// `raw` must be a valid, non-null pointer returned by a cuGraph sampling
358    /// C API. Ownership is transferred to this struct.
359    pub(crate) unsafe fn from_raw(
360        raw: *mut cugraph_sample_result_t,
361        device: i32,
362        resource_handle_lease: Option<ResourceHandleLease>,
363    ) -> Result<Self> {
364        let raw = require_non_null(raw, "SampleResult::from_raw", "sample_result")?;
365        Ok(Self {
366            raw,
367            device,
368            _resource_handle_lease: resource_handle_lease,
369            _nosend_nosync: std::marker::PhantomData,
370        })
371    }
372
373    /// Sampled major vertices, or `None` when the output is compressed.
374    pub fn majors(&self) -> Result<Option<VertexIds<'_>>> {
375        self.optional_raw_view(cugraph_sample_result_get_majors, "SampleResult::majors")
376            .and_then(|raw| raw.map(VertexIds::from_raw).transpose())
377    }
378
379    /// Sampled minor vertices.
380    pub fn minors(&self) -> Result<VertexIds<'_>> {
381        VertexIds::from_raw(self.required_raw_view(
382            cugraph_sample_result_get_minors,
383            "SampleResult::minors",
384            "minors",
385        )?)
386    }
387
388    /// Major offsets for compressed outputs.
389    pub fn major_offsets(&self) -> Result<Option<SizeValues<'_>>> {
390        self.optional_raw_view(cugraph_sample_result_get_major_offsets, "SampleResult::major_offsets")
391            .and_then(|raw| raw.map(SizeValues::from_raw).transpose())
392    }
393
394    /// Optional start labels.
395    pub fn start_labels(&self) -> Result<Option<IntegerValues<'_>>> {
396        self.optional_raw_view(cugraph_sample_result_get_start_labels, "SampleResult::start_labels")
397            .and_then(|raw| raw.map(IntegerValues::from_raw).transpose())
398    }
399
400    /// Optional sampled edge ids.
401    pub fn edge_ids(&self) -> Result<Option<IntegerValues<'_>>> {
402        self.optional_raw_view(cugraph_sample_result_get_edge_id, "SampleResult::edge_ids")
403            .and_then(|raw| raw.map(IntegerValues::from_raw).transpose())
404    }
405
406    /// Optional sampled edge type ids.
407    pub fn edge_types(&self) -> Result<Option<IntegerValues<'_>>> {
408        self.optional_raw_view(cugraph_sample_result_get_edge_type, "SampleResult::edge_types")
409            .and_then(|raw| raw.map(IntegerValues::from_raw).transpose())
410    }
411
412    /// Optional sampled edge weights.
413    pub fn edge_weights(&self) -> Result<Option<ScoreValues<'_>>> {
414        self.optional_raw_view(cugraph_sample_result_get_edge_weight, "SampleResult::edge_weights")
415            .and_then(|raw| raw.map(ScoreValues::from_raw).transpose())
416    }
417
418    /// Optional sampled edge start times.
419    pub fn edge_start_times(&self) -> Result<Option<IntegerValues<'_>>> {
420        self.optional_raw_view(cugraph_sample_result_get_edge_start_time, "SampleResult::edge_start_times")
421            .and_then(|raw| raw.map(IntegerValues::from_raw).transpose())
422    }
423
424    /// Optional sampled edge end times.
425    pub fn edge_end_times(&self) -> Result<Option<IntegerValues<'_>>> {
426        self.optional_raw_view(cugraph_sample_result_get_edge_end_time, "SampleResult::edge_end_times")
427            .and_then(|raw| raw.map(IntegerValues::from_raw).transpose())
428    }
429
430    /// Optional hop ids.
431    pub fn hops(&self) -> Result<Option<IntegerValues<'_>>> {
432        self.optional_raw_view(cugraph_sample_result_get_hop, "SampleResult::hops")
433            .and_then(|raw| raw.map(IntegerValues::from_raw).transpose())
434    }
435
436    /// Optional label-hop offsets.
437    pub fn label_hop_offsets(&self) -> Result<Option<SizeValues<'_>>> {
438        self.optional_raw_view(cugraph_sample_result_get_label_hop_offsets, "SampleResult::label_hop_offsets")
439            .and_then(|raw| raw.map(SizeValues::from_raw).transpose())
440    }
441
442    /// Optional label-type-hop offsets.
443    pub fn label_type_hop_offsets(&self) -> Result<Option<SizeValues<'_>>> {
444        self.optional_raw_view(cugraph_sample_result_get_label_type_hop_offsets, "SampleResult::label_type_hop_offsets")
445            .and_then(|raw| raw.map(SizeValues::from_raw).transpose())
446    }
447
448    /// Optional renumber map.
449    pub fn renumber_map(&self) -> Result<Option<VertexIds<'_>>> {
450        self.optional_raw_view(cugraph_sample_result_get_renumber_map, "SampleResult::renumber_map")
451            .and_then(|raw| raw.map(VertexIds::from_raw).transpose())
452    }
453
454    /// Optional renumber-map offsets.
455    pub fn renumber_map_offsets(&self) -> Result<Option<SizeValues<'_>>> {
456        self.optional_raw_view(cugraph_sample_result_get_renumber_map_offsets, "SampleResult::renumber_map_offsets")
457            .and_then(|raw| raw.map(SizeValues::from_raw).transpose())
458    }
459
460    /// Optional edge-renumber map.
461    pub fn edge_renumber_map(&self) -> Result<Option<IntegerValues<'_>>> {
462        self.optional_raw_view(cugraph_sample_result_get_edge_renumber_map, "SampleResult::edge_renumber_map")
463            .and_then(|raw| raw.map(IntegerValues::from_raw).transpose())
464    }
465
466    /// Optional edge-renumber-map offsets.
467    pub fn edge_renumber_map_offsets(&self) -> Result<Option<SizeValues<'_>>> {
468        self.optional_raw_view(
469            cugraph_sample_result_get_edge_renumber_map_offsets,
470            "SampleResult::edge_renumber_map_offsets",
471        )
472        .and_then(|raw| raw.map(SizeValues::from_raw).transpose())
473    }
474
475    fn optional_raw_view(
476        &self,
477        getter: unsafe extern "C" fn(
478            *const cugraph_sample_result_t,
479        ) -> *mut cugraph_sys::cugraph_type_erased_device_array_view_t,
480        operation: &'static str,
481    ) -> Result<Option<RawDeviceArrayView<'_>>> {
482        // SAFETY: `self.raw` is owned by this wrapper and remains live for
483        // the duration of this call.
484        let view_ptr = unsafe { getter(self.raw) };
485        if view_ptr.is_null() {
486            return Ok(None);
487        }
488
489        // SAFETY: sample-result getters allocate view metadata for the caller.
490        // The backing device arrays are owned by `self.raw`.
491        unsafe { RawDeviceArrayView::from_owned_view(view_ptr, Some(self.device), operation, "view").map(Some) }
492    }
493
494    fn required_raw_view(
495        &self,
496        getter: unsafe extern "C" fn(
497            *const cugraph_sample_result_t,
498        ) -> *mut cugraph_sys::cugraph_type_erased_device_array_view_t,
499        operation: &'static str,
500        pointer_name: &'static str,
501    ) -> Result<RawDeviceArrayView<'_>> {
502        self.optional_raw_view(getter, operation)?.ok_or_else(|| {
503            Error::new(ErrorKind::InvalidState, operation)
504                .with_message(format!("native sample result returned null {pointer_name}"))
505        })
506    }
507}
508
509impl Drop for SampleResult {
510    fn drop(&mut self) {
511        drop_device_resource(&mut self.raw, self.device, "SampleResult::drop", DeviceResourceKind::SampleResult);
512    }
513}
514
515type RandomWalkFn = unsafe extern "C" fn(
516    *const cugraph_sys::cugraph_resource_handle_t,
517    *mut cugraph_sys::cugraph_rng_state_t,
518    *mut cugraph_sys::cugraph_graph_t,
519    *const cugraph_sys::cugraph_type_erased_device_array_view_t,
520    usize,
521    *mut *mut cugraph_random_walk_result_t,
522    *mut *mut cugraph_error_t,
523) -> cugraph_sys::cugraph_error_code_t;
524
525pub(crate) fn uniform_random_walks<V: CugraphElement>(
526    handle: &EnteredContext<'_>,
527    rng_state: &mut RngState,
528    graph: &mut Graph<Forward>,
529    start_vertices: &DeviceArrayView<'_, V>,
530    params: &RandomWalkParams,
531) -> Result<RandomWalkResult> {
532    run_random_walks(
533        handle,
534        rng_state,
535        graph,
536        start_vertices,
537        params.max_length.get(),
538        cugraph_uniform_random_walks,
539        "uniform_random_walks",
540    )
541}
542
543pub(crate) fn biased_random_walks<V: CugraphElement>(
544    handle: &EnteredContext<'_>,
545    rng_state: &mut RngState,
546    graph: &mut Graph<Forward>,
547    start_vertices: &DeviceArrayView<'_, V>,
548    params: &RandomWalkParams,
549) -> Result<RandomWalkResult> {
550    run_random_walks(
551        handle,
552        rng_state,
553        graph,
554        start_vertices,
555        params.max_length.get(),
556        cugraph_biased_random_walks,
557        "biased_random_walks",
558    )
559}
560
561fn run_random_walks<V: CugraphElement>(
562    handle: &EnteredContext<'_>,
563    rng_state: &mut RngState,
564    graph: &mut Graph<Forward>,
565    start_vertices: &DeviceArrayView<'_, V>,
566    max_length: usize,
567    function: RandomWalkFn,
568    operation: &'static str,
569) -> Result<RandomWalkResult> {
570    graph.validate_vertex_input_dtype(V::data_type(), operation, "start_vertices")?;
571    let mut result: *mut cugraph_random_walk_result_t = ptr::null_mut();
572    let mut error: *mut cugraph_error_t = ptr::null_mut();
573    let graph_ptr = graph.as_ptr(handle, operation)?;
574
575    // SAFETY: `graph`, `rng_state`, and `start_vertices` are live cuGraph
576    // objects, and `result`/`error` are valid out-pointers.
577    unsafe {
578        let code = function(
579            handle.handle_ptr(),
580            rng_state.as_ptr(),
581            graph_ptr,
582            start_vertices.as_ptr(),
583            max_length,
584            &mut result,
585            &mut error,
586        );
587        check_error(code, error, handle.device(), operation)?;
588    }
589
590    // SAFETY: on success, cuGraph transfers ownership of the result pointer.
591    unsafe { RandomWalkResult::from_raw(result, handle.device(), graph.has_weights(), handle.runtime_bound_lease()) }
592}
593
594pub(crate) fn node2vec_random_walks<V: CugraphElement>(
595    handle: &EnteredContext<'_>,
596    rng_state: &mut RngState,
597    graph: &mut Graph<Forward>,
598    start_vertices: &DeviceArrayView<'_, V>,
599    params: &Node2VecRandomWalkParams,
600) -> Result<RandomWalkResult> {
601    graph.validate_vertex_input_dtype(V::data_type(), "node2vec_random_walks", "start_vertices")?;
602    let mut result: *mut cugraph_random_walk_result_t = ptr::null_mut();
603    let mut error: *mut cugraph_error_t = ptr::null_mut();
604    let graph_ptr = graph.as_ptr(handle, "node2vec_random_walks")?;
605
606    // SAFETY: `graph`, `rng_state`, and `start_vertices` are live cuGraph
607    // objects, and `result`/`error` are valid out-pointers.
608    unsafe {
609        let code = cugraph_sys::cugraph_node2vec_random_walks(
610            handle.handle_ptr(),
611            rng_state.as_ptr(),
612            graph_ptr,
613            start_vertices.as_ptr(),
614            params.max_length.get(),
615            params.p.get(),
616            params.q.get(),
617            &mut result,
618            &mut error,
619        );
620        check_error(code, error, handle.device(), "node2vec_random_walks")?;
621    }
622
623    // SAFETY: on success, cuGraph transfers ownership of the result pointer.
624    unsafe { RandomWalkResult::from_raw(result, handle.device(), graph.has_weights(), handle.runtime_bound_lease()) }
625}
626
627pub(crate) fn select_random_vertices<T, V>(
628    handle: &EnteredContext<'_>,
629    graph: &Graph<T>,
630    rng_state: &mut RngState,
631    num_vertices: usize,
632) -> Result<DeviceArray<V>>
633where
634    V: CugraphElement,
635{
636    graph.validate_vertex_input_dtype(V::data_type(), "select_random_vertices", "vertices")?;
637    let mut vertices: *mut cugraph_type_erased_device_array_t = ptr::null_mut();
638    let mut error: *mut cugraph_error_t = ptr::null_mut();
639    let graph_ptr = graph.as_ptr(handle, "select_random_vertices")?;
640
641    // SAFETY: `graph` and `rng_state` are live cuGraph objects, and
642    // `vertices`/`error` are valid out-pointers.
643    unsafe {
644        let code = cugraph_select_random_vertices(
645            handle.handle_ptr(),
646            graph_ptr,
647            rng_state.as_ptr(),
648            num_vertices,
649            &mut vertices,
650            &mut error,
651        );
652        check_error(code, error, handle.device(), "select_random_vertices")?;
653    }
654
655    // SAFETY: on success, cuGraph transfers ownership of the device array.
656    unsafe {
657        DeviceArray::from_raw(
658            vertices,
659            handle.device(),
660            "select_random_vertices",
661            "vertices",
662            handle.runtime_bound_lease(),
663        )
664    }
665}
666
667fn fan_out_view(params: &NeighborSamplingParams) -> Result<HostArrayView<'_>> {
668    if params.fan_out.is_empty() {
669        return Err(Error::new(ErrorKind::InvalidInput, "NeighborSamplingParams::fan_out")
670            .with_message("fan_out must contain at least one hop"));
671    }
672
673    let byte_len = params.fan_out.len() * std::mem::size_of::<i32>();
674    // SAFETY: `fan_out` is a contiguous `Vec<i32>` and the byte slice is
675    // borrowed for no longer than `params`.
676    let bytes = unsafe { std::slice::from_raw_parts(params.fan_out.as_ptr().cast::<u8>(), byte_len) };
677    HostArrayView::new(bytes, params.fan_out.len(), DataType::Int32)
678}
679
680pub(crate) fn homogeneous_uniform_neighbor_sample<V: CugraphElement>(
681    handle: &EnteredContext<'_>,
682    rng_state: &mut RngState,
683    graph: &mut Graph<Forward>,
684    start_vertices: &DeviceArrayView<'_, V>,
685    starting_vertex_label_offsets: Option<&DeviceArrayView<'_, usize>>,
686    params: &NeighborSamplingParams,
687) -> Result<SampleResult> {
688    graph.validate_vertex_input_dtype(V::data_type(), "homogeneous_uniform_neighbor_sample", "start_vertices")?;
689    let fan_out = fan_out_view(params)?;
690    let options = NativeSamplingOptions::new(&params.options)?;
691    let mut result: *mut cugraph_sample_result_t = ptr::null_mut();
692    let mut error: *mut cugraph_error_t = ptr::null_mut();
693    let label_offsets = starting_vertex_label_offsets.map(|v| v.as_ptr()).unwrap_or(ptr::null_mut());
694    let graph_ptr = graph.as_ptr(handle, "homogeneous_uniform_neighbor_sample")?;
695
696    // SAFETY: all input pointers are either null or live cuGraph view handles,
697    // and `result`/`error` are valid out-pointers.
698    unsafe {
699        let code = cugraph_homogeneous_uniform_neighbor_sample(
700            handle.handle_ptr(),
701            rng_state.as_ptr(),
702            graph_ptr,
703            start_vertices.as_ptr(),
704            label_offsets,
705            fan_out.as_ptr(),
706            options.as_ptr(),
707            to_bool_t(params.do_expensive_check),
708            &mut result,
709            &mut error,
710        );
711        check_error(code, error, handle.device(), "homogeneous_uniform_neighbor_sample")?;
712    }
713
714    // SAFETY: on success, cuGraph transfers ownership of the result pointer.
715    unsafe { SampleResult::from_raw(result, handle.device(), handle.runtime_bound_lease()) }
716}
717
718pub(crate) fn homogeneous_biased_neighbor_sample<V: CugraphElement>(
719    handle: &EnteredContext<'_>,
720    rng_state: &mut RngState,
721    graph: &mut Graph<Forward>,
722    start_vertices: &DeviceArrayView<'_, V>,
723    starting_vertex_label_offsets: Option<&DeviceArrayView<'_, usize>>,
724    params: &NeighborSamplingParams,
725) -> Result<SampleResult> {
726    graph.validate_vertex_input_dtype(V::data_type(), "homogeneous_biased_neighbor_sample", "start_vertices")?;
727    let fan_out = fan_out_view(params)?;
728    let options = NativeSamplingOptions::new(&params.options)?;
729    let mut result: *mut cugraph_sample_result_t = ptr::null_mut();
730    let mut error: *mut cugraph_error_t = ptr::null_mut();
731    let label_offsets = starting_vertex_label_offsets.map(|v| v.as_ptr()).unwrap_or(ptr::null_mut());
732    let graph_ptr = graph.as_ptr(handle, "homogeneous_biased_neighbor_sample")?;
733
734    // SAFETY: all input pointers are either null or live cuGraph view handles,
735    // and `result`/`error` are valid out-pointers. Edge biases are null because
736    // the upstream C constructors for independent edge-property views are not
737    // exposed in this header; native cuGraph then uses graph weights as biases
738    // when present.
739    unsafe {
740        let code = cugraph_homogeneous_biased_neighbor_sample(
741            handle.handle_ptr(),
742            rng_state.as_ptr(),
743            graph_ptr,
744            ptr::null(),
745            start_vertices.as_ptr(),
746            label_offsets,
747            fan_out.as_ptr(),
748            options.as_ptr(),
749            to_bool_t(params.do_expensive_check),
750            &mut result,
751            &mut error,
752        );
753        check_error(code, error, handle.device(), "homogeneous_biased_neighbor_sample")?;
754    }
755
756    // SAFETY: on success, cuGraph transfers ownership of the result pointer.
757    unsafe { SampleResult::from_raw(result, handle.device(), handle.runtime_bound_lease()) }
758}
759
760pub(crate) fn homogeneous_uniform_temporal_neighbor_sample<V, Tm>(
761    handle: &EnteredContext<'_>,
762    rng_state: &mut RngState,
763    graph: &mut Graph<Forward>,
764    temporal_property_name: &str,
765    start_vertices: &DeviceArrayView<'_, V>,
766    starting_vertex_times: Option<&DeviceArrayView<'_, Tm>>,
767    starting_vertex_label_offsets: Option<&DeviceArrayView<'_, usize>>,
768    params: &NeighborSamplingParams,
769) -> Result<SampleResult>
770where
771    V: CugraphElement,
772    Tm: CugraphElement,
773{
774    graph.validate_vertex_input_dtype(
775        V::data_type(),
776        "homogeneous_uniform_temporal_neighbor_sample",
777        "start_vertices",
778    )?;
779    if let Some(times) = starting_vertex_times
780        && times.len() != start_vertices.len()
781    {
782        return Err(Error::new(ErrorKind::InvalidInput, "homogeneous_uniform_temporal_neighbor_sample")
783            .with_message("starting_vertex_times length must match start_vertices length")
784            .with_context("start_vertices_len", start_vertices.len().to_string())
785            .with_context("starting_vertex_times_len", times.len().to_string()));
786    }
787
788    let temporal_property_name = CString::new(temporal_property_name).map_err(|_| {
789        Error::new(ErrorKind::InvalidInput, "homogeneous_uniform_temporal_neighbor_sample")
790            .with_message("temporal_property_name must not contain a NUL byte")
791    })?;
792    let fan_out = fan_out_view(params)?;
793    let options = NativeSamplingOptions::new(&params.options)?;
794    let mut result: *mut cugraph_sample_result_t = ptr::null_mut();
795    let mut error: *mut cugraph_error_t = ptr::null_mut();
796    let times = starting_vertex_times.map(|v| v.as_ptr()).unwrap_or(ptr::null_mut());
797    let label_offsets = starting_vertex_label_offsets.map(|v| v.as_ptr()).unwrap_or(ptr::null_mut());
798    let graph_ptr = graph.as_ptr(handle, "homogeneous_uniform_temporal_neighbor_sample")?;
799
800    // SAFETY: all input pointers are either null or live cuGraph view handles,
801    // `temporal_property_name` is a live NUL-terminated C string, and
802    // `result`/`error` are valid out-pointers.
803    unsafe {
804        let code = cugraph_homogeneous_uniform_temporal_neighbor_sample(
805            handle.handle_ptr(),
806            rng_state.as_ptr(),
807            graph_ptr,
808            temporal_property_name.as_ptr(),
809            start_vertices.as_ptr(),
810            times,
811            label_offsets,
812            fan_out.as_ptr(),
813            options.as_ptr(),
814            to_bool_t(params.do_expensive_check),
815            &mut result,
816            &mut error,
817        );
818        check_error(code, error, handle.device(), "homogeneous_uniform_temporal_neighbor_sample")?;
819    }
820
821    // SAFETY: on success, cuGraph transfers ownership of the result pointer.
822    unsafe { SampleResult::from_raw(result, handle.device(), handle.runtime_bound_lease()) }
823}
824
825pub(crate) fn homogeneous_biased_temporal_neighbor_sample<V, Tm>(
826    handle: &EnteredContext<'_>,
827    rng_state: &mut RngState,
828    graph: &mut Graph<Forward>,
829    temporal_property_name: &str,
830    start_vertices: &DeviceArrayView<'_, V>,
831    starting_vertex_times: Option<&DeviceArrayView<'_, Tm>>,
832    starting_vertex_label_offsets: Option<&DeviceArrayView<'_, usize>>,
833    params: &NeighborSamplingParams,
834) -> Result<SampleResult>
835where
836    V: CugraphElement,
837    Tm: CugraphElement,
838{
839    graph.validate_vertex_input_dtype(
840        V::data_type(),
841        "homogeneous_biased_temporal_neighbor_sample",
842        "start_vertices",
843    )?;
844    if let Some(times) = starting_vertex_times
845        && times.len() != start_vertices.len()
846    {
847        return Err(Error::new(ErrorKind::InvalidInput, "homogeneous_biased_temporal_neighbor_sample")
848            .with_message("starting_vertex_times length must match start_vertices length")
849            .with_context("start_vertices_len", start_vertices.len().to_string())
850            .with_context("starting_vertex_times_len", times.len().to_string()));
851    }
852
853    let temporal_property_name = CString::new(temporal_property_name).map_err(|_| {
854        Error::new(ErrorKind::InvalidInput, "homogeneous_biased_temporal_neighbor_sample")
855            .with_message("temporal_property_name must not contain a NUL byte")
856    })?;
857    let fan_out = fan_out_view(params)?;
858    let options = NativeSamplingOptions::new(&params.options)?;
859    let mut result: *mut cugraph_sample_result_t = ptr::null_mut();
860    let mut error: *mut cugraph_error_t = ptr::null_mut();
861    let times = starting_vertex_times.map(|v| v.as_ptr()).unwrap_or(ptr::null_mut());
862    let label_offsets = starting_vertex_label_offsets.map(|v| v.as_ptr()).unwrap_or(ptr::null_mut());
863    let graph_ptr = graph.as_ptr(handle, "homogeneous_biased_temporal_neighbor_sample")?;
864
865    // SAFETY: all input pointers are either null or live cuGraph view handles,
866    // `temporal_property_name` is a live NUL-terminated C string, and
867    // `result`/`error` are valid out-pointers. Edge biases are null for the
868    // same reason as the non-temporal biased wrapper.
869    unsafe {
870        let code = cugraph_homogeneous_biased_temporal_neighbor_sample(
871            handle.handle_ptr(),
872            rng_state.as_ptr(),
873            graph_ptr,
874            temporal_property_name.as_ptr(),
875            ptr::null(),
876            start_vertices.as_ptr(),
877            times,
878            label_offsets,
879            fan_out.as_ptr(),
880            options.as_ptr(),
881            to_bool_t(params.do_expensive_check),
882            &mut result,
883            &mut error,
884        );
885        check_error(code, error, handle.device(), "homogeneous_biased_temporal_neighbor_sample")?;
886    }
887
888    // SAFETY: on success, cuGraph transfers ownership of the result pointer.
889    unsafe { SampleResult::from_raw(result, handle.device(), handle.runtime_bound_lease()) }
890}
891
892#[cfg(test)]
893mod tests {
894    use std::ptr;
895
896    use super::*;
897    use crate::error::ErrorKind;
898
899    #[test]
900    fn random_walk_result_from_raw_rejects_null() {
901        // SAFETY: this exercises the constructor's null guard only.
902        let err = match unsafe { RandomWalkResult::from_raw(ptr::null_mut(), 0, false, None) } {
903            Err(err) => err,
904            Ok(_) => panic!("null must be rejected"),
905        };
906        assert_eq!(err.kind(), ErrorKind::Unexpected);
907        assert_eq!(err.operation(), "RandomWalkResult::from_raw");
908    }
909
910    #[test]
911    fn sample_result_from_raw_rejects_null() {
912        // SAFETY: this exercises the constructor's null guard only.
913        let err = match unsafe { SampleResult::from_raw(ptr::null_mut(), 0, None) } {
914            Err(err) => err,
915            Ok(_) => panic!("null must be rejected"),
916        };
917        assert_eq!(err.kind(), ErrorKind::Unexpected);
918        assert_eq!(err.operation(), "SampleResult::from_raw");
919    }
920
921    #[test]
922    fn neighbor_sampling_fan_out_rejects_empty() {
923        let params = NeighborSamplingParams::new(Vec::new());
924
925        let err = fan_out_view(&params).expect_err("empty fan_out must be rejected");
926
927        assert_eq!(err.kind(), ErrorKind::InvalidInput);
928        assert_eq!(err.operation(), "NeighborSamplingParams::fan_out");
929    }
930}
cugraph/sampling/mod.rs

cugraph/sampling/
mod.rs
