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