weavatrix_search_vector/storage/
mapped_access.rs

1use super::format::checksum;
2use super::mapped_validation::validate_sections;
3use super::{HEADER_LEN, MappedVectorIndex, SnapshotValidation};
4use crate::config::{DistanceMetric, IndexConfig};
5use crate::error::SearchError;
6use crate::mmap::Mapping;
7use crate::simd::DistanceKernel;
8use crate::vector::{distance, splitmix64, squared_norm};
9use std::path::Path;
10
11impl MappedVectorIndex {
12    /// Opens and fully validates a read-only memory-mapped snapshot.
13    ///
14    /// # Errors
15    ///
16    /// Returns a typed storage, version, integrity, or config error.
17    pub fn open(path: impl AsRef<Path>) -> Result<Self, SearchError> {
18        Self::open_with_validation(path, SnapshotValidation::Full)
19    }
20
21    /// Opens a read-only snapshot with an explicit validation policy.
22    ///
23    /// # Errors
24    ///
25    /// Returns a typed storage, version, integrity, or config error.
26    pub fn open_with_validation(
27        path: impl AsRef<Path>,
28        validation: SnapshotValidation,
29    ) -> Result<Self, SearchError> {
30        let mapping =
31            Mapping::open(path.as_ref()).map_err(|error| SearchError::storage("map", &error))?;
32        let header = super::Header::parse(mapping.bytes())?;
33        if header.file_len != mapping.len() {
34            return Err(SearchError::CorruptSnapshot(
35                "header file length does not match mapped file",
36            ));
37        }
38        if validation == SnapshotValidation::Full {
39            let actual = checksum(&mapping.bytes()[HEADER_LEN..]);
40            if actual != header.checksum {
41                return Err(SearchError::CorruptSnapshot(
42                    "snapshot payload checksum does not match",
43                ));
44            }
45        }
46        let graphs = validate_sections(&mapping, &header, validation)?;
47        Ok(Self {
48            mapping,
49            header,
50            graphs,
51            distance_kernel: DistanceKernel::detect(),
52        })
53    }
54
55    #[must_use]
56    pub fn len(&self) -> usize {
57        self.header.count
58    }
59
60    #[must_use]
61    pub fn is_empty(&self) -> bool {
62        self.header.count == 0
63    }
64
65    #[must_use]
66    pub fn dimensions(&self) -> usize {
67        self.header.config.dimensions
68    }
69
70    #[must_use]
71    pub fn config(&self) -> &IndexConfig {
72        &self.header.config
73    }
74
75    /// Returns the mapped file length. Query scratch and the small graph-layout
76    /// table are excluded.
77    #[must_use]
78    pub fn mapped_bytes(&self) -> usize {
79        self.mapping.len()
80    }
81
82    /// Iterates stable keys in ascending order.
83    #[must_use]
84    pub fn keys(&self) -> impl ExactSizeIterator<Item = u64> + '_ {
85        self.key_slice().iter().copied()
86    }
87
88    /// Returns the normalized mapped vector for `key`.
89    #[must_use]
90    pub fn vector(&self, key: u64) -> Option<&[f32]> {
91        let index = self.key_slice().binary_search(&key).ok()?;
92        Some(self.vector_at(index))
93    }
94
95    pub(super) fn key_slice(&self) -> &[u64] {
96        self.mapping
97            .u64_slice(self.header.keys_offset, self.len())
98            .expect("validated key section")
99    }
100
101    pub(super) fn vector_slice(&self) -> &[f32] {
102        self.mapping
103            .f32_slice(self.header.vectors_offset, self.len() * self.dimensions())
104            .expect("validated vector section")
105    }
106
107    pub(super) fn routing_codes(&self) -> &[u16] {
108        self.mapping
109            .u16_slice(self.header.routing_codes_offset, self.len())
110            .expect("validated routing-code section")
111    }
112
113    pub(super) fn routing_nodes(&self) -> &[u32] {
114        self.mapping
115            .u32_slice(self.header.routing_nodes_offset, self.len())
116            .expect("validated routing-node section")
117    }
118
119    pub(super) fn vector_at(&self, index: usize) -> &[f32] {
120        let start = index * self.dimensions();
121        &self.vector_slice()[start..start + self.dimensions()]
122    }
123
124    pub(super) fn query_squared_norm(&self, query: &[f32]) -> Result<f32, SearchError> {
125        if query.len() != self.dimensions() {
126            return Err(SearchError::DimensionMismatch {
127                expected: self.dimensions(),
128                actual: query.len(),
129                vector: None,
130            });
131        }
132        let norm = squared_norm(query, None)?;
133        if self.header.config.metric == DistanceMetric::Cosine && norm == 0.0 {
134            return Err(SearchError::ZeroVector { vector: None });
135        }
136        Ok(norm)
137    }
138
139    pub(super) fn distance_query(
140        &self,
141        index: usize,
142        query: &[f32],
143        query_squared_norm: f32,
144    ) -> f32 {
145        let vector = self.vector_at(index);
146        let vector_squared_norm = if self.header.config.metric == DistanceMetric::Cosine {
147            1.0
148        } else {
149            squared_norm(vector, Some(index)).expect("validated mapped vector")
150        };
151        distance(
152            self.distance_kernel,
153            self.header.config.metric,
154            vector,
155            vector_squared_norm,
156            query,
157            query_squared_norm,
158        )
159    }
160
161    pub(super) fn routing_signs(&self) -> impl Iterator<Item = u16> + '_ {
162        (0..self.dimensions()).map(|dimension| {
163            let mixed =
164                splitmix64(u64::try_from(dimension).unwrap_or(u64::MAX) ^ 0xa076_1d64_78bd_642f);
165            u16::try_from(mixed & u64::from(u16::MAX)).expect("masked routing signs fit u16")
166        })
167    }
168}
weavatrix_search_vector/storage/mapped_access.rs

weavatrix_search_vector/storage/
mapped_access.rs
