Skip to main content

ursa_engine/graph/
functions.rs

1//! The graph functions in SQL: which of `ursa_core::catalog`'s functions a
2//! query can call, how a call's SQL arguments become a catalog [`Call`], and
3//! the wording of every message about them.
4//!
5//! The catalog is the single source of truth for names, parameters, defaults,
6//! accepted values and output types, shared with the Python library, so a call
7//! means the same thing in both. This module only adds SQL's choices:
8//!
9//! - **Which functions.** Every algorithm, every statistic and the neighbour
10//!   aggregation. The layouts (`layout_fr`, `layout_random`,
11//!   `layout_circle`) are left out ([`in_sql`]), hidden until the
12//!   visualization work, as they are from the Python library's public API. A
13//!   call fails as an unknown function would, and no message or listing names
14//!   them.
15//! - **Node values.** The neighbour aggregation (`neighbors_agg`) reads one
16//!   value per node. SQL gives them as a query of `(id, value)` rows,
17//!   `values => (SELECT id, age FROM people)` ([`VALUES`], a parameter only
18//!   SQL has), or as an expression over the graph's node table, `attribute =>
19//!   age` (the catalog's `attribute`). Either is planned as a child of the
20//!   graph node, as a traversal's node ids are.
21//! - **No `dtype`.** SQL has `CAST`; a call never narrows its own column.
22//! - **`weight =>`** is an expression over the edge table's columns, planned by
23//!   the relation planner as a column of the node's edge plan, not a literal
24//!   resolved here.
25//! - **Messages**, worded for SQL: `pagerank(damping =>) must be a number; got
26//!   'high'`.
27
28use std::fmt;
29
30use datafusion::error::{DataFusionError, Result};
31use ursa_core::catalog::{
32    self, Arg, Call, ColumnType, Expected, Family, FunctionSpec, OutputSpec, ParamError,
33    ParamErrorKind, ParamSpec, ParamType, Shape, Value,
34};
35
36/// The multi-algorithm table function: `algorithms(g, degree(), c => louvain())`.
37pub const ALGORITHMS: &str = "algorithms";
38
39/// The function catalog as a table: `SELECT * FROM ursa_functions()`.
40pub const URSA_FUNCTIONS: &str = "ursa_functions";
41
42/// Whether SQL can call `spec` (see the module docs). Everything the engine
43/// reads from `ursa_core::catalog` goes through this predicate, so a function
44/// it rejects is unknown to SQL: not callable, not suggested, not listed by
45/// `ursa_functions()`.
46fn in_sql(spec: &FunctionSpec) -> bool {
47    spec.family != Family::Layout
48}
49
50/// The catalog function called `name`, if SQL can call it.
51pub fn lookup(name: &str) -> Option<&'static FunctionSpec> {
52    catalog::lookup(name).filter(|s| in_sql(s))
53}
54
55/// Every catalog function SQL can call, in catalog order.
56pub fn specs() -> impl Iterator<Item = &'static FunctionSpec> {
57    catalog::FUNCTIONS.iter().filter(|s| in_sql(s))
58}
59
60/// The names of the graph table functions: the catalog's, the traversals,
61/// then `algorithms`.
62pub fn names() -> impl Iterator<Item = &'static str> {
63    specs()
64        .map(|s| s.name)
65        .chain(super::traversal::TRAVERSALS.iter().map(|t| t.name))
66        .chain([ALGORITHMS])
67}
68
69/// `describe`'s count of the edges left out of the graph because an endpoint
70/// is not a key of its vertex table (#311). SQL's own column: a Python graph
71/// has no vertex table, so it would always be 0 there.
72pub const N_DROPPED_EDGES: OutputSpec = OutputSpec {
73    name: "n_dropped_edges",
74    ty: ColumnType::Int64,
75    nullable: false,
76};
77
78/// The columns a call to `spec` returns in SQL, besides `id`: the catalog's,
79/// then for `describe` [`N_DROPPED_EDGES`]. An output's `field` indexes this
80/// list; past the catalog's own, the engine computes it (`kernels`).
81pub fn outputs(spec: &FunctionSpec) -> Vec<OutputSpec> {
82    let mut outputs = spec.outputs.to_vec();
83    if spec.name == "describe" {
84        outputs.push(N_DROPPED_EDGES);
85    }
86    outputs
87}
88
89/// Whether `spec` returns one row per node (an algorithm), as
90/// opposed to one row for the whole graph (a statistic).
91pub fn is_per_node(spec: &FunctionSpec) -> bool {
92    spec.shape == Shape::PerNode
93}
94
95/// `neighbors_agg(values => (SELECT id, value FROM ...))`: the node values as
96/// a query. SQL's alternative to the catalog's `attribute`, which in SQL is an
97/// expression over the graph's node table.
98pub const VALUES: ParamSpec = ParamSpec {
99    name: "values",
100    ty: ParamType::Attribute,
101    default: Value::Null,
102    doc: "A query of two columns, node ids then values: the value of each node. Give this or \
103          attribute.",
104};
105
106/// The SQL description of the catalog's `attribute` parameter.
107const ATTRIBUTE_DOC: &str =
108    "An expression over the graph's node table, the value of each node. Give this or values.";
109
110/// `neighbors_agg`'s `agg` in SQL, where `min` and `max` also take strings,
111/// dates and times.
112const AGG_DOC: &str = "The aggregation; mean and sum need numbers, and min and max take numbers, \
113     strings, dates or times. count and n_unique return integers.";
114
115/// The parameters a SQL call can name: all but `dtype`, in the catalog's
116/// order, after [`VALUES`] for a function that reads node values.
117pub fn sql_params(spec: &FunctionSpec) -> impl Iterator<Item = &'static ParamSpec> {
118    let params: &'static [ParamSpec] = spec.params;
119    let values = spec.reads_attribute().then_some(&VALUES);
120    values
121        .into_iter()
122        .chain(params.iter().filter(|p| p.ty != ParamType::Dtype))
123}
124
125/// A parameter's one-line description, as SQL reads it.
126pub fn sql_doc(p: &ParamSpec) -> &'static str {
127    match (p.ty, p.name) {
128        (ParamType::Attribute, "attribute") => ATTRIBUTE_DOC,
129        (ParamType::Choice(_), "agg") => AGG_DOC,
130        _ => p.doc,
131    }
132}
133
134/// How SQL writes a value of parameter `name`, of type `ty`: `number`,
135/// `string`, `query` ...
136pub fn sql_type(name: &str, ty: ParamType) -> &'static str {
137    match ty {
138        ParamType::Float => "number",
139        ParamType::OptionalFloat => "number or NULL",
140        ParamType::Count => "integer",
141        ParamType::Seed => "integer or NULL",
142        ParamType::Bool => "boolean",
143        ParamType::Choice(_) | ParamType::Direction(_) | ParamType::Dtype => "string",
144        ParamType::Attribute if name == VALUES.name => "query",
145        ParamType::Weight | ParamType::Attribute => "expression",
146    }
147}
148
149/// A literal argument as written: a (possibly signed) number, a string, a
150/// boolean or `NULL`.
151#[derive(Debug, Clone, PartialEq)]
152pub enum Literal {
153    /// The number's text, sign included, so it parses exactly as Python
154    /// parses the same text.
155    Number(String),
156    String(String),
157    Boolean(bool),
158    Null,
159}
160
161impl fmt::Display for Literal {
162    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
163        match self {
164            Literal::Number(n) => write!(f, "{n}"),
165            Literal::String(s) => write!(f, "'{}'", s.replace('\'', "''")),
166            Literal::Boolean(b) => write!(f, "{b}"),
167            Literal::Null => write!(f, "NULL"),
168        }
169    }
170}
171
172impl Literal {
173    /// The catalog's view of the literal: an integer if the text is one, else
174    /// a float if it is one.
175    fn to_arg(&self) -> Arg {
176        match self {
177            Literal::Number(n) => n
178                .parse::<i128>()
179                .map(Arg::Int)
180                .or_else(|_| n.parse::<f64>().map(Arg::Float))
181                .unwrap_or(Arg::Other),
182            Literal::String(s) => Arg::Str(s.clone()),
183            Literal::Boolean(b) => Arg::Bool(*b),
184            Literal::Null => Arg::Null,
185        }
186    }
187}
188
189/// Resolve the literal arguments of a call to `spec` (names already
190/// normalized; `weight` already taken out by the caller) to a [`Call`].
191pub fn resolve(spec: &'static FunctionSpec, args: &[(String, Literal)]) -> Result<Call> {
192    if spec.narrows() {
193        if let Some((name, _)) = args.iter().find(|(n, _)| n == "dtype") {
194            let column = spec.outputs.first().map_or(spec.name, |o| o.name);
195            return Err(plan(format!(
196                "{}() has no parameter '{name}' in SQL; cast its column instead, for example \
197                 CAST({column} AS REAL)",
198                spec.name
199            )));
200        }
201    }
202    spec.resolve(args.iter().map(|(n, l)| (n.as_str(), l.to_arg())))
203        .map(|resolved| resolved.call)
204        .map_err(|e| param_error(spec, &e, args))
205}
206
207/// A catalog [`ParamError`] in SQL's words.
208fn param_error(spec: &FunctionSpec, e: &ParamError, args: &[(String, Literal)]) -> DataFusionError {
209    let (f, p) = (spec.name, e.param.as_str());
210    let given = || {
211        args.iter()
212            .find(|(n, _)| n == p)
213            .map_or(Literal::Null, |(_, l)| l.clone())
214    };
215    match e.kind {
216        ParamErrorKind::Unknown => unknown_param(spec, p),
217        ParamErrorKind::Duplicate => plan(format!("{f}() was given '{p}' more than once")),
218        ParamErrorKind::Expected(want) => plan(format!(
219            "{f}({p} =>) must be {}; got {}",
220            expected_text(want),
221            given()
222        )),
223        ParamErrorKind::NotOneOf(options) => {
224            let options: Vec<String> = options.iter().map(|o| format!("'{o}'")).collect();
225            plan(format!(
226                "{f}({p} => {}) is not supported; use one of {}",
227                given(),
228                options.join(", ")
229            ))
230        }
231        ParamErrorKind::Expression => plan(format!(
232            "{f}({p} =>) takes an expression over the edge table's columns; got {}",
233            given()
234        )),
235    }
236}
237
238/// What a parameter expected, in SQL's words.
239pub(crate) fn expected_text(want: Expected) -> &'static str {
240    match want {
241        Expected::Number => "a number",
242        Expected::FiniteNumber => "a finite number",
243        Expected::Count => "a non-negative integer below 2^32",
244        Expected::Seed => "a non-negative integer or NULL",
245        Expected::NumberOrNull => "a number or NULL",
246        Expected::Bool => "a boolean",
247    }
248}
249
250/// "pagerank() has no parameter 'dampng'; did you mean 'damping'? ...".
251pub fn unknown_param(spec: &FunctionSpec, name: &str) -> DataFusionError {
252    let names: Vec<&str> = sql_params(spec).map(|p| p.name).collect();
253    if names.is_empty() {
254        return plan(format!(
255            "{}() takes no parameters besides the graph; got '{name}'",
256            spec.name
257        ));
258    }
259    let hint = did_you_mean(name, names.iter().copied())
260        .map(|s| format!("; did you mean '{s}'?"))
261        .unwrap_or_else(|| ".".to_string());
262    plan(format!(
263        "{}() has no parameter '{name}'{hint} Its parameters are: {}",
264        spec.name,
265        names.join(", ")
266    ))
267}
268
269/// A catalog value as SQL writes it: `NULL`, `'out'`, `30`, `0.85`, `1e-6`.
270pub fn sql_value(value: &Value) -> String {
271    match value {
272        Value::Null => "NULL".to_string(),
273        Value::Bool(b) => b.to_string(),
274        Value::UInt(u) => u.to_string(),
275        Value::Float(f) => format!("{f:?}"),
276        Value::Str(s) => format!("'{s}'"),
277    }
278}
279
280/// The parameters of `spec` at their defaults: `damping => 0.85, ...`. Node
281/// values are shown as the query they take, and `attribute`, its
282/// alternative, is left out.
283fn defaults(spec: &FunctionSpec) -> Vec<String> {
284    sql_params(spec)
285        .filter_map(|p| match (p.ty, p.name) {
286            (ParamType::Attribute, "values") => Some(format!("{} => {VALUES_QUERY}", p.name)),
287            (ParamType::Attribute, _) => None,
288            _ => Some(format!("{} => {}", p.name, sql_value(&p.default))),
289        })
290        .collect()
291}
292
293/// The form of a `values =>` argument, for messages.
294pub const VALUES_QUERY: &str = "(SELECT id, value FROM ...)";
295
296/// `pagerank(<graph>, damping => 0.85, ...)`: the call with every parameter at
297/// its default, for messages.
298pub fn signature(spec: &FunctionSpec) -> String {
299    let mut parts = vec!["<graph>".to_string()];
300    parts.extend(defaults(spec));
301    format!("{}({})", spec.name, parts.join(", "))
302}
303
304/// `pagerank(damping => 0.85, ...)`: a call as written inside `algorithms()`.
305pub fn nested_signature(spec: &FunctionSpec) -> String {
306    format!("{}({})", spec.name, defaults(spec).join(", "))
307}
308
309/// The call as SQL would spell it, every literal parameter shown, and the
310/// weight when there is one: for `EXPLAIN`.
311pub fn display_call(call: &Call, weight: Option<&str>) -> String {
312    let values = call.values();
313    let parts: Vec<String> = sql_params(call.spec())
314        .filter_map(|p| {
315            if p.ty == ParamType::Weight {
316                return weight.map(|w| format!("{} => {w}", p.name));
317            }
318            values
319                .iter()
320                .find(|(q, _)| q.name == p.name)
321                .map(|(_, v)| format!("{} => {}", p.name, sql_value(v)))
322        })
323        .collect();
324    format!("{}({})", call.name(), parts.join(", "))
325}
326
327/// A plan error.
328pub(crate) fn plan(message: String) -> DataFusionError {
329    DataFusionError::Plan(message)
330}
331
332/// The closest of `candidates` to `name`, if one is close enough to be a typo
333/// (the graph DDL's rule, so every "did you mean" agrees).
334pub fn did_you_mean<'a>(
335    name: &str,
336    candidates: impl IntoIterator<Item = &'a str>,
337) -> Option<&'a str> {
338    super::closest_name(name, candidates)
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use ursa_core::Direction;
345
346    fn call(name: &str, args: &[(&str, Literal)]) -> Result<Call> {
347        let args: Vec<(String, Literal)> = args
348            .iter()
349            .map(|(n, v)| (n.to_string(), v.clone()))
350            .collect();
351        resolve(lookup(name).unwrap(), &args)
352    }
353
354    fn num(n: &str) -> Literal {
355        Literal::Number(n.into())
356    }
357
358    fn s(v: &str) -> Literal {
359        Literal::String(v.into())
360    }
361
362    fn message(r: Result<Call>) -> String {
363        crate::user_message(&r.unwrap_err())
364    }
365
366    #[test]
367    fn defaults_match_python() {
368        assert_eq!(
369            call("pagerank", &[]).unwrap(),
370            Call::PageRank {
371                damping: 0.85,
372                max_iter: 30,
373                tol: 1e-6,
374                direction: Direction::Out
375            }
376        );
377        assert_eq!(
378            call("degree", &[]).unwrap(),
379            Call::Degree {
380                direction: Direction::Out
381            }
382        );
383        assert_eq!(
384            call("connected_components", &[]).unwrap(),
385            Call::ConnectedComponents { strong: false }
386        );
387        assert_eq!(
388            call("louvain", &[]).unwrap(),
389            Call::Louvain {
390                resolution: 1.0,
391                seed: None
392            }
393        );
394        assert_eq!(call("triangle_count", &[]).unwrap(), Call::TriangleCount);
395        assert_eq!(
396            call("diameter", &[]).unwrap(),
397            Call::Diameter {
398                approximate: true,
399                direction: Direction::Out
400            }
401        );
402    }
403
404    #[test]
405    fn reads_arguments() {
406        assert_eq!(
407            call(
408                "pagerank",
409                &[
410                    ("damping", num("0.9")),
411                    ("max_iter", num("100")),
412                    ("tol", num("1e-10")),
413                    ("direction", s("both"))
414                ]
415            )
416            .unwrap(),
417            Call::PageRank {
418                damping: 0.9,
419                max_iter: 100,
420                tol: 1e-10,
421                direction: Direction::Both
422            }
423        );
424        // An integer is a number, as Python's float parameters take an int.
425        assert!(call("pagerank", &[("damping", num("1"))]).is_ok());
426        assert_eq!(
427            call(
428                "louvain",
429                &[("seed", num("7")), ("resolution", num("-0.5"))]
430            )
431            .unwrap(),
432            Call::Louvain {
433                resolution: -0.5,
434                seed: Some(7)
435            }
436        );
437        assert_eq!(
438            call("louvain", &[("seed", Literal::Null)]).unwrap(),
439            call("louvain", &[]).unwrap()
440        );
441        assert_eq!(
442            call("connected_components", &[("mode", s("strong"))]).unwrap(),
443            Call::ConnectedComponents { strong: true }
444        );
445        assert_eq!(
446            call("diameter", &[("approximate", Literal::Boolean(false))]).unwrap(),
447            Call::Diameter {
448                approximate: false,
449                direction: Direction::Out
450            }
451        );
452        assert_eq!(
453            call("betweenness", &[("sample", Literal::Null)]).unwrap(),
454            call("betweenness", &[]).unwrap()
455        );
456    }
457
458    #[test]
459    fn explains_bad_arguments() {
460        assert_eq!(
461            message(call("pagerank", &[("dampng", num("0.9"))])),
462            "pagerank() has no parameter 'dampng'; did you mean 'damping'? Its parameters are: \
463             damping, max_iter, tol, weight, direction"
464        );
465        assert_eq!(
466            message(call("pagerank", &[("color", num("1"))])),
467            "pagerank() has no parameter 'color'. Its parameters are: damping, max_iter, tol, \
468             weight, direction"
469        );
470        assert_eq!(
471            message(call("triangle_count", &[("direction", s("both"))])),
472            "triangle_count() takes no parameters besides the graph; got 'direction'"
473        );
474        assert_eq!(
475            message(call(
476                "pagerank",
477                &[("damping", num("0.9")), ("damping", num("0.8"))]
478            )),
479            "pagerank() was given 'damping' more than once"
480        );
481        assert_eq!(
482            message(call("pagerank", &[("damping", s("high"))])),
483            "pagerank(damping =>) must be a number; got 'high'"
484        );
485        assert_eq!(
486            message(call("pagerank", &[("damping", num("1e999"))])),
487            "pagerank(damping =>) must be a finite number; got 1e999"
488        );
489        assert_eq!(
490            message(call("pagerank", &[("max_iter", num("-1"))])),
491            "pagerank(max_iter =>) must be a non-negative integer below 2^32; got -1"
492        );
493        assert_eq!(
494            message(call("pagerank", &[("max_iter", num("2.5"))])),
495            "pagerank(max_iter =>) must be a non-negative integer below 2^32; got 2.5"
496        );
497        assert_eq!(
498            message(call("pagerank", &[("max_iter", num("4294967296"))])),
499            "pagerank(max_iter =>) must be a non-negative integer below 2^32; got 4294967296"
500        );
501        assert_eq!(
502            message(call("pagerank", &[("max_iter", Literal::Null)])),
503            "pagerank(max_iter =>) must be a non-negative integer below 2^32; got NULL"
504        );
505        assert_eq!(
506            message(call("louvain", &[("seed", num("-3"))])),
507            "louvain(seed =>) must be a non-negative integer or NULL; got -3"
508        );
509        assert_eq!(
510            message(call("betweenness", &[("sample", s("half"))])),
511            "betweenness(sample =>) must be a number or NULL; got 'half'"
512        );
513        assert_eq!(
514            message(call("diameter", &[("approximate", num("1"))])),
515            "diameter(approximate =>) must be a boolean; got 1"
516        );
517        assert_eq!(
518            message(call("pagerank", &[("direction", s("in"))])),
519            "pagerank(direction => 'in') is not supported; use one of 'out', 'both'"
520        );
521        assert_eq!(
522            message(call("degree", &[("direction", s("sideways"))])),
523            "degree(direction => 'sideways') is not supported; use one of 'out', 'in', 'both'"
524        );
525        assert_eq!(
526            message(call("connected_components", &[("mode", num("1"))])),
527            "connected_components(mode => 1) is not supported; use one of 'weak', 'strong'"
528        );
529        assert_eq!(
530            message(call("closeness", &[("dtype", s("f32"))])),
531            "closeness() has no parameter 'dtype' in SQL; cast its column instead, for example \
532             CAST(closeness AS REAL)"
533        );
534        assert_eq!(
535            message(call("degree", &[("dtype", s("f32"))])),
536            "degree() has no parameter 'dtype'. Its parameters are: direction"
537        );
538    }
539
540    #[test]
541    fn suggests_names() {
542        assert_eq!(did_you_mean("pagerenk", names()), Some("pagerank"));
543        assert_eq!(did_you_mean("closenes", names()), Some("closeness"));
544        assert_eq!(did_you_mean("algorithm", names()), Some("algorithms"));
545        assert_eq!(
546            did_you_mean("shortest_paths", names()),
547            Some("shortest_path")
548        );
549        assert_eq!(did_you_mean("random_walks", names()), Some("random_walk"));
550        assert_eq!(did_you_mean("dijkstra", names()), None);
551        assert_eq!(did_you_mean("neighbor_agg", names()), Some("neighbors_agg"));
552    }
553
554    #[test]
555    fn neighbors_agg_takes_values_or_an_attribute() {
556        let spec = lookup("neighbors_agg").unwrap();
557        let params: Vec<&str> = sql_params(spec).map(|p| p.name).collect();
558        assert_eq!(params, ["values", "direction", "attribute", "agg"]);
559        assert_eq!(
560            signature(spec),
561            "neighbors_agg(<graph>, values => (SELECT id, value FROM ...), direction => 'out', \
562             agg => 'mean')"
563        );
564        assert_eq!(
565            call(
566                "neighbors_agg",
567                &[("agg", s("n_unique")), ("direction", s("both"))]
568            )
569            .unwrap(),
570            Call::NeighborAgg {
571                direction: Direction::Both,
572                agg: ursa_core::algo::AggKind::NUnique
573            }
574        );
575        assert_eq!(
576            message(call("neighbors_agg", &[("agg", s("median"))])),
577            "neighbors_agg(agg => 'median') is not supported; use one of 'mean', 'sum', 'min', \
578             'max', 'count', 'n_unique'"
579        );
580    }
581
582    #[test]
583    fn layouts_are_hidden() {
584        for name in ["layout_fr", "layout_random", "layout_circle"] {
585            assert!(catalog::lookup(name).is_some(), "{name} is in the catalog");
586            assert!(lookup(name).is_none(), "{name} is hidden from SQL");
587            assert!(names().all(|n| n != name));
588            assert!(specs().all(|s| s.name != name));
589        }
590        // No suggestion leads to a hidden name.
591        for typo in [
592            "layout_f",
593            "layout_frr",
594            "layout_randm",
595            "layout_circl",
596            "layout",
597        ] {
598            assert_eq!(did_you_mean(typo, names()), None, "{typo}");
599        }
600    }
601
602    #[test]
603    fn signatures_show_defaults() {
604        assert_eq!(
605            signature(lookup("pagerank").unwrap()),
606            "pagerank(<graph>, damping => 0.85, max_iter => 30, tol => 1e-6, weight => NULL, \
607             direction => 'out')"
608        );
609        assert_eq!(
610            nested_signature(lookup("louvain").unwrap()),
611            "louvain(weight => NULL, resolution => 1.0, seed => NULL)"
612        );
613    }
614
615    #[test]
616    fn explain_spells_the_call() {
617        let pagerank = call("pagerank", &[]).unwrap();
618        assert_eq!(
619            display_call(&pagerank, None),
620            "pagerank(damping => 0.85, max_iter => 30, tol => 1e-6, direction => 'out')"
621        );
622        assert_eq!(
623            display_call(&pagerank, Some("w")),
624            "pagerank(damping => 0.85, max_iter => 30, tol => 1e-6, weight => w, \
625             direction => 'out')"
626        );
627        assert_eq!(
628            display_call(&call("louvain", &[]).unwrap(), None),
629            "louvain(resolution => 1.0, seed => NULL)"
630        );
631        assert_eq!(
632            display_call(&call("triangle_count", &[]).unwrap(), None),
633            "triangle_count()"
634        );
635    }
636}