tests for index

cozo-core/src/cozoscript.pest

@@ -14,14 +14,14 @@ imperative_script = {SOI ~ imperative_stmt+ ~ EOI}
 sys_script = {SOI ~ "::" ~ (compact_op | list_relations_op | list_relation_op | remove_relations_op | trigger_relation_op |
                     trigger_relation_show_op | rename_relations_op | running_op | kill_op | explain_op | access_level_op | index_op) ~ EOI}
 index_op = {"index" ~ (index_create | index_drop)}
-index_create = {"create" ~ compound_ident ~ ":" ~ ident ~ "{" ~ ident+ ~ "}"}
+index_create = {"create" ~ compound_ident ~ ":" ~ ident ~ "{" ~ (ident ~ ",")* ~ ident? ~ "}"}
 index_drop = {"drop" ~ compound_ident ~ ":" ~ ident }
 compact_op = {"compact"}
 running_op = {"running"}
 kill_op = {"kill" ~ expr}
 explain_op = {"explain" ~ "{" ~ query_script_inner_no_bracket ~ "}"}
 list_relations_op = {"relations"}
-list_relation_op = {"columns" ~ compound_ident}
+list_relation_op = {"columns" ~ compound_or_index_ident}
 remove_relations_op = {"remove" ~ (compound_ident ~ ",")* ~ compound_ident }
 rename_relations_op = {"rename" ~ (rename_pair ~ ",")* ~ rename_pair }
 access_level_op = {"access_level" ~ access_level ~ (compound_ident ~ ",")* ~ compound_ident}
@@ -46,8 +46,9 @@ var = @{(XID_START | "_") ~ (XID_CONTINUE | "_")*}
 param = @{"$" ~ (XID_CONTINUE | "_")*}
 ident = @{XID_START ~ ("_" | XID_CONTINUE)*}
 underscore_ident = @{("_" | XID_START) ~ ("_" | XID_CONTINUE)*}
-relation_ident = @{"*" ~ (compound_ident | underscore_ident)}
+relation_ident = @{"*" ~ (compound_or_index_ident | underscore_ident)}
-compound_ident = @{ident ~ ("." ~ ident)?}
+compound_ident = @{ident ~ ("." ~ ident)*}
+compound_or_index_ident = @{ident ~ ("." ~ ident)* ~ (":" ~ ident)?}
 
 rule = {rule_head ~ ":=" ~ rule_body ~ ";"?}
 const_rule = {rule_head ~ "<-" ~ expr ~ ";"?}

cozo-core/src/parse/sys.rs

@@ -10,7 +10,7 @@ use std::collections::BTreeMap;
 use std::sync::Arc;
 
 use itertools::Itertools;
-use miette::{miette, Diagnostic, Result};
+use miette::{miette, Diagnostic, Result, ensure};
 use thiserror::Error;
 
 use crate::data::program::InputProgram;
