1use std::{
4 cell::{Cell, RefCell},
5 marker::PhantomData,
6 ptr,
7};
8
9use cugraph_sys::{
10 cugraph_error_t, cugraph_graph_create_with_times_sg, cugraph_graph_t, cugraph_type_erased_device_array_view_t,
11 cugraph_type_erased_device_array_view_type,
12};
13
14use super::{Graph, Transposition};
15use crate::{
16 array::{DeviceArrayView, RawDeviceArrayView},
17 error::{Error, ErrorKind, Result, check_error, require_non_null},
18 handle::EnteredContext,
19 types::{CugraphElement, DataType, GraphProperties, to_bool_t},
20};
21
22#[must_use]
60#[derive(Debug)]
61pub struct GraphBuilder<'a> {
62 properties: GraphProperties,
63 src: *mut cugraph_type_erased_device_array_view_t,
64 src_len: usize,
65 dst: *mut cugraph_type_erased_device_array_view_t,
66 dst_len: usize,
67 weights: Option<(*mut cugraph_type_erased_device_array_view_t, usize)>,
68 edge_ids: Option<(*mut cugraph_type_erased_device_array_view_t, usize)>,
69 edge_type_ids: Option<(*mut cugraph_type_erased_device_array_view_t, usize)>,
70 vertices: Option<(*mut cugraph_type_erased_device_array_view_t, usize)>,
71 edge_start_time_ids: Option<(*mut cugraph_type_erased_device_array_view_t, usize)>,
72 edge_end_time_ids: Option<(*mut cugraph_type_erased_device_array_view_t, usize)>,
73 pub(super) vertex_dtype: DataType,
74 pub(super) weight_dtype: DataType,
75 renumber: bool,
76 drop_self_loops: bool,
77 drop_multi_edges: bool,
78 symmetrize: bool,
79 do_expensive_check: bool,
80 _lifetime: PhantomData<&'a ()>,
81}
82
83impl Clone for GraphBuilder<'_> {
84 fn clone(&self) -> Self {
85 *self
86 }
87}
88
89impl Copy for GraphBuilder<'_> {}
90
91impl<'a> GraphBuilder<'a> {
92 pub fn new<V: CugraphElement>(src: &'a DeviceArrayView<'a, V>, dst: &'a DeviceArrayView<'a, V>) -> Self {
116 Self {
117 properties: GraphProperties::default(),
118 src: src.as_ptr(),
119 src_len: src.len(),
120 dst: dst.as_ptr(),
121 dst_len: dst.len(),
122 weights: None,
123 edge_ids: None,
124 edge_type_ids: None,
125 vertices: None,
126 edge_start_time_ids: None,
127 edge_end_time_ids: None,
128 vertex_dtype: V::data_type(),
129 weight_dtype: DataType::Float32,
130 renumber: false,
131 drop_self_loops: false,
132 drop_multi_edges: false,
133 symmetrize: false,
134 do_expensive_check: false,
135 _lifetime: PhantomData,
136 }
137 }
138
139 pub fn properties(mut self, properties: GraphProperties) -> Self {
143 self.properties = properties;
144 self
145 }
146
147 pub fn weights<W: CugraphElement>(mut self, weights: &'a DeviceArrayView<'a, W>) -> Self {
151 self.weights = Some((weights.as_ptr(), weights.len()));
152 self.weight_dtype = W::data_type();
153 self
154 }
155
156 pub fn edge_ids<E: CugraphElement>(mut self, edge_ids: &'a DeviceArrayView<'a, E>) -> Self {
160 self.edge_ids = Some((edge_ids.as_ptr(), edge_ids.len()));
161 self
162 }
163
164 pub fn edge_type_ids<E: CugraphElement>(mut self, edge_type_ids: &'a DeviceArrayView<'a, E>) -> Self {
168 self.edge_type_ids = Some((edge_type_ids.as_ptr(), edge_type_ids.len()));
169 self
170 }
171
172 pub fn vertices<V: CugraphElement>(mut self, vertices: &'a DeviceArrayView<'a, V>) -> Self {
176 self.vertices = Some((vertices.as_ptr(), vertices.len()));
177 self
178 }
179
180 pub fn edge_start_time_ids<E: CugraphElement>(mut self, times: &'a DeviceArrayView<'a, E>) -> Self {
184 self.edge_start_time_ids = Some((times.as_ptr(), times.len()));
185 self
186 }
187
188 pub fn edge_end_time_ids<E: CugraphElement>(mut self, times: &'a DeviceArrayView<'a, E>) -> Self {
192 self.edge_end_time_ids = Some((times.as_ptr(), times.len()));
193 self
194 }
195
196 pub fn renumber(mut self) -> Self {
200 self.renumber = true;
201 self
202 }
203
204 pub fn drop_self_loops(mut self) -> Self {
206 self.drop_self_loops = true;
207 self
208 }
209
210 pub fn drop_multi_edges(mut self) -> Self {
212 self.drop_multi_edges = true;
213 self
214 }
215
216 pub fn symmetrize(mut self) -> Self {
221 self.symmetrize = true;
222 self
223 }
224
225 pub fn do_expensive_check(mut self) -> Self {
227 self.do_expensive_check = true;
228 self
229 }
230
231 pub fn build<T: Transposition>(&self, handle: &EnteredContext<'_>) -> Result<Graph<T>> {
246 self.build_inner(handle, T::store_transposed())
247 }
248
249 pub(crate) fn build_inner<T>(&self, handle: &EnteredContext<'_>, store_transposed: bool) -> Result<Graph<T>> {
251 self.validate()?;
252
253 let edge_id_dtype = self.edge_ids.map(|(ptr, _)| Self::view_dtype(ptr)).transpose()?;
254
255 let c_props = self.properties.into();
256 let vertices_ptr = self.vertices.map(|(p, _)| p as *const _).unwrap_or(ptr::null());
257 let weights_ptr = self.weights.map(|(p, _)| p as *const _).unwrap_or(ptr::null());
258 let edge_ids_ptr = self.edge_ids.map(|(p, _)| p as *const _).unwrap_or(ptr::null());
259 let edge_type_ids_ptr = self.edge_type_ids.map(|(p, _)| p as *const _).unwrap_or(ptr::null());
260 let start_times_ptr = self.edge_start_time_ids.map(|(p, _)| p as *const _).unwrap_or(ptr::null());
261 let end_times_ptr = self.edge_end_time_ids.map(|(p, _)| p as *const _).unwrap_or(ptr::null());
262
263 let mut graph: *mut cugraph_graph_t = ptr::null_mut();
264 let mut error: *mut cugraph_error_t = ptr::null_mut();
265
266 unsafe {
270 let code = cugraph_graph_create_with_times_sg(
271 handle.handle_ptr(),
272 &c_props,
273 vertices_ptr,
274 self.src,
275 self.dst,
276 weights_ptr,
277 edge_ids_ptr,
278 edge_type_ids_ptr,
279 start_times_ptr,
280 end_times_ptr,
281 to_bool_t(store_transposed),
282 to_bool_t(self.renumber),
283 to_bool_t(self.drop_self_loops),
284 to_bool_t(self.drop_multi_edges),
285 to_bool_t(self.symmetrize),
286 to_bool_t(self.do_expensive_check),
287 &mut graph,
288 &mut error,
289 );
290 check_error(code, error, handle.device(), "GraphBuilder::build")?;
291 }
292
293 let graph = require_non_null(graph, "GraphBuilder::build", "graph")?;
294
295 Ok(Graph {
296 raw: graph,
297 device: handle.device(),
298 properties: self.properties,
299 vertex_dtype: self.vertex_dtype,
300 weight_dtype: self.weight_dtype,
301 has_weights: self.weights.is_some(),
302 edge_id_dtype,
303 pending_input_completion: None,
304 resource_handle_leases: RefCell::new(handle.runtime_bound_lease().into_iter().collect()),
305 ordered_binding: Cell::new(handle.runtime_bound_lease().map(|lease| lease.runtime().binding_id())),
306 untracked_resource_handle_use: Cell::new(false),
307 _nosend_nosync: PhantomData,
308 _transposition: PhantomData,
309 })
310 }
311
312 pub(crate) fn from_raw_views(
314 properties: GraphProperties,
315 src: &'a RawDeviceArrayView<'a>,
316 dst: &'a RawDeviceArrayView<'a>,
317 ) -> Result<Self> {
318 Ok(Self {
319 properties,
320 src: src.as_ptr(),
321 src_len: src.len(),
322 dst: dst.as_ptr(),
323 dst_len: dst.len(),
324 weights: None,
325 edge_ids: None,
326 edge_type_ids: None,
327 vertices: None,
328 edge_start_time_ids: None,
329 edge_end_time_ids: None,
330 vertex_dtype: src.dtype()?,
331 weight_dtype: DataType::Float32,
332 renumber: false,
333 drop_self_loops: false,
334 drop_multi_edges: false,
335 symmetrize: false,
336 do_expensive_check: false,
337 _lifetime: PhantomData,
338 })
339 }
340
341 pub(crate) fn weights_raw(mut self, raw: &'a RawDeviceArrayView<'a>) -> Result<Self> {
343 self.weights = Some((raw.as_ptr(), raw.len()));
344 self.weight_dtype = raw.dtype()?;
345 Ok(self)
346 }
347
348 pub(crate) fn edge_ids_raw(mut self, raw: &'a RawDeviceArrayView<'a>) -> Result<Self> {
350 self.edge_ids = Some((raw.as_ptr(), raw.len()));
351 Ok(self)
352 }
353
354 pub(crate) fn renumber_flag(mut self, renumber: bool) -> Self {
357 self.renumber = renumber;
358 self
359 }
360
361 fn validate(&self) -> Result<()> {
364 let edge_count = self.src_len;
365
366 if self.dst_len != edge_count {
367 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
368 .with_message("src and dst must have the same number of elements")
369 .with_context("src_len", edge_count.to_string())
370 .with_context("dst_len", self.dst_len.to_string()));
371 }
372
373 Self::check_optional_len(self.weights, edge_count, "weights")?;
374 Self::check_optional_len(self.edge_ids, edge_count, "edge_ids")?;
375 Self::check_optional_len(self.edge_type_ids, edge_count, "edge_type_ids")?;
376 Self::check_optional_len(self.edge_start_time_ids, edge_count, "edge_start_time_ids")?;
377 Self::check_optional_len(self.edge_end_time_ids, edge_count, "edge_end_time_ids")?;
378
379 if self.symmetrize && self.edge_ids.is_some() {
381 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
382 .with_message("symmetrize cannot be combined with edge_ids"));
383 }
384 if self.symmetrize && self.edge_type_ids.is_some() {
385 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
386 .with_message("symmetrize cannot be combined with edge_type_ids"));
387 }
388
389 let has_start = self.edge_start_time_ids.is_some();
391 let has_end = self.edge_end_time_ids.is_some();
392 if has_start != has_end {
393 let missing = if has_start { "edge_end_time_ids" } else { "edge_start_time_ids" };
394 let present = if has_start { "edge_start_time_ids" } else { "edge_end_time_ids" };
395 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build").with_message(format!(
396 "{present} was provided but {missing} was not; both must be provided or both omitted"
397 )));
398 }
399
400 self.validate_dtypes()?;
402
403 Ok(())
404 }
405
406 fn validate_dtypes(&self) -> Result<()> {
410 let src_dtype = Self::view_dtype(self.src)?;
411 let dst_dtype = Self::view_dtype(self.dst)?;
412
413 if src_dtype != dst_dtype {
415 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
416 .with_message("src and dst must have the same dtype")
417 .with_context("src_dtype", src_dtype.to_string())
418 .with_context("dst_dtype", dst_dtype.to_string()));
419 }
420
421 if !src_dtype.is_supported_index_type() {
423 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
424 .with_message(format!("src/dst dtype must be Int32 or Int64, got {src_dtype}")));
425 }
426
427 if let Some((ptr, _)) = self.weights {
429 let dtype = Self::view_dtype(ptr)?;
430 if !dtype.is_float() {
431 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
432 .with_message(format!("weights dtype must be Float32 or Float64, got {dtype}")));
433 }
434 }
435
436 if let Some((ptr, _)) = self.edge_ids {
438 let dtype = Self::view_dtype(ptr)?;
439 if !dtype.is_supported_index_type() {
440 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
441 .with_message(format!("edge_ids dtype must be Int32 or Int64, got {dtype}")));
442 }
443 if dtype != src_dtype {
444 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
445 .with_message(
446 "edge_ids must have the same dtype as src/dst to prevent silent truncation by the C API",
447 )
448 .with_context("edge_ids_dtype", dtype.to_string())
449 .with_context("src_dst_dtype", src_dtype.to_string()));
450 }
451 }
452
453 if let Some((ptr, _)) = self.edge_type_ids {
455 let dtype = Self::view_dtype(ptr)?;
456 if dtype != DataType::Int32 {
457 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
458 .with_message(format!("edge_type_ids dtype must be Int32, got {dtype}")));
459 }
460 }
461
462 if let Some((ptr, _)) = self.edge_start_time_ids {
464 let dtype = Self::view_dtype(ptr)?;
465 if !dtype.is_supported_index_type() {
466 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
467 .with_message(format!("edge_start_time_ids dtype must be Int32 or Int64, got {dtype}")));
468 }
469 }
470 if let Some((ptr, _)) = self.edge_end_time_ids {
471 let dtype = Self::view_dtype(ptr)?;
472 if !dtype.is_supported_index_type() {
473 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
474 .with_message(format!("edge_end_time_ids dtype must be Int32 or Int64, got {dtype}")));
475 }
476 }
477
478 if let (Some((start_ptr, _)), Some((end_ptr, _))) = (self.edge_start_time_ids, self.edge_end_time_ids) {
480 let start_dtype = Self::view_dtype(start_ptr)?;
481 let end_dtype = Self::view_dtype(end_ptr)?;
482 if start_dtype != end_dtype {
483 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
484 .with_message("edge_start_time_ids and edge_end_time_ids must have the same dtype")
485 .with_context("start_dtype", start_dtype.to_string())
486 .with_context("end_dtype", end_dtype.to_string()));
487 }
488 }
489
490 if let Some((ptr, _)) = self.vertices {
496 let dtype = Self::view_dtype(ptr)?;
497 if !dtype.is_supported_index_type() {
498 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
499 .with_message(format!("vertices dtype must be Int32 or Int64, got {dtype}")));
500 }
501 if dtype != src_dtype {
502 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
503 .with_message(
504 "vertices must have the same dtype as src/dst to prevent silent truncation by the C API",
505 )
506 .with_context("vertices_dtype", dtype.to_string())
507 .with_context("src_dst_dtype", src_dtype.to_string()));
508 }
509 }
510
511 Ok(())
512 }
513
514 fn view_dtype(ptr: *mut cugraph_type_erased_device_array_view_t) -> Result<DataType> {
518 let c_type = unsafe { cugraph_type_erased_device_array_view_type(ptr) };
521 DataType::try_from(c_type).map_err(|err| err.with_operation("view_dtype"))
522 }
523
524 fn check_optional_len(
526 array: Option<(*mut cugraph_type_erased_device_array_view_t, usize)>,
527 expected: usize,
528 name: &'static str,
529 ) -> Result<()> {
530 if let Some((_, len)) = array
531 && len != expected
532 {
533 return Err(Error::new(ErrorKind::InvalidInput, "GraphBuilder::build")
534 .with_message(format!("{name} must have the same number of elements as src/dst"))
535 .with_context("src_len", expected.to_string())
536 .with_context(name, len.to_string()));
537 }
538 Ok(())
539 }
540}