uqa_storage/sqlite/vector_index/
brute_force.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Persistent brute-force vector index and exact search contract.
8
9use super::{
10    blob_to_vector, cosine_similarity, decode_doc_id, encode_doc_id, i64_to_usize, params,
11    select_top_k_scored, usize_to_u64, validate_persisted_ordinal_sequence,
12    validate_vector_ordinal_count, validate_vector_values, vector_to_blob, Arc, DocId,
13    EncodedDocVectors, ManagedConnection, Payload, PostingEntry, PostingList, SQLiteError,
14    SQLiteResult, StorageBackendResult, VectorIndex,
15};
16
17#[derive(Clone)]
18pub struct SQLiteVectorIndex {
19    pub(super) conn: ManagedConnection,
20    pub(super) table: String,
21    pub(super) field: String,
22    pub(super) dimensions: u32,
23}
24
25impl SQLiteVectorIndex {
26    pub fn new(
27        conn: ManagedConnection,
28        table: impl Into<String>,
29        field: impl Into<String>,
30        dimensions: u32,
31    ) -> Self {
32        Self {
33            conn,
34            table: table.into(),
35            field: field.into(),
36            dimensions,
37        }
38    }
39
40    pub(super) fn load_all(&self) -> SQLiteResult<Vec<(DocId, Vec<f32>)>> {
41        self.load_all_with_ordinals().map(|rows| {
42            rows.into_iter()
43                .map(|(doc_id, _, vector)| (doc_id, vector))
44                .collect()
45        })
46    }
47
48    pub(super) fn load_all_with_ordinals(&self) -> SQLiteResult<Vec<(DocId, u32, Vec<f32>)>> {
49        self.conn
50            .with(|connection| self.load_all_with_ordinals_from(connection))
51    }
52
53    pub(super) fn load_all_with_ordinals_from(
54        &self,
55        connection: &rusqlite::Connection,
56    ) -> SQLiteResult<Vec<(DocId, u32, Vec<f32>)>> {
57        let mut stmt = connection.prepare(
58            "SELECT doc_id, vector_ordinal, vector FROM _vectors
59                 WHERE table_name = ?1 AND field = ?2
60                 ORDER BY doc_id, vector_ordinal",
61        )?;
62        let rows = stmt.query_map(params![self.table, self.field], |r| {
63            Ok((
64                r.get::<_, i64>(0)?,
65                r.get::<_, i64>(1)?,
66                r.get::<_, Vec<u8>>(2)?,
67            ))
68        })?;
69        let mut out = Vec::new();
70        for row in rows {
71            let (doc_id, ordinal, blob) = row?;
72            let ordinal = u32::try_from(ordinal).map_err(|_| {
73                SQLiteError::StorageBackend(format!(
74                    "invalid vector ordinal {ordinal} for {}.{}",
75                    self.table, self.field
76                ))
77            })?;
78            let vector = blob_to_vector(&blob)?;
79            self.validate_dimensions_sqlite(&vector)?;
80            out.push((decode_doc_id(doc_id)?, ordinal, vector));
81        }
82        validate_persisted_ordinal_sequence(&out)?;
83        Ok(out)
84    }
85
86    pub(super) fn validate_dimensions_sqlite(&self, vector: &[f32]) -> SQLiteResult<()> {
87        validate_vector_values(self.dimensions, vector).map_err(|error| {
88            SQLiteError::StorageBackend(format!(
89                "invalid vector for {}.{}: {error}",
90                self.table, self.field
91            ))
92        })
93    }
94
95    pub(super) fn validate_dimensions(&self, vector: &[f32]) -> StorageBackendResult<()> {
96        Ok(self.validate_dimensions_sqlite(vector)?)
97    }
98
99    pub(super) fn stage_doc_vectors(
100        &self,
101        doc_id: DocId,
102        vectors: &[Vec<f32>],
103    ) -> SQLiteResult<EncodedDocVectors> {
104        for vector in vectors {
105            self.validate_dimensions_sqlite(vector)?;
106        }
107        validate_vector_ordinal_count(usize_to_u64("vector count", vectors.len())?)?;
108        let doc_id = encode_doc_id(doc_id)?;
109        let encoded = vectors
110            .iter()
111            .enumerate()
112            .map(|(ordinal, vector)| {
113                let ordinal = u32::try_from(ordinal).map_err(|_| {
114                    SQLiteError::StorageBackend(
115                        "vector ordinal exceeds the u32 index format".into(),
116                    )
117                })?;
118                Ok((i64::from(ordinal), vector_to_blob(vector)?))
119            })
120            .collect::<SQLiteResult<Vec<_>>>()?;
121        Ok((doc_id, encoded))
122    }
123}
124
125impl VectorIndex for SQLiteVectorIndex {
126    fn dimensions(&self) -> u32 {
127        self.dimensions
128    }
129
130    fn index_kind(&self) -> &'static str {
131        "sqlite-bruteforce"
132    }
133
134    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
135        self.add_many(doc_id, vec![vector])
136    }
137
138    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
139        let (doc_id, encoded_vectors) = self.stage_doc_vectors(doc_id, &vectors)?;
140        self.conn.with_mut(|conn| {
141            let tx = conn.savepoint()?;
142            tx.execute(
143                "DELETE FROM _vectors
144                 WHERE table_name = ?1 AND field = ?2 AND doc_id = ?3",
145                params![self.table, self.field, doc_id],
146            )?;
147            let mut stmt = tx.prepare(
148                "INSERT INTO _vectors (table_name, field, doc_id, vector_ordinal, vector)
149                 VALUES (?1, ?2, ?3, ?4, ?5)",
150            )?;
151            for (ordinal, vector) in &encoded_vectors {
152                stmt.execute(params![self.table, self.field, doc_id, ordinal, vector,])?;
153            }
154            drop(stmt);
155            tx.commit()?;
156            Ok(())
157        })?;
158        Ok(())
159    }
160
161    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
162        let doc_id = encode_doc_id(doc_id)?;
163        self.conn.with(|c| {
164            c.execute(
165                "DELETE FROM _vectors
166                 WHERE table_name = ?1 AND field = ?2 AND doc_id = ?3",
167                params![self.table, self.field, doc_id],
168            )?;
169            Ok(())
170        })?;
171        Ok(())
172    }
173
174    fn clear(&mut self) -> StorageBackendResult<()> {
175        self.conn.with(|c| {
176            c.execute(
177                "DELETE FROM _vectors WHERE table_name = ?1 AND field = ?2",
178                params![self.table, self.field],
179            )?;
180            Ok(())
181        })?;
182        Ok(())
183    }
184
185    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
186        self.validate_dimensions(query)?;
187        if k == 0 {
188            return Ok(PostingList::new());
189        }
190        let entries = self.load_all()?;
191        if entries.is_empty() {
192            return Ok(PostingList::new());
193        }
194        let mut best_by_doc: std::collections::BTreeMap<DocId, f32> =
195            std::collections::BTreeMap::new();
196        for (doc_id, vector) in &entries {
197            let sim = cosine_similarity(query, vector);
198            best_by_doc
199                .entry(*doc_id)
200                .and_modify(|best| {
201                    if sim > *best {
202                        *best = sim;
203                    }
204                })
205                .or_insert(sim);
206        }
207        let mut scored: Vec<(DocId, f32)> = best_by_doc.into_iter().collect();
208        select_top_k_scored(&mut scored, k);
209        scored.sort_by_key(|(id, _)| *id);
210        let entries: Vec<PostingEntry> = scored
211            .into_iter()
212            .map(|(doc_id, sim)| PostingEntry::new(doc_id, Payload::with_score(f64::from(sim))))
213            .collect();
214        Ok(PostingList::from_sorted_unchecked(entries))
215    }
216
217    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
218        self.validate_dimensions(query)?;
219        if !threshold.is_finite() {
220            return Err(crate::StorageBackendError::Other(format!(
221                "vector similarity threshold must be finite, got {threshold}"
222            )));
223        }
224        let entries = self.load_all()?;
225        let mut best_by_doc: std::collections::BTreeMap<DocId, f32> =
226            std::collections::BTreeMap::new();
227        for (doc_id, vector) in &entries {
228            let sim = cosine_similarity(query, vector);
229            if sim >= threshold {
230                best_by_doc
231                    .entry(*doc_id)
232                    .and_modify(|best| {
233                        if sim > *best {
234                            *best = sim;
235                        }
236                    })
237                    .or_insert(sim);
238            }
239        }
240        let mut out: Vec<PostingEntry> = best_by_doc
241            .into_iter()
242            .map(|(doc_id, sim)| PostingEntry::new(doc_id, Payload::with_score(f64::from(sim))))
243            .collect();
244        out.sort_by_key(|e| e.doc_id);
245        Ok(PostingList::from_sorted_unchecked(out))
246    }
247
248    fn count(&self) -> StorageBackendResult<usize> {
249        Ok(self.conn.with(|c| {
250            let n: i64 = c.query_row(
251                "SELECT COUNT(*) FROM _vectors WHERE table_name = ?1 AND field = ?2",
252                params![self.table, self.field],
253                |r| r.get(0),
254            )?;
255            i64_to_usize("vector count", n)
256        })?)
257    }
258
259    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
260        Ok(Arc::new(self.clone()))
261    }
262}
uqa_storage/sqlite/vector_index/brute_force.rs

uqa_storage/sqlite/vector_index/
brute_force.rs
