Skip to main content

weavatrix_search_vector/storage/
index_io.rs

1use super::decoder::decode_snapshot;
2use super::writer::{write_snapshot, write_snapshot_stream};
3use super::{FORMAT_VERSION, HEADER_LEN, Header, MappedVectorIndex, SnapshotMetadata};
4use crate::error::SearchError;
5use crate::hnsw::VectorIndex;
6use std::io::{Cursor, Read, Seek, Write};
7use std::path::Path;
8
9impl VectorIndex {
10    /// Persists normalized vectors, routing data, and every HNSW replica.
11    ///
12    /// The snapshot is written through a temporary sibling and flushed before
13    /// replacement.
14    ///
15    /// # Errors
16    ///
17    /// Returns a typed storage, allocation, or capacity error.
18    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), SearchError> {
19        write_snapshot(self, path.as_ref())
20    }
21
22    /// Loads a persisted index into owned memory without rebuilding HNSW.
23    ///
24    /// # Errors
25    ///
26    /// Returns a typed storage, version, integrity, allocation, or config
27    /// error.
28    pub fn load(path: impl AsRef<Path>) -> Result<Self, SearchError> {
29        MappedVectorIndex::open(path)?.to_owned()
30    }
31
32    /// Serializes the canonical snapshot into an owned byte buffer.
33    ///
34    /// # Errors
35    ///
36    /// Returns a typed storage, allocation, or capacity error.
37    pub fn to_bytes(&self) -> Result<Vec<u8>, SearchError> {
38        let mut cursor = Cursor::new(Vec::new());
39        write_snapshot_stream(self, &mut cursor)?;
40        Ok(cursor.into_inner())
41    }
42
43    /// Restores an owned index directly from a canonical snapshot buffer.
44    ///
45    /// # Errors
46    ///
47    /// Returns a typed version, integrity, allocation, or config error.
48    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SearchError> {
49        decode_snapshot(bytes)
50    }
51
52    /// Writes a canonical snapshot to a seekable stream.
53    ///
54    /// The caller should provide an empty or truncated stream.
55    ///
56    /// # Errors
57    ///
58    /// Returns a typed stream, allocation, or capacity error.
59    pub fn write_to<W: Write + Seek>(&self, writer: &mut W) -> Result<(), SearchError> {
60        write_snapshot_stream(self, writer)
61    }
62
63    /// Reads a complete canonical snapshot from a stream.
64    ///
65    /// # Errors
66    ///
67    /// Returns a typed stream, version, integrity, allocation, or config
68    /// error.
69    pub fn read_from<R: Read>(reader: &mut R) -> Result<Self, SearchError> {
70        let mut bytes = Vec::new();
71        reader
72            .read_to_end(&mut bytes)
73            .map_err(|error| SearchError::storage("read snapshot stream", &error))?;
74        Self::from_bytes(&bytes)
75    }
76
77    /// Reads snapshot metadata without decoding vectors or graphs.
78    ///
79    /// # Errors
80    ///
81    /// Returns a typed storage, version, or header error.
82    pub fn read_metadata(path: impl AsRef<Path>) -> Result<SnapshotMetadata, SearchError> {
83        let mut file = std::fs::File::open(path)
84            .map_err(|error| SearchError::storage("open snapshot metadata", &error))?;
85        let mut header = [0_u8; HEADER_LEN];
86        file.read_exact(&mut header)
87            .map_err(|error| SearchError::storage("read snapshot metadata", &error))?;
88        SnapshotMetadata::from_bytes(&header)
89    }
90}
91
92impl SnapshotMetadata {
93    /// Parses the fixed snapshot header from a full or header-only buffer.
94    ///
95    /// # Errors
96    ///
97    /// Returns a typed version or structural error.
98    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SearchError> {
99        let header = Header::parse(bytes)?;
100        Ok(Self {
101            format_version: FORMAT_VERSION,
102            vector_count: header.count,
103            serialized_bytes: header.file_len,
104            config: header.config,
105        })
106    }
107}