point lookup optimization

cozo-core/src/data/expr.rs

@@ -309,7 +309,7 @@ impl Expr {
             }
         }
     }
-    pub(crate) fn eval(&self, bindings: &Tuple) -> Result<DataValue> {
+    pub(crate) fn eval(&self, bindings: impl AsRef<[DataValue]>) -> Result<DataValue> {
         match self {
             Expr::Binding { var, tuple_pos, .. } => match tuple_pos {
                 None => {
@@ -327,29 +327,29 @@ impl Expr {
                     #[diagnostic(code(eval::unbound))]
                     struct TupleTooShortError(String, usize, usize, #[label] SourceSpan);
 
-                    Ok(bindings
+                    Ok(bindings.as_ref()
                         .get(*i)
                         .ok_or_else(|| {
-                            TupleTooShortError(var.name.to_string(), *i, bindings.len(), var.span)
+                            TupleTooShortError(var.name.to_string(), *i, bindings.as_ref().len(), var.span)
                         })?
                         .clone())
                 }
             },
             Expr::Const { val, .. } => Ok(val.clone()),
             Expr::Apply { op, args, .. } => {
-                let args: Box<[DataValue]> = args.iter().map(|v| v.eval(bindings)).try_collect()?;
+                let args: Box<[DataValue]> = args.iter().map(|v| v.eval(bindings.as_ref())).try_collect()?;
                 Ok((op.inner)(&args)
                     .map_err(|err| EvalRaisedError(self.span(), err.to_string()))?)
             }
             Expr::Cond { clauses, .. } => {
                 for (cond, val) in clauses {
-                    let cond_val = cond.eval(bindings)?;
+                    let cond_val = cond.eval(bindings.as_ref())?;
                     let cond_val = cond_val
                         .get_bool()
                         .ok_or_else(|| PredicateTypeError(cond.span(), cond_val))?;
 
                     if cond_val {
-                        return val.eval(bindings);
+                        return val.eval(bindings.as_ref());
                     }
                 }
                 Ok(DataValue::Null)
@@ -359,17 +359,17 @@ impl Expr {
                     Ok(DataValue::Null)
                 } else {
                     for item in clauses.iter().take(clauses.len() - 1) {
-                        let res = item.eval(bindings);
+                        let res = item.eval(bindings.as_ref());
                         if res.is_ok() {
                             return res;
                         }
                     }
-                    clauses[clauses.len() - 1].eval(bindings)
+                    clauses[clauses.len() - 1].eval(bindings.as_ref())
                 }
             }
         }
     }
-    pub(crate) fn eval_pred(&self, bindings: &Tuple) -> Result<bool> {
+    pub(crate) fn eval_pred(&self, bindings: impl AsRef<[DataValue]>) -> Result<bool> {
         match self.eval(bindings)? {
             DataValue::Bool(b) => Ok(b),
             v => {

cozo-core/src/data/tuple.rs

@@ -18,29 +18,30 @@ pub(crate) type TupleIter<'a> = Box<dyn Iterator<Item = Result<Tuple>> + 'a>;
 
 pub(crate) trait TupleT {
     fn encode_as_key(&self, prefix: RelationId) -> Vec<u8>;
-    fn decode_from_key(key: &[u8]) -> Self;
 }
 
-impl TupleT for  Tuple {
+impl<T> TupleT for T where T: AsRef<[DataValue]> {
     fn encode_as_key(&self, prefix: RelationId) -> Vec<u8> {
-        let len = self.len();
+        let len = self.as_ref().len();
         let mut ret = Vec::with_capacity(4 + 4 * len + 10 * len);
         let prefix_bytes = prefix.0.to_be_bytes();
         ret.extend(prefix_bytes);
-        for val in self.iter() {
+        for val in self.as_ref().iter() {
             ret.encode_datavalue(val);
         }
         ret
     }
-    fn decode_from_key(key: &[u8]) -> Self {
-        let mut remaining = &key[ENCODED_KEY_MIN_LEN..];
-        let mut ret = vec![];
-        while !remaining.is_empty() {
-            let (val, next) = DataValue::decode_from_key(remaining);
-            ret.push(val);
-            remaining = next;
-        }
-        ret
+}
+
+pub(crate) fn decode_tuple_from_key(key: &[u8]) -> Tuple {
+    let mut remaining = &key[ENCODED_KEY_MIN_LEN..];
+    let mut ret = vec![];
+    while !remaining.is_empty() {
+        let (val, next) = DataValue::decode_from_key(remaining);
+        ret.push(val);
+        remaining = next;
     }
+    ret
 }
+
 pub(crate) const ENCODED_KEY_MIN_LEN: usize = 8;

cozo-core/src/query/relation.rs

@@ -770,6 +770,72 @@ impl StoredRA {
             .map(|(a, _)| left_join_indices[a])
             .collect_vec();
 
+        let key_len = self.storage.metadata.keys.len();
+        if left_to_prefix_indices.len() >= key_len {
+            // We can use point lookups, finally
+            // TODO can use this branch if no values are touched!
+            return if self.storage.metadata.non_keys.is_empty() {
+                let it = left_iter
+                    .map_ok(move |tuple| -> Result<Option<Tuple>> {
+                        let prefix = left_to_prefix_indices
+                            .iter()
+                            .map(|i| tuple[*i].clone())
+                            .collect_vec();
+                        let filters = self.filters.clone();
+                        let key = &prefix[0..key_len];
+                        for p in filters.iter() {
+                            if !p.eval_pred(key)? {
+                                return Ok(None);
+                            }
+                        }
+                        if self.storage.exists(tx, key)? {
+                            let mut ret = tuple.clone();
+                            ret.extend_from_slice(key);
+                            Ok(Some(ret))
+                        } else {
+                            Ok(None)
+                        }
+                    })
+                    .flatten_ok()
+                    .filter_map(invert_option_err);
+                Ok(if eliminate_indices.is_empty() {
+                    Box::new(it)
+                } else {
+                    Box::new(it.map_ok(move |t| eliminate_from_tuple(t, &eliminate_indices)))
+                })
+            } else {
+                let it = left_iter
+                    .map_ok(move |tuple| -> Result<Option<Tuple>> {
+                        let prefix = left_to_prefix_indices
+                            .iter()
+                            .map(|i| tuple[*i].clone())
+                            .collect_vec();
+                        let filters = self.filters.clone();
+                        let key = &prefix[0..key_len];
+                        match self.storage.get(tx, key)? {
+                            None => Ok(None),
+                            Some(found) => {
+                                for p in filters.iter() {
+                                    if !p.eval_pred(&found)? {
+                                        return Ok(None);
+                                    }
+                                }
+                                let mut ret = tuple.clone();
+                                ret.extend(found);
+                                Ok(Some(ret))
+                            }
+                        }
+                    })
+                    .flatten_ok()
+                    .filter_map(invert_option_err);
+                Ok(if eliminate_indices.is_empty() {
+                    Box::new(it)
+                } else {
+                    Box::new(it.map_ok(move |t| eliminate_from_tuple(t, &eliminate_indices)))
+                })
+            };
+        }
+
         let mut skip_range_check = false;
         let it = left_iter
             .map_ok(move |tuple| {

cozo-core/src/runtime/relation.rs

@@ -19,7 +19,7 @@ use thiserror::Error;
 use crate::data::memcmp::MemCmpEncoder;
 use crate::data::relation::StoredRelationMetadata;
 use crate::data::symb::Symbol;
-use crate::data::tuple::{Tuple, TupleT, ENCODED_KEY_MIN_LEN};
+use crate::data::tuple::{Tuple, TupleT, ENCODED_KEY_MIN_LEN, decode_tuple_from_key};
 use crate::data::value::DataValue;
 use crate::parse::SourceSpan;
 use crate::runtime::transact::SessionTx;
@@ -225,6 +225,19 @@ impl RelationHandle {
         tx.tx.range_scan_tuple(&lower, &upper)
     }
 
+    pub(crate) fn get<'a>(&self, tx: &'a SessionTx<'_>, key: &[DataValue]) -> Result<Option<Tuple>> {
+        let key_data = key.encode_as_key(self.id);
+        Ok(match tx.tx.get(&key_data, false)? {
+            None => None,
+            Some(val_data) => Some(decode_tuple_from_kv(&key_data, &val_data)),
+        })
+    }
+
+    pub(crate) fn exists<'a>(&self, tx: &'a SessionTx<'_>, key: &[DataValue]) -> Result<bool> {
+        let key_data = key.encode_as_key(self.id);
+        Ok(tx.tx.exists(&key_data, false)?)
+    }
+
     pub(crate) fn scan_prefix<'a>(
         &self,
         tx: &'a SessionTx<'_>,
@@ -261,7 +274,7 @@ impl RelationHandle {
 /// in trait [`StoreTx`](crate::StoreTx).
 #[inline]
 pub fn decode_tuple_from_kv(key: &[u8], val: &[u8]) -> Tuple {
-    let mut tup = Tuple::decode_from_key(key);
+    let mut tup = decode_tuple_from_key(key);
     if !val.is_empty() {
         let vals: Vec<DataValue> = rmp_serde::from_slice(&val[ENCODED_KEY_MIN_LEN..]).unwrap();
         tup.extend(vals);