Skip to main content

ursa_engine/graph/
traversal.rs

1//! The traversals in SQL: `hop`, `shortest_path` and `random_walk`.
2//!
3//! A traversal returns rows of its own (reached pairs, path edges, walk
4//! steps), not one row per node, and some of its parameters are sets of node
5//! ids rather than literals. So the traversals are not in
6//! `ursa_core::catalog`, whose specs, [`ursa_core::catalog::Call`]s and
7//! [`ursa_core::catalog::Evaluator`] describe per-node columns and one-row
8//! statistics; they are specified here, with the same names, defaults, checks
9//! and output columns as the Python library's `ur.hop(edges, n,
10//! direction).from_(seeds)`, `ur.shortest_path(edges, source, target, weight,
11//! direction)` and `ur.random_walk(edges, start, steps, walks_per_node,
12//! seed)` (`python/ursa/_graph.py`, `ursa-plan`'s `result.rs`). The parity
13//! fixture records the Python signatures and outputs, and
14//! `tests/parity.rs` holds these specs to them.
15//!
16//! The literal parameters use the catalog's [`ParamType`]s and are checked by
17//! the catalog's rules (a count is a non-negative integer below 2^32, a seed a
18//! non-negative integer or NULL, a direction one of the listed names). The
19//! node-id parameters are planned by the relation planner as child plans of
20//! the traversal node (see `relation`), and their values are cast to the
21//! graph's id type: literals as they are planned, a subquery's values when
22//! the node runs ([`not_node_ids`] words the error for both).
23//!
24//! The kernels are `ursa_core::algo`'s, called exactly as `ursa-plan` calls
25//! them, so the rows are the Python library's, in the same order. One
26//! deliberate difference: SQL's `hop` names its columns `(seed, id,
27//! distance)`, where Python's `EdgeFrame` has `(src, dst)` and no distance
28//! (#328); the parity test compares `seed` and `id` with `src` and `dst`.
29
30use 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/// What a traversal parameter takes.
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum TraversalParamType {
51    /// Node ids: a list literal (`[0, 1]`), a scalar subquery returning one
52    /// column (`(SELECT id FROM ...)`), or one id.
53    Nodes,
54    /// One node id: a literal or a scalar subquery returning at most one row.
55    Node,
56    /// A literal or a weight expression, as a catalog function's parameter.
57    Param(ParamType),
58}
59
60/// One parameter of a traversal, besides the graph.
61#[derive(Debug, Clone, Copy)]
62pub struct TraversalParam {
63    pub name: &'static str,
64    pub ty: TraversalParamType,
65    /// The value when the argument is not given; `None` if it must be given.
66    pub default: Option<Value>,
67    /// One line on what it does.
68    pub doc: &'static str,
69}
70
71impl TraversalParam {
72    /// Whether the parameter is a node id or a set of them.
73    pub fn takes_nodes(&self) -> bool {
74        matches!(
75            self.ty,
76            TraversalParamType::Nodes | TraversalParamType::Node
77        )
78    }
79
80    /// Whether the parameter is the weight expression.
81    pub fn is_weight(&self) -> bool {
82        self.ty == TraversalParamType::Param(ParamType::Weight)
83    }
84
85    /// The parameter as a signature shows it: its default, or a placeholder
86    /// for a required one.
87    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/// The type of a traversal's output column.
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99pub enum OutputType {
100    /// A node id, of the graph's id type (`Int64` or `Utf8`).
101    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/// One output column of a traversal. None is nullable.
117#[derive(Debug, Clone, Copy)]
118pub struct TraversalOutput {
119    pub name: &'static str,
120    pub ty: OutputType,
121}
122
123/// One traversal.
124#[derive(Debug)]
125pub struct TraversalSpec {
126    pub name: &'static str,
127    /// One line on what it returns.
128    pub summary: &'static str,
129    /// The NetworkX call it matches, as `docs/semantics.md` pins it.
130    pub networkx: Option<&'static str>,
131    /// The parameters besides the graph, in SQL's order: the node ids first.
132    pub params: &'static [TraversalParam],
133    pub outputs: &'static [TraversalOutput],
134}
135
136impl TraversalSpec {
137    /// The parameter called `name`.
138    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    /// The parameters that take node ids, in order: each is one child plan of
144    /// the traversal node.
145    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    /// The directions the traversal accepts; empty for one without a
151    /// `direction` parameter.
152    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    /// Whether the traversal takes a `weight`.
163    pub fn weighted(&self) -> bool {
164        self.params.iter().any(TraversalParam::is_weight)
165    }
166
167    /// The output schema over a graph whose node ids are `id_type`.
168    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    /// `hop(<graph>, seeds => [<id>, ...], n => 1, direction => 'out')`: the
178    /// call with every parameter at its default, for messages.
179    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
215/// Every traversal, in the Python API's order.
216pub 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
323/// The traversal called `name`.
324pub fn lookup(name: &str) -> Option<&'static TraversalSpec> {
325    TRAVERSALS.iter().find(|t| t.name == name)
326}
327
328/// A validated traversal call: which traversal, with every literal parameter
329/// set. Its node ids and weight are the node's child plans.
330#[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    /// The literal parameter `name`'s value.
361    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    /// The call as SQL would spell it, for `EXPLAIN`: the node arguments and
379    /// the weight as written, then every literal parameter.
380    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
401/// "hop() has no parameter 'm'; did you mean 'n'? ...".
402pub 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
414/// "random_walk() needs 'steps': random_walk(<graph>, ...)".
415pub 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
424/// "hop(seeds =>) must be node ids of graph 'karate', which are integers; got
425/// 'a'": a node-id argument that is not of the graph's id type and does not
426/// cast to it. `got` is the offending value as SQL writes it, or its type.
427pub 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
455/// Whether values of type `from` may be node ids of a graph whose ids have
456/// type `id_type` (#271).
457///
458/// A node id is an integer or a string, whatever the graph's id type: on a
459/// graph with integer ids a string must then parse exactly as an integer
460/// (`'7'`, checked when the value is cast), and on a graph with string ids an
461/// integer becomes its decimal string, which loses nothing. Every other type
462/// is refused, even where Arrow would cast it: a float (`1.5`, or `2.0`), a
463/// boolean or a date as a seed is almost certainly a mistake, and casting it
464/// would quietly pick a node. NULL is accepted; it matches no node.
465pub 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
474/// Whether `text` is exactly an integer, as a string seed on a graph with
475/// integer ids must be (#271): digits, after an optional sign. Arrow's cast
476/// also takes surrounding whitespace (`' 1'`); a seed does not.
477pub 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
482/// The row of the first value in `values` that the node-id rule (#271)
483/// refuses before any cast: any non-null value of a type node ids may not
484/// have, or a string that is not exactly an integer on a graph with integer
485/// ids. `None` when every value passes (it may still fail the cast: an integer
486/// string out of range).
487pub 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
500/// Check one literal against a parameter, by the catalog's rules.
501fn 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
546/// Resolve a traversal's literal arguments (names normalized and known, none
547/// repeated; node ids and the weight already taken out by the caller) to a
548/// [`Traversal`].
549pub 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
600/// Run `traversal` over a graph. `nodes` holds, per node parameter in order,
601/// its ids as dense indexes (`None` for an id not in the graph or a null);
602/// `weights` is one `f64` per edge row for a weighted shortest path.
603pub 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    // An id not in the graph reaches nothing (#225), as in ursa-plan.
612    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                    // The cost to this edge's end: the settled distance when
644                    // weighted, the edge count otherwise.
645                    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}