1use 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
36pub const ALGORITHMS: &str = "algorithms";
38
39pub const URSA_FUNCTIONS: &str = "ursa_functions";
41
42fn in_sql(spec: &FunctionSpec) -> bool {
47 spec.family != Family::Layout
48}
49
50pub fn lookup(name: &str) -> Option<&'static FunctionSpec> {
52 catalog::lookup(name).filter(|s| in_sql(s))
53}
54
55pub fn specs() -> impl Iterator<Item = &'static FunctionSpec> {
57 catalog::FUNCTIONS.iter().filter(|s| in_sql(s))
58}
59
60pub 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
69pub const N_DROPPED_EDGES: OutputSpec = OutputSpec {
73 name: "n_dropped_edges",
74 ty: ColumnType::Int64,
75 nullable: false,
76};
77
78pub 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
89pub fn is_per_node(spec: &FunctionSpec) -> bool {
92 spec.shape == Shape::PerNode
93}
94
95pub 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
106const ATTRIBUTE_DOC: &str =
108 "An expression over the graph's node table, the value of each node. Give this or values.";
109
110const 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
115pub 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
125pub 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
134pub 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#[derive(Debug, Clone, PartialEq)]
152pub enum Literal {
153 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 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
189pub 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
207fn 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
238pub(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
250pub 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
269pub 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
280fn 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
293pub const VALUES_QUERY: &str = "(SELECT id, value FROM ...)";
295
296pub 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
304pub fn nested_signature(spec: &FunctionSpec) -> String {
306 format!("{}({})", spec.name, defaults(spec).join(", "))
307}
308
309pub 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
327pub(crate) fn plan(message: String) -> DataFusionError {
329 DataFusionError::Plan(message)
330}
331
332pub 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 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 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}