Skip to main content

ursa_engine/graph/
catalog_table.rs

1//! `ursa_functions()`: the graph functions, as a table.
2//!
3//! One row per function SQL can call: every algorithm, statistic and
4//! aggregation of `ursa_core::catalog` that SQL offers (not the layouts,
5//! hidden until the visualization work; see [`functions`]), then the
6//! traversals, then `algorithms`. The list columns are Arrow lists and
7//! structs, not strings, so they stay usable both ways:
8//!
9//! - from SQL: `unnest(parameters)`, `array_has(directions, 'in')`,
10//!   `parameters[1]['default']`;
11//! - as JSON (`ursa --format json`), where they are real arrays and objects an agent
12//!   reads without parsing a format of ours.
13//!
14//! Columns:
15//!
16//! | column       | type                                                     |
17//! |--------------|----------------------------------------------------------|
18//! | `name`       | the function's name                                      |
19//! | `family`     | `algorithm`, `statistic`, `aggregation`, `traversal` or `combinator` |
20//! | `summary`    | one line on what it computes                             |
21//! | `networkx`   | the NetworkX call it matches, or NULL                    |
22//! | `parameters` | list of {name, type, default, choices, doc}, call order  |
23//! | `directions` | the `direction =>` values it accepts; empty if none     |
24//! | `outputs`    | list of {name, type, nullable}: its columns besides `id` |
25//!
26//! A parameter's `type` is how SQL writes it (`number`, `number or NULL`,
27//! `integer`, `integer or NULL`, `boolean`, `string`, `expression`, `query`
28//! for `neighbors_agg`'s values, and for a traversal's node ids `node ids` or
29//! `node id`), and its `default` is the
30//! literal SQL would write (`0.85`, `'out'`, `NULL`), or `required`. A
31//! traversal's node-id output columns have the graph's id type, `Int64 or
32//! Utf8`.
33
34use std::sync::Arc;
35
36use datafusion::arrow::array::{
37    ArrayRef, BooleanBuilder, ListBuilder, StringBuilder, StructBuilder,
38};
39use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef};
40use datafusion::arrow::record_batch::RecordBatch;
41use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
42use datafusion::datasource::MemTable;
43use datafusion::error::Result;
44use ursa_core::catalog::Family;
45
46use super::functions::{self, plan, sql_value, ALGORITHMS, URSA_FUNCTIONS};
47use super::traversal::{OutputType, TraversalParamType, TRAVERSALS};
48
49/// The `ursa_functions()` table function.
50#[derive(Debug)]
51pub struct UrsaFunctions;
52
53impl TableFunctionImpl for UrsaFunctions {
54    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
55        if !args.exprs().is_empty() {
56            return Err(plan(format!(
57                "{URSA_FUNCTIONS}() takes no arguments: SELECT * FROM {URSA_FUNCTIONS}()"
58            )));
59        }
60        let batch = functions_batch()?;
61        Ok(Arc::new(MemTable::try_new(
62            batch.schema(),
63            vec![vec![batch]],
64        )?))
65    }
66}
67
68fn parameter_fields() -> Fields {
69    Fields::from(vec![
70        Field::new("name", DataType::Utf8, false),
71        Field::new("type", DataType::Utf8, false),
72        Field::new("default", DataType::Utf8, false),
73        Field::new(
74            "choices",
75            DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
76            false,
77        ),
78        Field::new("doc", DataType::Utf8, false),
79    ])
80}
81
82fn output_fields() -> Fields {
83    Fields::from(vec![
84        Field::new("name", DataType::Utf8, false),
85        Field::new("type", DataType::Utf8, false),
86        Field::new("nullable", DataType::Boolean, false),
87    ])
88}
89
90/// The schema of `ursa_functions()`.
91pub fn schema() -> SchemaRef {
92    let list = |item: DataType| DataType::List(Arc::new(Field::new("item", item, true)));
93    Arc::new(Schema::new(vec![
94        Field::new("name", DataType::Utf8, false),
95        Field::new("family", DataType::Utf8, false),
96        Field::new("summary", DataType::Utf8, false),
97        Field::new("networkx", DataType::Utf8, true),
98        Field::new(
99            "parameters",
100            list(DataType::Struct(parameter_fields())),
101            false,
102        ),
103        Field::new("directions", list(DataType::Utf8), false),
104        Field::new("outputs", list(DataType::Struct(output_fields())), false),
105    ]))
106}
107
108fn family_name(family: Family) -> &'static str {
109    match family {
110        Family::Algorithm => "algorithm",
111        Family::Layout => "layout",
112        Family::Aggregation => "aggregation",
113        Family::Statistic => "statistic",
114    }
115}
116
117/// The summary of the `algorithms` row.
118const ALGORITHMS_SUMMARY: &str = "Several per-node functions over one index, one row per node: \
119     algorithms(g, degree(direction => 'both'), community => louvain(seed => 1)). Each argument \
120     is a call with named parameters; its column is named by its key or after the function.";
121
122/// The rows of `ursa_functions()`.
123pub fn functions_batch() -> Result<RecordBatch> {
124    let mut name = StringBuilder::new();
125    let mut family = StringBuilder::new();
126    let mut summary = StringBuilder::new();
127    let mut networkx = StringBuilder::new();
128    let mut parameters =
129        ListBuilder::new(StructBuilder::from_fields(parameter_fields(), 0)).with_field(Arc::new(
130            Field::new("item", DataType::Struct(parameter_fields()), true),
131        ));
132    let mut directions = ListBuilder::new(StringBuilder::new());
133    let mut outputs = ListBuilder::new(StructBuilder::from_fields(output_fields(), 0)).with_field(
134        Arc::new(Field::new("item", DataType::Struct(output_fields()), true)),
135    );
136
137    for spec in functions::specs() {
138        name.append_value(spec.name);
139        family.append_value(family_name(spec.family));
140        summary.append_value(spec.summary);
141        networkx.append_option(spec.networkx);
142        let params = parameters.values();
143        for p in functions::sql_params(spec) {
144            params
145                .field_builder::<StringBuilder>(0)
146                .expect("name")
147                .append_value(p.name);
148            params
149                .field_builder::<StringBuilder>(1)
150                .expect("type")
151                .append_value(functions::sql_type(p.name, p.ty));
152            params
153                .field_builder::<StringBuilder>(2)
154                .expect("default")
155                .append_value(sql_value(&p.default));
156            let choices = params
157                .field_builder::<ListBuilder<Box<dyn datafusion::arrow::array::ArrayBuilder>>>(3)
158                .expect("choices");
159            for c in p.ty.choices() {
160                choices
161                    .values()
162                    .as_any_mut()
163                    .downcast_mut::<StringBuilder>()
164                    .expect("choices are strings")
165                    .append_value(c);
166            }
167            choices.append(true);
168            params
169                .field_builder::<StringBuilder>(4)
170                .expect("doc")
171                .append_value(functions::sql_doc(p));
172            params.append(true);
173        }
174        parameters.append(true);
175        for d in spec.directions() {
176            directions.values().append_value(d);
177        }
178        directions.append(true);
179        let outs = outputs.values();
180        for o in spec.outputs {
181            outs.field_builder::<StringBuilder>(0)
182                .expect("name")
183                .append_value(o.name);
184            outs.field_builder::<StringBuilder>(1)
185                .expect("type")
186                .append_value(o.ty.data_type().to_string());
187            outs.field_builder::<BooleanBuilder>(2)
188                .expect("nullable")
189                .append_value(o.nullable);
190            outs.append(true);
191        }
192        outputs.append(true);
193    }
194
195    for spec in TRAVERSALS {
196        name.append_value(spec.name);
197        family.append_value("traversal");
198        summary.append_value(spec.summary);
199        networkx.append_option(spec.networkx);
200        let params = parameters.values();
201        for p in spec.params {
202            let (ty, choices): (&str, &[&str]) = match &p.ty {
203                TraversalParamType::Nodes => ("node ids", &[]),
204                TraversalParamType::Node => ("node id", &[]),
205                TraversalParamType::Param(ty) => (functions::sql_type(p.name, *ty), ty.choices()),
206            };
207            params
208                .field_builder::<StringBuilder>(0)
209                .expect("name")
210                .append_value(p.name);
211            params
212                .field_builder::<StringBuilder>(1)
213                .expect("type")
214                .append_value(ty);
215            params
216                .field_builder::<StringBuilder>(2)
217                .expect("default")
218                .append_value(p.default.as_ref().map_or("required".to_string(), sql_value));
219            let choice_list = params
220                .field_builder::<ListBuilder<Box<dyn datafusion::arrow::array::ArrayBuilder>>>(3)
221                .expect("choices");
222            for c in choices {
223                choice_list
224                    .values()
225                    .as_any_mut()
226                    .downcast_mut::<StringBuilder>()
227                    .expect("choices are strings")
228                    .append_value(c);
229            }
230            choice_list.append(true);
231            params
232                .field_builder::<StringBuilder>(4)
233                .expect("doc")
234                .append_value(p.doc);
235            params.append(true);
236        }
237        parameters.append(true);
238        for d in spec.directions() {
239            directions.values().append_value(d);
240        }
241        directions.append(true);
242        let outs = outputs.values();
243        for o in spec.outputs {
244            let ty = match o.ty {
245                OutputType::Id => "Int64 or Utf8".to_string(),
246                other => other.data_type(&DataType::Null).to_string(),
247            };
248            outs.field_builder::<StringBuilder>(0)
249                .expect("name")
250                .append_value(o.name);
251            outs.field_builder::<StringBuilder>(1)
252                .expect("type")
253                .append_value(ty);
254            outs.field_builder::<BooleanBuilder>(2)
255                .expect("nullable")
256                .append_value(false);
257            outs.append(true);
258        }
259        outputs.append(true);
260    }
261
262    name.append_value(ALGORITHMS);
263    family.append_value("combinator");
264    summary.append_value(ALGORITHMS_SUMMARY);
265    networkx.append_null();
266    parameters.append(true);
267    directions.append(true);
268    outputs.append(true);
269
270    let columns: Vec<ArrayRef> = vec![
271        Arc::new(name.finish()),
272        Arc::new(family.finish()),
273        Arc::new(summary.finish()),
274        Arc::new(networkx.finish()),
275        Arc::new(parameters.finish()),
276        Arc::new(directions.finish()),
277        Arc::new(outputs.finish()),
278    ];
279    Ok(RecordBatch::try_new(schema(), columns)?)
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn one_row_per_function() {
288        let batch = functions_batch().unwrap();
289        let expected = functions::names().count();
290        assert_eq!(batch.num_rows(), expected);
291        assert_eq!(batch.schema(), schema());
292    }
293}