weavatrix_search_vector/
multi_io.rs

1use crate::error::SearchError;
2use crate::hnsw::VectorIndex;
3use crate::multi::{MultiVectorIndex, MultiVectorKey};
4use std::io::{Read, Write};
5use std::path::Path;
6
7const MAGIC: &[u8; 8] = b"WVMULT03";
8const VERSION: u32 = 1;
9const HEADER_LEN: usize = 56;
10const HEADER_LEN_U32: u32 = 56;
11const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
12const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
13
14impl MultiVectorIndex {
15    /// Serializes the graph and caller-visible `(key, vector_id)` identities.
16    ///
17    /// # Errors
18    ///
19    /// Returns a capacity or allocation error.
20    pub fn to_bytes(&self) -> Result<Vec<u8>, SearchError> {
21        let index = self.index.to_bytes()?;
22        let identity_bytes = self
23            .identities
24            .len()
25            .checked_mul(16)
26            .ok_or(SearchError::CapacityOverflow)?;
27        let file_len = HEADER_LEN
28            .checked_add(index.len())
29            .and_then(|value| value.checked_add(identity_bytes))
30            .ok_or(SearchError::CapacityOverflow)?;
31        let mut bytes = Vec::new();
32        bytes
33            .try_reserve_exact(file_len)
34            .map_err(|_| SearchError::AllocationFailed)?;
35        bytes.extend_from_slice(MAGIC);
36        bytes.resize(HEADER_LEN, 0);
37        bytes.extend_from_slice(&index);
38        for identity in &self.identities {
39            bytes.extend_from_slice(&identity.key.to_le_bytes());
40            bytes.extend_from_slice(&identity.vector_id.to_le_bytes());
41        }
42        let payload_checksum = checksum(&bytes[HEADER_LEN..]);
43        put_u32(&mut bytes, 8, VERSION);
44        put_u32(&mut bytes, 12, HEADER_LEN_U32);
45        put_u64(&mut bytes, 16, payload_checksum);
46        put_u64(&mut bytes, 24, to_u64(file_len)?);
47        put_u64(&mut bytes, 32, to_u64(index.len())?);
48        put_u64(&mut bytes, 40, to_u64(self.identities.len())?);
49        Ok(bytes)
50    }
51
52    /// Restores and validates a graph and its multi-vector identities.
53    ///
54    /// # Errors
55    ///
56    /// Returns a format, integrity, graph, or allocation error.
57    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SearchError> {
58        if bytes.len() < HEADER_LEN || bytes.get(..8) != Some(MAGIC.as_slice()) {
59            return Err(SearchError::CorruptSnapshot(
60                "multi-vector snapshot header is invalid",
61            ));
62        }
63        let version = read_u32(bytes, 8)?;
64        if version != VERSION {
65            return Err(SearchError::UnsupportedSnapshotVersion(version));
66        }
67        if read_u32(bytes, 12)? as usize != HEADER_LEN {
68            return Err(SearchError::CorruptSnapshot(
69                "multi-vector snapshot header length does not match",
70            ));
71        }
72        if to_usize(read_u64(bytes, 24)?)? != bytes.len() {
73            return Err(SearchError::CorruptSnapshot(
74                "multi-vector snapshot file length does not match",
75            ));
76        }
77        if checksum(&bytes[HEADER_LEN..]) != read_u64(bytes, 16)? {
78            return Err(SearchError::CorruptSnapshot(
79                "multi-vector snapshot checksum does not match",
80            ));
81        }
82        let index_len = to_usize(read_u64(bytes, 32)?)?;
83        let count = to_usize(read_u64(bytes, 40)?)?;
84        let index_end = HEADER_LEN
85            .checked_add(index_len)
86            .ok_or(SearchError::CapacityOverflow)?;
87        let index = VectorIndex::from_bytes(bytes.get(HEADER_LEN..index_end).ok_or(
88            SearchError::CorruptSnapshot("multi-vector graph is truncated"),
89        )?)?;
90        if index.len() != count
91            || !index
92                .keys()
93                .enumerate()
94                .all(|(position, key)| key == position as u64)
95        {
96            return Err(SearchError::CorruptSnapshot(
97                "multi-vector graph identities do not match",
98            ));
99        }
100        let identity_len = count.checked_mul(16).ok_or(SearchError::CapacityOverflow)?;
101        let identity_bytes = bytes
102            .get(index_end..)
103            .filter(|value| value.len() == identity_len)
104            .ok_or(SearchError::CorruptSnapshot(
105                "multi-vector identities are truncated",
106            ))?;
107        let mut identities = Vec::new();
108        identities
109            .try_reserve_exact(count)
110            .map_err(|_| SearchError::AllocationFailed)?;
111        for pair in identity_bytes.chunks_exact(16) {
112            let mut key = [0_u8; 8];
113            key.copy_from_slice(&pair[..8]);
114            let mut vector_id = [0_u8; 8];
115            vector_id.copy_from_slice(&pair[8..]);
116            identities.push(MultiVectorKey {
117                key: u64::from_le_bytes(key),
118                vector_id: u64::from_le_bytes(vector_id),
119            });
120        }
121        if identities
122            .windows(2)
123            .any(|pair| (pair[0].key, pair[0].vector_id) >= (pair[1].key, pair[1].vector_id))
124        {
125            return Err(SearchError::CorruptSnapshot(
126                "multi-vector identities are not strictly ordered",
127            ));
128        }
129        Ok(Self { index, identities })
130    }
131
132    /// Writes a complete multi-vector snapshot to a stream.
133    ///
134    /// # Errors
135    ///
136    /// Returns a serialization or stream error.
137    pub fn write_to(&self, mut writer: impl Write) -> Result<(), SearchError> {
138        writer
139            .write_all(&self.to_bytes()?)
140            .map_err(|error| SearchError::storage("write multi-vector snapshot", &error))
141    }
142
143    /// Reads a complete multi-vector snapshot from a stream.
144    ///
145    /// # Errors
146    ///
147    /// Returns a stream, format, or integrity error.
148    pub fn read_from(mut reader: impl Read) -> Result<Self, SearchError> {
149        let mut bytes = Vec::new();
150        reader
151            .read_to_end(&mut bytes)
152            .map_err(|error| SearchError::storage("read multi-vector snapshot", &error))?;
153        Self::from_bytes(&bytes)
154    }
155
156    /// Atomically saves a complete multi-vector snapshot.
157    ///
158    /// # Errors
159    ///
160    /// Returns a serialization or filesystem error.
161    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), SearchError> {
162        crate::atomic_file::atomic_write(path.as_ref(), &self.to_bytes()?)
163    }
164
165    /// Loads a complete multi-vector snapshot.
166    ///
167    /// # Errors
168    ///
169    /// Returns a filesystem, format, or integrity error.
170    pub fn load(path: impl AsRef<Path>) -> Result<Self, SearchError> {
171        let bytes = std::fs::read(path.as_ref())
172            .map_err(|error| SearchError::storage("read multi-vector snapshot", &error))?;
173        Self::from_bytes(&bytes)
174    }
175}
176
177fn checksum(bytes: &[u8]) -> u64 {
178    bytes.iter().fold(FNV_OFFSET, |hash, byte| {
179        (hash ^ u64::from(*byte)).wrapping_mul(FNV_PRIME)
180    })
181}
182
183fn read_u32(bytes: &[u8], offset: usize) -> Result<u32, SearchError> {
184    Ok(u32::from_le_bytes(
185        bytes
186            .get(offset..offset + 4)
187            .ok_or(SearchError::CorruptSnapshot(
188                "multi-vector header is truncated",
189            ))?
190            .try_into()
191            .expect("checked four-byte header value"),
192    ))
193}
194
195fn read_u64(bytes: &[u8], offset: usize) -> Result<u64, SearchError> {
196    Ok(u64::from_le_bytes(
197        bytes
198            .get(offset..offset + 8)
199            .ok_or(SearchError::CorruptSnapshot(
200                "multi-vector header is truncated",
201            ))?
202            .try_into()
203            .expect("checked eight-byte header value"),
204    ))
205}
206
207fn put_u32(bytes: &mut [u8], offset: usize, value: u32) {
208    bytes[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
209}
210
211fn put_u64(bytes: &mut [u8], offset: usize, value: u64) {
212    bytes[offset..offset + 8].copy_from_slice(&value.to_le_bytes());
213}
214
215fn to_u64(value: usize) -> Result<u64, SearchError> {
216    u64::try_from(value).map_err(|_| SearchError::CapacityOverflow)
217}
218
219fn to_usize(value: u64) -> Result<usize, SearchError> {
220    usize::try_from(value).map_err(|_| SearchError::CapacityOverflow)
221}
weavatrix_search_vector/multi_io.rs

weavatrix_search_vector/
multi_io.rs
