diff --git a/cozo-core/src/query/relation.rs b/cozo-core/src/query/relation.rs index a00529db..ff8580ca 100644 --- a/cozo-core/src/query/relation.rs +++ b/cozo-core/src/query/relation.rs @@ -1448,7 +1448,7 @@ pub(crate) struct InnerJoin { pub(crate) joiner: Joiner, pub(crate) to_eliminate: BTreeSet, pub(crate) span: SourceSpan, - mat_right_cache: RefCell>>>, + mat_right_cache: RefCell>>, 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>>, + materialized: Rc>, eliminate_indices: BTreeSet, left_join_indices: Vec, right_invert_indices: Vec, right_idx: usize, - prefix: Vec, + prefix: Tuple, left: TupleIter<'a>, - left_cache: Vec, + left_cache: Tuple, } impl<'a> CachedMaterializedIterator<'a> { - fn advance_right(&mut self) -> Option<&Vec> { + 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], + mat: &[Tuple], left_join_indices: &[usize], - left_tuple: &[DataValue], -) -> (Vec, 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, diff --git a/cozo-core/src/runtime/db.rs b/cozo-core/src/runtime/db.rs index 45b14424..2519eb6e 100644 --- a/cozo-core/src/runtime/db.rs +++ b/cozo-core/src/runtime/db.rs @@ -644,8 +644,7 @@ impl<'s, S: Storage<'s>> Db { 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!([])); + } +} diff --git a/cozo-core/src/runtime/in_mem.rs b/cozo-core/src/runtime/in_mem.rs index f57ca976..60e99796 100644 --- a/cozo-core/src/runtime/in_mem.rs +++ b/cozo-core/src/runtime/in_mem.rs @@ -36,9 +36,6 @@ impl Debug for StoredRelationId { pub(crate) struct InMemRelation { mem_db: Rc>>>>>, epoch_size: Arc, - // total: Rc>>, - // current: Rc>>, - // prev: Rc>>, 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 \ No newline at end of file