uqa_storage/vector_index/
retained.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Exact readers retain canonical vectors and their allowance independently of the live session.
8
9use std::sync::Arc;
10
11use uqa_core::{
12    memory::{Budgeted, BudgetedVec},
13    DocId, Payload, PostingEntry, PostingList,
14};
15
16use crate::read_control::StorageReadControl;
17use crate::vector_index::{
18    cosine_similarity, select_top_k_scored, validate_vector_values, VectorIndex,
19};
20use crate::{StorageBackendError, StorageBackendResult};
21
22mod builder;
23pub use builder::RetainedVectorIndexBuilder;
24
25pub(crate) type VectorEntries = Vec<(DocId, u32, Vec<f32>)>;
26
27/// An immutable exact index whose vector buffers and read workspace share the original retention allowance.
28#[derive(Clone)]
29pub struct RetainedVectorIndex {
30    entries: Arc<Budgeted<VectorEntries>>,
31    dimensions: u32,
32    index_kind: &'static str,
33    control: StorageReadControl,
34}
35
36impl RetainedVectorIndex {
37    pub(crate) fn from_entries(
38        entries: Budgeted<VectorEntries>,
39        dimensions: u32,
40        index_kind: &'static str,
41        control: &StorageReadControl,
42    ) -> StorageBackendResult<Self> {
43        control.check()?;
44        Ok(Self {
45            entries: entries.into_shared()?,
46            dimensions,
47            index_kind,
48            control: control.clone(),
49        })
50    }
51
52    fn scores(
53        &self,
54        query: &[f32],
55        threshold: Option<f32>,
56    ) -> StorageBackendResult<BudgetedVec<(DocId, f32)>> {
57        let mut scores = BudgetedVec::new(self.control.memory());
58        for vectors in self.entries.chunk_by(|a, b| a.0 == b.0) {
59            self.control.check()?;
60            let mut best = cosine_similarity(query, &vectors[0].2);
61            for (_, _, vector) in &vectors[1..] {
62                self.control.check()?;
63                let score = cosine_similarity(query, vector);
64                if score.total_cmp(&best).is_gt() {
65                    best = score;
66                }
67            }
68            if threshold.is_none_or(|threshold| best >= threshold) {
69                scores.push((vectors[0].0, best))?;
70            }
71        }
72        Ok(scores)
73    }
74
75    fn postings(&self, scores: &[(DocId, f32)]) -> StorageBackendResult<PostingList> {
76        let mut postings = BudgetedVec::new(self.control.memory());
77        postings.reserve(scores.len())?;
78        for &(doc_id, score) in scores {
79            self.control.check()?;
80            postings.push(PostingEntry::new(
81                doc_id,
82                Payload::with_score(f64::from(score)),
83            ))?;
84        }
85        let (postings, _memory) = postings.into_parts();
86        Ok(PostingList::from_sorted_unchecked(postings))
87    }
88}
89
90impl VectorIndex for RetainedVectorIndex {
91    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
92        self.control.check()?;
93        Ok(self
94            .entries
95            .binary_search_by_key(&doc_id, |entry| entry.0)
96            .is_ok())
97    }
98
99    fn dimensions(&self) -> u32 {
100        self.dimensions
101    }
102    fn index_kind(&self) -> &'static str {
103        self.index_kind
104    }
105    fn add(&mut self, _: DocId, _: Vec<f32>) -> StorageBackendResult<()> {
106        Err(read_only())
107    }
108    fn add_many(&mut self, _: DocId, _: Vec<Vec<f32>>) -> StorageBackendResult<()> {
109        Err(read_only())
110    }
111    fn delete(&mut self, _: DocId) -> StorageBackendResult<()> {
112        Err(read_only())
113    }
114    fn clear(&mut self) -> StorageBackendResult<()> {
115        Err(read_only())
116    }
117    fn initialize(&mut self) -> StorageBackendResult<()> {
118        Err(read_only())
119    }
120    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
121        self.control.check()?;
122        validate_vector_values(self.dimensions, query)?;
123        if k == 0 {
124            return Ok(PostingList::new());
125        }
126        // The tuple drops score storage before its reservation on every return path.
127        let mut scores = self.scores(query, None)?.into_parts();
128        select_top_k_scored(&mut scores.0, k);
129        scores.0.sort_unstable_by_key(|(doc_id, _)| *doc_id);
130        self.postings(&scores.0)
131    }
132    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
133        self.control.check()?;
134        validate_vector_values(self.dimensions, query)?;
135        if !threshold.is_finite() {
136            return Err(StorageBackendError::Other(format!(
137                "vector similarity threshold must be finite, got {threshold}"
138            )));
139        }
140        let scores = self.scores(query, Some(threshold))?;
141        self.postings(&scores)
142    }
143    fn count(&self) -> StorageBackendResult<usize> {
144        self.control.check()?;
145        Ok(self.entries.len())
146    }
147    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
148        Ok(Arc::new(self.clone()))
149    }
150}
151
152fn read_only() -> StorageBackendError {
153    StorageBackendError::Other("cannot write a retained vector snapshot".into())
154}
155
156#[cfg(test)]
157mod tests;
uqa_storage/vector_index/retained.rs

uqa_storage/vector_index/
retained.rs
