gqlite 1.5.0 → 1.7.0

data/ext/gqlitedb/src/store/sqlite.rs

@@ -3,14 +3,29 @@ use std::{cell::RefCell, path::PathBuf};
 use ccutils::pool::{self, Pool};
 
 use crate::{
+  error::StoreError,
   prelude::*,
   store::{
-    sqlbase::{SqlBindingValue, SqlResultValue},
     TransactionBoxable,
+    sqlbase::{SqlBindingValue, SqlResultValue},
   },
   utils::hex,
 };
 
+// Initialize sqlite-vec extension
+fn init_sqlite_vec()
+{
+  static INIT: std::sync::Once = std::sync::Once::new();
+
+  INIT.call_once(|| unsafe {
+    use rusqlite::ffi::sqlite3_auto_extension;
+    use sqlite_vec::sqlite3_vec_init;
+
+    #[allow(clippy::missing_transmute_annotations)]
+    sqlite3_auto_extension(Some(std::mem::transmute(sqlite3_vec_init as *const ())));
+  });
+}
+
 ccutils::assert_impl_all!(Store: Sync, Send);
 
 use askama::Template;
@@ -84,14 +99,12 @@ impl Drop for TransactionBase
   fn drop(&mut self)
   {
     if *self.active.borrow()
+      && let Err(e) = self.connection.execute("ROLLBACK", ())
     {
-      if let Err(e) = self.connection.execute("ROLLBACK", ())
-      {
-        println!(
-          "Rollback failed with error {:?}, future use of the connection are likely to fail.",
-          e
-        );
-      }
+      log::error!(
+        "Rollback failed with error {:?}, future use of the connection are likely to fail.",
+        e
+      );
     }
   }
 }
@@ -135,7 +148,7 @@ impl super::WriteTransaction for WriteTransaction
   }
 }
 
-trait GetConnection
+pub(crate) trait GetConnection
 {
   fn get_connection(&self) -> &rusqlite::Connection;
 }
@@ -277,6 +290,16 @@ mod templates
 //  ___) | || (_) | | |  __/
 // |____/ \__\___/|_|  \___|
 
+#[derive(Debug, serde::Serialize, serde::Deserialize)]
+struct VectorIndexMetadata
+{
+  name: String,
+  label: Option<String>,
+  property: Option<String>,
+  dimension: usize,
+  metric: store::VectorMetric,
+}
+
 type TransactionBox = store::TransactionBox<ReadTransaction, WriteTransaction>;
 
 pub(crate) struct Store
