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