uqa_storage/sqlite/vector_index/ivf/
lifecycle.rs1use std::sync::Arc;
10
11use uqa_core::{DocId, PostingList};
12
13use super::writing::drop_metadata;
14use super::SQLiteIVFIndex;
15use crate::sqlite::vector_index::SQLiteVectorIndex;
16use crate::sqlite::{ManagedConnection, Result as SQLiteResult};
17use crate::vector_index::{IVFIndexParams, VectorIndex};
18use crate::StorageBackendResult;
19
20impl SQLiteIVFIndex {
21 pub fn new(
22 conn: ManagedConnection,
23 table: impl Into<String>,
24 field: impl Into<String>,
25 dimensions: u32,
26 ) -> Self {
27 Self::from_params(conn, table, field, dimensions, IVFIndexParams::default())
28 }
29
30 pub fn with_params(
31 conn: ManagedConnection,
32 table: impl Into<String>,
33 field: impl Into<String>,
34 dimensions: u32,
35 nlist: usize,
36 nprobe: usize,
37 train_threshold: usize,
38 ) -> Self {
39 Self::from_params(
40 conn,
41 table,
42 field,
43 dimensions,
44 IVFIndexParams {
45 nlist: nlist.max(1),
46 nprobe: nprobe.max(1),
47 train_threshold: train_threshold.max(1),
48 },
49 )
50 }
51
52 fn from_params(
53 conn: ManagedConnection,
54 table: impl Into<String>,
55 field: impl Into<String>,
56 dimensions: u32,
57 params: IVFIndexParams,
58 ) -> Self {
59 Self {
60 persistent: SQLiteVectorIndex::new(conn, table, field, dimensions),
61 params,
62 }
63 }
64
65 pub fn open_existing(
66 conn: ManagedConnection,
67 table: impl Into<String>,
68 field: impl Into<String>,
69 dimensions: u32,
70 nlist: usize,
71 nprobe: usize,
72 train_threshold: usize,
73 ) -> Self {
74 Self::with_params(
75 conn,
76 table,
77 field,
78 dimensions,
79 nlist,
80 nprobe,
81 train_threshold,
82 )
83 }
84
85 pub fn drop_metadata(conn: &ManagedConnection, table: &str, field: &str) -> SQLiteResult<()> {
86 drop_metadata(conn, table, field)
87 }
88}
89
90impl VectorIndex for SQLiteIVFIndex {
91 fn dimensions(&self) -> u32 {
92 self.persistent.dimensions
93 }
94
95 fn index_kind(&self) -> &'static str {
96 "sqlite-ivf"
97 }
98
99 fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
100 self.replace_document(doc_id, vec![vector])
101 }
102
103 fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
104 self.replace_document(doc_id, vectors)
105 }
106
107 fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
108 self.delete_document(doc_id)
109 }
110
111 fn clear(&mut self) -> StorageBackendResult<()> {
112 self.clear_index()
113 }
114
115 fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
116 self.search_top_k(query, k)
117 }
118
119 fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
120 self.persistent.search_threshold(query, threshold)
121 }
122
123 fn count(&self) -> StorageBackendResult<usize> {
124 self.persistent.count()
125 }
126
127 fn initialize(&mut self) -> StorageBackendResult<()> {
128 self.initialize_metadata()
129 }
130
131 fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
132 Ok(Arc::new(self.clone()))
133 }
134}