@@ -150,20 +150,29 @@ pub(crate) fn parse_sys(
         Rule::index_op => {
             let inner = inner.into_inner().next().unwrap();
             match inner.as_rule() {
-                Rule::index_drop => {
+                Rule::index_create => {
+                    let span = inner.extract_span();
                     let mut inner = inner.into_inner();
                     let rel = inner.next().unwrap();
                     let name = inner.next().unwrap();
                     let cols = inner
                         .map(|p| Symbol::new(p.as_str(), p.extract_span()))
                         .collect_vec();
+
+
+                    #[derive(Debug, Diagnostic, Error)]
+                    #[error("index must have at least one column specified")]
+                    #[diagnostic(code(parser::empty_index))]
+                    struct EmptyIndex(#[label] SourceSpan);
+
+                    ensure!(!cols.is_empty(), EmptyIndex(span));
                     SysOp::CreateIndex(
                         Symbol::new(rel.as_str(), rel.extract_span()),
                         Symbol::new(name.as_str(), name.extract_span()),
                         cols,
                     )
                 }
-                Rule::index_create => {
+                Rule::index_drop => {
                     let mut inner = inner.into_inner();
                     let rel = inner.next().unwrap();
                     let name = inner.next().unwrap();

cozo-core/src/query/compile.rs

@@ -316,11 +316,10 @@ impl<'a> SessionTx<'a> {
                                 })
                                 .collect_vec();
 
-                            let final_joiner_vars = right_vars
+                            let mut final_joiner_vars = vec![];
-                                .iter()
+                            for idx in mapper.iter() {
-                                .take(store.metadata.keys.len())
+                                final_joiner_vars.push(right_vars[*idx].clone());
-                                .cloned()
+                            }
-                                .collect_vec();
 
                             let middle = RelAlgebra::relation(
                                 middle_vars,

cozo-core/src/runtime/db.rs

@@ -594,7 +594,7 @@ impl<'s, S: Storage<'s>> Db<S> {
         tx.commit_tx()?;
         Ok(())
     }
-    fn transact(&'s self) -> Result<SessionTx<'_>> {
+    pub(crate) fn transact(&'s self) -> Result<SessionTx<'_>> {
         let ret = SessionTx {
             store_tx: Box::new(self.db.transact(false)?),
             temp_store_tx: self.temp_db.transact(true)?,
@@ -603,7 +603,7 @@ impl<'s, S: Storage<'s>> Db<S> {
         };
         Ok(ret)
     }
-    fn transact_write(&'s self) -> Result<SessionTx<'_>> {
+    pub(crate) fn transact_write(&'s self) -> Result<SessionTx<'_>> {
         let ret = SessionTx {
             store_tx: Box::new(self.db.transact(true)?),
             temp_store_tx: self.temp_db.transact(true)?,
@@ -1623,7 +1623,11 @@ impl<'s, S: Storage<'s>> Db<S> {
             let n_dependents = meta.metadata.non_keys.len();
             let arity = n_keys + n_dependents;
             let name = meta.name;
-            let access_level = meta.access_level.to_string();
+            let access_level = if name.contains(':') {
+                "index".to_string()
+            } else {
+                meta.access_level.to_string()
+            };
             rows.push(vec![
                 json!(name),
                 json!(arity),

cozo-core/src/runtime/relation.rs

@@ -670,10 +670,10 @@ impl<'a> SessionTx<'a> {
             ));
         }
 
-        for key in rel_handle.metadata.keys.iter() {
+        'outer: for key in rel_handle.metadata.keys.iter() {
             for col in cols.iter() {
                 if col.name == key.name {
-                    break;
+                    continue 'outer;
                 }
             }
             col_defs.push(key.clone());

cozo-core/src/runtime/tests.rs

@@ -10,6 +10,7 @@
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
+use itertools::Itertools;
 use log::debug;
 use serde_json::json;
 
@@ -473,10 +474,178 @@ fn test_callback() {
     assert_eq!(collected[1].1.rows.len(), 2);
     assert_eq!(collected[1].1.rows[0].len(), 3);
     assert_eq!(collected[1].2.rows.len(), 1);
-    assert_eq!(collected[1].2.rows[0].len(), 3);
+    assert_eq!(
+        collected[1].2.rows[0],
+        vec![DataValue::from(1), DataValue::from(2), DataValue::from(3)]
+    );
     assert_eq!(collected[2].0, CallbackOp::Rm);
     assert_eq!(collected[2].1.rows.len(), 2);
     assert_eq!(collected[2].1.rows[0].len(), 2);
     assert_eq!(collected[2].2.rows.len(), 1);
     assert_eq!(collected[2].2.rows[0].len(), 3);
 }
+
+#[test]
+fn test_index() {
+    let db = new_cozo_mem().unwrap();
+    db.run_script(
+        ":create friends {fr: Int, to: Int => data: Any}",
+        Default::default(),
+    )
+    .unwrap();
+
+    db.run_script(
+        r"?[fr, to, data] <- [[1,2,3],[4,5,6]] :put friends {fr, to => data}",
+        Default::default(),
+    )
+    .unwrap();
+
+    assert!(db
+        .run_script("::index create friends:rev {to, no}", Default::default())
+        .is_err());
+    db.run_script("::index create friends:rev {to, data}", Default::default())
+        .unwrap();
+
+    db.run_script(
+        r"?[fr, to, data] <- [[1,2,5],[6,5,7]] :put friends {fr, to => data}",
+        Default::default(),
+    )
+    .unwrap();
+    db.run_script(
+        r"?[fr, to] <- [[4,5]] :rm friends {fr, to}",
+        Default::default(),
+    )
+    .unwrap();
+
+    let rels_data = db
+        .export_relations(["friends", "friends:rev"].into_iter())
+        .unwrap();
+    assert_eq!(
+        rels_data["friends"].clone().into_json()["rows"],
+        json!([[1, 2, 5], [6, 5, 7]])
+    );
+    assert_eq!(
+        rels_data["friends:rev"].clone().into_json()["rows"],
+        json!([[2, 5, 1], [5, 7, 6]])
+    );
+
+    let rels = db.run_script("::relations", Default::default()).unwrap();
+    assert_eq!(rels.rows[1][0], DataValue::from("friends:rev"));
+    assert_eq!(rels.rows[1][1], DataValue::from(3));
+    assert_eq!(rels.rows[1][2], DataValue::from("index"));
+
+    let cols = db
+        .run_script("::columns friends:rev", Default::default())
+        .unwrap();
+    assert_eq!(cols.rows.len(), 3);
+
+    let res = db
+        .run_script(
+            "?[fr, data] := *friends:rev{to: 2, fr, data}",
+            Default::default(),
+        )
+        .unwrap();
+    assert_eq!(res.into_json()["rows"], json!([[1, 5]]));
+
+    let res = db
+        .run_script(
+            "?[fr, data] := *friends{to: 2, fr, data}",
+            Default::default(),
+        )
+        .unwrap();
+    assert_eq!(res.into_json()["rows"], json!([[1, 5]]));
+
+    let expl = db
+        .run_script(
+            "::explain { ?[fr, data] := *friends{to: 2, fr, data} }",
+            Default::default(),
+        )
+        .unwrap();
+    let joins = expl.into_json()["rows"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|row| row.as_array().unwrap()[5].clone())
+        .collect_vec();
+    assert!(joins.contains(&json!(":friends:rev")));
+}
+
+#[test]
+fn test_index_short() {
+    let db = new_cozo_mem().unwrap();
+    db.run_script(
+        ":create friends {fr: Int, to: Int => data: Any}",
+        Default::default(),
+    )
+    .unwrap();
+
+    db.run_script(
+        r"?[fr, to, data] <- [[1,2,3],[4,5,6]] :put friends {fr, to => data}",
+        Default::default(),
+    )
+    .unwrap();
+
+    db.run_script("::index create friends:rev {to}", Default::default())
+        .unwrap();
+
+    db.run_script(
+        r"?[fr, to, data] <- [[1,2,5],[6,5,7]] :put friends {fr, to => data}",
+        Default::default(),
+    )
+    .unwrap();
+    db.run_script(
+        r"?[fr, to] <- [[4,5]] :rm friends {fr, to}",
+        Default::default(),
+    )
+    .unwrap();
+
+    let rels_data = db
+        .export_relations(["friends", "friends:rev"].into_iter())
+        .unwrap();
+    assert_eq!(
+        rels_data["friends"].clone().into_json()["rows"],
+        json!([[1, 2, 5], [6, 5, 7]])
+    );
+    assert_eq!(
+        rels_data["friends:rev"].clone().into_json()["rows"],
+        json!([[2, 1], [5, 6]])
+    );
+
+    let rels = db.run_script("::relations", Default::default()).unwrap();
+    assert_eq!(rels.rows[1][0], DataValue::from("friends:rev"));
+    assert_eq!(rels.rows[1][1], DataValue::from(2));
+    assert_eq!(rels.rows[1][2], DataValue::from("index"));
+
+    let cols = db
+        .run_script("::columns friends:rev", Default::default())
+        .unwrap();
+    assert_eq!(cols.rows.len(), 2);
+
+    let expl = db
+        .run_script(
+            "::explain { ?[fr, data] := *friends{to: 2, fr, data} }",
+            Default::default(),
+        )
+        .unwrap()
+        .into_json();
+
+    for row in expl["rows"].as_array().unwrap() {
+        println!("{}", row);
+    }
+
+    let joins = expl["rows"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|row| row.as_array().unwrap()[5].clone())
+        .collect_vec();
+    assert!(joins.contains(&json!(":friends:rev")));
+
+    let res = db
+        .run_script(
+            "?[fr, data] := *friends{to: 2, fr, data}",
+            Default::default(),
+        )
+        .unwrap();
+    assert_eq!(res.into_json()["rows"], json!([[1, 5]]));
+}