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`, and `neighbors_agg`'s column a type that depends on its `agg` and
33//! its values ([`NEIGHBORS_AGG_TYPE`]).
34
35use std::sync::Arc;
36
37/// The output type `ursa_functions()` gives `neighbors_agg`, whose column's
38/// type depends on `agg` and the values (`kernels::neighbors_agg_type`).
39pub const NEIGHBORS_AGG_TYPE: &str =
40    "Float64; UInt32 for count and n_unique; the values' type for min and max of strings, \
41     dates or times";
42
43use datafusion::arrow::array::{
44    ArrayRef, BooleanBuilder, ListBuilder, StringBuilder, StructBuilder,
45};
46use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef};
47use datafusion::arrow::record_batch::RecordBatch;
48use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
49use datafusion::datasource::MemTable;
50use datafusion::error::Result;
51use ursa_core::catalog::Family;
52
53use super::functions::{self, plan, sql_value, ALGORITHMS, URSA_FUNCTIONS};
54use super::traversal::{OutputType, TraversalParamType, TRAVERSALS};
55
56/// The `ursa_functions()` table function.
57#[derive(Debug)]
58pub struct UrsaFunctions;
59
60impl TableFunctionImpl for UrsaFunctions {
61    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
62        if !args.exprs().is_empty() {
63            return Err(plan(format!(
64                "{URSA_FUNCTIONS}() takes no arguments: SELECT * FROM {URSA_FUNCTIONS}()"
65            )));
66        }
67        let batch = functions_batch()?;
68        Ok(Arc::new(MemTable::try_new(
69            batch.schema(),
70            vec![vec![batch]],
71        )?))
72    }
73}
74
75fn parameter_fields() -> Fields {
76    Fields::from(vec![
77        Field::new("name", DataType::Utf8, false),
78        Field::new("type", DataType::Utf8, false),
79        Field::new("default", DataType::Utf8, false),
80        Field::new(
81            "choices",
82            DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
83            false,
84        ),
85        Field::new("doc", DataType::Utf8, false),
86    ])
87}
88
89fn output_fields() -> Fields {
90    Fields::from(vec![
91        Field::new("name", DataType::Utf8, false),
92        Field::new("type", DataType::Utf8, false),
93        Field::new("nullable", DataType::Boolean, false),
94    ])
95}
96
97/// The schema of `ursa_functions()`.
98pub fn schema() -> SchemaRef {
99    let list = |item: DataType| DataType::List(Arc::new(Field::new("item", item, true)));
100    Arc::new(Schema::new(vec![
101        Field::new("name", DataType::Utf8, false),
102        Field::new("family", DataType::Utf8, false),
103        Field::new("summary", DataType::Utf8, false),
104        Field::new("networkx", DataType::Utf8, true),
105        Field::new(
106            "parameters",
107            list(DataType::Struct(parameter_fields())),
108            false,
109        ),
110        Field::new("directions", list(DataType::Utf8), false),
111        Field::new("outputs", list(DataType::Struct(output_fields())), false),
112    ]))
113}
114
115fn family_name(family: Family) -> &'static str {
116    match family {
117        Family::Algorithm => "algorithm",
118        Family::Layout => "layout",
119        Family::Aggregation => "aggregation",
120        Family::Statistic => "statistic",
121    }
122}
123
124/// The summary of the `algorithms` row.
125const ALGORITHMS_SUMMARY: &str = "Several per-node functions over one index, one row per node: \
126     algorithms(g, degree(direction => 'both'), community => louvain(seed => 1)). Each argument \
127     is a call with named parameters; its column is named by its key or after the function.";
128
129/// The rows of `ursa_functions()`.
130pub fn functions_batch() -> Result<RecordBatch> {
131    let mut name = StringBuilder::new();
132    let mut family = StringBuilder::new();
133    let mut summary = StringBuilder::new();
134    let mut networkx = StringBuilder::new();
135    let mut parameters =
136        ListBuilder::new(StructBuilder::from_fields(parameter_fields(), 0)).with_field(Arc::new(
137            Field::new("item", DataType::Struct(parameter_fields()), true),
138        ));
139    let mut directions = ListBuilder::new(StringBuilder::new());
140    let mut outputs = ListBuilder::new(StructBuilder::from_fields(output_fields(), 0)).with_field(
141        Arc::new(Field::new("item", DataType::Struct(output_fields()), true)),
142    );
143
144    for spec in functions::specs() {
145        name.append_value(spec.name);
146        family.append_value(family_name(spec.family));
147        summary.append_value(spec.summary);
148        networkx.append_option(spec.networkx);
149        let params = parameters.values();
150        for p in functions::sql_params(spec) {
151            params
152                .field_builder::<StringBuilder>(0)
153                .expect("name")
154                .append_value(p.name);
155            params
156                .field_builder::<StringBuilder>(1)
157                .expect("type")
158                .append_value(functions::sql_type(p.name, p.ty));
159            params
160                .field_builder::<StringBuilder>(2)
161                .expect("default")
162                .append_value(sql_value(&p.default));
163            let choices = params
164                .field_builder::<ListBuilder<Box<dyn datafusion::arrow::array::ArrayBuilder>>>(3)
165                .expect("choices");
166            for c in p.ty.choices() {
167                choices
168                    .values()
169                    .as_any_mut()
170                    .downcast_mut::<StringBuilder>()
171                    .expect("choices are strings")
172                    .append_value(c);
173            }
174            choices.append(true);
175            params
176                .field_builder::<StringBuilder>(4)
177                .expect("doc")
178                .append_value(functions::sql_doc(p));
179            params.append(true);
180        }
181        parameters.append(true);
182        for d in spec.directions() {
183            directions.values().append_value(d);
184        }
185        directions.append(true);
186        let outs = outputs.values();
187        for o in functions::outputs(spec) {
188            outs.field_builder::<StringBuilder>(0)
189                .expect("name")
190                .append_value(o.name);
191            // A neighbour aggregation's type depends on `agg` and the values.
192            let ty = if spec.reads_attribute() {
193                NEIGHBORS_AGG_TYPE.to_string()
194            } else {
195                o.ty.data_type().to_string()
196            };
197            outs.field_builder::<StringBuilder>(1)
198                .expect("type")
199                .append_value(ty);
200            outs.field_builder::<BooleanBuilder>(2)
201                .expect("nullable")
202                .append_value(o.nullable);
203            outs.append(true);
204        }
205        outputs.append(true);
206    }
207
208    for spec in TRAVERSALS {
209        name.append_value(spec.name);
210        family.append_value("traversal");
211        summary.append_value(spec.summary);
212        networkx.append_option(spec.networkx);
213        let params = parameters.values();
214        for p in spec.params {
215            let (ty, choices): (&str, &[&str]) = match &p.ty {
216                TraversalParamType::Nodes => ("node ids", &[]),
217                TraversalParamType::Node => ("node id", &[]),
218                TraversalParamType::Param(ty) => (functions::sql_type(p.name, *ty), ty.choices()),
219            };
220            params
221                .field_builder::<StringBuilder>(0)
222                .expect("name")
223                .append_value(p.name);
224            params
225                .field_builder::<StringBuilder>(1)
226                .expect("type")
227                .append_value(ty);
228            params
229                .field_builder::<StringBuilder>(2)
230                .expect("default")
231                .append_value(p.default.as_ref().map_or("required".to_string(), sql_value));
232            let choice_list = params
233                .field_builder::<ListBuilder<Box<dyn datafusion::arrow::array::ArrayBuilder>>>(3)
234                .expect("choices");
235            for c in choices {
236                choice_list
237                    .values()
238                    .as_any_mut()
239                    .downcast_mut::<StringBuilder>()
240                    .expect("choices are strings")
241                    .append_value(c);
242            }
243            choice_list.append(true);
244            params
245                .field_builder::<StringBuilder>(4)
246                .expect("doc")
247                .append_value(p.doc);
248            params.append(true);
249        }
250        parameters.append(true);
251        for d in spec.directions() {
252            directions.values().append_value(d);
253        }
254        directions.append(true);
255        let outs = outputs.values();
256        for o in spec.outputs {
257            let ty = match o.ty {
258                OutputType::Id => "Int64 or Utf8".to_string(),
259                other => other.data_type(&DataType::Null).to_string(),
260            };
261            outs.field_builder::<StringBuilder>(0)
262                .expect("name")
263                .append_value(o.name);
264            outs.field_builder::<StringBuilder>(1)
265                .expect("type")
266                .append_value(ty);
267            outs.field_builder::<BooleanBuilder>(2)
268                .expect("nullable")
269                .append_value(false);
270            outs.append(true);
271        }
272        outputs.append(true);
273    }
274
275    name.append_value(ALGORITHMS);
276    family.append_value("combinator");
277    summary.append_value(ALGORITHMS_SUMMARY);
278    networkx.append_null();
279    parameters.append(true);
280    directions.append(true);
281    outputs.append(true);
282
283    let columns: Vec<ArrayRef> = vec![
284        Arc::new(name.finish()),
285        Arc::new(family.finish()),
286        Arc::new(summary.finish()),
287        Arc::new(networkx.finish()),
288        Arc::new(parameters.finish()),
289        Arc::new(directions.finish()),
290        Arc::new(outputs.finish()),
291    ];
292    Ok(RecordBatch::try_new(schema(), columns)?)
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    #[test]
300    fn one_row_per_function() {
301        let batch = functions_batch().unwrap();
302        let expected = functions::names().count();
303        assert_eq!(batch.num_rows(), expected);
304        assert_eq!(batch.schema(), schema());
305    }
306}