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}