@@ -292,7 +315,11 @@ impl Store
   {
     let path: PathBuf = path.as_ref().into();
     let connection = Pool::new(
-      move || Ok(rusqlite::Connection::open(&path)?),
+      move || {
+        init_sqlite_vec();
+        let conn = rusqlite::Connection::open(&path)?;
+        Ok(conn)
+      },
       pool::Options::default().minimum_pool_size(1).pool_size(3),
     )?;
     let s = Self { connection };
@@ -304,10 +331,12 @@ impl Store
     let id = uuid::Uuid::new_v4().as_u128();
     let connection = Pool::new(
       move || {
-        Ok(rusqlite::Connection::open_with_flags(
+        init_sqlite_vec();
+        let conn = rusqlite::Connection::open_with_flags(
           format!("file:{}?mode=memory&cache=shared", id),
           rusqlite::OpenFlags::default(),
-        )?)
+        )?;
+        Ok(conn)
       },
       pool::Options::default().minimum_pool_size(1).pool_size(3),
     )?;
@@ -329,6 +358,69 @@ impl Store
     }
     match (from.major, from.minor)
     {
+      (1, 2) =>
+      {
+        // Migrate vector index metadata from the 1.2 custom-table format:
+        //   _vector_indices_{graph}  (SQL table with columns index_name, label, property, dimension, metric)
+        // to the 1.3 unified metadata format:
+        //   {graph}/vector_indices   -> Vec<VectorIndexMetadata>
+
+        let graphs: Vec<String> = self.get_metadata_value_json_or_default(transaction, "graphs")?;
+
+        for graph_name in &graphs
+        {
+          let metadata_table = format!("_vector_indices_{}", graph_name);
+          if !self.check_if_table_exists(transaction, &metadata_table)?
+          {
+            continue;
+          }
+
+          type LegacyIndexMetadata = (String, Option<String>, Option<String>, usize, String);
+          let rows: Vec<LegacyIndexMetadata> = {
+            let conn = transaction.get_connection();
+            let mut stmt = conn.prepare(&format!(
+              "SELECT index_name, label, property, dimension, metric FROM {}",
+              metadata_table
+            ))?;
+            stmt
+              .query_map([], |row| {
+                Ok((
+                  row.get::<_, String>(0)?,
+                  row.get::<_, Option<String>>(1)?,
+                  row.get::<_, Option<String>>(2)?,
+                  row.get::<_, i32>(3)? as usize,
+                  row.get::<_, String>(4)?,
+                ))
+              })?
+              .collect::<Result<Vec<_>, _>>()?
+          };
+
+          let new_list: Vec<VectorIndexMetadata> = rows
+            .into_iter()
+            .map(
+              |(name, label, property, dimension, metric_str)| VectorIndexMetadata {
+                name,
+                label,
+                property,
+                dimension,
+                metric: store::VectorMetric::from_debug_str(&metric_str),
+              },
+            )
+            .collect();
+
+          self.set_metadata_value_json(
+            transaction,
+            format!("{}/vector_indices", graph_name),
+            &new_list,
+          )?;
+
+          transaction
+            .get_connection()
+            .execute(&format!("DROP TABLE IF EXISTS {}", metadata_table), ())?;
+        }
+
+        Ok(())
+      }
       (1, 0) | (1, 1) =>
       {
         // uuid function is needed for upgrade
@@ -390,6 +482,10 @@ impl super::sqlbase::SqlMetaDataQueries for Store
   {
     Ok(include_str!("../../templates/sql/sqlite/metadata_set.sql").to_string())
   }
+  fn metadata_delete_query() -> Result<String>
+  {
+    Ok(include_str!("../../templates/sql/sqlite/metadata_delete.sql").to_string())
+  }
 }
 
 impl super::sqlbase::SqlStore for Store
@@ -518,6 +614,472 @@ impl super::sqlbase::SqlStore for Store
       |row| Ok(f(row)),
     )?
   }
