uqa-storage 0.1.6

Document store, inverted index, IVF/HNSW vectors, B-tree, R*Tree, catalog
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Atomic vector and IVF-metadata mutations.

use rusqlite::params;
use uqa_core::DocId;

use super::metadata::{encode_metadata, EncodedIVFMetadata};
use super::writing::write_metadata;
use super::SQLiteIVFIndex;
use crate::sqlite::vector_index::codec::encode_doc_id;
use crate::sqlite::SQLiteError;
use crate::StorageBackendResult;

impl SQLiteIVFIndex {
    pub(super) fn replace_document(
        &self,
        doc_id: DocId,
        vectors: Vec<Vec<f32>>,
    ) -> StorageBackendResult<()> {
        let (encoded_doc_id, encoded_vectors) =
            self.persistent.stage_doc_vectors(doc_id, &vectors)?;
        let mut prospective = self.persistent.load_all_with_ordinals()?;
        prospective.retain(|(stored_doc_id, _, _)| *stored_doc_id != doc_id);
        for (ordinal, vector) in vectors.into_iter().enumerate() {
            prospective.push((
                doc_id,
                u32::try_from(ordinal).map_err(|_| {
                    SQLiteError::StorageBackend(
                        "vector ordinal exceeds the u32 index format".into(),
                    )
                })?,
                vector,
            ));
        }
        prospective.sort_by_key(|(doc_id, ordinal, _)| (*doc_id, *ordinal));
        let snapshot = self.metadata_for_entries(&prospective)?;
        let metadata = encode_metadata(self.params, &snapshot)?;
        self.apply_replacement(encoded_doc_id, &encoded_vectors, &metadata)
    }

    pub(super) fn delete_document(&self, doc_id: DocId) -> StorageBackendResult<()> {
        let encoded_doc_id = encode_doc_id(doc_id)?;
        let mut prospective = self.persistent.load_all_with_ordinals()?;
        let previous_len = prospective.len();
        prospective.retain(|(stored_doc_id, _, _)| *stored_doc_id != doc_id);
        if prospective.len() == previous_len {
            return Ok(());
        }
        let snapshot = self.metadata_for_entries(&prospective)?;
        let metadata = encode_metadata(self.params, &snapshot)?;
        self.apply_delete(encoded_doc_id, &metadata)
    }

    pub(super) fn clear_index(&self) -> StorageBackendResult<()> {
        self.persistent.conn.with_mut(|conn| {
            let tx = conn.savepoint()?;
            for table in [
                "_vectors",
                "_ivf_assignments",
                "_ivf_centroids",
                "_ivf_indexes",
            ] {
                tx.execute(
                    &format!("DELETE FROM {table} WHERE table_name = ?1 AND field = ?2"),
                    params![self.persistent.table, self.persistent.field],
                )?;
            }
            tx.commit()?;
            Ok(())
        })?;
        Ok(())
    }

    fn apply_replacement(
        &self,
        doc_id: i64,
        vectors: &[(i64, Vec<u8>)],
        metadata: &EncodedIVFMetadata,
    ) -> StorageBackendResult<()> {
        self.persistent.conn.with_mut(|conn| {
            let tx = conn.savepoint()?;
            tx.execute(
                "DELETE FROM _vectors
                  WHERE table_name = ?1 AND field = ?2 AND doc_id = ?3",
                params![self.persistent.table, self.persistent.field, doc_id],
            )?;
            let mut insert = tx.prepare(
                "INSERT INTO _vectors
                    (table_name, field, doc_id, vector_ordinal, vector)
                 VALUES (?1, ?2, ?3, ?4, ?5)",
            )?;
            for (ordinal, vector) in vectors {
                insert.execute(params![
                    self.persistent.table,
                    self.persistent.field,
                    doc_id,
                    ordinal,
                    vector
                ])?;
            }
            drop(insert);
            write_metadata(&tx, &self.persistent, metadata)?;
            tx.commit()?;
            Ok(())
        })?;
        Ok(())
    }

    fn apply_delete(&self, doc_id: i64, metadata: &EncodedIVFMetadata) -> StorageBackendResult<()> {
        self.persistent.conn.with_mut(|conn| {
            let tx = conn.savepoint()?;
            tx.execute(
                "DELETE FROM _vectors
                  WHERE table_name = ?1 AND field = ?2 AND doc_id = ?3",
                params![self.persistent.table, self.persistent.field, doc_id],
            )?;
            write_metadata(&tx, &self.persistent, metadata)?;
            tx.commit()?;
            Ok(())
        })?;
        Ok(())
    }
}