1use 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
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 fn is_per_node(spec: &FunctionSpec) -> bool {
72 spec.shape == Shape::PerNode
73}
74
75pub 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
86const ATTRIBUTE_DOC: &str =
88 "An expression over the graph's node table, the value of each node. Give this or values.";
89
90pub 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
100pub 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
108pub 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#[derive(Debug, Clone, PartialEq)]
126pub enum Literal {
127 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 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
163pub 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
181fn 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
212pub(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
224pub 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
243pub 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
254fn 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
267pub const VALUES_QUERY: &str = "(SELECT id, value FROM ...)";
269
270pub 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
278pub fn nested_signature(spec: &FunctionSpec) -> String {
280 format!("{}({})", spec.name, defaults(spec).join(", "))
281}
282
283pub 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
301pub(crate) fn plan(message: String) -> DataFusionError {
303 DataFusionError::Plan(message)
304}
305
306pub 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 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 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}