Skip to main content

uqa_storage/sqlite/vector_index/ivf/
lifecycle.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Construction and public vector-index contract for `SQLite` IVF.
8
9use 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}