Skip to main content

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 centroid artifact paths
7//!
8//! The workflow is:
9//! 1. During initial accumulation, segments store flat vectors.
10//! 2. A manual build trains the first coarse-centroid ANN generation.
11//! 3. A manual retrain stages and atomically publishes a replacement generation.
12//! 4. On index open, metadata loads the currently published artifacts.
13
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::io::Write;
17use std::path::Path;
18
19use crate::dsl::{BinaryIndexType, Schema, VectorIndexType};
20use crate::error::{Error, Result};
21
22/// Metadata file name at index level
23pub const INDEX_META_FILENAME: &str = "metadata.json";
24/// Temp file for atomic writes (write here, then rename to INDEX_META_FILENAME)
25const INDEX_META_TMP_FILENAME: &str = "metadata.json.tmp";
26
27/// Current metadata.json format version written by this build.
28///
29/// `load` refuses metadata stamped with a newer version: serde_json silently
30/// drops fields it does not know about, so loading newer metadata would
31/// misread index state and the next save would destructively rewrite the
32/// unknown fields away.
33pub const INDEX_META_FORMAT_VERSION: u32 = 4;
34
35/// Index-level centroids/codebooks are deliberately bounded before they are
36/// read or decoded. Besides limiting ordinary corruption damage, the matching
37/// bincode limit prevents a tiny forged collection length from requesting an
38/// effectively unbounded allocation.
39pub(crate) const MAX_TRAINED_ARTIFACT_BYTES: usize = 512 * 1024 * 1024;
40
41/// State of vector index for a field
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
43pub enum VectorIndexState {
44    /// Accumulating vectors - using Flat (brute-force) search
45    #[default]
46    Flat,
47    /// Index structures built - using ANN search
48    Built {
49        /// Total vector count when training happened
50        vector_count: usize,
51        /// Number of clusters used
52        num_clusters: usize,
53    },
54}
55
56fn default_true() -> bool {
57    true
58}
59
60/// Per-segment metadata stored in index metadata
61/// This allows merge decisions without loading segment files
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SegmentMetaInfo {
64    /// Number of documents in this segment
65    pub num_docs: u32,
66    /// Parent segment IDs that were merged to produce this segment (empty for fresh segments)
67    pub ancestors: Vec<String>,
68    /// Merge generation: 0 for fresh segments, max(parent generations) + 1 for merged segments
69    pub generation: u32,
70    /// Whether this segment has been reordered via Recursive Graph Bisection (BP).
71    /// Fresh segments and block-copy merges are not reordered. Only segments that have
72    /// been explicitly reordered (via background optimizer or reorder command) are marked true.
73    #[serde(default)]
74    pub reordered: bool,
75    /// Whether the last BP reorder pass ran to natural convergence. False when
76    /// a wall-clock BP budget ended the pass early — the segment is ordered
77    /// better than before, and a later warm-started pass can deepen it.
78    /// Old metadata (field absent) deserializes as converged.
79    #[serde(default = "default_true")]
80    pub bp_converged: bool,
81    /// Number of consecutive budget-exhausted BP rewrites in this segment's
82    /// current reordered lineage. Carried across replacement IDs so the
83    /// optimizer can impose a hard follow-up bound instead of rewriting forever.
84    #[serde(default)]
85    pub bp_unconverged_passes: u32,
86}
87
88/// Per-field vector index metadata
89#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
90#[serde(tag = "kind", content = "index", rename_all = "snake_case")]
91pub enum VectorFieldIndexType {
92    Float(VectorIndexType),
93    Binary(BinaryIndexType),
94}
95
96impl From<VectorIndexType> for VectorFieldIndexType {
97    fn from(value: VectorIndexType) -> Self {
98        Self::Float(value)
99    }
100}
101
102impl From<BinaryIndexType> for VectorFieldIndexType {
103    fn from(value: BinaryIndexType) -> Self {
104        Self::Binary(value)
105    }
106}
107
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct FieldVectorMeta {
110    /// Field ID
111    pub field_id: u32,
112    /// Configured index type (target type when built)
113    pub index_type: VectorFieldIndexType,
114    /// Current state
115    pub state: VectorIndexState,
116    /// Path to centroids file (relative to index dir)
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub centroids_file: Option<String>,
119    /// Legacy: path to a trained IVF-PQ codebook. Always `None` for current
120    /// formats; kept so pre-removal metadata deserializes into an actionable
121    /// error instead of dropping the field.
122    #[serde(skip_serializing_if = "Option::is_none")]
123    pub codebook_file: Option<String>,
124}
125
126/// Unified index metadata - single source of truth for index state
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct IndexMetadata {
129    /// Version for compatibility
130    pub version: u32,
131    /// Index schema
132    pub schema: Schema,
133    /// Segment metadata: segment_id -> info (doc count, etc.)
134    /// Using HashMap allows O(1) lookup and stores doc counts for merge decisions
135    #[serde(default)]
136    pub segment_metas: HashMap<String, SegmentMetaInfo>,
137    /// Per-field vector index metadata
138    #[serde(default)]
139    pub vector_fields: HashMap<u32, FieldVectorMeta>,
140    /// Aggregate vector count recorded by all built vector fields.
141    ///
142    /// The per-field `VectorIndexState::Built::vector_count` values are the
143    /// source of truth. This cached aggregate is refreshed whenever a field is
144    /// marked built, rather than being overwritten with whichever field was
145    /// trained last.
146    #[serde(default)]
147    pub total_vectors: usize,
148}
149
150impl IndexMetadata {
151    /// Create new metadata with schema
152    pub fn new(schema: Schema) -> Self {
153        Self {
154            version: INDEX_META_FORMAT_VERSION,
155            schema,
156            segment_metas: HashMap::new(),
157            vector_fields: HashMap::new(),
158            total_vectors: 0,
159        }
160    }
161
162    /// Get segment IDs as a sorted Vec (deterministic ordering)
163    pub fn segment_ids(&self) -> Vec<String> {
164        let mut ids: Vec<String> = self.segment_metas.keys().cloned().collect();
165        ids.sort();
166        ids
167    }
168
169    /// Add a fresh segment (gen=0, no ancestors, not reordered)
170    pub fn add_segment(&mut self, segment_id: String, num_docs: u32) {
171        self.segment_metas.insert(
172            segment_id,
173            SegmentMetaInfo {
174                num_docs,
175                ancestors: Vec::new(),
176                generation: 0,
177                reordered: false,
178                bp_converged: true,
179                bp_unconverged_passes: 0,
180            },
181        );
182    }
183
184    /// Add a merged segment with lineage info
185    pub fn add_merged_segment(
186        &mut self,
187        segment_id: String,
188        num_docs: u32,
189        ancestors: Vec<String>,
190        generation: u32,
191        reordered: bool,
192        bp_converged: bool,
193    ) {
194        self.add_segment_meta(
195            segment_id,
196            SegmentMetaInfo {
197                num_docs,
198                ancestors,
199                generation,
200                reordered,
201                bp_converged,
202                bp_unconverged_passes: 0,
203            },
204        );
205    }
206
207    /// Insert fully constructed lifecycle metadata. Merge/reorder code uses
208    /// this to carry bounded BP lineage; ordinary callers use the safer
209    /// constructors above, which start a fresh lineage.
210    pub(crate) fn add_segment_meta(&mut self, segment_id: String, info: SegmentMetaInfo) {
211        self.segment_metas.insert(segment_id, info);
212    }
213
214    /// Remove a segment
215    pub fn remove_segment(&mut self, segment_id: &str) {
216        self.segment_metas.remove(segment_id);
217    }
218
219    /// Check if segment exists
220    pub fn has_segment(&self, segment_id: &str) -> bool {
221        self.segment_metas.contains_key(segment_id)
222    }
223
224    /// Get segment doc count
225    pub fn segment_doc_count(&self, segment_id: &str) -> Option<u32> {
226        self.segment_metas.get(segment_id).map(|m| m.num_docs)
227    }
228
229    /// Check if a field has been built
230    pub fn is_field_built(&self, field_id: u32) -> bool {
231        self.vector_fields
232            .get(&field_id)
233            .map(|f| matches!(f.state, VectorIndexState::Built { .. }))
234            .unwrap_or(false)
235    }
236
237    /// Get field metadata
238    pub fn get_field_meta(&self, field_id: u32) -> Option<&FieldVectorMeta> {
239        self.vector_fields.get(&field_id)
240    }
241
242    /// Initialize field metadata (called when field is first seen)
243    pub fn init_field(&mut self, field_id: u32, index_type: impl Into<VectorFieldIndexType>) {
244        let index_type = index_type.into();
245        self.vector_fields
246            .entry(field_id)
247            .or_insert(FieldVectorMeta {
248                field_id,
249                index_type,
250                state: VectorIndexState::Flat,
251                centroids_file: None,
252                codebook_file: None,
253            });
254    }
255
256    /// Mark field as built with trained structures
257    pub fn mark_field_built(
258        &mut self,
259        field_id: u32,
260        vector_count: usize,
261        num_clusters: usize,
262        centroids_file: String,
263        codebook_file: Option<String>,
264    ) {
265        if let Some(field) = self.vector_fields.get_mut(&field_id) {
266            field.state = VectorIndexState::Built {
267                vector_count,
268                num_clusters,
269            };
270            field.centroids_file = Some(centroids_file);
271            field.codebook_file = codebook_file;
272            self.refresh_total_vectors();
273        }
274    }
275
276    /// Refresh the cached aggregate from the authoritative per-field states.
277    ///
278    /// Saturation keeps this infallible metadata helper safe even if it is
279    /// called after loading externally modified metadata with impossible
280    /// counts.
281    pub(crate) fn refresh_total_vectors(&mut self) {
282        self.total_vectors = self
283            .vector_fields
284            .values()
285            .filter_map(|field| match field.state {
286                VectorIndexState::Built { vector_count, .. } => Some(vector_count),
287                VectorIndexState::Flat => None,
288            })
289            .fold(0usize, usize::saturating_add);
290    }
291
292    /// Check if field should be built based on threshold
293    pub fn should_build_field(&self, field_id: u32, threshold: usize) -> bool {
294        // Don't build if already built
295        if self.is_field_built(field_id) {
296            return false;
297        }
298        // Build if we have enough vectors
299        self.total_vectors >= threshold
300    }
301
302    /// Load from directory
303    ///
304    /// If `metadata.json` is missing but `metadata.json.tmp` exists (crash
305    /// between write and rename), recovers from the temp file.
306    pub async fn load<D: crate::directories::Directory>(dir: &D) -> Result<Self> {
307        let path = Path::new(INDEX_META_FILENAME);
308        match dir.open_read(path).await {
309            Ok(slice) => {
310                let bytes = slice.read_bytes().await?;
311                Self::deserialize_versioned(bytes.as_slice())
312            }
313            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
314                // Try recovering from temp file (crash between write and rename)
315                let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
316                let slice = dir.open_read(tmp_path).await?;
317                let bytes = slice.read_bytes().await?;
318                let meta = Self::deserialize_versioned(bytes.as_slice())?;
319                log::warn!("Recovered metadata from temp file (previous crash during save)");
320                Ok(meta)
321            }
322            Err(e) => Err(Error::Io(e)),
323        }
324    }
325
326    /// Deserialize only the current format. Vector artifacts are intentionally
327    /// rebuilt when the ANN format changes; silently accepting older metadata
328    /// would mix incompatible segment and global-codebook generations.
329    fn deserialize_versioned(bytes: &[u8]) -> Result<Self> {
330        let meta: Self =
331            serde_json::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))?;
332        crate::dsl::reject_removed_vector_index_types(&meta.schema).map_err(Error::Schema)?;
333        if meta.version != INDEX_META_FORMAT_VERSION {
334            return Err(Error::Corruption(format!(
335                "metadata.json format version {} is incompatible with required version {}; \
336                 rebuild and republish the index with this Hermes version",
337                meta.version, INDEX_META_FORMAT_VERSION
338            )));
339        }
340        Ok(meta)
341    }
342
343    /// Save to directory (atomic: write temp file, then rename)
344    ///
345    /// Uses write-then-rename so a crash mid-write won't corrupt the
346    /// existing metadata file. On POSIX, rename is atomic.
347    pub async fn save<D: crate::directories::DirectoryWriter>(&self, dir: &D) -> Result<()> {
348        let bytes = self.serialize_to_bytes()?;
349        Self::save_bytes(dir, &bytes).await
350    }
351
352    /// Serialize metadata to bytes (cheap, no I/O).
353    /// Useful when you need to release a lock before doing disk I/O.
354    pub fn serialize_to_bytes(&self) -> Result<Vec<u8>> {
355        serde_json::to_vec_pretty(self).map_err(|e| Error::Serialization(e.to_string()))
356    }
357
358    /// Write pre-serialized metadata bytes to directory (atomic rename + fsync).
359    ///
360    /// The fsync ensures durability: without it, a power failure after rename
361    /// could lose the metadata update on systems with volatile write caches.
362    pub async fn save_bytes<D: crate::directories::DirectoryWriter>(
363        dir: &D,
364        bytes: &[u8],
365    ) -> Result<()> {
366        let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
367        let final_path = Path::new(INDEX_META_FILENAME);
368        // Metadata is tiny, but `DirectoryWriter::write` does not guarantee
369        // the file contents themselves are fsynced. Finish the streaming
370        // writer first (filesystem implementations call `File::sync_all`),
371        // then atomically publish the durable temp file by rename.
372        let mut writer = dir.streaming_writer(tmp_path).await.map_err(Error::Io)?;
373        writer.write_all(bytes).map_err(Error::Io)?;
374        writer.finish().map_err(Error::Io)?;
375        // Rename is the logical commit point: after it succeeds, readers can
376        // observe the new generation and callers must publish the matching
377        // in-memory/tracker state. Directory fsync only strengthens crash
378        // durability. It cannot safely turn an already-visible rename into a
379        // reported pre-commit failure, because cleanup could then delete files
380        // referenced by the metadata now on disk.
381        dir.rename(tmp_path, final_path).await.map_err(Error::Io)?;
382        if let Err(error) = dir.sync().await {
383            log::error!(
384                "[metadata] directory fsync failed after committed rename: {}. \
385                 Continuing with the renamed generation; crash durability is not guaranteed",
386                error,
387            );
388        }
389        Ok(())
390    }
391
392    /// Fallible schema-aware loader used for lifecycle publication.
393    #[cfg_attr(not(feature = "native"), allow(dead_code))]
394    pub(crate) async fn try_load_trained_from_fields<D: crate::directories::Directory>(
395        vector_fields: &HashMap<u32, FieldVectorMeta>,
396        schema: &Schema,
397        dir: &D,
398    ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
399        Self::load_trained_from_fields_impl(vector_fields, schema, dir).await
400    }
401
402    /// Load and validate the complete trained-artifact set described by a
403    /// `vector_fields` snapshot.
404    ///
405    /// This is intentionally all-or-nothing. A `Built` field is a durable
406    /// promise that every artifact required by its configured index exists and
407    /// is compatible with the schema. Returning a partial map would let some
408    /// segment builders publish ANN data while another field was silently
409    /// unusable, and would make the same index behave differently after a
410    /// restart.
411    async fn load_trained_from_fields_impl<D: crate::directories::Directory>(
412        vector_fields: &HashMap<u32, FieldVectorMeta>,
413        schema: &Schema,
414        dir: &D,
415    ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
416        use std::sync::Arc;
417
418        let mut centroids = rustc_hash::FxHashMap::default();
419        let mut binary_quantizers = rustc_hash::FxHashMap::default();
420
421        let mut built_fields: Vec<_> = vector_fields
422            .iter()
423            .filter(|(_, meta)| matches!(meta.state, VectorIndexState::Built { .. }))
424            .collect();
425        built_fields.sort_unstable_by_key(|(field_id, _)| **field_id);
426
427        log::debug!(
428            "[trained] loading trained structures, dense_vector_fields={:?}",
429            vector_fields.keys().collect::<Vec<_>>()
430        );
431
432        for (field_id, field_meta) in built_fields {
433            log::debug!(
434                "[trained] field {} state={:?} centroids_file={:?} codebook_file={:?}",
435                field_id,
436                field_meta.state,
437                field_meta.centroids_file,
438                field_meta.codebook_file,
439            );
440            if field_meta.field_id != *field_id {
441                return Err(Error::Corruption(format!(
442                    "trained vector metadata key {field_id} contains field_id {}",
443                    field_meta.field_id
444                )));
445            }
446
447            let expected_clusters = match field_meta.state {
448                VectorIndexState::Built { num_clusters, .. } if num_clusters > 0 => num_clusters,
449                VectorIndexState::Built { .. } => {
450                    return Err(Error::Corruption(format!(
451                        "trained vector metadata field {field_id} has zero clusters"
452                    )));
453                }
454                VectorIndexState::Flat => unreachable!("built_fields contains only Built entries"),
455            };
456
457            let centroids_file = field_meta.centroids_file.as_deref().ok_or_else(|| {
458                Error::Corruption(format!(
459                    "trained vector metadata field {field_id} is Built but has no centroids_file"
460                ))
461            })?;
462            match field_meta.index_type {
463                VectorFieldIndexType::Float(VectorIndexType::IvfPq) => {
464                    return Err(Error::Corruption(format!(
465                        "field {field_id} was trained as IVF-PQ, which is no longer \
466                         supported; recreate the index with `ivf_tq` and reindex \
467                         (docs/turboquant-quantization.md)"
468                    )));
469                }
470                VectorFieldIndexType::Float(index_type @ VectorIndexType::IvfTq) => {
471                    let entry = schema
472                        .get_field_entry(crate::dsl::Field(*field_id))
473                        .ok_or_else(|| {
474                            Error::Corruption(format!(
475                                "trained vector metadata references missing field {field_id}"
476                            ))
477                        })?;
478                    let schema_config = entry
479                        .dense_vector_config
480                        .as_ref()
481                        .filter(|_| entry.field_type == crate::dsl::FieldType::DenseVector)
482                        .ok_or_else(|| {
483                            Error::Corruption(format!(
484                                "trained vector metadata field {field_id} is not a float dense field"
485                            ))
486                        })?;
487                    if schema_config.index_type != index_type {
488                        return Err(Error::Corruption(format!(
489                            "trained vector metadata field {field_id} uses {index_type:?}, schema requires {:?}",
490                            schema_config.index_type
491                        )));
492                    }
493                    let c: crate::structures::CoarseCentroids =
494                        load_trained_artifact(dir, *field_id, "centroids", centroids_file).await?;
495                    let expected_dim = schema_config.dim;
496                    let actual_clusters = c.num_clusters as usize;
497                    let expected_values =
498                        actual_clusters.checked_mul(expected_dim).ok_or_else(|| {
499                            Error::Corruption(format!(
500                                "trained centroid dimensions overflow for field {field_id}"
501                            ))
502                        })?;
503                    if actual_clusters == 0
504                        || actual_clusters > expected_clusters
505                        || c.dim == 0
506                        || c.dim != expected_dim
507                        || c.centroids.len() != expected_values
508                        || c.centroids.iter().any(|value| !value.is_finite())
509                    {
510                        return Err(Error::Corruption(format!(
511                            "trained centroids for field {field_id} do not match metadata/schema"
512                        )));
513                    }
514                    c.validate_routing(schema_config.ivf_routing)
515                        .map_err(|error| {
516                            Error::Corruption(format!(
517                                "invalid trained centroid routing for field {field_id}: {error}"
518                            ))
519                        })?;
520                    // The TQ leaf codec is derived, never trained; ensure
521                    // `index_type` stays referenced for future variants.
522                    let _ = index_type;
523                    if field_meta.codebook_file.is_some() {
524                        return Err(Error::Corruption(format!(
525                            "trained IVF-TQ field {field_id} unexpectedly references a codebook file"
526                        )));
527                    }
528                    centroids.insert(*field_id, Arc::new(c));
529                }
530                VectorFieldIndexType::Binary(BinaryIndexType::Ivf) => {
531                    let entry = schema
532                        .get_field_entry(crate::dsl::Field(*field_id))
533                        .ok_or_else(|| {
534                            Error::Corruption(format!(
535                                "trained vector metadata references missing field {field_id}"
536                            ))
537                        })?;
538                    let schema_config = entry
539                        .binary_dense_vector_config
540                        .as_ref()
541                        .filter(|config| {
542                            entry.field_type == crate::dsl::FieldType::BinaryDenseVector
543                                && config.index_type == BinaryIndexType::Ivf
544                        })
545                        .ok_or_else(|| {
546                            Error::Corruption(format!(
547                                "trained vector metadata field {field_id} is not a binary IVF field"
548                            ))
549                        })?;
550                    let quantizer: crate::structures::BinaryCoarseQuantizer =
551                        load_trained_artifact(dir, *field_id, "binary centroids", centroids_file)
552                            .await?;
553                    quantizer.validate().map_err(|error| {
554                        Error::Corruption(format!(
555                            "invalid binary coarse quantizer for field {field_id}: {error}"
556                        ))
557                    })?;
558                    let actual_clusters = quantizer.num_clusters as usize;
559                    if actual_clusters > expected_clusters
560                        || schema_config.dim != quantizer.dim_bits
561                    {
562                        return Err(Error::Corruption(format!(
563                            "binary coarse quantizer for field {field_id} does not match metadata/schema"
564                        )));
565                    }
566                    quantizer
567                        .validate_routing(schema_config.ivf_routing)
568                        .map_err(|error| {
569                            Error::Corruption(format!(
570                                "invalid binary centroid routing for field {field_id}: {error}"
571                            ))
572                        })?;
573                    binary_quantizers.insert(*field_id, Arc::new(quantizer));
574                }
575                unsupported => {
576                    return Err(Error::Corruption(format!(
577                        "field {field_id} is Built for {unsupported:?}, which has no global IVF artifacts"
578                    )));
579                }
580            }
581        }
582
583        if centroids.is_empty() && binary_quantizers.is_empty() {
584            Ok(None)
585        } else {
586            let trained = crate::segment::TrainedVectorStructures {
587                #[cfg(feature = "native")]
588                _ann_pins: Default::default(),
589                centroids,
590                binary_quantizers,
591            };
592            #[cfg(feature = "native")]
593            let trained = {
594                let mut trained = trained;
595                trained.pin_ann_structures(crate::segment::pin::pin_policy());
596                trained
597            };
598            Ok(Some(trained))
599        }
600    }
601}
602
603fn validate_trained_artifact_path(field_id: u32, kind: &str, filename: &str) -> Result<()> {
604    use std::path::Component;
605
606    let path = Path::new(filename);
607    if filename.is_empty()
608        || path.is_absolute()
609        || path.components().any(|component| {
610            matches!(
611                component,
612                Component::ParentDir | Component::RootDir | Component::Prefix(_)
613            )
614        })
615    {
616        return Err(Error::Corruption(format!(
617            "trained {kind} path for field {field_id} is not a safe relative path: '{filename}'"
618        )));
619    }
620    Ok(())
621}
622
623async fn load_trained_artifact<T, D>(
624    dir: &D,
625    field_id: u32,
626    kind: &str,
627    filename: &str,
628) -> Result<T>
629where
630    T: serde::de::DeserializeOwned,
631    D: crate::directories::Directory,
632{
633    validate_trained_artifact_path(field_id, kind, filename)?;
634    let path = Path::new(filename);
635    let file_size = dir.file_size(path).await.map_err(|error| {
636        Error::Corruption(format!(
637            "failed to stat trained {kind} '{filename}' for field {field_id}: {error}"
638        ))
639    })?;
640    validate_trained_artifact_size(field_id, kind, filename, file_size)?;
641    let slice = dir.open_read(path).await.map_err(|error| {
642        Error::Corruption(format!(
643            "failed to open trained {kind} '{filename}' for field {field_id}: {error}"
644        ))
645    })?;
646    validate_trained_artifact_size(field_id, kind, filename, slice.len())?;
647    let bytes = slice.read_bytes().await.map_err(|error| {
648        Error::Corruption(format!(
649            "failed to read trained {kind} '{filename}' for field {field_id}: {error}"
650        ))
651    })?;
652    let (artifact, consumed) = bincode::serde::decode_from_slice::<T, _>(
653        bytes.as_slice(),
654        bincode::config::standard().with_limit::<MAX_TRAINED_ARTIFACT_BYTES>(),
655    )
656    .map_err(|error| {
657        Error::Corruption(format!(
658            "failed to deserialize trained {kind} '{filename}' for field {field_id}: {error}"
659        ))
660    })?;
661    if consumed != bytes.len() {
662        return Err(Error::Corruption(format!(
663            "trained {kind} '{filename}' for field {field_id} has {} trailing bytes",
664            bytes.len() - consumed
665        )));
666    }
667    Ok(artifact)
668}
669
670fn validate_trained_artifact_size(
671    field_id: u32,
672    kind: &str,
673    filename: &str,
674    file_size: u64,
675) -> Result<()> {
676    if file_size > MAX_TRAINED_ARTIFACT_BYTES as u64 {
677        return Err(Error::Corruption(format!(
678            "trained {kind} '{filename}' for field {field_id} is {file_size} bytes, \
679             exceeding the {MAX_TRAINED_ARTIFACT_BYTES}-byte safety limit"
680        )));
681    }
682    Ok(())
683}
684
685#[cfg(test)]
686mod tests {
687    use super::*;
688    use crate::directories::DirectoryWriter;
689
690    #[derive(Clone, Default)]
691    struct SyncFailDirectory(crate::directories::RamDirectory);
692
693    #[async_trait::async_trait]
694    impl crate::directories::Directory for SyncFailDirectory {
695        async fn exists(&self, path: &Path) -> std::io::Result<bool> {
696            self.0.exists(path).await
697        }
698
699        async fn file_size(&self, path: &Path) -> std::io::Result<u64> {
700            self.0.file_size(path).await
701        }
702
703        async fn open_read(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
704            self.0.open_read(path).await
705        }
706
707        async fn read_range(
708            &self,
709            path: &Path,
710            range: std::ops::Range<u64>,
711        ) -> std::io::Result<crate::directories::OwnedBytes> {
712            self.0.read_range(path, range).await
713        }
714
715        async fn list_files(&self, prefix: &Path) -> std::io::Result<Vec<std::path::PathBuf>> {
716            self.0.list_files(prefix).await
717        }
718
719        async fn open_lazy(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
720            self.0.open_lazy(path).await
721        }
722    }
723
724    #[async_trait::async_trait]
725    impl crate::directories::DirectoryWriter for SyncFailDirectory {
726        async fn write(&self, path: &Path, data: &[u8]) -> std::io::Result<()> {
727            self.0.write(path, data).await
728        }
729
730        async fn delete(&self, path: &Path) -> std::io::Result<()> {
731            self.0.delete(path).await
732        }
733
734        async fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
735            self.0.rename(from, to).await
736        }
737
738        async fn sync(&self) -> std::io::Result<()> {
739            Err(std::io::Error::other("injected directory fsync failure"))
740        }
741
742        async fn streaming_writer(
743            &self,
744            path: &Path,
745        ) -> std::io::Result<Box<dyn crate::directories::StreamingWriter>> {
746            self.0.streaming_writer(path).await
747        }
748    }
749
750    fn test_schema() -> Schema {
751        Schema::default()
752    }
753
754    fn dense_schema(index_type: VectorIndexType) -> (Schema, crate::dsl::Field) {
755        let mut builder = crate::dsl::SchemaBuilder::default();
756        let config = match index_type {
757            VectorIndexType::IvfTq => crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1),
758            other => panic!("unsupported trained test index type: {other:?}"),
759        };
760        let field = builder.add_dense_vector_field_with_config("embedding", true, true, config);
761        (builder.build(), field)
762    }
763
764    fn test_centroids() -> crate::structures::CoarseCentroids {
765        crate::structures::CoarseCentroids {
766            num_clusters: 1,
767            dim: 2,
768            centroids: vec![0.25, 0.75],
769            version: 7,
770            soar_config: None,
771            routing_index: None,
772        }
773    }
774
775    async fn write_bincode(
776        directory: &crate::directories::RamDirectory,
777        filename: &str,
778        value: &impl serde::Serialize,
779    ) {
780        let bytes = bincode::serde::encode_to_vec(value, bincode::config::standard()).unwrap();
781        directory.write(Path::new(filename), &bytes).await.unwrap();
782    }
783
784    #[test]
785    fn test_metadata_init() {
786        let mut meta = IndexMetadata::new(test_schema());
787        assert_eq!(meta.total_vectors, 0);
788        assert!(meta.segment_metas.is_empty());
789        assert!(!meta.is_field_built(0));
790
791        meta.init_field(0, VectorIndexType::IvfTq);
792        assert!(!meta.is_field_built(0));
793        assert!(meta.vector_fields.contains_key(&0));
794    }
795
796    #[tokio::test]
797    async fn load_refuses_metadata_stamped_with_a_newer_format_version() {
798        let directory = crate::directories::RamDirectory::new();
799        let mut metadata = IndexMetadata::new(test_schema());
800        metadata.version = INDEX_META_FORMAT_VERSION + 1;
801        metadata.save(&directory).await.unwrap();
802
803        let error = IndexMetadata::load(&directory)
804            .await
805            .expect_err("metadata from a newer format version must be refused, not silently pruned")
806            .to_string();
807        assert!(error.contains("version 4"), "{error}");
808        assert!(error.contains("incompatible"), "{error}");
809    }
810
811    #[tokio::test]
812    async fn tmp_recovery_refuses_metadata_stamped_with_a_newer_format_version() {
813        let directory = crate::directories::RamDirectory::new();
814        let mut metadata = IndexMetadata::new(test_schema());
815        metadata.version = INDEX_META_FORMAT_VERSION + 1;
816        let bytes = metadata.serialize_to_bytes().unwrap();
817        // Simulate a crash between write and rename: only the temp file exists.
818        directory
819            .write(Path::new(INDEX_META_TMP_FILENAME), &bytes)
820            .await
821            .unwrap();
822
823        let error = IndexMetadata::load(&directory)
824            .await
825            .expect_err("temp-file recovery must apply the same version gate")
826            .to_string();
827        assert!(error.contains("version 4"), "{error}");
828    }
829
830    #[tokio::test]
831    async fn save_treats_post_rename_sync_failure_as_committed() {
832        let directory = SyncFailDirectory::default();
833        let mut metadata = IndexMetadata::new(test_schema());
834        metadata.add_segment("committed".to_string(), 7);
835
836        metadata.save(&directory).await.unwrap();
837
838        let loaded = IndexMetadata::load(&directory).await.unwrap();
839        assert_eq!(loaded.segment_doc_count("committed"), Some(7));
840    }
841
842    #[tokio::test]
843    async fn trained_artifacts_load_only_when_the_complete_built_set_is_valid() {
844        let mut builder = crate::dsl::SchemaBuilder::default();
845        let config = crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1);
846        let first = builder.add_dense_vector_field_with_config(
847            "first_embedding",
848            true,
849            true,
850            config.clone(),
851        );
852        let second =
853            builder.add_dense_vector_field_with_config("second_embedding", true, true, config);
854        let schema = builder.build();
855        let directory = crate::directories::RamDirectory::new();
856        let mut metadata = IndexMetadata::new(schema.clone());
857        metadata.init_field(first.0, VectorIndexType::IvfTq);
858        metadata.init_field(second.0, VectorIndexType::IvfTq);
859        metadata.mark_field_built(first.0, 10, 1, "field_0_centroids.bin".into(), None);
860        metadata.mark_field_built(second.0, 10, 1, "field_1_centroids.bin".into(), None);
861        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
862
863        let error = IndexMetadata::try_load_trained_from_fields(
864            &metadata.vector_fields,
865            &schema,
866            &directory,
867        )
868        .await
869        .err()
870        .expect("missing artifact must fail the complete load")
871        .to_string();
872        assert!(error.contains("field_1_centroids.bin"), "{error}");
873        assert!(error.contains("field 1"), "{error}");
874    }
875
876    #[tokio::test]
877    async fn index_open_fails_closed_when_built_artifact_is_missing() {
878        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
879        let directory = crate::directories::RamDirectory::new();
880        let mut metadata = IndexMetadata::new(schema);
881        metadata.init_field(field.0, VectorIndexType::IvfTq);
882        metadata.mark_field_built(field.0, 10, 1, "missing_centroids.bin".into(), None);
883        metadata.save(&directory).await.unwrap();
884
885        let error = match crate::index::Index::open(directory, crate::index::IndexConfig::default())
886            .await
887        {
888            Ok(_) => panic!("Index::open accepted a Built field with no artifact"),
889            Err(error) => error.to_string(),
890        };
891        assert!(error.contains("missing_centroids.bin"), "{error}");
892    }
893
894    #[tokio::test]
895    async fn ivf_tq_built_state_rejects_a_codebook_file() {
896        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
897        let directory = crate::directories::RamDirectory::new();
898        let mut metadata = IndexMetadata::new(schema.clone());
899        metadata.init_field(field.0, VectorIndexType::IvfTq);
900        metadata.mark_field_built(
901            field.0,
902            10,
903            1,
904            "field_0_centroids.bin".into(),
905            Some("field_0_codebook.bin".into()),
906        );
907        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
908
909        let error = IndexMetadata::try_load_trained_from_fields(
910            &metadata.vector_fields,
911            &schema,
912            &directory,
913        )
914        .await
915        .err()
916        .expect("IVF-TQ Built state with a codebook file must fail")
917        .to_string();
918        assert!(error.contains("codebook"), "{error}");
919    }
920
921    #[tokio::test]
922    async fn legacy_ivf_pq_trained_field_fails_with_actionable_error() {
923        // Simulates metadata written by a pre-removal version: the schema
924        // gate rejects `ivf_pq` fields, so build the raw field-state map
925        // directly against a current-format schema.
926        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
927        let directory = crate::directories::RamDirectory::new();
928        let mut metadata = IndexMetadata::new(schema.clone());
929        metadata.init_field(field.0, VectorIndexType::IvfTq);
930        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
931        // Overwrite the recorded type the way pre-removal metadata carries it
932        // (init_field never downgrades an existing entry).
933        metadata
934            .vector_fields
935            .get_mut(&field.0)
936            .expect("field initialized")
937            .index_type = VectorFieldIndexType::Float(VectorIndexType::IvfPq);
938        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
939
940        let error = IndexMetadata::try_load_trained_from_fields(
941            &metadata.vector_fields,
942            &schema,
943            &directory,
944        )
945        .await
946        .err()
947        .expect("legacy IVF-PQ trained state must fail loudly")
948        .to_string();
949        assert!(error.contains("no longer"), "{error}");
950        assert!(error.contains("ivf_tq"), "{error}");
951    }
952
953    #[tokio::test]
954    async fn requested_cluster_count_accepts_a_quality_clamped_artifact() {
955        let mut builder = crate::dsl::SchemaBuilder::default();
956        let field = builder.add_dense_vector_field_with_config(
957            "embedding",
958            true,
959            true,
960            crate::dsl::DenseVectorConfig::ivf_tq(2, Some(4), 1),
961        );
962        let schema = builder.build();
963        let directory = crate::directories::RamDirectory::new();
964        let mut metadata = IndexMetadata::new(schema.clone());
965        metadata.init_field(field.0, VectorIndexType::IvfTq);
966        metadata.mark_field_built(field.0, 1, 4, "field_0_centroids.bin".into(), None);
967        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
968
969        let trained = IndexMetadata::try_load_trained_from_fields(
970            &metadata.vector_fields,
971            &schema,
972            &directory,
973        )
974        .await
975        .unwrap()
976        .unwrap();
977        assert_eq!(trained.centroids[&field.0].num_clusters, 1);
978    }
979
980    #[tokio::test]
981    async fn trained_artifact_loader_rejects_trailing_data() {
982        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
983        let directory = crate::directories::RamDirectory::new();
984        let mut metadata = IndexMetadata::new(schema.clone());
985        metadata.init_field(field.0, VectorIndexType::IvfTq);
986        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
987        let mut bytes =
988            bincode::serde::encode_to_vec(test_centroids(), bincode::config::standard()).unwrap();
989        bytes.extend_from_slice(&[0xaa, 0xbb]);
990        directory
991            .write(Path::new("field_0_centroids.bin"), &bytes)
992            .await
993            .unwrap();
994
995        let error = IndexMetadata::try_load_trained_from_fields(
996            &metadata.vector_fields,
997            &schema,
998            &directory,
999        )
1000        .await
1001        .err()
1002        .expect("trailing artifact bytes must fail validation")
1003        .to_string();
1004        assert!(error.contains("trailing bytes"), "{error}");
1005    }
1006
1007    #[test]
1008    fn trained_artifact_size_limit_rejects_before_reading() {
1009        let error = validate_trained_artifact_size(
1010            3,
1011            "centroids",
1012            "field_3_centroids.bin",
1013            MAX_TRAINED_ARTIFACT_BYTES as u64 + 1,
1014        )
1015        .unwrap_err()
1016        .to_string();
1017        assert!(error.contains("exceeding"), "{error}");
1018        assert!(error.contains("field 3"), "{error}");
1019    }
1020
1021    #[tokio::test]
1022    async fn trained_artifact_decode_limit_rejects_forged_collection_length() {
1023        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1024        let directory = crate::directories::RamDirectory::new();
1025        let mut metadata = IndexMetadata::new(schema.clone());
1026        metadata.init_field(field.0, VectorIndexType::IvfTq);
1027        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1028
1029        // CoarseCentroids begins with num_clusters=1, dim=2, then the Vec
1030        // length. Bincode's standard varint marker 253 introduces a u64; this
1031        // tiny payload claims an impossible f32 vector and must hit the decode
1032        // limit before any large allocation is attempted.
1033        let mut bytes = vec![1, 2, 253];
1034        bytes.extend_from_slice(&u64::MAX.to_le_bytes());
1035        directory
1036            .write(Path::new("field_0_centroids.bin"), &bytes)
1037            .await
1038            .unwrap();
1039
1040        let error = IndexMetadata::try_load_trained_from_fields(
1041            &metadata.vector_fields,
1042            &schema,
1043            &directory,
1044        )
1045        .await
1046        .err()
1047        .expect("forged collection length must fail the bounded decoder")
1048        .to_string();
1049        assert!(error.contains("failed to deserialize"), "{error}");
1050    }
1051
1052    #[test]
1053    fn test_metadata_segments() {
1054        let mut meta = IndexMetadata::new(test_schema());
1055        meta.add_segment("abc123".to_string(), 50);
1056        meta.add_segment("def456".to_string(), 100);
1057        assert_eq!(meta.segment_metas.len(), 2);
1058        assert_eq!(meta.segment_doc_count("abc123"), Some(50));
1059        assert_eq!(meta.segment_doc_count("def456"), Some(100));
1060
1061        // Overwrites existing
1062        meta.add_segment("abc123".to_string(), 75);
1063        assert_eq!(meta.segment_metas.len(), 2);
1064        assert_eq!(meta.segment_doc_count("abc123"), Some(75));
1065
1066        meta.remove_segment("abc123");
1067        assert_eq!(meta.segment_metas.len(), 1);
1068        assert!(meta.has_segment("def456"));
1069        assert!(!meta.has_segment("abc123"));
1070    }
1071
1072    #[test]
1073    fn test_mark_field_built() {
1074        let mut meta = IndexMetadata::new(test_schema());
1075        meta.init_field(0, VectorIndexType::IvfTq);
1076        meta.total_vectors = 10000;
1077
1078        assert!(!meta.is_field_built(0));
1079
1080        meta.mark_field_built(0, 10000, 256, "field_0_centroids.bin".to_string(), None);
1081
1082        assert!(meta.is_field_built(0));
1083        let field = meta.get_field_meta(0).unwrap();
1084        assert_eq!(
1085            field.centroids_file.as_deref(),
1086            Some("field_0_centroids.bin")
1087        );
1088    }
1089
1090    #[test]
1091    fn total_vectors_is_aggregate_of_built_field_counts() {
1092        let mut meta = IndexMetadata::new(test_schema());
1093        meta.init_field(7, VectorIndexType::IvfTq);
1094        meta.init_field(3, VectorIndexType::IvfTq);
1095
1096        // Build in reverse field-id order to ensure the result is not tied to
1097        // HashMap or training iteration order.
1098        meta.mark_field_built(7, 400, 20, "field_7_centroids.bin".to_string(), None);
1099        assert_eq!(meta.total_vectors, 400);
1100        meta.mark_field_built(3, 250, 15, "field_3_centroids.bin".to_string(), None);
1101        assert_eq!(meta.total_vectors, 650);
1102
1103        // Rebuilding a field replaces its contribution; it does not add a
1104        // duplicate training snapshot.
1105        meta.mark_field_built(7, 425, 20, "field_7_centroids.bin".to_string(), None);
1106        assert_eq!(meta.total_vectors, 675);
1107    }
1108
1109    #[test]
1110    fn test_should_build_field() {
1111        let mut meta = IndexMetadata::new(test_schema());
1112        meta.init_field(0, VectorIndexType::IvfTq);
1113
1114        // Below threshold
1115        meta.total_vectors = 500;
1116        assert!(!meta.should_build_field(0, 1000));
1117
1118        // Above threshold
1119        meta.total_vectors = 1500;
1120        assert!(meta.should_build_field(0, 1000));
1121
1122        // Already built - should not build again
1123        meta.mark_field_built(0, 1500, 256, "centroids.bin".to_string(), None);
1124        assert!(!meta.should_build_field(0, 1000));
1125    }
1126
1127    #[test]
1128    fn test_serialization() {
1129        let mut meta = IndexMetadata::new(test_schema());
1130        meta.add_segment("seg1".to_string(), 100);
1131        meta.init_field(0, VectorIndexType::IvfTq);
1132        meta.total_vectors = 5000;
1133
1134        let json = serde_json::to_string_pretty(&meta).unwrap();
1135        let loaded: IndexMetadata = serde_json::from_str(&json).unwrap();
1136
1137        assert_eq!(loaded.segment_ids().len(), meta.segment_ids().len());
1138        assert_eq!(loaded.segment_doc_count("seg1"), Some(100));
1139        assert_eq!(loaded.total_vectors, meta.total_vectors);
1140        assert!(loaded.vector_fields.contains_key(&0));
1141    }
1142
1143    #[test]
1144    fn old_metadata_defaults_the_bp_retry_counter() {
1145        let mut meta = IndexMetadata::new(test_schema());
1146        meta.add_segment("legacy".to_string(), 10);
1147        let mut json = serde_json::to_value(&meta).unwrap();
1148        json["segment_metas"]["legacy"]
1149            .as_object_mut()
1150            .unwrap()
1151            .remove("bp_unconverged_passes");
1152
1153        let loaded: IndexMetadata = serde_json::from_value(json).unwrap();
1154        assert_eq!(loaded.segment_metas["legacy"].bp_unconverged_passes, 0);
1155    }
1156
1157    #[test]
1158    fn test_merged_segment_lineage() {
1159        let mut meta = IndexMetadata::new(test_schema());
1160        meta.add_segment("a".to_string(), 50);
1161        meta.add_segment("b".to_string(), 75);
1162
1163        // Fresh segments: gen=0, no ancestors
1164        assert_eq!(meta.segment_metas["a"].generation, 0);
1165        assert!(meta.segment_metas["a"].ancestors.is_empty());
1166
1167        // Merge a+b → c
1168        meta.add_merged_segment(
1169            "c".to_string(),
1170            125,
1171            vec!["a".to_string(), "b".to_string()],
1172            1,
1173            false,
1174            true,
1175        );
1176        assert_eq!(meta.segment_metas["c"].generation, 1);
1177        assert_eq!(meta.segment_metas["c"].ancestors, vec!["a", "b"]);
1178        assert_eq!(meta.segment_doc_count("c"), Some(125));
1179
1180        // Merge c+d → e (gen should be 2)
1181        meta.add_segment("d".to_string(), 30);
1182        meta.add_merged_segment(
1183            "e".to_string(),
1184            155,
1185            vec!["c".to_string(), "d".to_string()],
1186            2,
1187            false,
1188            true,
1189        );
1190        assert_eq!(meta.segment_metas["e"].generation, 2);
1191    }
1192}