hermes_core/segment/
vector_data.rs

1//! Vector index data structures shared between builder and reader
2
3use std::mem::size_of;
4
5use serde::{Deserialize, Serialize};
6
7/// Magic number for binary flat vector format ("FVD2" in little-endian)
8const FLAT_BINARY_MAGIC: u32 = 0x46564432;
9
10/// Binary header: magic(u32) + dim(u32) + num_vectors(u32)
11const FLAT_BINARY_HEADER_SIZE: usize = 3 * size_of::<u32>();
12/// Per-vector element size
13const FLOAT_SIZE: usize = size_of::<f32>();
14/// Per-doc_id entry: doc_id(u32) + ordinal(u16)
15const DOC_ID_ENTRY_SIZE: usize = size_of::<u32>() + size_of::<u16>();
16
17/// Flat vector data for brute-force search (accumulating state)
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct FlatVectorData {
20    pub dim: usize,
21    pub vectors: Vec<Vec<f32>>,
22    /// Document IDs with ordinals: (doc_id, ordinal) pairs
23    /// Ordinal tracks which vector in a multi-valued field
24    pub doc_ids: Vec<(u32, u16)>,
25}
26
27impl FlatVectorData {
28    /// Estimate memory usage
29    pub fn estimated_memory_bytes(&self) -> usize {
30        let vec_overhead = size_of::<Vec<f32>>();
31        let vectors_bytes: usize = self
32            .vectors
33            .iter()
34            .map(|v| v.capacity() * FLOAT_SIZE + vec_overhead)
35            .sum();
36        let doc_ids_bytes = self.doc_ids.capacity() * size_of::<(u32, u16)>();
37        vectors_bytes + doc_ids_bytes + vec_overhead * 2
38    }
39
40    /// Serialize to compact binary format.
41    ///
42    /// Format: header + vectors(dim*n*sizeof(f32)) + doc_ids(n*DOC_ID_ENTRY_SIZE)
43    /// Much more compact than JSON and avoids intermediate allocations.
44    pub fn to_binary_bytes(&self) -> Vec<u8> {
45        let num_vectors = self.doc_ids.len();
46        let total = Self::serialized_binary_size(self.dim, num_vectors);
47        let mut buf = Vec::with_capacity(total);
48
49        buf.extend_from_slice(&FLAT_BINARY_MAGIC.to_le_bytes());
50        buf.extend_from_slice(&(self.dim as u32).to_le_bytes());
51        buf.extend_from_slice(&(num_vectors as u32).to_le_bytes());
52
53        for vec in &self.vectors {
54            for &val in vec {
55                buf.extend_from_slice(&val.to_le_bytes());
56            }
57        }
58
59        for &(doc_id, ordinal) in &self.doc_ids {
60            buf.extend_from_slice(&doc_id.to_le_bytes());
61            buf.extend_from_slice(&ordinal.to_le_bytes());
62        }
63
64        buf
65    }
66
67    /// Compute the serialized size without actually serializing.
68    pub fn serialized_binary_size(index_dim: usize, num_vectors: usize) -> usize {
69        FLAT_BINARY_HEADER_SIZE
70            + num_vectors * index_dim * FLOAT_SIZE
71            + num_vectors * DOC_ID_ENTRY_SIZE
72    }
73
74    /// Stream directly from flat f32 storage to a writer (zero-buffer serialization).
75    ///
76    /// `flat_vectors` is contiguous storage of dim*n floats.
77    /// `original_dim` is the dimension in flat_vectors (may differ from index_dim for MRL).
78    pub fn serialize_binary_from_flat_streaming(
79        index_dim: usize,
80        flat_vectors: &[f32],
81        original_dim: usize,
82        doc_ids: &[(u32, u16)],
83        writer: &mut dyn std::io::Write,
84    ) -> std::io::Result<()> {
85        let num_vectors = doc_ids.len();
86
87        writer.write_all(&FLAT_BINARY_MAGIC.to_le_bytes())?;
88        writer.write_all(&(index_dim as u32).to_le_bytes())?;
89        writer.write_all(&(num_vectors as u32).to_le_bytes())?;
90
91        if index_dim == original_dim {
92            // No trimming — write all floats directly
93            // SAFETY: reinterpret f32 slice as bytes for efficient bulk write
94            let bytes: &[u8] = unsafe {
95                std::slice::from_raw_parts(
96                    flat_vectors.as_ptr() as *const u8,
97                    flat_vectors.len() * FLOAT_SIZE,
98                )
99            };
100            writer.write_all(bytes)?;
101        } else {
102            // Trim each vector to index_dim (matryoshka/MRL)
103            for i in 0..num_vectors {
104                let start = i * original_dim;
105                let slice = &flat_vectors[start..start + index_dim];
106                let bytes: &[u8] = unsafe {
107                    std::slice::from_raw_parts(slice.as_ptr() as *const u8, index_dim * FLOAT_SIZE)
108                };
109                writer.write_all(bytes)?;
110            }
111        }
112
113        for &(doc_id, ordinal) in doc_ids {
114            writer.write_all(&doc_id.to_le_bytes())?;
115            writer.write_all(&ordinal.to_le_bytes())?;
116        }
117
118        Ok(())
119    }
120
121    /// Deserialize from binary format.
122    pub fn from_binary_bytes(data: &[u8]) -> std::io::Result<Self> {
123        if data.len() < FLAT_BINARY_HEADER_SIZE {
124            return Err(std::io::Error::new(
125                std::io::ErrorKind::InvalidData,
126                "FlatVectorData binary too short",
127            ));
128        }
129
130        let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
131        if magic != FLAT_BINARY_MAGIC {
132            return Err(std::io::Error::new(
133                std::io::ErrorKind::InvalidData,
134                "Invalid FlatVectorData binary magic",
135            ));
136        }
137
138        let dim = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;
139        let num_vectors = u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize;
140
141        let vectors_start = FLAT_BINARY_HEADER_SIZE;
142        let vectors_byte_len = num_vectors * dim * FLOAT_SIZE;
143        let doc_ids_start = vectors_start + vectors_byte_len;
144        let doc_ids_byte_len = num_vectors * DOC_ID_ENTRY_SIZE;
145
146        if data.len() < doc_ids_start + doc_ids_byte_len {
147            return Err(std::io::Error::new(
148                std::io::ErrorKind::InvalidData,
149                "FlatVectorData binary truncated",
150            ));
151        }
152
153        let mut vectors = Vec::with_capacity(num_vectors);
154        for i in 0..num_vectors {
155            let mut vec = Vec::with_capacity(dim);
156            let base = vectors_start + i * dim * FLOAT_SIZE;
157            for j in 0..dim {
158                let off = base + j * FLOAT_SIZE;
159                vec.push(f32::from_le_bytes([
160                    data[off],
161                    data[off + 1],
162                    data[off + 2],
163                    data[off + 3],
164                ]));
165            }
166            vectors.push(vec);
167        }
168
169        let mut doc_ids = Vec::with_capacity(num_vectors);
170        for i in 0..num_vectors {
171            let off = doc_ids_start + i * DOC_ID_ENTRY_SIZE;
172            let doc_id =
173                u32::from_le_bytes([data[off], data[off + 1], data[off + 2], data[off + 3]]);
174            let ordinal = u16::from_le_bytes([data[off + 4], data[off + 5]]);
175            doc_ids.push((doc_id, ordinal));
176        }
177
178        Ok(FlatVectorData {
179            dim,
180            vectors,
181            doc_ids,
182        })
183    }
184}
185
186/// IVF-RaBitQ index data with embedded centroids and codebook
187#[derive(Debug, Clone, Serialize, Deserialize)]
188pub struct IVFRaBitQIndexData {
189    pub index: crate::structures::IVFRaBitQIndex,
190    pub centroids: crate::structures::CoarseCentroids,
191    pub codebook: crate::structures::RaBitQCodebook,
192}
193
194impl IVFRaBitQIndexData {
195    pub fn to_bytes(&self) -> std::io::Result<Vec<u8>> {
196        serde_json::to_vec(self)
197            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
198    }
199
200    pub fn from_bytes(data: &[u8]) -> std::io::Result<Self> {
201        serde_json::from_slice(data)
202            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
203    }
204}
205
206/// ScaNN index data with embedded centroids and codebook
207#[derive(Debug, Clone, Serialize, Deserialize)]
208pub struct ScaNNIndexData {
209    pub index: crate::structures::IVFPQIndex,
210    pub centroids: crate::structures::CoarseCentroids,
211    pub codebook: crate::structures::PQCodebook,
212}
213
214impl ScaNNIndexData {
215    pub fn to_bytes(&self) -> std::io::Result<Vec<u8>> {
216        serde_json::to_vec(self)
217            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
218    }
219
220    pub fn from_bytes(data: &[u8]) -> std::io::Result<Self> {
221        serde_json::from_slice(data)
222            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
223    }
224}
hermes_core/segment/vector_data.rs

hermes_core/segment/
vector_data.rs
