1use std::sync::Arc;
36
37pub 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#[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
97pub 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
124const 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
129pub 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 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}