add tests for limit/offset

cozo-core/src/query/relation.rs

@@ -1448,7 +1448,7 @@ pub(crate) struct InnerJoin {
     pub(crate) joiner: Joiner,
     pub(crate) to_eliminate: BTreeSet<Symbol>,
     pub(crate) span: SourceSpan,
-    mat_right_cache: RefCell<Rc<Vec<Vec<DataValue>>>>,
+    mat_right_cache: RefCell<Rc<Vec<Tuple>>>,
     cached: AtomicBool,
 }
 
@@ -1614,7 +1614,7 @@ impl InnerJoin {
         let left_cache = match left_iter.next() {
             None => return Ok(Box::new(iter::empty())),
             Some(Err(err)) => return Err(err),
-            Some(Ok(data)) => data.0,
+            Some(Ok(data)) => data,
         };
 
         let right_join_indices_set = BTreeSet::from_iter(right_join_indices.iter().cloned());
@@ -1642,7 +1642,7 @@ impl InnerJoin {
                             .iter()
                             .map(|i| tuple.0[*i].clone())
                             .collect_vec();
-                        cache.insert(stored_tuple);
+                        cache.insert(Tuple(stored_tuple));
                     }
                     Err(e) => return Err(e),
                 }
@@ -1659,7 +1659,7 @@ impl InnerJoin {
         let it = CachedMaterializedIterator {
             eliminate_indices,
             left: left_iter,
-            left_cache,
+            left_cache: left_cache,
             left_join_indices,
             materialized: cached_data,
             right_invert_indices,
@@ -1671,23 +1671,23 @@ impl InnerJoin {
 }
 
 struct CachedMaterializedIterator<'a> {
-    materialized: Rc<Vec<Vec<DataValue>>>,
+    materialized: Rc<Vec<Tuple>>,
     eliminate_indices: BTreeSet<usize>,
     left_join_indices: Vec<usize>,
     right_invert_indices: Vec<usize>,
     right_idx: usize,
-    prefix: Vec<DataValue>,
+    prefix: Tuple,
     left: TupleIter<'a>,
-    left_cache: Vec<DataValue>,
+    left_cache: Tuple,
 }
 
 impl<'a> CachedMaterializedIterator<'a> {
-    fn advance_right(&mut self) -> Option<&Vec<DataValue>> {
+    fn advance_right(&mut self) -> Option<&Tuple> {
         if self.right_idx == self.materialized.len() {
             None
         } else {
             let ret = &self.materialized[self.right_idx];
-            if ret.starts_with(&self.prefix) {
+            if ret.0.starts_with(&self.prefix.0) {
                 self.right_idx += 1;
                 Some(ret)
             } else {
@@ -1703,9 +1703,9 @@ impl<'a> CachedMaterializedIterator<'a> {
                     let data = data.clone();
                     let mut ret = self.left_cache.clone();
                     for i in &self.right_invert_indices {
-                        ret.push(data[*i].clone());
+                        ret.0.push(data.0[*i].clone());
                     }
-                    let tuple = eliminate_from_tuple(Tuple(ret), &self.eliminate_indices);
+                    let tuple = eliminate_from_tuple(ret, &self.eliminate_indices);
                     return Ok(Some(tuple));
                 }
                 None => {
@@ -1713,7 +1713,7 @@ impl<'a> CachedMaterializedIterator<'a> {
                     match next_left {
                         None => return Ok(None),
                         Some(l) => {
-                            let left_tuple = l?.0;
+                            let left_tuple = l?;
                             let (prefix, idx) = build_mat_range_iter(
                                 &self.materialized,
                                 &self.left_join_indices,
@@ -1732,14 +1732,16 @@ impl<'a> CachedMaterializedIterator<'a> {
 }
 
 fn build_mat_range_iter(
-    mat: &[Vec<DataValue>],
+    mat: &[Tuple],
     left_join_indices: &[usize],
-    left_tuple: &[DataValue],
-) -> (Vec<DataValue>, usize) {
-    let prefix = left_join_indices
-        .iter()
-        .map(|i| left_tuple[*i].clone())
-        .collect_vec();
+    left_tuple: &Tuple,
+) -> (Tuple, usize) {
+    let prefix = Tuple(
+        left_join_indices
+            .iter()
+            .map(|i| left_tuple.0[*i].clone())
+            .collect_vec(),
+    );
     let idx = match mat.binary_search(&prefix) {
         Ok(i) => i,
         Err(i) => i,

cozo-core/src/runtime/db.rs

@@ -644,8 +644,7 @@ impl<'s, S: Storage<'s>> Db<S> {
         match op {
             SysOp::Explain(prog) => {
                 let mut tx = self.transact()?;
-                let (stratified_program, _) =
-                    prog.to_normalized_program(&tx)?.stratify()?;
+                let (stratified_program, _) = prog.to_normalized_program(&tx)?.stratify()?;
                 let program = stratified_program.magic_sets_rewrite(&tx)?;
                 let (compiled, _) = tx.stratified_magic_compile(&program)?;
                 tx.commit_tx()?;
@@ -1082,3 +1081,48 @@ impl Poison {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use itertools::Itertools;
+    use serde_json::json;
+
+    use crate::new_cozo_mem;
+
+    #[test]
+    fn test_limit_offset() {
+        let db = new_cozo_mem().unwrap();
+        let res = db
+            .run_script("?[a] := a in [5,3,1,2,4] :limit 2", Default::default())
+            .unwrap()
+            .rows
+            .into_iter()
+            .flatten()
+            .collect_vec();
+        assert_eq!(json!(res), json!([3, 5]));
+        let res = db
+            .run_script("?[a] := a in [5,3,1,2,4] :limit 2 :offset 1", Default::default())
+            .unwrap()
+            .rows
+            .into_iter()
+            .flatten()
+            .collect_vec();
+        assert_eq!(json!(res), json!([1, 3]));
+        let res = db
+            .run_script("?[a] := a in [5,3,1,2,4] :limit 2 :offset 4", Default::default())
+            .unwrap()
+            .rows
+            .into_iter()
+            .flatten()
+            .collect_vec();
+        assert_eq!(json!(res), json!([4]));
+        let res = db
+            .run_script("?[a] := a in [5,3,1,2,4] :limit 2 :offset 5", Default::default())
+            .unwrap()
+            .rows
+            .into_iter()
+            .flatten()
+            .collect_vec();
+        assert_eq!(json!(res), json!([]));
+    }
+}

cozo-core/src/runtime/in_mem.rs

@@ -36,9 +36,6 @@ impl Debug for StoredRelationId {
 pub(crate) struct InMemRelation {
     mem_db: Rc<RefCell<Vec<Rc<RefCell<BTreeMap<Tuple, Tuple>>>>>>,
     epoch_size: Arc<AtomicU32>,
-    // total: Rc<RefCell<BTreeMap<Tuple, Tuple>>>,
-    // current: Rc<RefCell<BTreeMap<Tuple, Tuple>>>,
-    // prev: Rc<RefCell<BTreeMap<Tuple, Tuple>>>,
     pub(crate) id: StoredRelationId,
     pub(crate) rule_name: MagicSymbol,
     pub(crate) arity: usize,
@@ -99,6 +96,7 @@ impl InMemRelation {
                     if ma.is_none() {
                         tuple.0[i].clone()
                     } else {
+                        // placeholder for meet aggregation
                         DataValue::Guard
                     }
                 })
@@ -131,6 +129,7 @@ impl InMemRelation {
                         if aggr.is_some() {
                             Ok(tuple.0[i].clone())
                         } else {
+                            // placeholder for key part
                             Ok(DataValue::Guard)
                         }
                     })
@@ -163,6 +162,7 @@ impl InMemRelation {
         let mut epoch_map = epoch_map.borrow_mut();
 
         if should_skip {
+            // put guard, so that when iterating results, those with guards are ignored
             epoch_map.insert(tuple, Tuple(vec![DataValue::Guard]));
         } else {
             epoch_map.insert(tuple, Tuple::default());
@@ -198,6 +198,7 @@ impl InMemRelation {
                         k.0.iter()
                             .zip(v.0.iter())
                             .map(|(kel, vel)| {
+                                // merge meet aggregation kv
                                 if matches!(kel, DataValue::Guard) {
                                     vel.clone()
                                 } else {
@@ -227,12 +228,14 @@ impl InMemRelation {
                 if v.0.is_empty() {
                     Some(k.clone())
                 } else if v.0.last() == Some(&DataValue::Guard) {
+                    // ignore since we are using :offset
                     None
                 } else {
                     let combined =
                         k.0.iter()
                             .zip(v.0.iter())
                             .map(|(kel, vel)| {
+                                // merge kv parts of meet aggr
                                 if matches!(kel, DataValue::Guard) {
                                     vel.clone()
                                 } else {
@@ -273,6 +276,7 @@ impl InMemRelation {
                         k.0.iter()
                             .zip(v.0.iter())
                             .map(|(kel, vel)| {
+                                // merge kv parts of meet aggr
                                 if matches!(kel, DataValue::Guard) {
                                     vel.clone()
                                 } else {
@@ -311,3 +315,5 @@ impl InMemRelation {
         res.into_iter().map(Ok)
     }
 }
+
+// meet put