uqa-storage 0.5.2

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

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

//! Canonical tensor reads for compound index preparation.

use uqa_core::{
    memory::{Budgeted, BudgetedVec},
    DocId,
};

use super::super::{
    codec::{blob_to_vector, other_error, read_str, read_u64, vector_field_prefix},
    KeyValueRead,
};
use super::KeyValueVectorIndex;
use crate::StorageBackendResult;

pub(super) use crate::vector_index::retained::VectorEntries as Entries;

impl KeyValueVectorIndex {
    pub(in crate::key_value) fn visit_canonical_from(
        &self,
        read: &dyn KeyValueRead,
        mut visit: impl FnMut(DocId, u32, &[f32]) -> StorageBackendResult<()>,
    ) -> StorageBackendResult<()> {
        read.control().check()?;
        let mut decoder = CanonicalDecoder::new(self);
        read.visit_prefix(
            &vector_field_prefix(&self.table, &self.field)?,
            &mut |key, value| {
                read.control().check()?;
                let mut payload = read.control().memory().reserve(value.len())?;
                let (document, ordinal, vector) = decoder.decode(key, value)?;
                let capacity = vector
                    .capacity()
                    .checked_mul(size_of::<f32>())
                    .ok_or(uqa_core::memory::MemoryError::SizeOverflow)?;
                payload.grow(capacity.saturating_sub(value.len()))?;
                visit(document, ordinal, &vector)
            },
        )
    }

    pub(in crate::key_value) fn load_all_from(
        &self,
        read: &dyn KeyValueRead,
    ) -> StorageBackendResult<Budgeted<Entries>> {
        // Destroy row allocations before releasing either their slot or payload allowance.
        let mut output = (
            BudgetedVec::new(read.control().memory()),
            read.control().memory().reserve(0)?,
        );
        let (vectors, payload) = (&mut output.0, &mut output.1);
        let mut decoder = CanonicalDecoder::new(self);
        read.visit_prefix(
            &vector_field_prefix(&self.table, &self.field)?,
            &mut |key, value| {
                payload.grow(value.len())?;
                vectors.reserve(1)?;
                let entry = decoder.decode(key, value)?;
                let capacity = entry
                    .2
                    .capacity()
                    .checked_mul(size_of::<f32>())
                    .ok_or(uqa_core::memory::MemoryError::SizeOverflow)?;
                payload.grow(capacity.saturating_sub(value.len()))?;
                vectors.push(entry)?;
                Ok(())
            },
        )?;
        let (vectors, mut memory) = output.0.into_parts();
        memory.absorb(output.1);
        Ok(Budgeted::new(vectors, memory))
    }
}

pub(super) struct CanonicalDecoder<'a> {
    index: &'a KeyValueVectorIndex,
    current_doc: Option<DocId>,
    expected_ordinal: u32,
}

impl<'a> CanonicalDecoder<'a> {
    pub(super) fn new(index: &'a KeyValueVectorIndex) -> Self {
        Self {
            index,
            current_doc: None,
            expected_ordinal: 0,
        }
    }

    pub(super) fn decode(
        &mut self,
        key: &[u8],
        value: &[u8],
    ) -> StorageBackendResult<(DocId, u32, Vec<f32>)> {
        let mut offset = 1;
        let _table = read_str(key, &mut offset)?;
        let _field = read_str(key, &mut offset)?;
        let doc_id = read_u64(key, &mut offset)?;
        let ordinal = u32::try_from(read_u64(key, &mut offset)?)
            .map_err(|_| other_error("persisted vector ordinal exceeds u32 index format"))?;
        if offset != key.len() {
            return Err(other_error("persisted vector key has trailing bytes"));
        }
        if self.current_doc != Some(doc_id) {
            self.current_doc = Some(doc_id);
            self.expected_ordinal = 0;
        }
        if ordinal != self.expected_ordinal {
            return Err(other_error(format!("invalid persisted vector ordinal sequence for document {doc_id}: expected {}, found {ordinal}", self.expected_ordinal)));
        }
        self.expected_ordinal = self
            .expected_ordinal
            .checked_add(1)
            .ok_or_else(|| other_error("persisted vector ordinal sequence overflow"))?;
        let vector = blob_to_vector(value)?;
        self.index.validate_dimensions(&vector)?;
        Ok((doc_id, ordinal, vector))
    }
}