1use 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#[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
90pub 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
117const 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
122pub 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}