oneai_vector/
sqlite_vec.rs

1//! `VectorBackend` backed by [`sqlite-vec`] exact KNN.
2//!
3//! `sqlite-vec` ships only the C function `sqlite3_vec_init`; there is no Rust
4//! wrapper. We static-register it as a SQLite auto extension via
5//! `sqlite3_auto_extension` (one `Once` per process) so the `vec0` virtual
6//! table is available on every connection without `load_extension`. This was
7//! verified in the P0 spike — notably the rusqlite `bundled` feature alone
8//! suffices (no `load_extension` feature is needed) because the extension is
9//! linked into the same `SQLITE_CORE` build.
10//!
11//! Scale ceiling: `vec0` does **exact** brute-force KNN (not HNSW), so it is
12//! the right pick for mobile / small persistent stores (~10⁵ vectors). For
13//! desktop/server scale use [`UsearchBackend`](crate::UsearchBackend).
14//!
15//! Metadata filtering is **post-filter**: `vec0` KNN runs first over the full
16//! set, then non-matching rows are dropped. To stay correct under selective
17//! filters we over-fetch (`limit = top_k * OVER_FETCH_FACTOR`, capped) and
18//! post-filter down to `top_k`. This is a documented trade-off of the exact-
19//! KNN engine; for highly selective filters on large sets prefer `UsearchBackend`
20//! which supports true pre-filter via a search callback.
21
22use std::collections::HashMap;
23use std::sync::Once;
24
25use async_trait::async_trait;
26use oneai_core::traits::{Filter, Metadata, VectorBackend, VectorHit};
27use rusqlite::ffi::sqlite3_auto_extension;
28use tokio::sync::Mutex;
29
30/// Multiplier on `top_k` when a filter is present, so post-filtering still
31/// yields enough hits. Capped to keep KNN bounded.
32const OVER_FETCH_FACTOR: usize = 8;
33/// Hard cap on rows fetched per KNN query (post-filter trades correctness for
34/// memory/time at extreme scales).
35const MAX_FETCH: usize = 512;
36
37/// Registers `sqlite3_vec_init` as a SQLite auto extension, exactly once per
38/// process. Safe to call repeatedly.
39///
40/// The `transmute` here is verbatim the registration pattern published by the
41/// sqlite-vec crate itself (sqlite-vec-0.1.9/src/lib.rs:15) — `sqlite3_vec_init`
42/// is an untyped `extern "C"` fn and must be widened to SQLite's
43/// `sqlite3_callback` signature. `missing_transmute_annotations` is cosmetic
44/// on this canonical FFI shim, so it's suppressed locally.
45#[allow(clippy::missing_transmute_annotations)]
46fn register_vec_once() {
47    static ONCE: Once = Once::new();
48    ONCE.call_once(|| {
49        // SAFETY: `sqlite3_vec_init` is the canonical entry point exported by
50        // the sqlite-vec crate; registering it via `sqlite3_auto_extension`
51        // makes the `vec0` module available on every subsequently opened
52        // connection. The transmute is the documented registration pattern
53        // (the C signature is `int sqlite3_vec_init(sqlite3*, char**, const
54        // sqlite3_api_routines*)`, matching `sqlite3_auto_extension`'s
55        // expected `sqlite3_callback`). No mutable state of ours is touched.
56        unsafe {
57            sqlite3_auto_extension(Some(std::mem::transmute(
58                sqlite_vec::sqlite3_vec_init as *const (),
59            )));
60        }
61    });
62}
63
64/// Exact-KNN vector backend over a `vec0` virtual table.
65pub struct SqliteVecBackend {
66    conn: Mutex<rusqlite::Connection>,
67    dim: usize,
68    /// rowid ↔ string id (vec0 keys by integer rowid; we map to the trait's
69    /// string id).
70    rowid_to_id: Mutex<HashMap<i64, String>>,
71    id_to_rowid: Mutex<HashMap<String, i64>>,
72    /// metadata keyed by string id (vec0 stores vectors; metadata lives here).
73    meta: Mutex<HashMap<String, Metadata>>,
74    next_rowid: Mutex<i64>,
75}
76
77impl SqliteVecBackend {
78    /// Open (or create) a backend at `path` (`":memory:"` for in-RAM) with the
79    /// given fixed dimension. The `vec0` virtual table is created if absent.
80    pub fn open(path: &str, dim: usize) -> oneai_core::Result<Self> {
81        register_vec_once();
82        let conn = rusqlite::Connection::open(path)
83            .map_err(|e| oneai_core::OneAIError::Rag(format!("sqlite-vec open {path}: {e}")))?;
84        conn.execute_batch(&format!(
85            // The embedding column is fixed-dim float; the side `document_id`
86            // text column is unused (we map rowid↔id in Rust) but kept for
87            // compatibility with the vec0 schema requirement that a virtual
88            // table have a recognizable shape.
89            "CREATE VIRTUAL TABLE IF NOT EXISTS vec_chunks USING vec0(embedding float[{dim}], document_id text);",
90        ))
91        .map_err(|e| oneai_core::OneAIError::Rag(format!("create vec0: {e}")))?;
92        Ok(Self {
93            conn: Mutex::new(conn),
94            dim,
95            rowid_to_id: Mutex::new(HashMap::new()),
96            id_to_rowid: Mutex::new(HashMap::new()),
97            meta: Mutex::new(HashMap::new()),
98            next_rowid: Mutex::new(1),
99        })
100    }
101
102    /// In-RAM backend (shorthand for `open(":memory:", dim)`).
103    pub fn in_memory(dim: usize) -> oneai_core::Result<Self> {
104        Self::open(":memory:", dim)
105    }
106}
107
108#[async_trait]
109impl VectorBackend for SqliteVecBackend {
110    async fn upsert(
111        &self,
112        id: &str,
113        embedding: &[f32],
114        metadata: Metadata,
115    ) -> oneai_core::Result<()> {
116        if embedding.len() != self.dim {
117            return Err(oneai_core::OneAIError::Rag(format!(
118                "SqliteVecBackend: dim mismatch (got {}, expected {})",
119                embedding.len(),
120                self.dim
121            )));
122        }
123        // vec0 stores embeddings as a JSON array string, e.g. "[0.1,0.2,...]".
124        let emb_json = serde_json::to_string(embedding)
125            .map_err(|e| oneai_core::OneAIError::Rag(format!("embed json: {e}")))?;
126        let conn = self.conn.lock().await;
127        let mut id_to_rowid = self.id_to_rowid.lock().await;
128        let mut rowid_to_id = self.rowid_to_id.lock().await;
129        let mut meta = self.meta.lock().await;
130        let mut next_rowid = self.next_rowid.lock().await;
131
132        if let Some(&rowid) = id_to_rowid.get(id) {
133            // Replace existing row by rowid.
134            conn.execute(
135                "DELETE FROM vec_chunks WHERE rowid = ?1",
136                rusqlite::params![rowid],
137            )
138            .map_err(|e| oneai_core::OneAIError::Rag(format!("vec0 delete-row: {e}")))?;
139            conn.execute(
140                "INSERT INTO vec_chunks(rowid, embedding, document_id) VALUES (?1, ?2, ?3)",
141                rusqlite::params![rowid, emb_json, id],
142            )
143            .map_err(|e| oneai_core::OneAIError::Rag(format!("vec0 insert: {e}")))?;
144        } else {
145            let rowid = *next_rowid;
146            *next_rowid += 1;
147            conn.execute(
148                "INSERT INTO vec_chunks(rowid, embedding, document_id) VALUES (?1, ?2, ?3)",
149                rusqlite::params![rowid, emb_json, id],
150            )
151            .map_err(|e| oneai_core::OneAIError::Rag(format!("vec0 insert: {e}")))?;
152            id_to_rowid.insert(id.to_string(), rowid);
153            rowid_to_id.insert(rowid, id.to_string());
154        }
155        meta.insert(id.to_string(), metadata);
156        Ok(())
157    }
158
159    async fn search(
160        &self,
161        query: &[f32],
162        top_k: usize,
163        filter: Option<&Filter>,
164    ) -> oneai_core::Result<Vec<VectorHit>> {
165        if query.len() != self.dim {
166            return Err(oneai_core::OneAIError::Rag(format!(
167                "SqliteVecBackend: query dim mismatch (got {}, expected {})",
168                query.len(),
169                self.dim
170            )));
171        }
172        let q_json = serde_json::to_string(query)
173            .map_err(|e| oneai_core::OneAIError::Rag(format!("query json: {e}")))?;
174        let conn = self.conn.lock().await;
175        let rowid_to_id = self.rowid_to_id.lock().await;
176        let meta = self.meta.lock().await;
177
178        let limit = match filter {
179            Some(f) if !f.metadata_eq.is_empty() && !f.metadata_in.is_empty() => (top_k
180                .saturating_mul(OVER_FETCH_FACTOR))
181            .min(MAX_FETCH)
182            .max(top_k),
183            _ => top_k.max(1),
184        };
185        let sql = "SELECT rowid, distance FROM vec_chunks WHERE embedding MATCH ?1 ORDER BY distance LIMIT ?2";
186        let mut stmt = conn
187            .prepare(sql)
188            .map_err(|e| oneai_core::OneAIError::Rag(format!("vec0 prepare: {e}")))?;
189        let rows = stmt
190            .query_map(rusqlite::params![q_json, limit as i64], |r| {
191                let rowid: i64 = r.get(0)?;
192                let dist: f32 = r.get(1)?;
193                Ok((rowid, dist))
194            })
195            .map_err(|e| oneai_core::OneAIError::Rag(format!("vec0 query: {e}")))?;
196
197        let mut hits: Vec<VectorHit> = Vec::new();
198        for r in rows {
199            let (rowid, dist) =
200                r.map_err(|e| oneai_core::OneAIError::Rag(format!("vec0 row: {e}")))?;
201            let Some(id) = rowid_to_id.get(&rowid) else {
202                continue;
203            };
204            let metadata = meta.get(id).cloned().unwrap_or_default();
205            if let Some(f) = filter {
206                if !f.matches(&metadata) {
207                    continue;
208                }
209            }
210            // cosine distance from sqlite-vec is in [0,2]; similarity = 1 - dist.
211            let score = (1.0 - dist).clamp(-1.0, 1.0);
212            hits.push(VectorHit {
213                id: id.clone(),
214                score,
215                metadata,
216            });
217            if filter.is_none() && hits.len() >= top_k {
218                break;
219            }
220        }
221        // Sort defensively (KNN ORDER BY distance already orders, but post-
222        // filtering may have removed rows; re-sort by similarity desc).
223        hits.sort_by(|a, b| {
224            b.score
225                .partial_cmp(&a.score)
226                .unwrap_or(std::cmp::Ordering::Equal)
227        });
228        hits.truncate(top_k);
229        Ok(hits)
230    }
231
232    async fn delete(&self, id: &str) -> oneai_core::Result<()> {
233        let conn = self.conn.lock().await;
234        let mut id_to_rowid = self.id_to_rowid.lock().await;
235        let mut rowid_to_id = self.rowid_to_id.lock().await;
236        let mut meta = self.meta.lock().await;
237        if let Some(rowid) = id_to_rowid.remove(id) {
238            conn.execute(
239                "DELETE FROM vec_chunks WHERE rowid = ?1",
240                rusqlite::params![rowid],
241            )
242            .map_err(|e| oneai_core::OneAIError::Rag(format!("vec0 delete: {e}")))?;
243            rowid_to_id.remove(&rowid);
244            meta.remove(id);
245        }
246        Ok(())
247    }
248
249    fn dimension(&self) -> usize {
250        self.dim
251    }
252}
253
254// NOTE: `vec0` computes distance natively (no `cosine` helper needed) and
255// metadata is held in a side `HashMap` (no `meta_to_json`/`meta_from_json`
256// in the vec0 table) — so the shared crate helpers are intentionally unused
257// by this backend.
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    #[tokio::test]
264    async fn knn_upsert_search_filter_delete() {
265        let be = SqliteVecBackend::in_memory(4).unwrap();
266        be.upsert(
267            "a",
268            &[0.1, 0.2, 0.3, 0.4],
269            Metadata::from([("k".into(), "x".into())]),
270        )
271        .await
272        .unwrap();
273        be.upsert(
274            "b",
275            &[0.11, 0.21, 0.31, 0.41],
276            Metadata::from([("k".into(), "y".into())]),
277        )
278        .await
279        .unwrap();
280        be.upsert(
281            "c",
282            &[0.9, 0.8, 0.7, 0.6],
283            Metadata::from([("k".into(), "x".into())]),
284        )
285        .await
286        .unwrap();
287
288        let hits = be.search(&[0.1, 0.2, 0.3, 0.4], 3, None).await.unwrap();
289        assert_eq!(hits[0].id, "a");
290        assert_eq!(hits[1].id, "b");
291        assert!(hits[0].score > 0.99);
292
293        let f = Filter::new().with_eq("k", "x");
294        let hits = be.search(&[0.1, 0.2, 0.3, 0.4], 5, Some(&f)).await.unwrap();
295        assert_eq!(hits.len(), 2);
296        assert!(hits.iter().all(|h| h.metadata["k"] == "x"));
297
298        be.delete("b").await.unwrap();
299        let hits = be
300            .search(&[0.11, 0.21, 0.31, 0.41], 10, None)
301            .await
302            .unwrap();
303        assert!(!hits.iter().any(|h| h.id == "b"));
304    }
305
306    #[tokio::test]
307    async fn upsert_replaces_same_id() {
308        let be = SqliteVecBackend::in_memory(4).unwrap();
309        be.upsert("a", &[0.1, 0.2, 0.3, 0.4], Metadata::new())
310            .await
311            .unwrap();
312        be.upsert("a", &[0.9, 0.8, 0.7, 0.6], Metadata::new())
313            .await
314            .unwrap();
315        let hits = be.search(&[0.9, 0.8, 0.7, 0.6], 1, None).await.unwrap();
316        assert_eq!(hits[0].id, "a");
317        assert!(hits[0].score > 0.99);
318        // Only one row for 'a'.
319        let hits = be.search(&[0.1, 0.2, 0.3, 0.4], 10, None).await.unwrap();
320        assert_eq!(hits.iter().filter(|h| h.id == "a").count(), 1);
321    }
322}
oneai_vector/sqlite_vec.rs

oneai_vector/
sqlite_vec.rs
