gqlite 1.5.0 → 1.7.0

data/ext/db-index/src/hnsw/metric.rs

@@ -0,0 +1,244 @@
+use super::kernels::cosine_distance;
+use super::kernels::cosine_distance_qi8;
+use super::kernels::inner_product_distance;
+use super::kernels::inner_product_distance_qi8;
+use super::kernels::l2_sq_qi8;
+use super::kernels::Kernel;
+use super::scalar::Scalar;
+use super::vector::Dense;
+use super::vector::Qi8;
+use super::vector::Qi8Ref;
+use super::vector::VectorFamily;
+use std::marker::PhantomData;
+
+pub trait Metric: Clone + Send + Sync + 'static
+{
+  type Family: VectorFamily;
+  fn distance<'a, 'b>(
+    &self,
+    a: <Self::Family as VectorFamily>::Ref<'a>,
+    b: <Self::Family as VectorFamily>::Ref<'b>,
+  ) -> f32;
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct L2<S: Scalar = f32>(PhantomData<S>);
+
+impl<S: Scalar> L2<S>
+{
+  pub fn new() -> Self
+  {
+    Self(PhantomData)
+  }
+}
+
+impl<S: Scalar> Metric for L2<S>
+{
+  type Family = Dense<S>;
+
+  fn distance(&self, a: &[S], b: &[S]) -> f32
+  {
+    debug_assert_eq!(a.len(), b.len());
+    <S as Kernel>::l2_sq(a, b)
+  }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct InnerProduct<S: Scalar = f32>(PhantomData<S>);
+
+impl<S: Scalar> InnerProduct<S>
+{
+  pub fn new() -> Self
+  {
+    Self(PhantomData)
+  }
+}
+
+impl<S: Scalar> Metric for InnerProduct<S>
+{
+  type Family = Dense<S>;
+
+  fn distance(&self, a: &[S], b: &[S]) -> f32
+  {
+    debug_assert_eq!(a.len(), b.len());
+    inner_product_distance::<S>(a, b)
+  }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct Cosine<S: Scalar = f32>(PhantomData<S>);
+
+impl<S: Scalar> Cosine<S>
+{
+  pub fn new() -> Self
+  {
+    Self(PhantomData)
+  }
+}
+
+impl<S: Scalar> Metric for Cosine<S>
+{
+  type Family = Dense<S>;
+
+  fn distance(&self, a: &[S], b: &[S]) -> f32
+  {
+    debug_assert_eq!(a.len(), b.len());
+    cosine_distance::<S>(a, b)
+  }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct L2Qi8;
+
+impl L2Qi8
+{
+  pub fn new() -> Self
+  {
+    Self
+  }
+}
+
+impl Metric for L2Qi8
+{
+  type Family = Qi8;
+
+  fn distance<'a, 'b>(&self, a: Qi8Ref<'a>, b: Qi8Ref<'b>) -> f32
+  {
+    l2_sq_qi8(a.data, a.scale, a.zero_point, b.data, b.scale, b.zero_point)
+  }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct InnerProductQi8;
+
+impl InnerProductQi8
+{
+  pub fn new() -> Self
+  {
+    Self
+  }
+}
+
+impl Metric for InnerProductQi8
+{
+  type Family = Qi8;
+
+  fn distance<'a, 'b>(&self, a: Qi8Ref<'a>, b: Qi8Ref<'b>) -> f32
+  {
+    inner_product_distance_qi8(a.data, a.scale, a.zero_point, b.data, b.scale, b.zero_point)
+  }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct CosineQi8;
+
+impl CosineQi8
+{
+  pub fn new() -> Self
+  {
+    Self
+  }
+}
+
+impl Metric for CosineQi8
+{
+  type Family = Qi8;
+
+  fn distance<'a, 'b>(&self, a: Qi8Ref<'a>, b: Qi8Ref<'b>) -> f32
+  {
+    cosine_distance_qi8(a.data, a.scale, a.zero_point, b.data, b.scale, b.zero_point)
+  }
+}
+
+pub fn normalize_cosine_in_place<S: Scalar>(vector: &mut [S])
+{
+  let mut norm_sq = 0.0_f32;
+  for &v in vector.iter()
+  {
+    let v = v.to_f32();
+    norm_sq += v * v;
+  }
+  if norm_sq == 0.0
+  {
+    return;
+  }
+  let inv_norm = norm_sq.sqrt().recip();
+  for v in vector.iter_mut()
+  {
+    *v = S::from_f32(v.to_f32() * inv_norm);
+  }
+}
+
+#[cfg(test)]
+mod tests
+{
+  use super::*;
+  use approx::assert_relative_eq;
+  use rand_09::prelude::Rng;
+  use rand_09::rngs::StdRng;
+  use rand_09::SeedableRng;
+
+  fn l2_ref(a: &[f32], b: &[f32]) -> f32
+  {
+    a.iter()
+      .zip(b.iter())
+      .map(|(a, b)| {
+        let d = a - b;
+        d * d
+      })
+      .sum()
+  }
+
+  fn ip_ref(a: &[f32], b: &[f32]) -> f32
+  {
+    1.0 - a.iter().zip(b.iter()).map(|(a, b)| a * b).sum::<f32>()
+  }
+
+  #[test]
+  fn l2_matches_scalar_across_dims()
+  {
+    let mut rng = StdRng::seed_from_u64(123);
+    let dims = [
+      1usize, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 255,
+    ];
+    let metric = L2::<f32>::new();
+    for &dim in &dims
+    {
+      for _ in 0..100
+      {
+        let a: Vec<f32> = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect();
+        let b: Vec<f32> = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect();
+        assert_relative_eq!(
+          metric.distance(&a, &b),
+          l2_ref(&a, &b),
+          epsilon = 1e-3,
+          max_relative = 1e-3
+        );
+      }
+    }
+  }
+
+  #[test]
+  fn inner_product_matches_scalar_across_dims()
+  {
+    let mut rng = StdRng::seed_from_u64(456);
+    let dims = [
+      1usize, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 255,
+    ];
+    let metric = InnerProduct::<f32>::new();
+    for &dim in &dims
+    {
+      for _ in 0..100
+      {
+        let a: Vec<f32> = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect();
+        let b: Vec<f32> = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect();
+        assert_relative_eq!(
+          metric.distance(&a, &b),
+          ip_ref(&a, &b),
+          epsilon = 1e-3,
+          max_relative = 1e-3
+        );
+      }
+    }
+  }
+}

data/ext/db-index/src/hnsw/scalar.rs

@@ -0,0 +1,72 @@
+use super::kernels::Kernel;
+use bytemuck::Pod;
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[repr(u8)]
+pub enum Dtype
+{
+  F32 = 0,
+  F16 = 1,
+  BF16 = 2,
+  QI8 = 3,
+}
+
+pub trait Scalar: Kernel + Pod
+{
+  const DTYPE: Dtype;
+  fn to_f32(self) -> f32;
+  fn from_f32(v: f32) -> Self;
+}
+
+impl Scalar for f32
+{
+  const DTYPE: Dtype = Dtype::F32;
+
+  #[inline]
+  fn to_f32(self) -> f32
+  {
+    self
+  }
+
+  #[inline]
+  fn from_f32(v: f32) -> Self
+  {
+    v
+  }
+}
+
+impl Scalar for half::f16
+{
+  const DTYPE: Dtype = Dtype::F16;
+
+  #[inline]
+  fn to_f32(self) -> f32
+  {
+    self.to_f32()
+  }
+
+  #[inline]
+  fn from_f32(v: f32) -> Self
+  {
+    Self::from_f32(v)
+  }
+}
+
+impl Scalar for half::bf16
+{
+  const DTYPE: Dtype = Dtype::BF16;
+
+  #[inline]
+  fn to_f32(self) -> f32
+  {
+    self.to_f32()
+  }
+
+  #[inline]
+  fn from_f32(v: f32) -> Self
+  {
+    Self::from_f32(v)
+  }
+}

data/ext/db-index/src/hnsw/simple_store.rs

@@ -0,0 +1,140 @@
+use super::error::Result;
+use super::graph_store::{GraphStoreRead, GraphStoreWrite};
+use crate::hnsw::id::NodeId;
+use std::collections::HashMap;
+use std::hash::Hash;
+
+/// Reference implementation of a graph store.
+///
+/// A simple, single-threaded, growable store with no locks or atomics.
+/// Provides conformance and testing baseline for the algorithm.
+pub struct SimpleGraphStore<K>
+{
+  levels: Vec<Option<i32>>, // None == removed or never set
+  keys: Vec<Option<K>>,
+  links: Vec<Vec<Vec<NodeId>>>, // [node][level] -> neighbor ids
+  by_key: HashMap<K, NodeId>,
+  entry: Option<(NodeId, i32)>,
+  next_node_id: usize, // monotonic, never decremented
+}
+
+impl<K> SimpleGraphStore<K>
+{
+  pub fn new() -> Self
+  {
+    Self {
+      levels: Vec::new(),
+      keys: Vec::new(),
+      links: Vec::new(),
+      by_key: HashMap::new(),
+      entry: None,
+      next_node_id: 0,
+    }
+  }
+}
+
+impl<K: Clone + Eq + Hash> GraphStoreRead<K> for SimpleGraphStore<K>
+{
+  type Error = std::convert::Infallible;
+
+  fn node_count(&self) -> Result<usize, Self::Error>
+  {
+    Ok(self.next_node_id)
+  }
+
+  fn entry_point(&self) -> Result<Option<(NodeId, i32)>, Self::Error>
+  {
+    Ok(self.entry)
+  }
+
+  fn contains(&self, n: NodeId) -> Result<bool, Self::Error>
+  {
+    Ok(self.levels.get(n.as_usize()).is_some_and(|l| l.is_some()))
+  }
+
+  fn key_of(&self, n: NodeId) -> Result<Option<K>, Self::Error>
+  {
+    Ok(self.keys.get(n.as_usize()).and_then(|k| k.clone()))
+  }
+
+  fn node_of(&self, key: &K) -> Result<Option<NodeId>, Self::Error>
+  {
+    Ok(self.by_key.get(key).copied())
+  }
+
+  fn level(&self, n: NodeId) -> Result<i32, Self::Error>
+  {
+    Ok(
+      self
+        .levels
+        .get(n.as_usize())
+        .copied()
+        .flatten()
+        .expect("level of absent node"),
+    )
+  }
+
+  fn neighbors(&self, n: NodeId, level: usize, out: &mut Vec<NodeId>) -> Result<(), Self::Error>
+  {
+    out.clear();
+    if let Some(per_level) = self.links.get(n.as_usize())
+    {
+      if let Some(list) = per_level.get(level)
+      {
+        out.extend_from_slice(list);
+      }
+    }
+    Ok(())
+  }
+}
+
+impl<K: Clone + Eq + Hash> GraphStoreWrite<K> for SimpleGraphStore<K>
+{
+  fn insert_node(&mut self, key: K, level: i32) -> Result<NodeId, Self::Error>
+  {
+    let id = NodeId::from_u32(self.next_node_id as u32);
+    self.next_node_id += 1;
+    self.levels.push(Some(level));
+    self.keys.push(Some(key.clone()));
+    self.links.push(vec![Vec::new(); (level as usize) + 1]);
+    self.by_key.insert(key, id);
+    Ok(id)
+  }
+
+  fn set_neighbors(&mut self, n: NodeId, level: usize, nbrs: &[NodeId]) -> Result<(), Self::Error>
+  {
+    let per_level = &mut self.links[n.as_usize()];
+    if per_level.len() <= level
+    {
+      per_level.resize(level + 1, Vec::new());
+    }
+    per_level[level] = nbrs.to_vec();
+    Ok(())
+  }
+
+  fn set_entry_point(&mut self, ep: Option<(NodeId, i32)>) -> Result<(), Self::Error>
+  {
+    self.entry = ep;
+    Ok(())
+  }
+
+  fn remove_node(&mut self, n: NodeId) -> Result<(), Self::Error>
+  {
+    let i = n.as_usize();
+    if let Some(k) = self.keys[i].take()
+    {
+      self.by_key.remove(&k);
+    }
+    self.levels[i] = None;
+    self.links[i].clear();
+    Ok(())
+  }
+}
+
+impl<K> Default for SimpleGraphStore<K>
+{
+  fn default() -> Self
+  {
+    Self::new()
+  }
+}