uqa_storage/ivf_index/
canonical.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Raw tensor reads remain independent of training, centroids and probe selection.
8
9use super::IVFIndex;
10use crate::{
11    read_control::StorageReadControl,
12    vector_index::{copy_vector, ordinal_count, VectorRead},
13    StorageBackendResult,
14};
15use std::ops::Bound::{Excluded, Unbounded};
16use uqa_core::{memory::BudgetedVec, DocId};
17
18impl VectorRead for IVFIndex {
19    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
20        control.check()
21    }
22    fn dimensions(&self) -> u32 {
23        self.dimensions
24    }
25    fn next_document_after(
26        &self,
27        after: Option<DocId>,
28        control: &StorageReadControl,
29    ) -> StorageBackendResult<Option<DocId>> {
30        control.check()?;
31        Ok(self
32            .vectors
33            .lock()
34            .range((
35                after.map_or(Unbounded, |document| Excluded((document, u32::MAX))),
36                Unbounded,
37            ))
38            .next()
39            .map(|(key, _)| key.0))
40    }
41    fn document_vector_count(
42        &self,
43        document: DocId,
44        control: &StorageReadControl,
45    ) -> StorageBackendResult<u64> {
46        ordinal_count(
47            self.vectors
48                .lock()
49                .range((document, 0)..=(document, u32::MAX))
50                .map(|(key, _)| key.1),
51            control,
52        )
53    }
54    fn read_vector(
55        &self,
56        document: DocId,
57        ordinal: u32,
58        control: &StorageReadControl,
59    ) -> StorageBackendResult<Option<BudgetedVec<f32>>> {
60        control.check()?;
61        let vectors = self.vectors.lock();
62        vectors
63            .get(&(document, ordinal))
64            .map(|vector| copy_vector(&vector.raw_vector, control))
65            .transpose()
66    }
67}
uqa_storage/ivf_index/canonical.rs

uqa_storage/ivf_index/
canonical.rs
