change algo relation format

README.md

@@ -21,4 +21,4 @@
   * [ ] connected components
   * [ ] label propagation
   * [ ] louvain modularity
-* [ ] direct loading of data
+* [x] direct loading of data

src/algo/mod.rs

@@ -44,17 +44,17 @@ impl MagicAlgoRuleArg {
         stores: &'a BTreeMap<MagicSymbol, DerivedRelStore>,
     ) -> Result<TupleIter<'a>> {
         Ok(match self {
-            MagicAlgoRuleArg::InMem(s) => {
+            MagicAlgoRuleArg::InMem(s, _) => {
                 let store = stores
                     .get(s)
                     .ok_or_else(|| anyhow!("rule not found: {:?}", s))?;
                 Box::new(store.scan_all())
             }
-            MagicAlgoRuleArg::Stored(s) => {
+            MagicAlgoRuleArg::Stored(s, _) => {
                 let view_rel = tx.get_view_rel(s)?;
                 Box::new(view_rel.scan_all()?)
             }
-            MagicAlgoRuleArg::Triple(attr, dir) => match dir {
+            MagicAlgoRuleArg::Triple(attr, _, dir) => match dir {
                 TripleDir::Fwd => {
                     if attr.with_history {
                         Box::new(

src/cozoscript.pest

@@ -34,8 +34,10 @@ head_arg = {aggr_arg | var}
 aggr_arg = {ident ~ "(" ~ var ~ ("," ~ expr)* ~ ")"}
 algo_arg = _{algo_rel | algo_opt_pair}
 algo_opt_pair = {ident ~ "=" ~ expr}
-algo_rel = {ident | view_ident | algo_triple_rel}
-algo_triple_rel = { rev_triple_marker? ~ "[" ~ ident ~ "]"}
+algo_rel = {algo_rule_rel | algo_view_rel | algo_triple_rel}
+algo_rule_rel = {ident ~ "[" ~ (var ~ ",")* ~ var? ~ "]"}
+algo_view_rel = {view_ident ~ "[" ~ (var ~ ",")* ~ var? ~ "]"}
+algo_triple_rel = { "[" ~ var ~ rev_triple_marker? ~ ident ~ var ~ "]"}
 rev_triple_marker = {"<"}
 
 rule_body = {(disjunction ~ ",")* ~ disjunction?}

src/data/program.rs

@@ -75,16 +75,16 @@ pub(crate) enum TripleDir {
 
 #[derive(Debug, Clone)]
 pub(crate) enum AlgoRuleArg {
-    InMem(Symbol),
-    Stored(Symbol),
-    Triple(Attribute, TripleDir),
+    InMem(Symbol, Vec<Symbol>),
+    Stored(Symbol, Vec<Symbol>),
+    Triple(Attribute, Vec<Symbol>, TripleDir),
 }
 
 #[derive(Debug, Clone)]
 pub(crate) enum MagicAlgoRuleArg {
-    InMem(MagicSymbol),
-    Stored(Symbol),
-    Triple(Attribute, TripleDir),
+    InMem(MagicSymbol, Vec<Symbol>),
+    Stored(Symbol, Vec<Symbol>),
+    Triple(Attribute, Vec<Symbol>, TripleDir),
 }
 
 #[derive(Debug, Clone)]

src/parse/cozoscript/query.rs

@@ -4,24 +4,24 @@ use std::str::FromStr;
 use anyhow::{anyhow, Result};
 use itertools::Itertools;
 use lazy_static::lazy_static;
-use pest::prec_climber::{Assoc, Operator, PrecClimber};
 use pest::Parser;
+use pest::prec_climber::{Assoc, Operator, PrecClimber};
 use serde_json::{json, Map};
 
 use crate::data::json::JsonValue;
+use crate::parse::cozoscript::{CozoScriptParser, Pair, Pairs, Rule};
 use crate::parse::cozoscript::number::parse_int;
 use crate::parse::cozoscript::schema::parsed_schema_to_json;
 use crate::parse::cozoscript::string::parse_string;
-use crate::parse::cozoscript::tx::parsed_tx_to_json;
-use crate::parse::cozoscript::{CozoScriptParser, Pair, Pairs, Rule};
 use crate::parse::cozoscript::sys::parsed_db_op_to_enum;
+use crate::parse::cozoscript::tx::parsed_tx_to_json;
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub(crate) enum ScriptType {
     Query,
     Schema,
     Tx,
-    Sys
+    Sys,
 }
 
 pub(crate) fn parse_query_to_json(src: &str) -> Result<(ScriptType, JsonValue)> {
@@ -40,7 +40,7 @@ pub(crate) fn parse_query_to_json(src: &str) -> Result<(ScriptType, JsonValue)>
             let opts = parsed_db_op_to_enum(parsed.into_inner())?;
             let opts = serde_json::to_value(&opts)?;
             (ScriptType::Sys, opts)
-        },
+        }
         _ => unreachable!(),
     })
 }
@@ -238,23 +238,35 @@ fn parse_algo_rule(src: Pair<'_>) -> Result<JsonValue> {
             Rule::algo_rel => {
                 let inner = nxt.into_inner().next().unwrap();
                 match inner.as_rule() {
-                    Rule::ident => algo_rels.push(json!({"rule": inner.as_str()})),
-                    Rule::view_ident => {
-                        algo_rels.push(json!({"view": inner.as_str().strip_prefix(':').unwrap()}))
+                    Rule::algo_rule_rel => {
+                        let mut els = inner.into_inner();
+                        let name = els.next().unwrap().as_str();
+                        let args = els.map(|v| v.as_str()).collect_vec();
+                        algo_rels.push(json!({"rule": name, "rel_args": args}));
+                    }
+                    Rule::algo_view_rel => {
+                        let mut els = inner.into_inner();
+                        let name = els.next().unwrap().as_str().strip_prefix(':').unwrap();
+                        let args = els.map(|v| v.as_str()).collect_vec();
+                        algo_rels.push(json!({"view": name, "rel_args": args}));
                     }
                     Rule::algo_triple_rel => {
-                        let mut inner = inner.into_inner();
+                        let mut els = inner.into_inner();
+                        let fst = els.next().unwrap().as_str();
+                        let mdl = els.next().unwrap();
                         let mut backward = false;
-                        let ident = inner.next().unwrap();
-                        let ident = match ident.as_rule() {
+                        let ident = match mdl.as_rule() {
                             Rule::rev_triple_marker => {
                                 backward = true;
-                                inner.next().unwrap()
+                                els.next().unwrap().as_str()
+                            }
+                            Rule::ident => {
+                                mdl.as_str()
                             }
-                            _ => ident,
+                            _ => unreachable!()
                         };
-                        let ident = ident.as_str();
-                        algo_rels.push(json!({"triple": ident, "backward": backward}));
+                        let snd = els.next().unwrap().as_str();
+                        algo_rels.push(json!({"triple": ident, "backward": backward, "rel_args": [fst, snd]}))
                     }
                     _ => unreachable!(),
                 }

src/parse/query.rs

@@ -147,16 +147,31 @@ impl SessionTx {
                     .as_array()
                     .ok_or_else(|| anyhow!("'relations' field must be an array"))?
                 {
+                    let args: Vec<Symbol> = rel_def
+                        .get("rel_args")
+                        .ok_or_else(|| anyhow!("field 'rel_args' required in {}", rel_def))?
+                        .as_array()
+                        .ok_or_else(|| anyhow!("field 'rel_args' must be an array in {}", rel_def))?
+                        .iter()
+                        .map(|v| -> Result<Symbol> {
+                            let s = v.as_str().ok_or_else(|| {
+                                anyhow!("element of 'rel_args' must be string, got {}", v)
+                            })?;
+                            let s = Symbol::from(s);
+                            s.validate_query_var()?;
+                            Ok(s)
+                        })
+                        .try_collect()?;
                     if let Some(rule_name) = rel_def.get("rule") {
                         let rule_name = rule_name
                             .as_str()
                             .ok_or_else(|| anyhow!("'rule' must be a string, got {}", rule_name))?;
-                        relations.push(AlgoRuleArg::InMem(Symbol::from(rule_name)));
+                        relations.push(AlgoRuleArg::InMem(Symbol::from(rule_name), args));
                     } else if let Some(view_name) = rel_def.get("view") {
                         let view_name = view_name
                             .as_str()
                             .ok_or_else(|| anyhow!("'view' must be a string, got {}", view_name))?;
-                        relations.push(AlgoRuleArg::Stored(Symbol::from(view_name)));
+                        relations.push(AlgoRuleArg::Stored(Symbol::from(view_name), args));
                     } else if let Some(triple_name) = rel_def.get("triple") {
                         let attr = self.parse_triple_atom_attr(triple_name)?;
                         // let triple_name = triple_name.as_str().ok_or_else(|| {
@@ -168,7 +183,7 @@ impl SessionTx {
                             Some(JsonValue::Bool(false)) => TripleDir::Fwd,
                             d => bail!("'backward' must be a boolean, got {}", d.unwrap()),
                         };
-                        relations.push(AlgoRuleArg::Triple(attr, dir));
+                        relations.push(AlgoRuleArg::Triple(attr, args, dir));
                     }
                 }
                 if let Some(opts) = algo_rule.get("options") {

src/query/magic.rs

@@ -265,7 +265,7 @@ impl NormalFormProgram {
                 }
                 NormalFormAlgoOrRules::Algo(algo) => {
                     for rel in algo.rule_args.iter() {
-                        if let AlgoRuleArg::InMem(r) = rel {
+                        if let AlgoRuleArg::InMem(r, _args) = rel {
                             downstream_rules.insert(r.clone());
                         }
                     }
@@ -304,14 +304,15 @@ impl NormalFormProgram {
                                 .rule_args
                                 .iter()
                                 .map(|r| match r {
-                                    AlgoRuleArg::InMem(m) => {
-                                        MagicAlgoRuleArg::InMem(MagicSymbol::Muggle {
-                                            inner: m.clone(),
-                                        })
+                                    AlgoRuleArg::InMem(m, args) => MagicAlgoRuleArg::InMem(
+                                        MagicSymbol::Muggle { inner: m.clone() },
+                                        args.clone(),
+                                    ),
+                                    AlgoRuleArg::Stored(s, args) => {
+                                        MagicAlgoRuleArg::Stored(s.clone(), args.clone())
                                     }
-                                    AlgoRuleArg::Stored(s) => MagicAlgoRuleArg::Stored(s.clone()),
-                                    AlgoRuleArg::Triple(t, d) => {
-                                        MagicAlgoRuleArg::Triple(t.clone(), *d)
+                                    AlgoRuleArg::Triple(t, args, d) => {
+                                        MagicAlgoRuleArg::Triple(t.clone(), args.clone(), *d)
                                     }
                                 })
                                 .collect_vec(),

src/query/stratify.rs

@@ -123,11 +123,11 @@ fn convert_normal_form_program_to_graph(
                 let mut ret: BTreeMap<&Symbol, bool> = BTreeMap::default();
                 for rel in &algo.rule_args {
                     match rel {
-                        AlgoRuleArg::InMem(r) => {
+                        AlgoRuleArg::InMem(r, _args) => {
                             ret.insert(r, true);
                         }
-                        AlgoRuleArg::Stored(_) => {}
-                        AlgoRuleArg::Triple(_, _) => {}
+                        AlgoRuleArg::Stored(_, _) => {}
+                        AlgoRuleArg::Triple(_, _, _) => {}
                     }
                 }
                 (k, ret)

tests/air_routes.rs

@@ -92,7 +92,7 @@ fn air_routes() -> Result<()> {
     let deg_centrality_time = Instant::now();
     let res = db.run_script(
         r#"
-        deg_centrality <- degree_centrality!(:flies_to);
+        deg_centrality <- degree_centrality!(:flies_to[?a, ?b]);
         ?[?total, ?out, ?in] := deg_centrality[?node, ?total, ?out, ?in];
         :order -?total;
         :limit 10;
@@ -114,7 +114,7 @@ fn air_routes() -> Result<()> {
         r#"
         flies_to[?a, ?b] := [?r route.src ?ac], [?r route.dst ?bc],
                             [?ac airport.iata ?a], [?bc airport.iata ?b];
-        deg_centrality <- degree_centrality!(flies_to);
+        deg_centrality <- degree_centrality!(flies_to[?a, ?b]);
         ?[?node, ?total, ?out, ?in] := deg_centrality[?node, ?total, ?out, ?in];
         :order -?total;
         :limit 10;