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
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
125
126
127
128
129
130
131
132

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

//! Evaluated vectors move into one charged immutable corpus before publication.

use super::{RetainedVectorIndex, VectorEntries};
use crate::{
    read_control::StorageReadControl, vector_index::validate_vector_values, StorageBackendError,
    StorageBackendResult,
};
use uqa_core::{
    memory::{Budgeted, BudgetedVec, MemoryError, MemoryReservation},
    ordering::sort_by_with_control,
    DocId,
};

/// Append selected document tensors without copying their buffers or reserving their payload twice. The caller converts inputs under this builder's allowance and transfers their reservations. Finalization orders identities, validates unique tensor ordinals and retains the same allowance for readers.
pub struct RetainedVectorIndexBuilder {
    entries: BudgetedVec<(DocId, u32, Vec<f32>)>,
    payload: MemoryReservation,
    dimensions: u32,
    control: StorageReadControl,
}

impl RetainedVectorIndexBuilder {
    pub fn new(dimensions: u32, control: &StorageReadControl) -> Self {
        Self {
            entries: BudgetedVec::new(control.memory()),
            payload: control.memory().empty_reservation(),
            dimensions,
            control: control.clone(),
        }
    }

    pub fn dimensions(&self) -> u32 {
        self.dimensions
    }

    pub fn add_document(
        &mut self,
        id: DocId,
        vectors: Budgeted<Vec<Vec<f32>>>,
    ) -> StorageBackendResult<()> {
        self.control.check()?;
        u32::try_from(vectors.len().saturating_sub(1)).map_err(|_| {
            StorageBackendError::Other("vector ordinal exceeds u32 index format".into())
        })?;
        // Drop unadopted vector buffers before releasing their reservation on every error.
        let mut pending = vectors.into_parts();
        if !pending.1.budget().shares_allowance(self.control.memory()) {
            return Err(StorageBackendError::Other(
                "retained vectors must use the original index allowance".into(),
            ));
        }
        let mut bytes = 0_usize;
        for vector in &pending.0 {
            self.control.check()?;
            bytes = vector
                .capacity()
                .checked_mul(size_of::<f32>())
                .and_then(|payload| bytes.checked_add(payload))
                .ok_or(MemoryError::SizeOverflow)?;
        }
        let required = pending
            .0
            .capacity()
            .checked_mul(size_of::<Vec<f32>>())
            .and_then(|headers| headers.checked_add(bytes))
            .ok_or(MemoryError::SizeOverflow)?;
        if pending.1.bytes() < required {
            return Err(StorageBackendError::Other(
                "retained vector reservation does not cover its buffers".into(),
            ));
        }
        for vector in &pending.0 {
            self.control.check()?;
            validate_vector_values(self.dimensions, vector)?;
        }
        self.entries.reserve(pending.0.len())?;
        let original = self.entries.len();
        let append = (|| -> StorageBackendResult<()> {
            for (ordinal, vector) in pending.0.drain(..).enumerate() {
                self.control.check()?;
                let ordinal = u32::try_from(ordinal).expect("validated vector ordinal count");
                self.entries.push((id, ordinal, vector))?;
            }
            self.control.check()
        })();
        if let Err(error) = append {
            self.entries.truncate(original);
            return Err(error);
        }
        self.payload.absorb(pending.1.split(bytes));
        // The now-empty outer buffer drops with the rest of its reservation.
        Ok(())
    }

    pub fn finish(mut self) -> StorageBackendResult<RetainedVectorIndex> {
        sort_by_with_control(
            &mut self.entries,
            &mut || self.control.check(),
            |left, right, _| Ok((left.0, left.1).cmp(&(right.0, right.1))),
        )?;
        let mut previous: Option<(DocId, u32)> = None;
        for (id, ordinal, _) in self.entries.iter() {
            self.control.check()?;
            let expected = match previous {
                Some((previous_id, previous_ordinal)) if previous_id == *id => {
                    previous_ordinal.checked_add(1)
                }
                _ => Some(0),
            };
            if expected != Some(*ordinal) {
                return Err(StorageBackendError::Other(
                    "retained vector input repeats a document or tensor ordinal".into(),
                ));
            }
            previous = Some((*id, *ordinal));
        }
        let (entries, mut memory): (VectorEntries, _) = self.entries.into_parts();
        memory.absorb(self.payload);
        RetainedVectorIndex::from_entries(
            Budgeted::new(entries, memory),
            self.dimensions,
            "memory-bruteforce",
            &self.control,
        )
    }
}