hermes_core/index/
metadata.rs

1//! Unified index metadata - segments list + vector index state
2//!
3//! This module manages all index-level metadata in a single `metadata.json` file:
4//! - List of committed segments
5//! - Vector index state per field (Flat/Built)
6//! - Trained centroids/codebooks paths
7//!
8//! The workflow is:
9//! 1. During accumulation: segments store Flat vectors, state is Flat
10//! 2. When threshold crossed: train ONCE, update state to Built
11//! 3. On index open: load metadata, skip re-training if already built
12
13use serde::{Deserialize, Serialize};
14use std::collections::HashMap;
15use std::path::Path;
16
17use crate::dsl::VectorIndexType;
18use crate::error::{Error, Result};
19use crate::schema::Schema;
20
21/// Metadata file name at index level
22pub const INDEX_META_FILENAME: &str = "metadata.json";
23/// Temp file for atomic writes (write here, then rename to INDEX_META_FILENAME)
24const INDEX_META_TMP_FILENAME: &str = "metadata.json.tmp";
25
26/// State of vector index for a field
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
28pub enum VectorIndexState {
29    /// Accumulating vectors - using Flat (brute-force) search
30    #[default]
31    Flat,
32    /// Index structures built - using ANN search
33    Built {
34        /// Total vector count when training happened
35        vector_count: usize,
36        /// Number of clusters used
37        num_clusters: usize,
38    },
39}
40
41/// Per-segment metadata stored in index metadata
42/// This allows merge decisions without loading segment files
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct SegmentMetaInfo {
45    /// Number of documents in this segment
46    pub num_docs: u32,
47}
48
49/// Per-field vector index metadata
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct FieldVectorMeta {
52    /// Field ID
53    pub field_id: u32,
54    /// Configured index type (target type when built)
55    pub index_type: VectorIndexType,
56    /// Current state
57    pub state: VectorIndexState,
58    /// Path to centroids file (relative to index dir)
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub centroids_file: Option<String>,
61    /// Path to codebook file (relative to index dir, for ScaNN)
62    #[serde(skip_serializing_if = "Option::is_none")]
63    pub codebook_file: Option<String>,
64}
65
66/// Unified index metadata - single source of truth for index state
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct IndexMetadata {
69    /// Version for compatibility
70    pub version: u32,
71    /// Index schema
72    pub schema: Schema,
73    /// Segment metadata: segment_id -> info (doc count, etc.)
74    /// Using HashMap allows O(1) lookup and stores doc counts for merge decisions
75    #[serde(default)]
76    pub segment_metas: HashMap<String, SegmentMetaInfo>,
77    /// Per-field vector index metadata
78    #[serde(default)]
79    pub vector_fields: HashMap<u32, FieldVectorMeta>,
80    /// Total vectors across all segments (updated on commit)
81    #[serde(default)]
82    pub total_vectors: usize,
83}
84
85impl IndexMetadata {
86    /// Create new metadata with schema
87    pub fn new(schema: Schema) -> Self {
88        Self {
89            version: 1,
90            schema,
91            segment_metas: HashMap::new(),
92            vector_fields: HashMap::new(),
93            total_vectors: 0,
94        }
95    }
96
97    /// Get segment IDs as a sorted Vec (deterministic ordering for doc_id_offset assignment)
98    pub fn segment_ids(&self) -> Vec<String> {
99        let mut ids: Vec<String> = self.segment_metas.keys().cloned().collect();
100        ids.sort();
101        ids
102    }
103
104    /// Add or update a segment with its doc count
105    pub fn add_segment(&mut self, segment_id: String, num_docs: u32) {
106        self.segment_metas
107            .insert(segment_id, SegmentMetaInfo { num_docs });
108    }
109
110    /// Remove a segment
111    pub fn remove_segment(&mut self, segment_id: &str) {
112        self.segment_metas.remove(segment_id);
113    }
114
115    /// Check if segment exists
116    pub fn has_segment(&self, segment_id: &str) -> bool {
117        self.segment_metas.contains_key(segment_id)
118    }
119
120    /// Get segment doc count
121    pub fn segment_doc_count(&self, segment_id: &str) -> Option<u32> {
122        self.segment_metas.get(segment_id).map(|m| m.num_docs)
123    }
124
125    /// Check if a field has been built
126    pub fn is_field_built(&self, field_id: u32) -> bool {
127        self.vector_fields
128            .get(&field_id)
129            .map(|f| matches!(f.state, VectorIndexState::Built { .. }))
130            .unwrap_or(false)
131    }
132
133    /// Get field metadata
134    pub fn get_field_meta(&self, field_id: u32) -> Option<&FieldVectorMeta> {
135        self.vector_fields.get(&field_id)
136    }
137
138    /// Initialize field metadata (called when field is first seen)
139    pub fn init_field(&mut self, field_id: u32, index_type: VectorIndexType) {
140        self.vector_fields
141            .entry(field_id)
142            .or_insert(FieldVectorMeta {
143                field_id,
144                index_type,
145                state: VectorIndexState::Flat,
146                centroids_file: None,
147                codebook_file: None,
148            });
149    }
150
151    /// Mark field as built with trained structures
152    pub fn mark_field_built(
153        &mut self,
154        field_id: u32,
155        vector_count: usize,
156        num_clusters: usize,
157        centroids_file: String,
158        codebook_file: Option<String>,
159    ) {
160        if let Some(field) = self.vector_fields.get_mut(&field_id) {
161            field.state = VectorIndexState::Built {
162                vector_count,
163                num_clusters,
164            };
165            field.centroids_file = Some(centroids_file);
166            field.codebook_file = codebook_file;
167        }
168    }
169
170    /// Check if field should be built based on threshold
171    pub fn should_build_field(&self, field_id: u32, threshold: usize) -> bool {
172        // Don't build if already built
173        if self.is_field_built(field_id) {
174            return false;
175        }
176        // Build if we have enough vectors
177        self.total_vectors >= threshold
178    }
179
180    /// Load from directory
181    ///
182    /// If `metadata.json` is missing but `metadata.json.tmp` exists (crash
183    /// between write and rename), recovers from the temp file.
184    pub async fn load<D: crate::directories::Directory>(dir: &D) -> Result<Self> {
185        let path = Path::new(INDEX_META_FILENAME);
186        match dir.open_read(path).await {
187            Ok(slice) => {
188                let bytes = slice.read_bytes().await?;
189                serde_json::from_slice(bytes.as_slice())
190                    .map_err(|e| Error::Serialization(e.to_string()))
191            }
192            Err(_) => {
193                // Try recovering from temp file (crash between write and rename)
194                let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
195                let slice = dir.open_read(tmp_path).await?;
196                let bytes = slice.read_bytes().await?;
197                let meta: Self = serde_json::from_slice(bytes.as_slice())
198                    .map_err(|e| Error::Serialization(e.to_string()))?;
199                log::warn!("Recovered metadata from temp file (previous crash during save)");
200                Ok(meta)
201            }
202        }
203    }
204
205    /// Save to directory (atomic: write temp file, then rename)
206    ///
207    /// Uses write-then-rename so a crash mid-write won't corrupt the
208    /// existing metadata file. On POSIX, rename is atomic.
209    pub async fn save<D: crate::directories::DirectoryWriter>(&self, dir: &D) -> Result<()> {
210        let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
211        let final_path = Path::new(INDEX_META_FILENAME);
212        let bytes =
213            serde_json::to_vec_pretty(self).map_err(|e| Error::Serialization(e.to_string()))?;
214        dir.write(tmp_path, &bytes).await.map_err(Error::Io)?;
215        dir.rename(tmp_path, final_path).await.map_err(Error::Io)?;
216        Ok(())
217    }
218
219    /// Load trained structures from index-level files for fields in Built state.
220    ///
221    /// Returns `Some(TrainedVectorStructures)` if at least one field has trained
222    /// centroids, `None` otherwise.
223    pub async fn load_trained_structures<D: crate::directories::Directory>(
224        &self,
225        dir: &D,
226    ) -> Option<crate::segment::TrainedVectorStructures> {
227        use std::sync::Arc;
228
229        let mut centroids = rustc_hash::FxHashMap::default();
230        let mut codebooks = rustc_hash::FxHashMap::default();
231
232        log::info!(
233            "[trained] loading trained structures, vector_fields={:?}",
234            self.vector_fields.keys().collect::<Vec<_>>()
235        );
236
237        for (field_id, field_meta) in &self.vector_fields {
238            log::info!(
239                "[trained] field {} state={:?} centroids_file={:?} codebook_file={:?}",
240                field_id,
241                field_meta.state,
242                field_meta.centroids_file,
243                field_meta.codebook_file,
244            );
245            if !matches!(field_meta.state, VectorIndexState::Built { .. }) {
246                log::info!("[trained] field {} skipped (not Built)", field_id);
247                continue;
248            }
249
250            // Load centroids
251            match &field_meta.centroids_file {
252                None => {
253                    log::warn!(
254                        "[trained] field {} is Built but has no centroids_file",
255                        field_id
256                    );
257                }
258                Some(file) => match dir.open_read(Path::new(file)).await {
259                    Err(e) => {
260                        log::warn!(
261                            "[trained] field {} failed to open centroids file '{}': {}",
262                            field_id,
263                            file,
264                            e
265                        );
266                    }
267                    Ok(slice) => match slice.read_bytes().await {
268                        Err(e) => {
269                            log::warn!(
270                                "[trained] field {} failed to read centroids file '{}': {}",
271                                field_id,
272                                file,
273                                e
274                            );
275                        }
276                        Ok(bytes) => {
277                            match serde_json::from_slice::<crate::structures::CoarseCentroids>(
278                                bytes.as_slice(),
279                            ) {
280                                Err(e) => {
281                                    log::warn!(
282                                        "[trained] field {} failed to deserialize centroids from '{}': {}",
283                                        field_id,
284                                        file,
285                                        e
286                                    );
287                                }
288                                Ok(c) => {
289                                    log::info!(
290                                        "[trained] field {} loaded centroids ({} clusters)",
291                                        field_id,
292                                        c.num_clusters
293                                    );
294                                    centroids.insert(*field_id, Arc::new(c));
295                                }
296                            }
297                        }
298                    },
299                },
300            }
301
302            // Load codebook (for ScaNN)
303            match &field_meta.codebook_file {
304                None => {} // optional, not all index types use codebooks
305                Some(file) => match dir.open_read(Path::new(file)).await {
306                    Err(e) => {
307                        log::warn!(
308                            "[trained] field {} failed to open codebook file '{}': {}",
309                            field_id,
310                            file,
311                            e
312                        );
313                    }
314                    Ok(slice) => match slice.read_bytes().await {
315                        Err(e) => {
316                            log::warn!(
317                                "[trained] field {} failed to read codebook file '{}': {}",
318                                field_id,
319                                file,
320                                e
321                            );
322                        }
323                        Ok(bytes) => {
324                            match serde_json::from_slice::<crate::structures::PQCodebook>(
325                                bytes.as_slice(),
326                            ) {
327                                Err(e) => {
328                                    log::warn!(
329                                        "[trained] field {} failed to deserialize codebook from '{}': {}",
330                                        field_id,
331                                        file,
332                                        e
333                                    );
334                                }
335                                Ok(c) => {
336                                    log::info!("[trained] field {} loaded codebook", field_id);
337                                    codebooks.insert(*field_id, Arc::new(c));
338                                }
339                            }
340                        }
341                    },
342                },
343            }
344        }
345
346        if centroids.is_empty() {
347            None
348        } else {
349            Some(crate::segment::TrainedVectorStructures {
350                centroids,
351                codebooks,
352            })
353        }
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360
361    fn test_schema() -> Schema {
362        Schema::default()
363    }
364
365    #[test]
366    fn test_metadata_init() {
367        let mut meta = IndexMetadata::new(test_schema());
368        assert_eq!(meta.total_vectors, 0);
369        assert!(meta.segment_metas.is_empty());
370        assert!(!meta.is_field_built(0));
371
372        meta.init_field(0, VectorIndexType::IvfRaBitQ);
373        assert!(!meta.is_field_built(0));
374        assert!(meta.vector_fields.contains_key(&0));
375    }
376
377    #[test]
378    fn test_metadata_segments() {
379        let mut meta = IndexMetadata::new(test_schema());
380        meta.add_segment("abc123".to_string(), 50);
381        meta.add_segment("def456".to_string(), 100);
382        assert_eq!(meta.segment_metas.len(), 2);
383        assert_eq!(meta.segment_doc_count("abc123"), Some(50));
384        assert_eq!(meta.segment_doc_count("def456"), Some(100));
385
386        // Overwrites existing
387        meta.add_segment("abc123".to_string(), 75);
388        assert_eq!(meta.segment_metas.len(), 2);
389        assert_eq!(meta.segment_doc_count("abc123"), Some(75));
390
391        meta.remove_segment("abc123");
392        assert_eq!(meta.segment_metas.len(), 1);
393        assert!(meta.has_segment("def456"));
394        assert!(!meta.has_segment("abc123"));
395    }
396
397    #[test]
398    fn test_mark_field_built() {
399        let mut meta = IndexMetadata::new(test_schema());
400        meta.init_field(0, VectorIndexType::IvfRaBitQ);
401        meta.total_vectors = 10000;
402
403        assert!(!meta.is_field_built(0));
404
405        meta.mark_field_built(0, 10000, 256, "field_0_centroids.bin".to_string(), None);
406
407        assert!(meta.is_field_built(0));
408        let field = meta.get_field_meta(0).unwrap();
409        assert_eq!(
410            field.centroids_file.as_deref(),
411            Some("field_0_centroids.bin")
412        );
413    }
414
415    #[test]
416    fn test_should_build_field() {
417        let mut meta = IndexMetadata::new(test_schema());
418        meta.init_field(0, VectorIndexType::IvfRaBitQ);
419
420        // Below threshold
421        meta.total_vectors = 500;
422        assert!(!meta.should_build_field(0, 1000));
423
424        // Above threshold
425        meta.total_vectors = 1500;
426        assert!(meta.should_build_field(0, 1000));
427
428        // Already built - should not build again
429        meta.mark_field_built(0, 1500, 256, "centroids.bin".to_string(), None);
430        assert!(!meta.should_build_field(0, 1000));
431    }
432
433    #[test]
434    fn test_serialization() {
435        let mut meta = IndexMetadata::new(test_schema());
436        meta.add_segment("seg1".to_string(), 100);
437        meta.init_field(0, VectorIndexType::IvfRaBitQ);
438        meta.total_vectors = 5000;
439
440        let json = serde_json::to_string_pretty(&meta).unwrap();
441        let loaded: IndexMetadata = serde_json::from_str(&json).unwrap();
442
443        assert_eq!(loaded.segment_ids().len(), meta.segment_ids().len());
444        assert_eq!(loaded.segment_doc_count("seg1"), Some(100));
445        assert_eq!(loaded.total_vectors, meta.total_vectors);
446        assert!(loaded.vector_fields.contains_key(&0));
447    }
448}
hermes_core/index/metadata.rs

hermes_core/index/
metadata.rs
