uqa_storage/key_value/
ivf_records.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Key/Value IVF wire records; numerical preparation belongs to the common index owner.
8
9use super::{
10    codec::{decode_u64_value, decode_value},
11    ivf_persistence::{metadata_from_snapshot, PersistedIVFMetadata, IVF_FORMAT_VERSION},
12    TAG_IVF_ASSIGNMENT, TAG_IVF_CENTROID, TAG_IVF_METADATA, TAG_VECTOR,
13};
14use crate::{
15    mvcc::{
16        IVFRecordHeader, IVFRecordKey, IVFRecordLayout, IVFRecordValue, VersionError, VersionResult,
17    },
18    read_control::StorageReadControl,
19};
20use uqa_core::{memory::BudgetedVec, DocId};
21
22use super::vector_records::{ordinal, tail, usize_value, vector_bytes, GUARD};
23pub struct KeyValueIVFRecords;
24
25impl IVFRecordLayout for KeyValueIVFRecords {
26    fn metadata_key(
27        &self,
28        key: &[u8],
29        control: &StorageReadControl,
30    ) -> VersionResult<Option<BudgetedVec<u8>>> {
31        control.cancellation().check()?;
32        let Some(tag) = key.first().copied() else {
33            return Ok(None);
34        };
35        let numbers = match tag {
36            TAG_IVF_METADATA => 0,
37            TAG_IVF_CENTROID => 1,
38            TAG_IVF_ASSIGNMENT => 2,
39            _ => return Ok(None),
40        };
41        let (end, _) = tail(key, tag, numbers)?;
42        let mut output = BudgetedVec::new(control.memory());
43        output.extend_from_slice(&key[..end])?;
44        output[0] = TAG_IVF_METADATA;
45        Ok(Some(output))
46    }
47    fn key(
48        &self,
49        metadata: &[u8],
50        address: IVFRecordKey,
51        control: &StorageReadControl,
52    ) -> VersionResult<BudgetedVec<u8>> {
53        control.cancellation().check()?;
54        tail(metadata, TAG_IVF_METADATA, 0)?;
55        let mut key = BudgetedVec::new(control.memory());
56        key.extend_from_slice(metadata)?;
57        match address {
58            IVFRecordKey::Structure => {
59                key[0] = GUARD;
60                key.push(0)?;
61            }
62            IVFRecordKey::Document(document) => {
63                key[0] = GUARD;
64                key.push(1)?;
65                key.extend_from_slice(&document.to_be_bytes())?;
66            }
67            IVFRecordKey::Vectors => key[0] = TAG_VECTOR,
68            IVFRecordKey::Centroids => key[0] = TAG_IVF_CENTROID,
69            IVFRecordKey::Assignments => key[0] = TAG_IVF_ASSIGNMENT,
70            IVFRecordKey::Centroid(centroid) => {
71                key[0] = TAG_IVF_CENTROID;
72                key.extend_from_slice(&(centroid as u64).to_be_bytes())?;
73            }
74            IVFRecordKey::Assignment(document, ordinal) => {
75                key[0] = TAG_IVF_ASSIGNMENT;
76                key.extend_from_slice(&document.to_be_bytes())?;
77                key.extend_from_slice(&u64::from(ordinal).to_be_bytes())?;
78            }
79        }
80        Ok(key)
81    }
82    fn header(
83        &self,
84        key: &[u8],
85        value: &[u8],
86        control: &StorageReadControl,
87    ) -> VersionResult<IVFRecordHeader> {
88        control.cancellation().check()?;
89        tail(key, TAG_IVF_METADATA, 0)?;
90        let meta: PersistedIVFMetadata = decode_value(value)?;
91        if meta.format_version != IVF_FORMAT_VERSION {
92            return Err(VersionError::InvalidEncoding(
93                "unsupported IVF record format",
94            ));
95        }
96        let params = crate::IVFIndexParams {
97            nlist: usize_value(meta.nlist)?,
98            nprobe: usize_value(meta.nprobe)?,
99            train_threshold: usize_value(meta.train_threshold)?,
100        }
101        .validate()?;
102        Ok(IVFRecordHeader {
103            dimensions: meta.dimensions,
104            params,
105            state: meta.state.into(),
106            trained_size: usize_value(meta.trained_size)?,
107            deletes_since_train: usize_value(meta.deletes_since_train)?,
108            vector_count: usize_value(meta.vector_count)?,
109            revision: Some(meta.revision),
110        })
111    }
112    fn vector_id(&self, key: &[u8], control: &StorageReadControl) -> VersionResult<(DocId, u32)> {
113        control.cancellation().check()?;
114        let (_, [document, order]) = tail(key, TAG_VECTOR, 2)?;
115        Ok((document, ordinal(order)?))
116    }
117    fn vector(
118        &self,
119        key: &[u8],
120        value: &[u8],
121        control: &StorageReadControl,
122    ) -> VersionResult<(DocId, u32, BudgetedVec<f32>)> {
123        let (document, order) = self.vector_id(key, control)?;
124        Ok((document, order, vector_bytes(value, control)?))
125    }
126    fn centroid(
127        &self,
128        key: &[u8],
129        value: &[u8],
130        control: &StorageReadControl,
131    ) -> VersionResult<(usize, BudgetedVec<f32>)> {
132        let (_, [centroid, _]) = tail(key, TAG_IVF_CENTROID, 1)?;
133        Ok((usize_value(centroid)?, vector_bytes(value, control)?))
134    }
135    fn assignment(
136        &self,
137        key: &[u8],
138        value: &[u8],
139        control: &StorageReadControl,
140    ) -> VersionResult<(DocId, u32, usize)> {
141        control.cancellation().check()?;
142        let (_, [document, order]) = tail(key, TAG_IVF_ASSIGNMENT, 2)?;
143        Ok((
144            document,
145            ordinal(order)?,
146            usize_value(decode_u64_value(value)?)?,
147        ))
148    }
149    fn encode(
150        &self,
151        key: &[u8],
152        template: &[u8],
153        value: IVFRecordValue<'_>,
154        control: &StorageReadControl,
155    ) -> VersionResult<BudgetedVec<u8>> {
156        control.cancellation().check()?;
157        let mut output = BudgetedVec::new(control.memory());
158        match value {
159            IVFRecordValue::Header { snapshot, revision } => {
160                let header = self.header(key, template, control)?;
161                let meta = metadata_from_snapshot(
162                    header.dimensions,
163                    header.params,
164                    snapshot,
165                    revision.ok_or(VersionError::InvalidEncoding(
166                        "missing IVF mutation counter",
167                    ))?,
168                )?;
169                return super::record_json::encode(&meta, control);
170            }
171            IVFRecordValue::Centroid(vector) => {
172                tail(key, TAG_IVF_CENTROID, 1)?;
173                for value in vector {
174                    control.cancellation().check()?;
175                    output.extend_from_slice(&value.to_le_bytes())?;
176                }
177            }
178            IVFRecordValue::Assignment(centroid) => {
179                tail(key, TAG_IVF_ASSIGNMENT, 2)?;
180                output.extend_from_slice(&(centroid as u64).to_be_bytes())?;
181            }
182        }
183        Ok(output)
184    }
185}
uqa_storage/key_value/ivf_records.rs

uqa_storage/key_value/
ivf_records.rs
