HNSW compilation

cozo-core/src/query/compile.rs

@@ -211,7 +211,40 @@ impl<'a> SessionTx<'a> {
                     ret = ret.join(right, prev_joiner_vars, right_joiner_vars, rule_app.span);
                 }
                 MagicAtom::HnswSearch(s) => {
-                    todo!("HNSW search")
+                    // already existing vars
+                    let mut prev_joiner_vars = vec![];
+                    // vars introduced by right and joined
+                    let mut right_joiner_vars = vec![];
+                    // used to split in case we need to join again
+                    let mut right_joiner_vars_pos = vec![];
+                    // vars introduced by right, regardless of joining
+                    let mut right_vars = vec![];
+                    // used for choosing indices
+                    let mut join_indices = vec![];
+
+                    for (i, var) in s.all_bindings().enumerate() {
+                        if seen_variables.contains(var) {
+                            prev_joiner_vars.push(var.clone());
+                            let rk = gen_symb(var.span);
+                            right_vars.push(rk.clone());
+                            right_joiner_vars.push(rk);
+                            right_joiner_vars_pos.push(i);
+                            join_indices.push(IndexPositionUse::Join)
+                        } else {
+                            seen_variables.insert(var.clone());
+                            right_vars.push(var.clone());
+                            if var.is_generated_ignored_symbol() {
+                                join_indices.push(IndexPositionUse::Ignored)
+                            } else {
+                                join_indices.push(IndexPositionUse::BindForLater)
+                            }
+                        }
+                    }
+
+                    // scan original relation
+                    let right = RelAlgebra::hnsw_search(right_vars, s.clone())?;
+                    debug_assert_eq!(prev_joiner_vars.len(), right_joiner_vars.len());
+                    ret = ret.join(right, prev_joiner_vars, right_joiner_vars, s.span);
                 }
                 MagicAtom::Relation(rel_app) => {
                     let store = self.get_relation(&rel_app.name, false)?;

cozo-core/src/query/ra.rs

@@ -17,7 +17,7 @@ use miette::{bail, Diagnostic, Result};
 use thiserror::Error;
 
 use crate::data::expr::{compute_bounds, eval_bytecode, eval_bytecode_pred, Bytecode, Expr};
-use crate::data::program::MagicSymbol;
+use crate::data::program::{HnswSearch, MagicSymbol};
 use crate::data::relation::{ColType, NullableColType};
 use crate::data::symb::Symbol;
 use crate::data::tuple::{Tuple, TupleIter};
@@ -53,17 +53,11 @@ impl RelAlgebra {
             RelAlgebra::Filter(i) => i.span,
             RelAlgebra::Unification(i) => i.span,
             RelAlgebra::StoredWithValidity(i) => i.span,
-            RelAlgebra::HnswSearch(_) => todo!(),
+            RelAlgebra::HnswSearch(i) => i.hnsw_search.span,
         }
     }
 }
 
-pub(crate) struct HnswSearchRA {
-    pub(crate) parent: Box<RelAlgebra>,
-    pub(crate) to_eliminate: BTreeSet<Symbol>,
-    pub(crate) span: SourceSpan,
-}
-
 pub(crate) struct UnificationRA {
     pub(crate) parent: Box<RelAlgebra>,
     pub(crate) binding: Symbol,
@@ -284,6 +278,11 @@ impl Debug for RelAlgebra {
                 .field(&r.storage.name)
                 .field(&r.filters)
                 .finish(),
+            RelAlgebra::HnswSearch(s) => f
+                .debug_tuple("HnswSearch")
+                .field(&bindings)
+                .field(&s.hnsw_search.idx_handle.name)
+                .finish(),
             RelAlgebra::StoredWithValidity(r) => f
                 .debug_tuple("StoredWithValidity")
                 .field(&bindings)
@@ -328,9 +327,6 @@ impl Debug for RelAlgebra {
                 .field(&r.binding)
                 .field(&r.expr)
                 .finish(),
-            RelAlgebra::HnswSearch(_) => {
-                todo!("HnswSearch")
-            }
         }
     }
 }
