1use std::sync::Arc;
31
32use datafusion::arrow::array::{Array, ArrayRef, AsArray, Float64Array, Int64Array};
33use datafusion::arrow::compute::cast;
34use datafusion::arrow::datatypes::{DataType, Field, Schema};
35use datafusion::arrow::record_batch::RecordBatch;
36use datafusion::error::{DataFusionError, Result};
37use ursa_core::algo::{
38 k_hop_distances_masked, random_walk_masked, shortest_path_masked,
39 shortest_path_weighted_with_cost_masked,
40};
41use ursa_core::catalog::{
42 direction_name, parse_direction, Expected, ParamType, Value, ALL_DIRECTIONS,
43};
44use ursa_core::{Direction, IdMap, Topology};
45
46use super::functions::{did_you_mean, expected_text, plan, sql_value, Literal};
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum TraversalParamType {
51 Nodes,
54 Node,
56 Param(ParamType),
58}
59
60#[derive(Debug, Clone, Copy)]
62pub struct TraversalParam {
63 pub name: &'static str,
64 pub ty: TraversalParamType,
65 pub default: Option<Value>,
67 pub doc: &'static str,
69}
70
71impl TraversalParam {
72 pub fn takes_nodes(&self) -> bool {
74 matches!(
75 self.ty,
76 TraversalParamType::Nodes | TraversalParamType::Node
77 )
78 }
79
80 pub fn is_weight(&self) -> bool {
82 self.ty == TraversalParamType::Param(ParamType::Weight)
83 }
84
85 pub fn shown_default(&self) -> String {
88 match (self.default, self.ty) {
89 (Some(v), _) => sql_value(&v),
90 (None, TraversalParamType::Nodes) => "[<id>, ...]".to_string(),
91 (None, TraversalParamType::Node) => "<id>".to_string(),
92 (None, TraversalParamType::Param(_)) => "<integer>".to_string(),
93 }
94 }
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99pub enum OutputType {
100 Id,
102 Int64,
103 Float64,
104}
105
106impl OutputType {
107 pub fn data_type(&self, id_type: &DataType) -> DataType {
108 match self {
109 OutputType::Id => id_type.clone(),
110 OutputType::Int64 => DataType::Int64,
111 OutputType::Float64 => DataType::Float64,
112 }
113 }
114}
115
116#[derive(Debug, Clone, Copy)]
118pub struct TraversalOutput {
119 pub name: &'static str,
120 pub ty: OutputType,
121}
122
123#[derive(Debug)]
125pub struct TraversalSpec {
126 pub name: &'static str,
127 pub summary: &'static str,
129 pub networkx: Option<&'static str>,
131 pub params: &'static [TraversalParam],
133 pub outputs: &'static [TraversalOutput],
134}
135
136impl TraversalSpec {
137 pub fn param(&self, name: &str) -> Option<&'static TraversalParam> {
139 let params: &'static [TraversalParam] = self.params;
140 params.iter().find(|p| p.name == name)
141 }
142
143 pub fn node_params(&self) -> impl Iterator<Item = &'static TraversalParam> {
146 let params: &'static [TraversalParam] = self.params;
147 params.iter().filter(|p| p.takes_nodes())
148 }
149
150 pub fn directions(&self) -> &'static [&'static str] {
153 self.params
154 .iter()
155 .find_map(|p| match p.ty {
156 TraversalParamType::Param(ParamType::Direction(d)) => Some(d),
157 _ => None,
158 })
159 .unwrap_or(&[])
160 }
161
162 pub fn weighted(&self) -> bool {
164 self.params.iter().any(TraversalParam::is_weight)
165 }
166
167 pub fn schema(&self, id_type: &DataType) -> Schema {
169 Schema::new(
170 self.outputs
171 .iter()
172 .map(|o| Field::new(o.name, o.ty.data_type(id_type), false))
173 .collect::<Vec<_>>(),
174 )
175 }
176
177 pub fn signature(&self) -> String {
180 let mut parts = vec!["<graph>".to_string()];
181 parts.extend(
182 self.params
183 .iter()
184 .map(|p| format!("{} => {}", p.name, p.shown_default())),
185 );
186 format!("{}({})", self.name, parts.join(", "))
187 }
188}
189
190const fn param(
191 name: &'static str,
192 ty: TraversalParamType,
193 default: Option<Value>,
194 doc: &'static str,
195) -> TraversalParam {
196 TraversalParam {
197 name,
198 ty,
199 default,
200 doc,
201 }
202}
203
204const fn out(name: &'static str, ty: OutputType) -> TraversalOutput {
205 TraversalOutput { name, ty }
206}
207
208const DIRECTION: TraversalParam = param(
209 "direction",
210 TraversalParamType::Param(ParamType::Direction(ALL_DIRECTIONS)),
211 Some(Value::Str("out")),
212 "Follow out-edges, in-edges, or both (the undirected view).",
213);
214
215pub static TRAVERSALS: &[TraversalSpec] = &[
217 TraversalSpec {
218 name: "hop",
219 summary: "The nodes within n steps of each seed: one (seed, id, distance) row per seed \
220 and node reached, distance the fewest steps from the seed to it. A seed is not \
221 its own row; a repeated seed repeats its rows.",
222 networkx: Some(
223 "nx.single_source_shortest_path_length(G, seed, cutoff=n), the seed left out \
224 (in: on G.reverse(); both: on U)",
225 ),
226 params: &[
227 param(
228 "seeds",
229 TraversalParamType::Nodes,
230 None,
231 "The nodes to start from: a list, a query of one column, or one id.",
232 ),
233 param(
234 "n",
235 TraversalParamType::Param(ParamType::Count),
236 Some(Value::UInt(1)),
237 "How many steps to take.",
238 ),
239 DIRECTION,
240 ],
241 outputs: &[
242 out("seed", OutputType::Id),
243 out("id", OutputType::Id),
244 out("distance", OutputType::Int64),
245 ],
246 },
247 TraversalSpec {
248 name: "shortest_path",
249 summary: "One shortest path from source to target: one row per edge, in order, with its \
250 0-based position (hop) and the cost from source to its end. No path, no rows.",
251 networkx: Some(
252 "nx.shortest_path(G, source, target) (weighted: nx.dijkstra_path(G, source, target, \
253 weight=...))",
254 ),
255 params: &[
256 param(
257 "source",
258 TraversalParamType::Node,
259 None,
260 "The node the path starts at: an id, or a query of one value.",
261 ),
262 param(
263 "target",
264 TraversalParamType::Node,
265 None,
266 "The node the path ends at: an id, or a query of one value.",
267 ),
268 param(
269 "weight",
270 TraversalParamType::Param(ParamType::Weight),
271 Some(Value::Null),
272 "A non-negative weight per edge for the least-cost path; null is the fewest \
273 edges.",
274 ),
275 DIRECTION,
276 ],
277 outputs: &[
278 out("src", OutputType::Id),
279 out("dst", OutputType::Id),
280 out("hop", OutputType::Int64),
281 out("cost", OutputType::Float64),
282 ],
283 },
284 TraversalSpec {
285 name: "random_walk",
286 summary: "Random walks along out-edges from each start node: one (walk_id, step, node) \
287 row per step. A walk stops early at a node with no out-edges.",
288 networkx: None,
289 params: &[
290 param(
291 "start",
292 TraversalParamType::Nodes,
293 None,
294 "The nodes to start walks at: a list, a query of one column, or one id.",
295 ),
296 param(
297 "steps",
298 TraversalParamType::Param(ParamType::Count),
299 None,
300 "How many steps each walk takes at most.",
301 ),
302 param(
303 "walks_per_node",
304 TraversalParamType::Param(ParamType::Count),
305 Some(Value::UInt(1)),
306 "How many walks start at each start node.",
307 ),
308 param(
309 "seed",
310 TraversalParamType::Param(ParamType::Seed),
311 Some(Value::Null),
312 "Seed for the deterministic random stream; null is a fixed default.",
313 ),
314 ],
315 outputs: &[
316 out("walk_id", OutputType::Int64),
317 out("step", OutputType::Int64),
318 out("node", OutputType::Id),
319 ],
320 },
321];
322
323pub fn lookup(name: &str) -> Option<&'static TraversalSpec> {
325 TRAVERSALS.iter().find(|t| t.name == name)
326}
327
328#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
331pub enum Traversal {
332 Hop {
333 n: u32,
334 direction: Direction,
335 },
336 ShortestPath {
337 direction: Direction,
338 },
339 RandomWalk {
340 steps: u32,
341 walks_per_node: u32,
342 seed: Option<u64>,
343 },
344}
345
346impl Traversal {
347 pub fn spec(&self) -> &'static TraversalSpec {
348 let name = match self {
349 Traversal::Hop { .. } => "hop",
350 Traversal::ShortestPath { .. } => "shortest_path",
351 Traversal::RandomWalk { .. } => "random_walk",
352 };
353 lookup(name).expect("every traversal has a spec")
354 }
355
356 pub fn name(&self) -> &'static str {
357 self.spec().name
358 }
359
360 fn value(&self, name: &str) -> Option<Value> {
362 let seed = |s: Option<u64>| s.map_or(Value::Null, Value::UInt);
363 Some(match (self, name) {
364 (Traversal::Hop { n, .. }, "n") => Value::UInt(u64::from(*n)),
365 (Traversal::Hop { direction, .. }, "direction")
366 | (Traversal::ShortestPath { direction }, "direction") => {
367 Value::Str(direction_name(*direction))
368 }
369 (Traversal::RandomWalk { steps, .. }, "steps") => Value::UInt(u64::from(*steps)),
370 (Traversal::RandomWalk { walks_per_node, .. }, "walks_per_node") => {
371 Value::UInt(u64::from(*walks_per_node))
372 }
373 (Traversal::RandomWalk { seed: s, .. }, "seed") => seed(*s),
374 _ => return None,
375 })
376 }
377
378 pub fn display(&self, nodes: &[String], weight: Option<&str>) -> String {
381 let mut nodes = nodes.iter();
382 let parts: Vec<String> = self
383 .spec()
384 .params
385 .iter()
386 .filter_map(|p| {
387 if p.takes_nodes() {
388 return nodes.next().map(|n| format!("{} => {n}", p.name));
389 }
390 if p.is_weight() {
391 return weight.map(|w| format!("{} => {w}", p.name));
392 }
393 self.value(p.name)
394 .map(|v| format!("{} => {}", p.name, sql_value(&v)))
395 })
396 .collect();
397 format!("{}({})", self.name(), parts.join(", "))
398 }
399}
400
401pub fn unknown_param(spec: &TraversalSpec, name: &str) -> DataFusionError {
403 let names: Vec<&str> = spec.params.iter().map(|p| p.name).collect();
404 let hint = did_you_mean(name, names.iter().copied())
405 .map(|s| format!("; did you mean '{s}'?"))
406 .unwrap_or_else(|| ".".to_string());
407 plan(format!(
408 "{}() has no parameter '{name}'{hint} Its parameters are: {}",
409 spec.name,
410 names.join(", ")
411 ))
412}
413
414pub fn missing_param(spec: &TraversalSpec, param: &TraversalParam) -> DataFusionError {
416 plan(format!(
417 "{}() needs '{}': {}",
418 spec.name,
419 param.name,
420 spec.signature()
421 ))
422}
423
424pub fn not_node_ids(
428 function: &str,
429 param: &TraversalParam,
430 graph: &str,
431 id_type: &DataType,
432 got: &str,
433) -> String {
434 let (integer, string) = match param.ty {
435 TraversalParamType::Node => (
436 "a node id of graph '{graph}', which is an integer",
437 "a node id of graph '{graph}', which is a string",
438 ),
439 _ => (
440 "node ids of graph '{graph}', which are integers",
441 "node ids of graph '{graph}', which are strings",
442 ),
443 };
444 let expected = match id_type {
445 DataType::Int64 => integer,
446 _ => string,
447 }
448 .replace("{graph}", graph);
449 format!(
450 "{function}({} =>) must be {expected}; got {got}",
451 param.name
452 )
453}
454
455pub fn is_node_id_type(from: &DataType, id_type: &DataType) -> bool {
466 debug_assert!(matches!(id_type, DataType::Int64 | DataType::Utf8));
467 match from {
468 DataType::Null | DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => true,
469 DataType::Dictionary(_, values) => is_node_id_type(values, id_type),
470 t => t.is_integer(),
471 }
472}
473
474pub fn is_integer_text(text: &str) -> bool {
478 let digits = text.strip_prefix(['-', '+']).unwrap_or(text);
479 !digits.is_empty() && digits.bytes().all(|b| b.is_ascii_digit())
480}
481
482pub fn first_refused_node_id(values: &ArrayRef, id_type: &DataType) -> Option<usize> {
488 let valid = |i: &usize| values.is_valid(*i);
489 if !is_node_id_type(values.data_type(), id_type) {
490 return (0..values.len()).find(valid);
491 }
492 if id_type != &DataType::Int64 || values.data_type().is_integer() {
493 return None;
494 }
495 let text = cast(values, &DataType::Utf8).ok()?;
496 let text = text.as_string::<i32>();
497 (0..text.len()).find(|i| valid(i) && !is_integer_text(text.value(*i)))
498}
499
500fn check(spec: &TraversalSpec, name: &str, ty: ParamType, literal: &Literal) -> Result<Value> {
502 let expected = |want: Expected| {
503 plan(format!(
504 "{}({name} =>) must be {}; got {literal}",
505 spec.name,
506 expected_text(want)
507 ))
508 };
509 let int = || match literal {
510 Literal::Number(n) => n.parse::<i128>().ok(),
511 _ => None,
512 };
513 match ty {
514 ParamType::Count => match int() {
515 Some(i) if (0..=i128::from(u32::MAX)).contains(&i) => Ok(Value::UInt(i as u64)),
516 _ => Err(expected(Expected::Count)),
517 },
518 ParamType::Seed => match (literal, int()) {
519 (Literal::Null, _) => Ok(Value::Null),
520 (_, Some(i)) if (0..=i128::from(u64::MAX)).contains(&i) => Ok(Value::UInt(i as u64)),
521 _ => Err(expected(Expected::Seed)),
522 },
523 ParamType::Direction(choices) => match literal {
524 Literal::String(s) if choices.contains(&s.as_str()) => Ok(Value::Str(
525 choices
526 .iter()
527 .find(|c| **c == s.as_str())
528 .expect("checked above"),
529 )),
530 _ => {
531 let options: Vec<String> = choices.iter().map(|o| format!("'{o}'")).collect();
532 Err(plan(format!(
533 "{}({name} => {literal}) is not supported; use one of {}",
534 spec.name,
535 options.join(", ")
536 )))
537 }
538 },
539 other => Err(DataFusionError::Internal(format!(
540 "{}({name}): no check for a parameter of type {other:?}",
541 spec.name
542 ))),
543 }
544}
545
546pub fn resolve(spec: &'static TraversalSpec, literals: &[(String, Literal)]) -> Result<Traversal> {
550 let mut values: Vec<(&'static str, Value)> = Vec::new();
551 for p in spec.params {
552 let TraversalParamType::Param(ty) = p.ty else {
553 continue;
554 };
555 if ty == ParamType::Weight {
556 continue;
557 }
558 let value = match literals.iter().find(|(n, _)| n == p.name) {
559 Some((_, literal)) => check(spec, p.name, ty, literal)?,
560 None => p.default.ok_or_else(|| missing_param(spec, p))?,
561 };
562 values.push((p.name, value));
563 }
564 let get = |name: &str| {
565 values
566 .iter()
567 .find(|(n, _)| *n == name)
568 .map(|(_, v)| *v)
569 .expect("every literal parameter has a value")
570 };
571 let count = |name: &str| match get(name) {
572 Value::UInt(u) => u as u32,
573 other => unreachable!("a count is checked: {other:?}"),
574 };
575 let seed = |name: &str| match get(name) {
576 Value::UInt(u) => Some(u),
577 _ => None,
578 };
579 let direction = |name: &str| match get(name) {
580 Value::Str(s) => parse_direction(s).expect("a direction is checked"),
581 other => unreachable!("a direction is checked: {other:?}"),
582 };
583 Ok(match spec.name {
584 "hop" => Traversal::Hop {
585 n: count("n"),
586 direction: direction("direction"),
587 },
588 "shortest_path" => Traversal::ShortestPath {
589 direction: direction("direction"),
590 },
591 "random_walk" => Traversal::RandomWalk {
592 steps: count("steps"),
593 walks_per_node: count("walks_per_node"),
594 seed: seed("seed"),
595 },
596 other => unreachable!("no traversal {other}"),
597 })
598}
599
600pub fn batch(
604 topology: &Topology,
605 ids: &IdMap,
606 traversal: &Traversal,
607 nodes: &[Vec<Option<u32>>],
608 weights: Option<&[f64]>,
609 schema: Arc<Schema>,
610) -> Result<RecordBatch> {
611 let known = |i: usize| -> Vec<u32> { nodes[i].iter().flatten().copied().collect() };
613 let columns: Vec<ArrayRef> = match *traversal {
614 Traversal::Hop { n, direction } => {
615 let (seed, reached, distance) =
616 k_hop_distances_masked(topology, None, &known(0), n, direction);
617 let distance: Vec<i64> = distance.into_iter().map(i64::from).collect();
618 vec![
619 ids.gather_user(&seed),
620 ids.gather_user(&reached),
621 Arc::new(Int64Array::from(distance)),
622 ]
623 }
624 Traversal::ShortestPath { direction } => {
625 let one = |i: usize| nodes[i].first().copied().flatten();
626 let route = match (one(0), one(1)) {
627 (Some(source), Some(target)) => match weights {
628 Some(w) => shortest_path_weighted_with_cost_masked(
629 topology, None, w, source, target, direction,
630 )
631 .map(|(path, costs)| (path, Some(costs))),
632 None => shortest_path_masked(topology, None, source, target, direction)
633 .map(|path| (path, None)),
634 },
635 _ => None,
636 };
637 let (mut src, mut dst, mut hop, mut cost) = (vec![], vec![], vec![], vec![]);
638 if let Some((path, costs)) = route {
639 for (i, pair) in path.windows(2).enumerate() {
640 src.push(pair[0]);
641 dst.push(pair[1]);
642 hop.push(i as i64);
643 cost.push(match &costs {
646 Some(c) => c[i + 1],
647 None => (i + 1) as f64,
648 });
649 }
650 }
651 vec![
652 ids.gather_user(&src),
653 ids.gather_user(&dst),
654 Arc::new(Int64Array::from(hop)),
655 Arc::new(Float64Array::from(cost)),
656 ]
657 }
658 Traversal::RandomWalk {
659 steps,
660 walks_per_node,
661 seed,
662 } => {
663 let walks = random_walk_masked(topology, None, &known(0), steps, walks_per_node, seed);
664 vec![
665 Arc::new(Int64Array::from(walks.walk_id)),
666 Arc::new(Int64Array::from(walks.step)),
667 ids.gather_user(&walks.node),
668 ]
669 }
670 };
671 RecordBatch::try_new(schema, columns)
672 .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
673}
674
675#[cfg(test)]
676mod tests {
677 use super::*;
678
679 #[test]
680 fn node_id_rule() {
681 for ok in ["0", "7", "-3", "+12", "9223372036854775808"] {
682 assert!(is_integer_text(ok), "{ok}");
683 }
684 for refused in ["", "-", "1.5", "2.0", " 1", "1 ", "1e3", "0x10", "one"] {
685 assert!(!is_integer_text(refused), "{refused}");
686 }
687 for id_type in [DataType::Int64, DataType::Utf8] {
688 for ok in [
689 DataType::Null,
690 DataType::Int8,
691 DataType::UInt64,
692 DataType::Utf8View,
693 ] {
694 assert!(is_node_id_type(&ok, &id_type), "{ok}");
695 }
696 for refused in [DataType::Float64, DataType::Boolean, DataType::Date32] {
697 assert!(!is_node_id_type(&refused, &id_type), "{refused}");
698 }
699 }
700 let strings: ArrayRef = Arc::new(datafusion::arrow::array::StringArray::from(vec![
701 Some("1"),
702 None,
703 Some("2.0"),
704 ]));
705 assert_eq!(first_refused_node_id(&strings, &DataType::Int64), Some(2));
706 assert_eq!(first_refused_node_id(&strings, &DataType::Utf8), None);
707 let floats: ArrayRef = Arc::new(Float64Array::from(vec![None, Some(1.0)]));
708 assert_eq!(first_refused_node_id(&floats, &DataType::Int64), Some(1));
709 }
710
711 fn call(name: &str, args: &[(&str, Literal)]) -> Result<Traversal> {
712 let args: Vec<(String, Literal)> = args
713 .iter()
714 .map(|(n, v)| (n.to_string(), v.clone()))
715 .collect();
716 resolve(lookup(name).unwrap(), &args)
717 }
718
719 fn num(n: &str) -> Literal {
720 Literal::Number(n.into())
721 }
722
723 fn message(r: Result<Traversal>) -> String {
724 crate::user_message(&r.unwrap_err())
725 }
726
727 #[test]
728 fn defaults_match_python() {
729 assert_eq!(
730 call("hop", &[]).unwrap(),
731 Traversal::Hop {
732 n: 1,
733 direction: Direction::Out
734 }
735 );
736 assert_eq!(
737 call("shortest_path", &[]).unwrap(),
738 Traversal::ShortestPath {
739 direction: Direction::Out
740 }
741 );
742 assert_eq!(
743 call("random_walk", &[("steps", num("4"))]).unwrap(),
744 Traversal::RandomWalk {
745 steps: 4,
746 walks_per_node: 1,
747 seed: None
748 }
749 );
750 }
751
752 #[test]
753 fn explains_bad_arguments() {
754 assert_eq!(
755 message(call("hop", &[("n", num("-1"))])),
756 "hop(n =>) must be a non-negative integer below 2^32; got -1"
757 );
758 assert_eq!(
759 message(call("hop", &[("n", num("1.5"))])),
760 "hop(n =>) must be a non-negative integer below 2^32; got 1.5"
761 );
762 assert_eq!(
763 message(call("hop", &[("direction", Literal::String("up".into()))])),
764 "hop(direction => 'up') is not supported; use one of 'out', 'in', 'both'"
765 );
766 assert_eq!(
767 message(call("random_walk", &[])),
768 "random_walk() needs 'steps': random_walk(<graph>, start => [<id>, ...], steps => \
769 <integer>, walks_per_node => 1, seed => NULL)"
770 );
771 assert_eq!(
772 message(call(
773 "random_walk",
774 &[("steps", num("2")), ("seed", Literal::String("x".into()))]
775 )),
776 "random_walk(seed =>) must be a non-negative integer or NULL; got 'x'"
777 );
778 assert_eq!(
779 crate::user_message(&unknown_param(lookup("hop").unwrap(), "m")),
780 "hop() has no parameter 'm'; did you mean 'n'? Its parameters are: seeds, n, \
781 direction"
782 );
783 }
784
785 #[test]
786 fn signatures_show_defaults() {
787 assert_eq!(
788 lookup("hop").unwrap().signature(),
789 "hop(<graph>, seeds => [<id>, ...], n => 1, direction => 'out')"
790 );
791 assert_eq!(
792 lookup("shortest_path").unwrap().signature(),
793 "shortest_path(<graph>, source => <id>, target => <id>, weight => NULL, direction \
794 => 'out')"
795 );
796 }
797
798 #[test]
799 fn explain_spells_the_call() {
800 let hop = call("hop", &[("n", num("2"))]).unwrap();
801 assert_eq!(
802 hop.display(&["[0]".into()], None),
803 "hop(seeds => [0], n => 2, direction => 'out')"
804 );
805 let path = call("shortest_path", &[]).unwrap();
806 assert_eq!(
807 path.display(&["0".into(), "33".into()], Some("weight")),
808 "shortest_path(source => 0, target => 33, weight => weight, direction => 'out')"
809 );
810 }
811}