+
+  fn create_vector_index_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+    spec: &store::VectorIndexSpec,
+  ) -> Result<()>
+  {
+    if !spec.dtype.is_float()
+    {
+      return Err(
+        crate::error::StoreError::InvalidFormat(format!(
+          "SQLite vector index only supports floating-point dtypes (F32 or F64); got {:?}",
+          spec.dtype
+        ))
+        .into(),
+      );
+    }
+
+    let vec_table = format!("vec_{}_{}", graph_name, spec.name);
+
+    // Check if index already exists
+    let mut indices: Vec<VectorIndexMetadata> = self
+      .get_metadata_value_json_or_default(transaction, format!("{}/vector_indices", graph_name))?;
+    if indices.iter().any(|m| m.name == spec.name)
+    {
+      return Err(
+        crate::error::StoreError::DuplicatedVectorIndex {
+          index_name: spec.name.clone(),
+        }
+        .into(),
+      );
+    }
+
+    transaction.get_connection().execute(
+      &format!(
+        "CREATE VIRTUAL TABLE IF NOT EXISTS {} USING vec0(embedding float[{}], node_id text)",
+        vec_table, spec.dimension
+      ),
+      (),
+    )?;
+
+    // Persist metadata for the new index
+    indices.push(VectorIndexMetadata {
+      name: spec.name.clone(),
+      label: spec.label.clone(),
+      property: spec.property.clone(),
+      dimension: spec.dimension,
+      metric: spec.metric,
+    });
+    self.set_metadata_value_json(
+      transaction,
+      format!("{}/vector_indices", graph_name),
+      &indices,
+    )?;
+
+    // Retroactively populate the index from existing nodes that match label+property
+    if let Some(label) = &spec.label
+      && let Some(property) = &spec.property
+    {
+      let nodes_table = format!("gqlite_{}_nodes", graph_name);
+      let sql = format!(
+        "SELECT hex(node_key), json_extract(properties, '$.{prop}') \
+         FROM {tbl} \
+         WHERE EXISTS ( \
+             SELECT 1 FROM json_each(labels) AS lbl WHERE lbl.value = ? \
+         ) \
+         AND json_extract(properties, '$.{prop}') IS NOT NULL",
+        prop = property,
+        tbl = nodes_table,
+      );
+      let rows: Vec<(String, String)> = {
+        let conn = transaction.get_connection();
+        match conn.prepare(&sql)
+        {
+          Ok(mut stmt) => stmt
+            .query_map(rusqlite::params![label], |row| {
+              Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
+            })?
+            .collect::<Result<Vec<_>, _>>()?,
+          Err(_) => Vec::new(),
+        }
+      };
+      for (node_key_hex, embedding_json) in rows
+      {
+        if let Some(vector) =
+          crate::store::vector_extract::extract_f32_vec_from_json(&embedding_json)
+          && !vector.is_empty()
+          && let Ok(uuid_val) = u128::from_str_radix(&node_key_hex, 16)
+        {
+          let key = graph::Key::new(uuid_val);
+          let tensor = value::Tensor::from_f32_slice(&vector);
+          let _ = self.upsert_vector_impl(transaction, graph_name, &spec.name, key, &tensor);
+        }
+      }
+    }
+
+    Ok(())
+  }
+
+  fn upsert_vector_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+    index_name: &str,
+    node_id: graph::Key,
+    tensor: &value::Tensor,
+  ) -> Result<()>
+  {
+    let f32_data: std::borrow::Cow<[f32]> = match tensor.dtype()
+    {
+      value::DType::F32 => match tensor.as_f32_slice()
+      {
+        Some(s) => std::borrow::Cow::Borrowed(s),
+        None => std::borrow::Cow::Owned(tensor.to_f32_vec()),
+      },
+      value::DType::F64 => std::borrow::Cow::Owned(tensor.to_f32_vec()),
+      other =>
+      {
+        return Err(
+          StoreError::InvalidFormat(format!(
+            "SQLite vector index only accepts F32 or F64 tensors; got {:?}",
+            other
+          ))
+          .into(),
+        );
+      }
+    };
+
+    let vec_table = format!("vec_{}_{}", graph_name, index_name);
+    let indices: Vec<VectorIndexMetadata> = self
+      .get_metadata_value_json_or_default(transaction, format!("{}/vector_indices", graph_name))?;
+    let dimension = indices
+      .iter()
+      .find(|m| m.name == index_name)
+      .map(|m| m.dimension)
+      .ok_or_else(|| {
+        StoreError::InvalidFormat(format!(
+          "vector index '{}' does not exist in graph '{}'",
+          index_name, graph_name
+        ))
+      })?;
+
+    if f32_data.len() != dimension
+    {
+      return Err(
+        StoreError::InvalidFormat(format!(
+          "Vector dimension {} does not match index dimension {}",
+          f32_data.len(),
+          dimension
+        ))
+        .into(),
+      );
+    }
+
+    let vector_bytes: Vec<u8> = f32_data.iter().flat_map(|f| f.to_le_bytes()).collect();
+
+    let node_id_hex = hex(node_id);
+
+    transaction.get_connection().execute(
+      &format!("DELETE FROM {} WHERE node_id = ?", vec_table),
+      rusqlite::params![&node_id_hex],
+    )?;
+
+    transaction.get_connection().execute(
+      &format!(
+        "INSERT INTO {} (embedding, node_id) VALUES (?, ?)",
+        vec_table
+      ),
+      rusqlite::params![&vector_bytes[..], &node_id_hex],
+    )?;
+
+    Ok(())
+  }
+
+  fn list_vector_indexes_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+  ) -> Result<Vec<store::VectorIndexSpec>>
+  {
+    let indices: Vec<VectorIndexMetadata> = self
+      .get_metadata_value_json_or_default(transaction, format!("{}/vector_indices", graph_name))?;
+
+    let specs = indices
+      .into_iter()
+      .map(|m| store::VectorIndexSpec {
+        name: m.name,
+        label: m.label,
+        property: m.property,
+        dimension: m.dimension,
+        metric: m.metric,
+        ..Default::default()
+      })
+      .collect();
+
+    Ok(specs)
+  }
+
+  fn vector_search_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+    index_name: &str,
+    query: &value::Tensor,
+    k: usize,
+  ) -> Result<Vec<store::VectorSearchResult>>
+  {
+    if k == 0
+    {
+      return Ok(Vec::new());
+    }
+
+    let f32_query: std::borrow::Cow<[f32]> = match query.dtype()
+    {
+      value::DType::F32 => match query.as_f32_slice()
+      {
+        Some(s) => std::borrow::Cow::Borrowed(s),
+        None => std::borrow::Cow::Owned(query.to_f32_vec()),
+      },
+      value::DType::F64 => std::borrow::Cow::Owned(query.to_f32_vec()),
+      other =>
+      {
+        return Err(
+          StoreError::InvalidFormat(format!(
+            "SQLite vector search only accepts F32 or F64 tensors; got {:?}",
+            other
+          ))
+          .into(),
+        );
+      }
+    };
+
+    let vec_table = format!("vec_{}_{}", graph_name, index_name);
+
+    let indices: Vec<VectorIndexMetadata> = self
+      .get_metadata_value_json_or_default(transaction, format!("{}/vector_indices", graph_name))?;
+    let dimension = indices
+      .iter()
+      .find(|m| m.name == index_name)
+      .map(|m| m.dimension)
+      .ok_or_else(|| {
+        StoreError::InvalidFormat(format!(
+          "Vector index '{}' does not exist in graph '{}'",
+          index_name, graph_name
+        ))
+      })?;
+
+    if f32_query.len() != dimension
+    {
+      return Err(
+        StoreError::InvalidFormat(format!(
+          "Vector dimension {} does not match index dimension {}",
+          f32_query.len(),
+          dimension
+        ))
+        .into(),
+      );
+    }
+
+    let mut check_stmt = transaction
+      .get_connection()
+      .prepare("SELECT 1 FROM sqlite_master WHERE type='table' AND name=?")?;
+    let table_exists = check_stmt.exists(rusqlite::params![&vec_table])?;
+
+    if !table_exists
+    {
+      return Err(
+        StoreError::InvalidFormat(format!(
+          "Vector index '{}' does not exist in graph '{}'",
+          index_name, graph_name
+        ))
+        .into(),
+      );
+    }
+
+    let query_json = format!(
+      "[{}]",
+      f32_query
+        .iter()
+        .map(|f| f.to_string())
+        .collect::<Vec<_>>()
+        .join(", ")
+    );
+
+    let mut stmt = transaction.get_connection().prepare(&format!(
+      "SELECT node_id, distance FROM {}
+         WHERE embedding MATCH ?
+         ORDER BY distance
+         LIMIT ?",
+      vec_table
+    ))?;
+
+    let results = stmt
+      .query_map(rusqlite::params![&query_json, k as i64], |row| {
+        let node_id_hex: String = row.get(0)?;
+        let distance: f32 = row.get(1)?;
+        Ok((node_id_hex, distance))
+      })?
+      .collect::<Result<Vec<_>, _>>()?;
+
+    let mut search_results = Vec::new();
+    for (node_id_hex, distance) in results
+    {
+      if node_id_hex.len() != 32
+      {
+        return Err(
+          StoreError::InvalidFormat(format!(
+            "Invalid node_id encoding: expected 32 hex chars, got {}",
+            node_id_hex.len()
+          ))
+          .into(),
+        );
+      }
+
+      let key_uuid = u128::from_str_radix(&node_id_hex, 16).map_err(|_| {
+        StoreError::InvalidFormat(format!(
+          "Invalid hex encoding in node_id: '{}'",
+          node_id_hex
+        ))
+      })?;
+      let key = graph::Key::new(key_uuid);
+      search_results.push(store::VectorSearchResult {
+        node_id: key,
+        score: distance,
+      });
+    }
+
+    Ok(search_results)
+  }
+
+  fn drop_vector_index_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+    index_name: &str,
+  ) -> Result<()>
+  {
+    let mut indices: Vec<VectorIndexMetadata> = self
+      .get_metadata_value_json_or_default(transaction, format!("{}/vector_indices", graph_name))?;
+
+    if !indices.iter().any(|m| m.name == index_name)
+    {
+      return Err(
+        StoreError::InvalidFormat(format!(
+          "Vector index '{}' does not exist in graph '{}'",
+          index_name, graph_name
+        ))
+        .into(),
+      );
+    }
+
+    let vec_table = format!("vec_{}_{}", graph_name, index_name);
+    transaction
+      .get_connection()
+      .execute(&format!("DROP TABLE IF EXISTS {}", vec_table), ())?;
+
+    indices.retain(|m| m.name != index_name);
+    self.set_metadata_value_json(
+      transaction,
+      format!("{}/vector_indices", graph_name),
+      &indices,
+    )?;
+
+    Ok(())
+  }
+
+  fn drop_graph_vector_indices_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+  ) -> Result<()>
+  {
+    let indices_key = format!("{}/vector_indices", graph_name);
+    let indices: Vec<VectorIndexMetadata> =
+      self.get_metadata_value_json_or_default(transaction, &indices_key)?;
+
+    for m in &indices
+    {
+      let vec_table = format!("vec_{}_{}", graph_name, m.name);
+      transaction
+        .get_connection()
+        .execute(&format!("DROP TABLE IF EXISTS {}", vec_table), ())?;
+    }
+
+    self.delete_metadata(transaction, &indices_key)?;
+    Ok(())
+  }
+
+  fn get_vector_indices_for_property_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+    label: &str,
+    property: &str,
+  ) -> Result<Vec<(String, usize, store::VectorMetric)>>
+  {
+    let indices: Vec<VectorIndexMetadata> = self
+      .get_metadata_value_json_or_default(transaction, format!("{}/vector_indices", graph_name))?;
+
+    let result = indices
+      .into_iter()
+      .filter(|m| m.label.as_deref() == Some(label) && m.property.as_deref() == Some(property))
+      .map(|m| (m.name, m.dimension, m.metric))
+      .collect();
+
+    Ok(result)
+  }
+
+  fn apply_vector_indices_for_node_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+    node: &crate::graph::Node,
+  ) -> Result<()>
+  {
+    for label in node.labels()
+    {
+      for property in node.properties().keys()
+      {
+        let indices =
+          self.get_vector_indices_for_property_impl(transaction, graph_name, label, property)?;
+
+        for (index_name, _dimension, _metric) in indices
+        {
+          if let Some(val) = node.properties().get(property)
+            && let Some(vector) = crate::store::vector_extract::extract_f32_vec(val)
+            && !vector.is_empty()
+          {
+            let tensor = value::Tensor::from_f32_slice(&vector);
+            self
+              .upsert_vector_impl(transaction, graph_name, &index_name, node.key(), &tensor)
+              .ok();
+          }
+        }
+      }
+    }
+
+    Ok(())
+  }
+
+  fn delete_vectors_for_nodes_impl(
+    &self,
+    transaction: &mut Self::TransactionBox,
+    graph_name: &str,
+    node_ids: &[String],
+  ) -> Result<()>
+  {
+    let indices: Vec<VectorIndexMetadata> = self
+      .get_metadata_value_json_or_default(transaction, format!("{}/vector_indices", graph_name))
+      .unwrap_or_default();
+
+    for m in &indices
+    {
+      let vec_table = format!("vec_{}_{}", graph_name, m.name);
+      for node_id in node_ids
+      {
+        let _ = transaction.get_connection().execute(
+          &format!("DELETE FROM {} WHERE node_id = ?", vec_table),
+          rusqlite::params![node_id],
+        );
+      }
+    }
+
+    Ok(())
+  }
 }
 
 impl super::sqlbase::SqlQueries for Store
@@ -578,7 +1140,7 @@ impl super::sqlbase::SqlQueries for Store
   }
 
   fn node_delete_query(graph_name: impl AsRef<str>, keys: impl AsRef<Vec<String>>)
-    -> Result<String>
+  -> Result<String>
   {
     Ok(
       templates::NodeDelete {
@@ -627,7 +1189,7 @@ impl super::sqlbase::SqlQueries for Store
     )
   }
   fn edge_delete_query(graph_name: impl AsRef<str>, keys: impl AsRef<Vec<String>>)
-    -> Result<String>
+  -> Result<String>
   {
     Ok(
       templates::EdgeDelete {