@@ -353,6 +349,9 @@ impl RelAlgebra {
             RelAlgebra::Stored(v) => {
                 v.fill_binding_indices_and_compile()?;
             }
+            RelAlgebra::HnswSearch(_) => {
+                todo!("fill_binding_indices_and_compile for HnswSearch")
+            }
             RelAlgebra::StoredWithValidity(v) => {
                 v.fill_binding_indices_and_compile()?;
             }
@@ -374,9 +373,6 @@ impl RelAlgebra {
                 r.left.fill_binding_indices_and_compile()?;
                 r.right.fill_binding_indices_and_compile()?;
             }
-            RelAlgebra::HnswSearch(_) => {
-                todo!()
-            }
         }
         Ok(())
     }
@@ -406,6 +402,12 @@ impl RelAlgebra {
             span,
         })
     }
+    pub(crate) fn hnsw_search(bindings: Vec<Symbol>, search: HnswSearch) -> Result<Self> {
+        Ok(Self::HnswSearch(HnswSearchRA {
+            bindings,
+            hnsw_search: search,
+        }))
+    }
     pub(crate) fn relation(
         bindings: Vec<Symbol>,
         storage: RelationHandle,
@@ -838,6 +840,12 @@ pub(crate) struct StoredRA {
     pub(crate) span: SourceSpan,
 }
 
+#[derive(Debug)]
+pub(crate) struct HnswSearchRA {
+    pub(crate) bindings: Vec<Symbol>,
+    pub(crate) hnsw_search: HnswSearch,
+}
+
 #[derive(Debug)]
 pub(crate) struct StoredWithValidityRA {
     pub(crate) bindings: Vec<Symbol>,
@@ -1573,7 +1581,7 @@ impl RelAlgebra {
             RelAlgebra::Filter(r) => r.do_eliminate_temp_vars(used),
             RelAlgebra::NegJoin(r) => r.do_eliminate_temp_vars(used),
             RelAlgebra::Unification(r) => r.do_eliminate_temp_vars(used),
-            RelAlgebra::HnswSearch(_) => {todo!()}
+            RelAlgebra::HnswSearch(_) => Ok(()),
         }
     }
 
@@ -1588,7 +1596,7 @@ impl RelAlgebra {
             RelAlgebra::Filter(r) => Some(&r.to_eliminate),
             RelAlgebra::NegJoin(r) => Some(&r.to_eliminate),
             RelAlgebra::Unification(u) => Some(&u.to_eliminate),
-            RelAlgebra::HnswSearch(_) => {todo!()}
+            RelAlgebra::HnswSearch(_) => None,
         }
     }
 
@@ -1618,7 +1626,7 @@ impl RelAlgebra {
                 bindings.push(u.binding.clone());
                 bindings
             }
-            RelAlgebra::HnswSearch(_) => {todo!()}
+            RelAlgebra::HnswSearch(s) => s.bindings.clone(),
         }
     }
     pub(crate) fn iter<'a>(
@@ -1637,7 +1645,9 @@ impl RelAlgebra {
             RelAlgebra::Filter(r) => r.iter(tx, delta_rule, stores),
             RelAlgebra::NegJoin(r) => r.iter(tx, delta_rule, stores),
             RelAlgebra::Unification(r) => r.iter(tx, delta_rule, stores),
-            RelAlgebra::HnswSearch(_) => {todo!()}
+            RelAlgebra::HnswSearch(_) => {
+                todo!()
+            }
         }
     }
 }
@@ -1817,6 +1827,7 @@ impl InnerJoin {
                     "stored_mat_join"
                 }
             }
+            RelAlgebra::HnswSearch(_) => "hnsw_search_join",
             RelAlgebra::StoredWithValidity(_) => {
                 let join_indices = self
                     .joiner
@@ -1840,9 +1851,6 @@ impl InnerJoin {
             RelAlgebra::NegJoin(_) => {
                 panic!("joining on NegJoin")
             }
-            RelAlgebra::HnswSearch(_) => {
-                todo!("joining on HnswSearch")
-            }
         }
     }
     pub(crate) fn iter<'a>(

cozo-core/src/runtime/tests.rs

@@ -824,7 +824,7 @@ fn test_vec_index() {
     .unwrap();
     db.run_script(
         r"
-        ?[key, v] := ~a:vec{k: 'a', v | query: [1,1,1,1,1,1,1,1]}
+        ?[v] := ~a:vec{k: 'a', v | query: [1,1,1,1,1,1,1,1]}
         ",
         Default::default(),
     )