gqlite 1.5.0 → 1.7.0

data/ext/gqlitedb/src/store/redb/index.rs

@@ -0,0 +1,274 @@
+use crate::graph;
+use crate::prelude::*;
+use crate::store::VectorMetric;
+use db_index::hnsw::metric::{Cosine, InnerProduct, L2};
+use db_index::hnsw::{Dense, Metric};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+pub(crate) type IndexRef = Arc<dyn Index>;
+
+/// Resolved HNSW graph construction and search parameters.
+#[derive(
+  Debug, Clone, Copy, PartialEq, Eq, smart_default::SmartDefault, Serialize, Deserialize,
+)]
+pub(crate) struct HnswParams
+{
+  #[default(16)]
+  pub m: usize,
+  #[default(200)]
+  pub ef_construction: usize,
+  #[default(50)]
+  pub ef_search: usize,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) enum IndexKind
+{
+  Vector
+  {
+    dimension: usize,
+    metric: VectorMetric,
+    params: HnswParams,
+  },
+}
+
+pub(crate) trait Index: Send + Sync
+{
+  fn name(&self) -> &str;
+  fn on_insert(&self, node: &graph::Node) -> Result<()>;
+  fn on_update(&self, old: &graph::Node, new: &graph::Node) -> Result<()>;
+  fn as_any(&self) -> &dyn std::any::Any;
+}
+
+pub(crate) struct IndexRegistry
+{
+  inner: Mutex<HashMap<String, IndexRef>>,
+}
+
+impl IndexRegistry
+{
+  pub(crate) fn new() -> Self
+  {
+    Self {
+      inner: Mutex::new(HashMap::new()),
+    }
+  }
+
+  pub(crate) fn register_index(&self, name: String, index: IndexRef)
+  {
+    let mut guard = self.inner.lock().unwrap();
+    guard.insert(name, index);
+  }
+
+  pub(crate) fn drop_index(&self, name: &str) -> Option<IndexRef>
+  {
+    let mut guard = self.inner.lock().unwrap();
+    guard.remove(name)
+  }
+
+  pub(crate) fn get_index(&self, name: &str) -> Option<IndexRef>
+  {
+    let guard = self.inner.lock().unwrap();
+    guard.get(name).cloned()
+  }
+
+  pub(crate) fn iter_indexes(&self) -> Vec<IndexRef>
+  {
+    let guard = self.inner.lock().unwrap();
+    guard.values().cloned().collect()
+  }
+}
+
+impl std::fmt::Debug for IndexRegistry
+{
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
+  {
+    let keys: Vec<String> = self.inner.lock().unwrap().keys().cloned().collect();
+    let keys_str = keys.join(", ");
+    f.write_fmt(format_args!("IndexRegistry {{ {} }} ", keys_str))
+  }
+}
+
+pub(crate) enum MetricWrapper
+{
+  Cosine(Cosine<f32>),
+  L2(L2<f32>),
+  DotProduct(InnerProduct<f32>),
+}
+
+impl Clone for MetricWrapper
+{
+  fn clone(&self) -> Self
+  {
+    match self
+    {
+      MetricWrapper::Cosine(m) => MetricWrapper::Cosine(m.clone()),
+      MetricWrapper::L2(m) => MetricWrapper::L2(m.clone()),
+      MetricWrapper::DotProduct(m) => MetricWrapper::DotProduct(m.clone()),
+    }
+  }
+}
+
+impl MetricWrapper
+{
+  pub(crate) fn new(metric: VectorMetric) -> Self
+  {
+    match metric
+    {
+      VectorMetric::Cosine => MetricWrapper::Cosine(Cosine::new()),
+      VectorMetric::L2 => MetricWrapper::L2(L2::new()),
+      VectorMetric::DotProduct => MetricWrapper::DotProduct(InnerProduct::new()),
+    }
+  }
+}
+
+impl Metric for MetricWrapper
+{
+  type Family = Dense<f32>;
+
+  fn distance<'a, 'b>(
+    &self,
+    a: <Self::Family as db_index::hnsw::vector::VectorFamily>::Ref<'a>,
+    b: <Self::Family as db_index::hnsw::vector::VectorFamily>::Ref<'b>,
+  ) -> f32
+  {
+    match self
+    {
+      MetricWrapper::Cosine(m) => m.distance(a, b),
+      MetricWrapper::L2(m) => m.distance(a, b),
+      MetricWrapper::DotProduct(m) => m.distance(a, b),
+    }
+  }
+}
+
+pub(crate) struct HnswVectorIndex
+{
+  name: String,
+  dimension: usize,
+  metric: VectorMetric,
+  params: HnswParams,
+}
+
+impl HnswVectorIndex
+{
+  pub(crate) fn new(
+    name: String,
+    dimension: usize,
+    metric: VectorMetric,
+    params: HnswParams,
+  ) -> Self
+  {
+    Self {
+      name,
+      dimension,
+      metric,
+      params,
+    }
+  }
+
+  /// Create a new HnswIndex ready for upsert or search operations.
+  /// This is called per transaction to build a fresh stateless index.
+  pub(crate) fn upsert_and_search_ready(&self) -> db_index::hnsw::HnswIndex<MetricWrapper>
+  {
+    db_index::hnsw::HnswIndex::new(
+      MetricWrapper::new(self.metric),
+      db_index::hnsw::HnswConfig::new(self.dimension, u32::MAX as usize)
+        .m(self.params.m)
+        .ef_construction(self.params.ef_construction)
+        .ef_search(self.params.ef_search),
+    )
+  }
+
+  #[cfg(test)]
+  pub(crate) fn get_hnsw_parameters(&self) -> HnswParams
+  {
+    self.params
+  }
+
+  pub(crate) fn dimension(&self) -> usize
+  {
+    self.dimension
+  }
+
+  pub(crate) fn ef_search(&self) -> usize
+  {
+    self.params.ef_search
+  }
+
+  pub(crate) fn encoding(&self) -> db_index::hnsw::store::VectorEncoding
+  {
+    // Default to F32 for now; revisit per WI #8.7's dtype measurement
+    db_index::hnsw::store::VectorEncoding::F32
+  }
+}
+
+impl Index for HnswVectorIndex
+{
+  fn name(&self) -> &str
+  {
+    &self.name
+  }
+
+  fn on_insert(&self, _node: &graph::Node) -> Result<()>
+  {
+    // No-op: real work happens through persistent stores in upsert_vector
+    Ok(())
+  }
+
+  fn on_update(&self, _old: &graph::Node, _new: &graph::Node) -> Result<()>
+  {
+    // No-op: real work happens through persistent stores in upsert_vector
+    Ok(())
+  }
+
+  fn as_any(&self) -> &dyn std::any::Any
+  {
+    self
+  }
+}
+
+pub(crate) fn create_index(kind: IndexKind, name: &str, _property: Option<String>) -> IndexRef
+{
+  match kind
+  {
+    IndexKind::Vector {
+      dimension,
+      metric,
+      params,
+    } => Arc::new(HnswVectorIndex::new(
+      name.to_string(),
+      dimension,
+      metric,
+      params,
+    )),
+  }
+}
+
+#[cfg(test)]
+mod tests
+{
+  use super::*;
+  use crate::store::VectorMetric;
+
+  #[test]
+  fn test_hnsw_parameters_are_stored_in_config()
+  {
+    let hnsw = HnswVectorIndex::new(
+      "idx".into(),
+      3,
+      VectorMetric::Cosine,
+      HnswParams {
+        m: 8,
+        ef_construction: 64,
+        ef_search: 24,
+      },
+    );
+
+    let params = hnsw.get_hnsw_parameters();
+    assert_eq!(params.m, 8);
+    assert_eq!(params.ef_construction, 64);
+    assert_eq!(params.ef_search, 24);
+  }
+}