uqa-storage 0.4.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

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

//! Raw tensor reads remain independent of training, centroids and probe selection.

use super::IVFIndex;
use crate::{
    read_control::StorageReadControl,
    vector_index::{copy_vector, ordinal_count, VectorRead},
    StorageBackendResult,
};
use std::ops::Bound::{Excluded, Unbounded};
use uqa_core::{memory::BudgetedVec, DocId};

impl VectorRead for IVFIndex {
    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
        control.check()
    }
    fn dimensions(&self) -> u32 {
        self.dimensions
    }
    fn next_document_after(
        &self,
        after: Option<DocId>,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<DocId>> {
        control.check()?;
        Ok(self
            .vectors
            .lock()
            .range((
                after.map_or(Unbounded, |document| Excluded((document, u32::MAX))),
                Unbounded,
            ))
            .next()
            .map(|(key, _)| key.0))
    }
    fn document_vector_count(
        &self,
        document: DocId,
        control: &StorageReadControl,
    ) -> StorageBackendResult<u64> {
        ordinal_count(
            self.vectors
                .lock()
                .range((document, 0)..=(document, u32::MAX))
                .map(|(key, _)| key.1),
            control,
        )
    }
    fn read_vector(
        &self,
        document: DocId,
        ordinal: u32,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<BudgetedVec<f32>>> {
        control.check()?;
        let vectors = self.vectors.lock();
        vectors
            .get(&(document, ordinal))
            .map(|vector| copy_vector(&vector.raw_vector, control))
            .transpose()
    }
}