uqa_storage/key_value/conformance/
vector_merging.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Concurrent vector generations agree byte-for-byte with the serial index owner.
8
9use std::sync::Arc;
10
11use super::{expect, expect_eq};
12use crate::{
13    key_value::{
14        codec::vector_field_prefix,
15        index_keys::{
16            hnsw_metadata_key, hnsw_node_prefix, ivf_assignment_prefix, ivf_centroid_prefix,
17            ivf_metadata_key,
18        },
19        KeyValueHNSWIndex, KeyValueIVFIndex,
20    },
21    HNSWIndexParams, IVFIndexParams, KeyValueStore, MemoryKeyValueStore, StorageBackendResult,
22    VectorIndex,
23};
24
25const FIELD: &str = "embedding\0日本語";
26#[derive(Clone, Copy, Debug)]
27pub enum VectorMergeKind {
28    IVF,
29    HNSW,
30}
31impl VectorMergeKind {
32    pub(super) fn open(
33        self,
34        store: Arc<dyn KeyValueStore>,
35        table: &str,
36        restore: bool,
37        alternate: bool,
38    ) -> StorageBackendResult<Box<dyn VectorIndex>> {
39        Ok(match self {
40            Self::IVF => {
41                let params = IVFIndexParams {
42                    nlist: if alternate { 3 } else { 2 },
43                    nprobe: 2,
44                    train_threshold: 3,
45                };
46                Box::new(if restore {
47                    KeyValueIVFIndex::restore(store, table, FIELD, 2, params)?
48                } else {
49                    KeyValueIVFIndex::create(store, table, FIELD, 2, params)?
50                })
51            }
52            Self::HNSW => {
53                let params = HNSWIndexParams {
54                    m: if alternate { 3 } else { 2 },
55                    ef_construction: 8,
56                    ef_search: 64,
57                    rebuild_threshold: 1,
58                    seed: 7,
59                };
60                Box::new(if restore {
61                    KeyValueHNSWIndex::restore(store, table, FIELD, 2, params)?
62                } else {
63                    KeyValueHNSWIndex::create(store, table, FIELD, 2, params)?
64                })
65            }
66        })
67    }
68}
69fn create(
70    store: Arc<dyn KeyValueStore>,
71    table: &str,
72    seed: u64,
73    kind: VectorMergeKind,
74) -> StorageBackendResult<Box<dyn VectorIndex>> {
75    let mut index = kind.open(store, table, false, false)?;
76    for document in 1..=seed {
77        index.add(document, vec![1.0, document as f32])?;
78    }
79    index.initialize()?;
80    Ok(index)
81}
82fn left_changes(index: &mut dyn VectorIndex) -> StorageBackendResult<()> {
83    index.add_many(11, vec![vec![1.0, 0.0], vec![0.0, 1.0]])?;
84    index.delete(1)?;
85    index.add(1, vec![0.25, 0.75])?;
86    index.add_many(11, vec![])?;
87    index.add_many(11, vec![vec![0.75, 0.25], vec![0.0, 1.0]])?;
88    index.add(13, vec![0.5, 0.5])
89}
90fn right_changes(index: &mut dyn VectorIndex) -> StorageBackendResult<()> {
91    index.add(12, vec![0.4, 0.6])?;
92    index.delete(2)
93}
94fn compare(
95    actual: &dyn KeyValueStore,
96    expected: &dyn KeyValueStore,
97    table: &str,
98    kind: VectorMergeKind,
99) -> StorageBackendResult<()> {
100    let derived = match kind {
101        VectorMergeKind::IVF => vec![
102            ivf_metadata_key(table, FIELD)?,
103            ivf_centroid_prefix(table, FIELD)?,
104            ivf_assignment_prefix(table, FIELD)?,
105        ],
106        VectorMergeKind::HNSW => vec![
107            hnsw_metadata_key(table, FIELD)?,
108            hnsw_node_prefix(table, FIELD)?,
109        ],
110    };
111    for prefix in std::iter::once(vector_field_prefix(table, FIELD)?).chain(derived) {
112        expect_eq(
113            &actual.scan_prefix(&prefix)?,
114            &expected.scan_prefix(&prefix)?,
115            "merged canonical vectors and complete vector generation equal serial execution",
116        )?;
117    }
118    Ok(())
119}
120fn serial(
121    table: &str,
122    seed: u64,
123    kind: VectorMergeKind,
124) -> StorageBackendResult<Arc<dyn KeyValueStore>> {
125    let store: Arc<dyn KeyValueStore> = Arc::new(MemoryKeyValueStore::new());
126    let mut index = create(store.clone(), table, seed, kind)?;
127    right_changes(&mut *index)?;
128    left_changes(&mut *index)?;
129    Ok(store)
130}
131
132/// Verify independent writers across initial training and retraining, ordered replacement/deletion and savepoint undo.
133pub fn verify_vector_document_merges(
134    a: &Arc<dyn KeyValueStore>,
135    b: &Arc<dyn KeyValueStore>,
136    kind: VectorMergeKind,
137) -> StorageBackendResult<()> {
138    for seed in [0, 2, 8] {
139        let table = format!("{kind:?}-merge\0{seed}");
140        let mut left = create(a.clone(), &table, seed, kind)?;
141        let mut right = kind.open(b.clone(), &table, true, false)?;
142        let baseline = left.snapshot()?;
143        a.begin_transaction()?;
144        b.begin_transaction()?;
145        a.savepoint("vector_branch")?;
146        left.add(99, vec![1.0, 0.0])?;
147        let discarded = left.snapshot()?;
148        a.rollback_to_savepoint("vector_branch")?;
149        a.release_savepoint("vector_branch")?;
150        left_changes(&mut *left)?;
151        let private = left.snapshot()?;
152        right_changes(&mut *right)?;
153        b.commit_transaction()?;
154        expect(
155            a.in_transaction(),
156            "second vector writer commits before the first writer ends",
157        )?;
158        expect_eq(
159            &left.count()?,
160            &private.count()?,
161            "private vector view survives other writer publication",
162        )?;
163        a.commit_transaction()?;
164        compare(
165            a.as_ref(),
166            serial(&table, seed, kind)?.as_ref(),
167            &table,
168            kind,
169        )?;
170        expect_eq(
171            &baseline.count()?,
172            &(seed as usize),
173            "old vector generation remains pinned",
174        )?;
175        expect_eq(
176            &discarded.count()?,
177            &(seed as usize + 1),
178            "discarded savepoint snapshot remains readable",
179        )?;
180        expect_eq(
181            &left.count()?,
182            &right.count()?,
183            "both live vector handles see the merged generation",
184        )?;
185    }
186    Ok(())
187}
188
189/// Verify complete merged generations after all prior file handles close.
190pub fn verify_vector_merge_reopen(
191    store: &Arc<dyn KeyValueStore>,
192    kind: VectorMergeKind,
193) -> StorageBackendResult<()> {
194    for seed in [0, 2, 8] {
195        let table = format!("{kind:?}-merge\0{seed}");
196        compare(
197            store.as_ref(),
198            serial(&table, seed, kind)?.as_ref(),
199            &table,
200            kind,
201        )?;
202        let index = kind.open(store.clone(), &table, true, false)?;
203        let mut documents = index
204            .search_knn(&[1.0, 0.0], 100)?
205            .iter()
206            .map(|posting| posting.doc_id)
207            .collect::<Vec<_>>();
208        documents.sort_unstable();
209        let mut expected = (3..=seed).chain([1, 11, 12, 13]).collect::<Vec<_>>();
210        expected.sort_unstable();
211        expect_eq(&documents, &expected, "reopened merged vector search")?;
212    }
213    Ok(())
214}
215
216/// Verify overlapping documents, rebuild/clear/drop and same-definition recreation retain their original conflicts.
217pub fn verify_vector_merge_conflicts(
218    a: &Arc<dyn KeyValueStore>,
219    b: &Arc<dyn KeyValueStore>,
220    kind: VectorMergeKind,
221) -> StorageBackendResult<()> {
222    for case in 0..6 {
223        for reverse in [false, true] {
224            let table = format!("{kind:?}-conflict-{case}-{reverse}");
225            let mut left = create(a.clone(), &table, 3, kind)?;
226            let mut right = kind.open(b.clone(), &table, true, false)?;
227            a.begin_transaction()?;
228            b.begin_transaction()?;
229            match case {
230                0 => left.add_many(11, vec![])?,
231                1 => left.delete(11)?,
232                _ => left.add(11, vec![1.0, 0.0])?,
233            }
234            match case {
235                0 | 1 => right.add(11, vec![0.0, 1.0])?,
236                2 => right.clear()?,
237                3 => right.initialize()?,
238                _ => {
239                    KeyValueIVFIndex::drop_metadata(b.as_ref(), &table, FIELD)?;
240                    right = kind.open(b.clone(), &table, false, case == 5)?;
241                    right.initialize()?;
242                }
243            }
244            let (winner, loser) = if reverse { (a, b) } else { (b, a) };
245            winner.commit_transaction()?;
246            let stable = winner.scan_prefix(b"")?;
247            expect(
248                loser.commit_transaction().is_err(),
249                "overlapping or structural vector change conflicts",
250            )?;
251            loser.rollback_transaction()?;
252            expect_eq(
253                &winner.scan_prefix(b"")?,
254                &stable,
255                "rejected vector writer publishes no partial state",
256            )?;
257        }
258    }
259    verify_catalog_lifecycle(a, b, kind)
260}
261
262fn verify_catalog_lifecycle(
263    a: &Arc<dyn KeyValueStore>,
264    b: &Arc<dyn KeyValueStore>,
265    kind: VectorMergeKind,
266) -> StorageBackendResult<()> {
267    use crate::{CatalogFacade, KeyValueCatalog, RelationIdentity, TableSchema, VectorFieldSchema};
268    let catalog = KeyValueCatalog::new(b.clone());
269    catalog.save_schema("public")?;
270    for case in 0..5 {
271        for reverse in [false, true] {
272            let table = format!("public.{kind:?}_lifecycle_{case}_{reverse}");
273            let destination = format!("{table}_renamed");
274            let schema = TableSchema {
275                relation: RelationIdentity::from_legacy_name(&table)
276                    .expect("valid fixture relation"),
277                security: crate::RelationSecurityRow::legacy("uqa"),
278                object_id: [1
279                    + case * 4
280                    + u8::from(reverse) * 2
281                    + u8::from(matches!(kind, VectorMergeKind::HNSW));
282                    16],
283                storage_generation: [case + 1; 16],
284                analyzer_json: serde_json::to_string(&uqa_analysis::whitespace_analyzer())?,
285                fts_fields: vec![],
286                vector_fields: vec![VectorFieldSchema {
287                    field: FIELD.into(),
288                    dimensions: 2,
289                }],
290                columns_json: "[]".into(),
291                constraints_json: String::new(),
292            };
293            catalog.save_table(&schema)?;
294            let mut left = create(a.clone(), &table, 3, kind)?;
295            a.begin_transaction()?;
296            b.begin_transaction()?;
297            left.add(11, vec![0.5, 0.5])?;
298            match case {
299                0 => catalog.drop_table_and_data(&table)?,
300                1 => catalog.purge_table_data(&table)?,
301                2 => catalog.rename_table_data(&table, &destination)?,
302                3 => catalog.drop_column_data(&table, FIELD)?,
303                _ => catalog.rename_column_data(&table, FIELD, "replacement")?,
304            }
305            if case <= 1 {
306                catalog.save_table(&schema)?;
307                create(b.clone(), &table, 3, kind)?;
308            }
309            let (winner, loser) = if reverse { (a, b) } else { (b, a) };
310            winner.commit_transaction()?;
311            let stable = winner.scan_prefix(b"")?;
312            expect(
313                loser.commit_transaction().is_err(),
314                "vector source and catalog lifecycle preserve conflicts in both orders",
315            )?;
316            loser.rollback_transaction()?;
317            expect_eq(
318                &winner.scan_prefix(b"")?,
319                &stable,
320                "vector catalog conflict preserves the winner including recreated names",
321            )?;
322        }
323    }
324    Ok(())
325}
uqa_storage/key_value/conformance/vector_merging.rs

uqa_storage/key_value/conformance/
vector_merging.rs
