Skip to main content

cugraph/structure/
selection.rs

1use crate::{
2    array::{DeviceArray, DeviceArrayView},
3    error::Result,
4    handle::EnteredContext,
5    types::CugraphElement,
6};
7
8/// Unified public vertex-selection input used across graph algorithms.
9///
10/// This type covers three distinct semantics:
11///
12/// - [`All`](Self::All) for algorithms that operate over every graph vertex
13/// - [`Device`](Self::Device) for an explicit device-resident subset
14/// - [`Host`](Self::Host) for a host-resident subset uploaded at the API boundary
15///
16/// # Examples
17///
18/// ```no_run
19/// use cugraph::{
20///     ResourceHandle, cudf_interop::device_array_view_from_slice,
21///     structure::VertexSelection,
22/// };
23///
24/// let runtime = std::sync::Arc::new(cudf::Runtime::new().unwrap());
25/// let handle =
26///     ResourceHandle::for_runtime(std::sync::Arc::clone(&runtime)).unwrap();
27///
28/// let device_vertices = runtime.series(&[0i32, 1]).unwrap();
29/// let column = device_vertices.as_column();
30/// let slice = column.device_slice::<i32>().unwrap();
31/// let device_view = device_array_view_from_slice::<i32>(&slice).unwrap();
32/// let selection = VertexSelection::from(&device_view);
33///
34/// let host_selection: VertexSelection<'_, i32> =
35///     VertexSelection::from(&[0i32, 1][..]);
36/// let all: VertexSelection<'_, i32> = VertexSelection::All;
37/// # let _ = (selection, host_selection, all);
38/// ```
39#[derive(Debug, Clone, Copy, Default)]
40pub enum VertexSelection<'a, V: CugraphElement = i32> {
41    /// Select all graph vertices.
42    #[default]
43    All,
44    /// Select an explicit subset that already lives on the GPU.
45    Device(&'a DeviceArrayView<'a, V>),
46    /// Select an explicit subset from host memory, uploaded on demand.
47    Host(&'a [V]),
48}
49
50impl<'a, V: CugraphElement> From<&'a DeviceArrayView<'a, V>> for VertexSelection<'a, V> {
51    fn from(view: &'a DeviceArrayView<'a, V>) -> Self {
52        Self::Device(view)
53    }
54}
55
56impl<'a, V: CugraphElement> From<&'a [V]> for VertexSelection<'a, V> {
57    fn from(slice: &'a [V]) -> Self {
58        Self::Host(slice)
59    }
60}
61
62impl<'a, V: CugraphElement> VertexSelection<'a, V> {
63    pub(crate) fn resolve(self, handle: &EnteredContext<'_>) -> Result<ResolvedVertexSelection<'a, V>> {
64        match self {
65            Self::All => Ok(ResolvedVertexSelection::All),
66            Self::Device(view) => Ok(ResolvedVertexSelection::Borrowed(view)),
67            Self::Host(slice) => {
68                let array = DeviceArray::from_slice(handle, slice)?;
69                Ok(ResolvedVertexSelection::Owned(array))
70            },
71        }
72    }
73}
74
75pub(crate) enum ResolvedVertexSelection<'a, V: CugraphElement> {
76    All,
77    Borrowed(&'a DeviceArrayView<'a, V>),
78    Owned(DeviceArray<V>),
79}
80
81impl<V: CugraphElement> ResolvedVertexSelection<'_, V> {
82    pub(crate) fn view(&self) -> Result<Option<DeviceArrayView<'_, V>>> {
83        match self {
84            Self::All => Ok(None),
85            Self::Borrowed(view) => Ok(Some(DeviceArrayView::from_raw_unchecked(view.raw_view()))),
86            Self::Owned(array) => array.view().map(Some),
87        }
88    }
89}