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    /// ScaNN global model generation encoded into every segment payload.
125    /// Absent for legacy IVF/TQ fields and older metadata.
126    #[serde(default, skip_serializing_if = "Option::is_none")]
127    pub artifact_generation: Option<u64>,
128    /// Content fingerprint of the exact ScaNN model encoded into every
129    /// segment payload. This prevents same-generation accidental mixing.
130    #[serde(default, skip_serializing_if = "Option::is_none")]
131    pub artifact_id: Option<u64>,
132}
133
134/// Unified index metadata - single source of truth for index state
135#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct IndexMetadata {
137    /// Version for compatibility
138    pub version: u32,
139    /// Monotonic publication identity for schema/vector generations. Segment
140    /// commits are already detected by their IDs; this also makes query-only
141    /// ALTERs (for example `nprobe`) visible to cached readers.
142    #[serde(default)]
143    pub publication_generation: u64,
144    /// Index schema
145    pub schema: Schema,
146    /// Segment metadata: segment_id -> info (doc count, etc.)
147    /// Using HashMap allows O(1) lookup and stores doc counts for merge decisions
148    #[serde(default)]
149    pub segment_metas: HashMap<String, SegmentMetaInfo>,
150    /// Per-field vector index metadata
151    #[serde(default)]
152    pub vector_fields: HashMap<u32, FieldVectorMeta>,
153    /// Aggregate vector count recorded by all built vector fields.
154    ///
155    /// The per-field `VectorIndexState::Built::vector_count` values are the
156    /// source of truth. This cached aggregate is refreshed whenever a field is
157    /// marked built, rather than being overwritten with whichever field was
158    /// trained last.
159    #[serde(default)]
160    pub total_vectors: usize,
161}
162
163impl IndexMetadata {
164    /// Create new metadata with schema
165    pub fn new(schema: Schema) -> Self {
166        Self {
167            version: INDEX_META_FORMAT_VERSION,
168            publication_generation: 0,
169            schema,
170            segment_metas: HashMap::new(),
171            vector_fields: HashMap::new(),
172            total_vectors: 0,
173        }
174    }
175
176    /// Get segment IDs as a sorted Vec (deterministic ordering)
177    pub fn segment_ids(&self) -> Vec<String> {
178        let mut ids: Vec<String> = self.segment_metas.keys().cloned().collect();
179        ids.sort();
180        ids
181    }
182
183    /// Add a fresh segment (gen=0, no ancestors, not reordered)
184    pub fn add_segment(&mut self, segment_id: String, num_docs: u32) {
185        self.segment_metas.insert(
186            segment_id,
187            SegmentMetaInfo {
188                num_docs,
189                ancestors: Vec::new(),
190                generation: 0,
191                reordered: false,
192                bp_converged: true,
193                bp_unconverged_passes: 0,
194            },
195        );
196    }
197
198    /// Add a merged segment with lineage info
199    pub fn add_merged_segment(
200        &mut self,
201        segment_id: String,
202        num_docs: u32,
203        ancestors: Vec<String>,
204        generation: u32,
205        reordered: bool,
206        bp_converged: bool,
207    ) {
208        self.add_segment_meta(
209            segment_id,
210            SegmentMetaInfo {
211                num_docs,
212                ancestors,
213                generation,
214                reordered,
215                bp_converged,
216                bp_unconverged_passes: 0,
217            },
218        );
219    }
220
221    /// Insert fully constructed lifecycle metadata. Merge/reorder code uses
222    /// this to carry bounded BP lineage; ordinary callers use the safer
223    /// constructors above, which start a fresh lineage.
224    pub(crate) fn add_segment_meta(&mut self, segment_id: String, info: SegmentMetaInfo) {
225        self.segment_metas.insert(segment_id, info);
226    }
227
228    /// Remove a segment
229    pub fn remove_segment(&mut self, segment_id: &str) {
230        self.segment_metas.remove(segment_id);
231    }
232
233    /// Check if segment exists
234    pub fn has_segment(&self, segment_id: &str) -> bool {
235        self.segment_metas.contains_key(segment_id)
236    }
237
238    /// Get segment doc count
239    pub fn segment_doc_count(&self, segment_id: &str) -> Option<u32> {
240        self.segment_metas.get(segment_id).map(|m| m.num_docs)
241    }
242
243    /// Check if a field has been built
244    pub fn is_field_built(&self, field_id: u32) -> bool {
245        self.vector_fields
246            .get(&field_id)
247            .map(|f| matches!(f.state, VectorIndexState::Built { .. }))
248            .unwrap_or(false)
249    }
250
251    /// Get field metadata
252    pub fn get_field_meta(&self, field_id: u32) -> Option<&FieldVectorMeta> {
253        self.vector_fields.get(&field_id)
254    }
255
256    /// Initialize field metadata (called when field is first seen)
257    pub fn init_field(&mut self, field_id: u32, index_type: impl Into<VectorFieldIndexType>) {
258        let index_type = index_type.into();
259        self.vector_fields
260            .entry(field_id)
261            .or_insert(FieldVectorMeta {
262                field_id,
263                index_type,
264                state: VectorIndexState::Flat,
265                centroids_file: None,
266                codebook_file: None,
267                artifact_generation: None,
268                artifact_id: None,
269            });
270    }
271
272    /// Mark field as built with trained structures
273    pub fn mark_field_built(
274        &mut self,
275        field_id: u32,
276        vector_count: usize,
277        num_clusters: usize,
278        centroids_file: String,
279        codebook_file: Option<String>,
280    ) {
281        if let Some(field) = self.vector_fields.get_mut(&field_id) {
282            field.state = VectorIndexState::Built {
283                vector_count,
284                num_clusters,
285            };
286            field.centroids_file = Some(centroids_file);
287            field.codebook_file = codebook_file;
288            field.artifact_generation = None;
289            field.artifact_id = None;
290            self.refresh_total_vectors();
291        }
292    }
293
294    /// Mark a ScaNN field built against one immutable global model. Reuses
295    /// `centroids_file` as the trained-artifact path for wire compatibility;
296    /// the explicit generation/fingerprint fields make its semantics loud.
297    pub fn mark_scann_field_built(
298        &mut self,
299        field_id: u32,
300        vector_count: usize,
301        num_leaves: usize,
302        artifact_file: String,
303        artifact_generation: u64,
304        artifact_id: u64,
305    ) -> Result<()> {
306        if artifact_generation == 0 || artifact_id == 0 {
307            return Err(Error::Corruption(format!(
308                "ScaNN field {field_id} cannot publish a zero generation or artifact fingerprint"
309            )));
310        }
311        let field = self.vector_fields.get_mut(&field_id).ok_or_else(|| {
312            Error::Corruption(format!(
313                "ScaNN field {field_id} must be initialized before it is marked built"
314            ))
315        })?;
316        if !matches!(
317            field.index_type,
318            VectorFieldIndexType::Float(VectorIndexType::Scann)
319                | VectorFieldIndexType::Binary(BinaryIndexType::Scann)
320        ) {
321            return Err(Error::Corruption(format!(
322                "field {field_id} is not configured as ScaNN"
323            )));
324        }
325        field.state = VectorIndexState::Built {
326            vector_count,
327            num_clusters: num_leaves,
328        };
329        field.centroids_file = Some(artifact_file);
330        field.codebook_file = None;
331        field.artifact_generation = Some(artifact_generation);
332        field.artifact_id = Some(artifact_id);
333        self.refresh_total_vectors();
334        Ok(())
335    }
336
337    /// Refresh the cached aggregate from the authoritative per-field states.
338    ///
339    /// Saturation keeps this infallible metadata helper safe even if it is
340    /// called after loading externally modified metadata with impossible
341    /// counts.
342    pub(crate) fn refresh_total_vectors(&mut self) {
343        self.total_vectors = self
344            .vector_fields
345            .values()
346            .filter_map(|field| match field.state {
347                VectorIndexState::Built { vector_count, .. } => Some(vector_count),
348                VectorIndexState::Flat => None,
349            })
350            .fold(0usize, usize::saturating_add);
351    }
352
353    /// Check if field should be built based on threshold
354    pub fn should_build_field(&self, field_id: u32, threshold: usize) -> bool {
355        // Don't build if already built
356        if self.is_field_built(field_id) {
357            return false;
358        }
359        // Build if we have enough vectors
360        self.total_vectors >= threshold
361    }
362
363    /// Load from directory
364    ///
365    /// If `metadata.json` is missing but `metadata.json.tmp` exists (crash
366    /// between write and rename), recovers from the temp file.
367    pub async fn load<D: crate::directories::Directory>(dir: &D) -> Result<Self> {
368        let path = Path::new(INDEX_META_FILENAME);
369        match dir.open_read(path).await {
370            Ok(slice) => {
371                let bytes = slice.read_bytes().await?;
372                Self::deserialize_versioned(bytes.as_slice())
373            }
374            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
375                // Try recovering from temp file (crash between write and rename)
376                let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
377                let slice = dir.open_read(tmp_path).await?;
378                let bytes = slice.read_bytes().await?;
379                let meta = Self::deserialize_versioned(bytes.as_slice())?;
380                log::warn!("Recovered metadata from temp file (previous crash during save)");
381                Ok(meta)
382            }
383            Err(e) => Err(Error::Io(e)),
384        }
385    }
386
387    /// Deserialize only the current format. Vector artifacts are intentionally
388    /// rebuilt when the ANN format changes; silently accepting older metadata
389    /// would mix incompatible segment and global-codebook generations.
390    fn deserialize_versioned(bytes: &[u8]) -> Result<Self> {
391        let meta: Self =
392            serde_json::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))?;
393        crate::dsl::reject_removed_vector_index_types(&meta.schema).map_err(Error::Schema)?;
394        if meta.version != INDEX_META_FORMAT_VERSION {
395            return Err(Error::Corruption(format!(
396                "metadata.json format version {} is incompatible with required version {}; \
397                 rebuild and republish the index with this Hermes version",
398                meta.version, INDEX_META_FORMAT_VERSION
399            )));
400        }
401        Ok(meta)
402    }
403
404    /// Save to directory (atomic: write temp file, then rename)
405    ///
406    /// Uses write-then-rename so a crash mid-write won't corrupt the
407    /// existing metadata file. On POSIX, rename is atomic.
408    pub async fn save<D: crate::directories::DirectoryWriter>(&self, dir: &D) -> Result<()> {
409        let bytes = self.serialize_to_bytes()?;
410        Self::save_bytes(dir, &bytes).await
411    }
412
413    /// Serialize metadata to bytes (cheap, no I/O).
414    /// Useful when you need to release a lock before doing disk I/O.
415    pub fn serialize_to_bytes(&self) -> Result<Vec<u8>> {
416        serde_json::to_vec_pretty(self).map_err(|e| Error::Serialization(e.to_string()))
417    }
418
419    /// Write pre-serialized metadata bytes to directory (atomic rename + fsync).
420    ///
421    /// The fsync ensures durability: without it, a power failure after rename
422    /// could lose the metadata update on systems with volatile write caches.
423    pub async fn save_bytes<D: crate::directories::DirectoryWriter>(
424        dir: &D,
425        bytes: &[u8],
426    ) -> Result<()> {
427        let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
428        let final_path = Path::new(INDEX_META_FILENAME);
429        // Metadata is tiny, but `DirectoryWriter::write` does not guarantee
430        // the file contents themselves are fsynced. Finish the streaming
431        // writer first (filesystem implementations call `File::sync_all`),
432        // then atomically publish the durable temp file by rename.
433        let mut writer = dir.streaming_writer(tmp_path).await.map_err(Error::Io)?;
434        writer.write_all(bytes).map_err(Error::Io)?;
435        writer.finish().map_err(Error::Io)?;
436        // Rename is the logical commit point: after it succeeds, readers can
437        // observe the new generation and callers must publish the matching
438        // in-memory/tracker state. Directory fsync only strengthens crash
439        // durability. It cannot safely turn an already-visible rename into a
440        // reported pre-commit failure, because cleanup could then delete files
441        // referenced by the metadata now on disk.
442        dir.rename(tmp_path, final_path).await.map_err(Error::Io)?;
443        if let Err(error) = dir.sync().await {
444            log::error!(
445                "[metadata] directory fsync failed after committed rename: {}. \
446                 Continuing with the renamed generation; crash durability is not guaranteed",
447                error,
448            );
449        }
450        Ok(())
451    }
452
453    /// Fallible schema-aware loader used for lifecycle publication.
454    #[cfg_attr(not(feature = "native"), allow(dead_code))]
455    pub(crate) async fn try_load_trained_from_fields<D: crate::directories::Directory>(
456        vector_fields: &HashMap<u32, FieldVectorMeta>,
457        schema: &Schema,
458        dir: &D,
459    ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
460        Self::load_trained_from_fields_impl(vector_fields, schema, dir).await
461    }
462
463    /// Load and validate the complete trained-artifact set described by a
464    /// `vector_fields` snapshot.
465    ///
466    /// This is intentionally all-or-nothing. A `Built` field is a durable
467    /// promise that every artifact required by its configured index exists and
468    /// is compatible with the schema. Returning a partial map would let some
469    /// segment builders publish ANN data while another field was silently
470    /// unusable, and would make the same index behave differently after a
471    /// restart.
472    async fn load_trained_from_fields_impl<D: crate::directories::Directory>(
473        vector_fields: &HashMap<u32, FieldVectorMeta>,
474        schema: &Schema,
475        dir: &D,
476    ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
477        use std::sync::Arc;
478
479        let mut centroids = rustc_hash::FxHashMap::default();
480        let mut binary_quantizers = rustc_hash::FxHashMap::default();
481        let mut scann_artifacts = rustc_hash::FxHashMap::default();
482
483        let mut built_fields: Vec<_> = vector_fields
484            .iter()
485            .filter(|(_, meta)| matches!(meta.state, VectorIndexState::Built { .. }))
486            .collect();
487        built_fields.sort_unstable_by_key(|(field_id, _)| **field_id);
488
489        log::debug!(
490            "[trained] index={} loading trained structures, dense_vector_fields={:?}",
491            schema.index_label(),
492            vector_fields.keys().collect::<Vec<_>>()
493        );
494
495        for (field_id, field_meta) in built_fields {
496            log::debug!(
497                "[trained] index={} field {} state={:?} centroids_file={:?} codebook_file={:?}",
498                schema.index_label(),
499                field_id,
500                field_meta.state,
501                field_meta.centroids_file,
502                field_meta.codebook_file,
503            );
504            if field_meta.field_id != *field_id {
505                return Err(Error::Corruption(format!(
506                    "trained vector metadata key {field_id} contains field_id {}",
507                    field_meta.field_id
508                )));
509            }
510
511            let expected_clusters = match field_meta.state {
512                VectorIndexState::Built { num_clusters, .. } if num_clusters > 0 => num_clusters,
513                VectorIndexState::Built { .. } => {
514                    return Err(Error::Corruption(format!(
515                        "trained vector metadata field {field_id} has zero clusters"
516                    )));
517                }
518                VectorIndexState::Flat => unreachable!("built_fields contains only Built entries"),
519            };
520
521            let centroids_file = field_meta.centroids_file.as_deref().ok_or_else(|| {
522                Error::Corruption(format!(
523                    "trained vector metadata field {field_id} is Built but has no centroids_file"
524                ))
525            })?;
526            match field_meta.index_type {
527                VectorFieldIndexType::Float(VectorIndexType::IvfPq) => {
528                    return Err(Error::Corruption(format!(
529                        "field {field_id} was trained as IVF-PQ, which is no longer \
530                         supported; recreate the index with `ivf_tq` and reindex \
531                         (docs/turboquant-quantization.md)"
532                    )));
533                }
534                VectorFieldIndexType::Float(index_type @ VectorIndexType::IvfTq) => {
535                    let entry = schema
536                        .get_field_entry(crate::dsl::Field(*field_id))
537                        .ok_or_else(|| {
538                            Error::Corruption(format!(
539                                "trained vector metadata references missing field {field_id}"
540                            ))
541                        })?;
542                    let schema_config = entry
543                        .dense_vector_config
544                        .as_ref()
545                        .filter(|_| entry.field_type == crate::dsl::FieldType::DenseVector)
546                        .ok_or_else(|| {
547                            Error::Corruption(format!(
548                                "trained vector metadata field {field_id} is not a float dense field"
549                            ))
550                        })?;
551                    if schema_config.index_type != index_type {
552                        return Err(Error::Corruption(format!(
553                            "trained vector metadata field {field_id} uses {index_type:?}, schema requires {:?}",
554                            schema_config.index_type
555                        )));
556                    }
557                    let c: crate::structures::CoarseCentroids =
558                        load_trained_artifact(dir, *field_id, "centroids", centroids_file).await?;
559                    let expected_dim = schema_config.dim;
560                    let actual_clusters = c.num_clusters as usize;
561                    let expected_values =
562                        actual_clusters.checked_mul(expected_dim).ok_or_else(|| {
563                            Error::Corruption(format!(
564                                "trained centroid dimensions overflow for field {field_id}"
565                            ))
566                        })?;
567                    if actual_clusters == 0
568                        || actual_clusters > expected_clusters
569                        || c.dim == 0
570                        || c.dim != expected_dim
571                        || c.centroids.len() != expected_values
572                        || c.centroids.iter().any(|value| !value.is_finite())
573                    {
574                        return Err(Error::Corruption(format!(
575                            "trained centroids for field {field_id} do not match metadata/schema"
576                        )));
577                    }
578                    if !crate::structures::is_ivf_tq_cosine_generation(c.version) {
579                        return Err(Error::Corruption(format!(
580                            "trained IVF-TQ centroids for field {field_id} use an \
581                             unsupported legacy generation; rebuild the index"
582                        )));
583                    }
584                    c.validate_routing(schema_config.ivf_routing)
585                        .map_err(|error| {
586                            Error::Corruption(format!(
587                                "invalid trained centroid routing for field {field_id}: {error}"
588                            ))
589                        })?;
590                    // The TQ leaf codec is derived, never trained; ensure
591                    // `index_type` stays referenced for future variants.
592                    let _ = index_type;
593                    if field_meta.codebook_file.is_some() {
594                        return Err(Error::Corruption(format!(
595                            "trained IVF-TQ field {field_id} unexpectedly references a codebook file"
596                        )));
597                    }
598                    centroids.insert(*field_id, Arc::new(c));
599                }
600                VectorFieldIndexType::Binary(BinaryIndexType::Ivf) => {
601                    let entry = schema
602                        .get_field_entry(crate::dsl::Field(*field_id))
603                        .ok_or_else(|| {
604                            Error::Corruption(format!(
605                                "trained vector metadata references missing field {field_id}"
606                            ))
607                        })?;
608                    let schema_config = entry
609                        .binary_dense_vector_config
610                        .as_ref()
611                        .filter(|config| {
612                            entry.field_type == crate::dsl::FieldType::BinaryDenseVector
613                                && config.index_type == BinaryIndexType::Ivf
614                        })
615                        .ok_or_else(|| {
616                            Error::Corruption(format!(
617                                "trained vector metadata field {field_id} is not a binary IVF field"
618                            ))
619                        })?;
620                    let quantizer: crate::structures::BinaryCoarseQuantizer =
621                        load_trained_artifact(dir, *field_id, "binary centroids", centroids_file)
622                            .await?;
623                    quantizer.validate().map_err(|error| {
624                        Error::Corruption(format!(
625                            "invalid binary coarse quantizer for field {field_id}: {error}"
626                        ))
627                    })?;
628                    let actual_clusters = quantizer.num_clusters as usize;
629                    if actual_clusters > expected_clusters
630                        || schema_config.dim != quantizer.dim_bits
631                    {
632                        return Err(Error::Corruption(format!(
633                            "binary coarse quantizer for field {field_id} does not match metadata/schema"
634                        )));
635                    }
636                    quantizer
637                        .validate_routing(schema_config.ivf_routing)
638                        .map_err(|error| {
639                            Error::Corruption(format!(
640                                "invalid binary centroid routing for field {field_id}: {error}"
641                            ))
642                        })?;
643                    binary_quantizers.insert(*field_id, Arc::new(quantizer));
644                }
645                VectorFieldIndexType::Float(VectorIndexType::Scann)
646                | VectorFieldIndexType::Binary(BinaryIndexType::Scann) => {
647                    if field_meta.codebook_file.is_some() {
648                        return Err(Error::Corruption(format!(
649                            "trained ScaNN field {field_id} unexpectedly references a separate codebook file"
650                        )));
651                    }
652                    let expected_generation = field_meta.artifact_generation.ok_or_else(|| {
653                        Error::Corruption(format!(
654                            "trained ScaNN field {field_id} has no artifact generation"
655                        ))
656                    })?;
657                    let expected_artifact_id = field_meta.artifact_id.ok_or_else(|| {
658                        Error::Corruption(format!(
659                            "trained ScaNN field {field_id} has no artifact fingerprint"
660                        ))
661                    })?;
662                    validate_trained_artifact_path(
663                        field_id.to_owned(),
664                        "ScaNN artifact",
665                        centroids_file,
666                    )?;
667                    let path = Path::new(centroids_file);
668                    let slice = dir.open_read(path).await.map_err(|error| {
669                        Error::Corruption(format!(
670                            "failed to open trained ScaNN artifact '{centroids_file}' for field {field_id}: {error}"
671                        ))
672                    })?;
673                    let raw = slice.read_bytes().await.map_err(|error| {
674                        Error::Corruption(format!(
675                            "failed to map trained ScaNN artifact '{centroids_file}' for field {field_id}: {error}"
676                        ))
677                    })?;
678                    let artifact =
679                        crate::segment::ScannTrainedArtifactBytes::open(raw).map_err(|error| {
680                            Error::Corruption(format!(
681                                "invalid trained ScaNN artifact for field {field_id}: {error}"
682                            ))
683                        })?;
684                    if artifact.generation() != expected_generation
685                        || artifact.artifact_id() != expected_artifact_id
686                        || artifact.config().num_leaves as usize != expected_clusters
687                    {
688                        return Err(Error::Corruption(format!(
689                            "trained ScaNN artifact for field {field_id} does not match metadata"
690                        )));
691                    }
692                    let entry = schema
693                        .get_field_entry(crate::dsl::Field(*field_id))
694                        .ok_or_else(|| {
695                            Error::Corruption(format!(
696                                "trained vector metadata references missing field {field_id}"
697                            ))
698                        })?;
699                    let schema_matches = match field_meta.index_type {
700                        VectorFieldIndexType::Float(VectorIndexType::Scann) => entry
701                            .dense_vector_config
702                            .as_ref()
703                            .filter(|_| entry.field_type == crate::dsl::FieldType::DenseVector)
704                            .is_some_and(|config| {
705                                config.index_type == VectorIndexType::Scann
706                                    && config.dim == artifact.config().dimension as usize
707                                    && scann_explicit_geometry_matches(
708                                        config.num_clusters,
709                                        config.tree_levels,
710                                        artifact.config().num_leaves as usize,
711                                        artifact.config().tree_levels,
712                                    )
713                                    && matches!(
714                                        artifact.config().encoding,
715                                        crate::structures::vector::scann::ScannEncoding::AsymmetricHash { .. }
716                                    )
717                            }),
718                        VectorFieldIndexType::Binary(BinaryIndexType::Scann) => entry
719                            .binary_dense_vector_config
720                            .as_ref()
721                            .filter(|_| {
722                                entry.field_type == crate::dsl::FieldType::BinaryDenseVector
723                            })
724                            .is_some_and(|config| {
725                                config.index_type == BinaryIndexType::Scann
726                                    && config.dim == artifact.config().dimension as usize
727                                    && scann_explicit_geometry_matches(
728                                        config.num_clusters,
729                                        config.tree_levels,
730                                        artifact.config().num_leaves as usize,
731                                        artifact.config().tree_levels,
732                                    )
733                                    && artifact.config().encoding
734                                        == crate::structures::vector::scann::ScannEncoding::BinaryHamming
735                            }),
736                        _ => false,
737                    };
738                    if !schema_matches {
739                        return Err(Error::Corruption(format!(
740                            "trained ScaNN artifact for field {field_id} does not match schema geometry/encoding"
741                        )));
742                    }
743                    scann_artifacts.insert(*field_id, Arc::new(artifact));
744                }
745                unsupported => {
746                    return Err(Error::Corruption(format!(
747                        "field {field_id} is Built for {unsupported:?}, which has no global IVF artifacts"
748                    )));
749                }
750            }
751        }
752
753        if centroids.is_empty() && binary_quantizers.is_empty() && scann_artifacts.is_empty() {
754            Ok(None)
755        } else {
756            let trained = crate::segment::TrainedVectorStructures {
757                #[cfg(feature = "native")]
758                _ann_pins: Default::default(),
759                centroids,
760                binary_quantizers,
761                scann_artifacts,
762            };
763            #[cfg(feature = "native")]
764            let trained = {
765                let mut trained = trained;
766                trained.pin_ann_structures(crate::segment::pin::pin_policy());
767                trained
768            };
769            Ok(Some(trained))
770        }
771    }
772}
773
774/// Omitted ScaNN geometry is autopilot: the trained artifact's resolved
775/// values are authoritative and are also pinned by `VectorIndexState::Built`.
776/// Explicit schema values remain strict compatibility requirements.
777fn scann_explicit_geometry_matches(
778    configured_leaves: Option<usize>,
779    configured_levels: Option<u8>,
780    resolved_leaves: usize,
781    resolved_levels: u8,
782) -> bool {
783    configured_leaves.is_none_or(|leaves| leaves == resolved_leaves)
784        && configured_levels.is_none_or(|levels| levels == resolved_levels)
785}
786
787fn validate_trained_artifact_path(field_id: u32, kind: &str, filename: &str) -> Result<()> {
788    use std::path::Component;
789
790    let path = Path::new(filename);
791    if filename.is_empty()
792        || path.is_absolute()
793        || path.components().any(|component| {
794            matches!(
795                component,
796                Component::ParentDir | Component::RootDir | Component::Prefix(_)
797            )
798        })
799    {
800        return Err(Error::Corruption(format!(
801            "trained {kind} path for field {field_id} is not a safe relative path: '{filename}'"
802        )));
803    }
804    Ok(())
805}
806
807async fn load_trained_artifact<T, D>(
808    dir: &D,
809    field_id: u32,
810    kind: &str,
811    filename: &str,
812) -> Result<T>
813where
814    T: serde::de::DeserializeOwned,
815    D: crate::directories::Directory,
816{
817    validate_trained_artifact_path(field_id, kind, filename)?;
818    let path = Path::new(filename);
819    let file_size = dir.file_size(path).await.map_err(|error| {
820        Error::Corruption(format!(
821            "failed to stat trained {kind} '{filename}' for field {field_id}: {error}"
822        ))
823    })?;
824    validate_trained_artifact_size(field_id, kind, filename, file_size)?;
825    let slice = dir.open_read(path).await.map_err(|error| {
826        Error::Corruption(format!(
827            "failed to open trained {kind} '{filename}' for field {field_id}: {error}"
828        ))
829    })?;
830    validate_trained_artifact_size(field_id, kind, filename, slice.len())?;
831    let bytes = slice.read_bytes().await.map_err(|error| {
832        Error::Corruption(format!(
833            "failed to read trained {kind} '{filename}' for field {field_id}: {error}"
834        ))
835    })?;
836    let (artifact, consumed) = bincode::serde::decode_from_slice::<T, _>(
837        bytes.as_slice(),
838        bincode::config::standard().with_limit::<MAX_TRAINED_ARTIFACT_BYTES>(),
839    )
840    .map_err(|error| {
841        Error::Corruption(format!(
842            "failed to deserialize trained {kind} '{filename}' for field {field_id}: {error}"
843        ))
844    })?;
845    if consumed != bytes.len() {
846        return Err(Error::Corruption(format!(
847            "trained {kind} '{filename}' for field {field_id} has {} trailing bytes",
848            bytes.len() - consumed
849        )));
850    }
851    Ok(artifact)
852}
853
854fn validate_trained_artifact_size(
855    field_id: u32,
856    kind: &str,
857    filename: &str,
858    file_size: u64,
859) -> Result<()> {
860    if file_size > MAX_TRAINED_ARTIFACT_BYTES as u64 {
861        return Err(Error::Corruption(format!(
862            "trained {kind} '{filename}' for field {field_id} is {file_size} bytes, \
863             exceeding the {MAX_TRAINED_ARTIFACT_BYTES}-byte safety limit"
864        )));
865    }
866    Ok(())
867}
868
869#[cfg(test)]
870mod tests {
871    use super::*;
872    use crate::directories::DirectoryWriter;
873
874    #[derive(Clone, Default)]
875    struct SyncFailDirectory(crate::directories::RamDirectory);
876
877    #[async_trait::async_trait]
878    impl crate::directories::Directory for SyncFailDirectory {
879        async fn exists(&self, path: &Path) -> std::io::Result<bool> {
880            self.0.exists(path).await
881        }
882
883        async fn file_size(&self, path: &Path) -> std::io::Result<u64> {
884            self.0.file_size(path).await
885        }
886
887        async fn open_read(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
888            self.0.open_read(path).await
889        }
890
891        async fn read_range(
892            &self,
893            path: &Path,
894            range: std::ops::Range<u64>,
895        ) -> std::io::Result<crate::directories::OwnedBytes> {
896            self.0.read_range(path, range).await
897        }
898
899        async fn list_files(&self, prefix: &Path) -> std::io::Result<Vec<std::path::PathBuf>> {
900            self.0.list_files(prefix).await
901        }
902
903        async fn open_lazy(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
904            self.0.open_lazy(path).await
905        }
906    }
907
908    #[async_trait::async_trait]
909    impl crate::directories::DirectoryWriter for SyncFailDirectory {
910        async fn write(&self, path: &Path, data: &[u8]) -> std::io::Result<()> {
911            self.0.write(path, data).await
912        }
913
914        async fn delete(&self, path: &Path) -> std::io::Result<()> {
915            self.0.delete(path).await
916        }
917
918        async fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
919            self.0.rename(from, to).await
920        }
921
922        async fn sync(&self) -> std::io::Result<()> {
923            Err(std::io::Error::other("injected directory fsync failure"))
924        }
925
926        async fn streaming_writer(
927            &self,
928            path: &Path,
929        ) -> std::io::Result<Box<dyn crate::directories::StreamingWriter>> {
930            self.0.streaming_writer(path).await
931        }
932    }
933
934    fn test_schema() -> Schema {
935        Schema::default()
936    }
937
938    fn dense_schema(index_type: VectorIndexType) -> (Schema, crate::dsl::Field) {
939        let mut builder = crate::dsl::SchemaBuilder::default();
940        let config = match index_type {
941            VectorIndexType::IvfTq => crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1),
942            other => panic!("unsupported trained test index type: {other:?}"),
943        };
944        let field = builder.add_dense_vector_field_with_config("embedding", true, true, config);
945        (builder.build(), field)
946    }
947
948    fn test_centroids() -> crate::structures::CoarseCentroids {
949        crate::structures::CoarseCentroids {
950            num_clusters: 1,
951            dim: 2,
952            centroids: vec![0.25, 0.75],
953            version: crate::structures::mark_ivf_tq_cosine_generation(7),
954            soar_config: None,
955            routing_index: None,
956        }
957    }
958
959    async fn write_bincode(
960        directory: &crate::directories::RamDirectory,
961        filename: &str,
962        value: &impl serde::Serialize,
963    ) {
964        let bytes = bincode::serde::encode_to_vec(value, bincode::config::standard()).unwrap();
965        directory.write(Path::new(filename), &bytes).await.unwrap();
966    }
967
968    #[test]
969    fn test_metadata_init() {
970        let mut meta = IndexMetadata::new(test_schema());
971        assert_eq!(meta.total_vectors, 0);
972        assert!(meta.segment_metas.is_empty());
973        assert!(!meta.is_field_built(0));
974
975        meta.init_field(0, VectorIndexType::IvfTq);
976        assert!(!meta.is_field_built(0));
977        assert!(meta.vector_fields.contains_key(&0));
978    }
979
980    #[tokio::test]
981    async fn load_refuses_metadata_stamped_with_a_newer_format_version() {
982        let directory = crate::directories::RamDirectory::new();
983        let mut metadata = IndexMetadata::new(test_schema());
984        metadata.version = INDEX_META_FORMAT_VERSION + 1;
985        metadata.save(&directory).await.unwrap();
986
987        let error = IndexMetadata::load(&directory)
988            .await
989            .expect_err("metadata from a newer format version must be refused, not silently pruned")
990            .to_string();
991        assert!(error.contains("version 4"), "{error}");
992        assert!(error.contains("incompatible"), "{error}");
993    }
994
995    #[tokio::test]
996    async fn tmp_recovery_refuses_metadata_stamped_with_a_newer_format_version() {
997        let directory = crate::directories::RamDirectory::new();
998        let mut metadata = IndexMetadata::new(test_schema());
999        metadata.version = INDEX_META_FORMAT_VERSION + 1;
1000        let bytes = metadata.serialize_to_bytes().unwrap();
1001        // Simulate a crash between write and rename: only the temp file exists.
1002        directory
1003            .write(Path::new(INDEX_META_TMP_FILENAME), &bytes)
1004            .await
1005            .unwrap();
1006
1007        let error = IndexMetadata::load(&directory)
1008            .await
1009            .expect_err("temp-file recovery must apply the same version gate")
1010            .to_string();
1011        assert!(error.contains("version 4"), "{error}");
1012    }
1013
1014    #[tokio::test]
1015    async fn save_treats_post_rename_sync_failure_as_committed() {
1016        let directory = SyncFailDirectory::default();
1017        let mut metadata = IndexMetadata::new(test_schema());
1018        metadata.add_segment("committed".to_string(), 7);
1019
1020        metadata.save(&directory).await.unwrap();
1021
1022        let loaded = IndexMetadata::load(&directory).await.unwrap();
1023        assert_eq!(loaded.segment_doc_count("committed"), Some(7));
1024    }
1025
1026    #[tokio::test]
1027    async fn trained_artifacts_load_only_when_the_complete_built_set_is_valid() {
1028        let mut builder = crate::dsl::SchemaBuilder::default();
1029        let config = crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1);
1030        let first = builder.add_dense_vector_field_with_config(
1031            "first_embedding",
1032            true,
1033            true,
1034            config.clone(),
1035        );
1036        let second =
1037            builder.add_dense_vector_field_with_config("second_embedding", true, true, config);
1038        let schema = builder.build();
1039        let directory = crate::directories::RamDirectory::new();
1040        let mut metadata = IndexMetadata::new(schema.clone());
1041        metadata.init_field(first.0, VectorIndexType::IvfTq);
1042        metadata.init_field(second.0, VectorIndexType::IvfTq);
1043        metadata.mark_field_built(first.0, 10, 1, "field_0_centroids.bin".into(), None);
1044        metadata.mark_field_built(second.0, 10, 1, "field_1_centroids.bin".into(), None);
1045        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1046
1047        let error = IndexMetadata::try_load_trained_from_fields(
1048            &metadata.vector_fields,
1049            &schema,
1050            &directory,
1051        )
1052        .await
1053        .err()
1054        .expect("missing artifact must fail the complete load")
1055        .to_string();
1056        assert!(error.contains("field_1_centroids.bin"), "{error}");
1057        assert!(error.contains("field 1"), "{error}");
1058    }
1059
1060    #[tokio::test]
1061    async fn index_open_fails_closed_when_built_artifact_is_missing() {
1062        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1063        let directory = crate::directories::RamDirectory::new();
1064        let mut metadata = IndexMetadata::new(schema);
1065        metadata.init_field(field.0, VectorIndexType::IvfTq);
1066        metadata.mark_field_built(field.0, 10, 1, "missing_centroids.bin".into(), None);
1067        metadata.save(&directory).await.unwrap();
1068
1069        let error = match crate::index::Index::open(directory, crate::index::IndexConfig::default())
1070            .await
1071        {
1072            Ok(_) => panic!("Index::open accepted a Built field with no artifact"),
1073            Err(error) => error.to_string(),
1074        };
1075        assert!(error.contains("missing_centroids.bin"), "{error}");
1076    }
1077
1078    #[tokio::test]
1079    async fn ivf_tq_built_state_rejects_a_codebook_file() {
1080        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1081        let directory = crate::directories::RamDirectory::new();
1082        let mut metadata = IndexMetadata::new(schema.clone());
1083        metadata.init_field(field.0, VectorIndexType::IvfTq);
1084        metadata.mark_field_built(
1085            field.0,
1086            10,
1087            1,
1088            "field_0_centroids.bin".into(),
1089            Some("field_0_codebook.bin".into()),
1090        );
1091        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1092
1093        let error = IndexMetadata::try_load_trained_from_fields(
1094            &metadata.vector_fields,
1095            &schema,
1096            &directory,
1097        )
1098        .await
1099        .err()
1100        .expect("IVF-TQ Built state with a codebook file must fail")
1101        .to_string();
1102        assert!(error.contains("codebook"), "{error}");
1103    }
1104
1105    #[tokio::test]
1106    async fn legacy_ivf_tq_centroid_generation_is_rejected_while_loading() {
1107        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1108        let directory = crate::directories::RamDirectory::new();
1109        let mut metadata = IndexMetadata::new(schema.clone());
1110        metadata.init_field(field.0, VectorIndexType::IvfTq);
1111        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1112        let mut legacy = test_centroids();
1113        legacy.version = 7;
1114        write_bincode(&directory, "field_0_centroids.bin", &legacy).await;
1115
1116        let error = IndexMetadata::try_load_trained_from_fields(
1117            &metadata.vector_fields,
1118            &schema,
1119            &directory,
1120        )
1121        .await
1122        .err()
1123        .expect("legacy IVF-TQ centroid state must fail while loading")
1124        .to_string();
1125        assert!(error.contains("unsupported legacy generation"), "{error}");
1126        assert!(error.contains("rebuild the index"), "{error}");
1127    }
1128
1129    #[tokio::test]
1130    async fn legacy_ivf_pq_trained_field_fails_with_actionable_error() {
1131        // Simulates metadata written by a pre-removal version: the schema
1132        // gate rejects `ivf_pq` fields, so build the raw field-state map
1133        // directly against a current-format schema.
1134        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1135        let directory = crate::directories::RamDirectory::new();
1136        let mut metadata = IndexMetadata::new(schema.clone());
1137        metadata.init_field(field.0, VectorIndexType::IvfTq);
1138        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1139        // Overwrite the recorded type the way pre-removal metadata carries it
1140        // (init_field never downgrades an existing entry).
1141        metadata
1142            .vector_fields
1143            .get_mut(&field.0)
1144            .expect("field initialized")
1145            .index_type = VectorFieldIndexType::Float(VectorIndexType::IvfPq);
1146        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1147
1148        let error = IndexMetadata::try_load_trained_from_fields(
1149            &metadata.vector_fields,
1150            &schema,
1151            &directory,
1152        )
1153        .await
1154        .err()
1155        .expect("legacy IVF-PQ trained state must fail loudly")
1156        .to_string();
1157        assert!(error.contains("no longer"), "{error}");
1158        assert!(error.contains("ivf_tq"), "{error}");
1159    }
1160
1161    #[tokio::test]
1162    async fn requested_cluster_count_accepts_a_quality_clamped_artifact() {
1163        let mut builder = crate::dsl::SchemaBuilder::default();
1164        let field = builder.add_dense_vector_field_with_config(
1165            "embedding",
1166            true,
1167            true,
1168            crate::dsl::DenseVectorConfig::ivf_tq(2, Some(4), 1),
1169        );
1170        let schema = builder.build();
1171        let directory = crate::directories::RamDirectory::new();
1172        let mut metadata = IndexMetadata::new(schema.clone());
1173        metadata.init_field(field.0, VectorIndexType::IvfTq);
1174        metadata.mark_field_built(field.0, 1, 4, "field_0_centroids.bin".into(), None);
1175        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1176
1177        let trained = IndexMetadata::try_load_trained_from_fields(
1178            &metadata.vector_fields,
1179            &schema,
1180            &directory,
1181        )
1182        .await
1183        .unwrap()
1184        .unwrap();
1185        assert_eq!(trained.centroids[&field.0].num_clusters, 1);
1186    }
1187
1188    #[tokio::test]
1189    async fn trained_artifact_loader_rejects_trailing_data() {
1190        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1191        let directory = crate::directories::RamDirectory::new();
1192        let mut metadata = IndexMetadata::new(schema.clone());
1193        metadata.init_field(field.0, VectorIndexType::IvfTq);
1194        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1195        let mut bytes =
1196            bincode::serde::encode_to_vec(test_centroids(), bincode::config::standard()).unwrap();
1197        bytes.extend_from_slice(&[0xaa, 0xbb]);
1198        directory
1199            .write(Path::new("field_0_centroids.bin"), &bytes)
1200            .await
1201            .unwrap();
1202
1203        let error = IndexMetadata::try_load_trained_from_fields(
1204            &metadata.vector_fields,
1205            &schema,
1206            &directory,
1207        )
1208        .await
1209        .err()
1210        .expect("trailing artifact bytes must fail validation")
1211        .to_string();
1212        assert!(error.contains("trailing bytes"), "{error}");
1213    }
1214
1215    #[test]
1216    fn trained_artifact_size_limit_rejects_before_reading() {
1217        let error = validate_trained_artifact_size(
1218            3,
1219            "centroids",
1220            "field_3_centroids.bin",
1221            MAX_TRAINED_ARTIFACT_BYTES as u64 + 1,
1222        )
1223        .unwrap_err()
1224        .to_string();
1225        assert!(error.contains("exceeding"), "{error}");
1226        assert!(error.contains("field 3"), "{error}");
1227    }
1228
1229    #[tokio::test]
1230    async fn trained_artifact_decode_limit_rejects_forged_collection_length() {
1231        let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1232        let directory = crate::directories::RamDirectory::new();
1233        let mut metadata = IndexMetadata::new(schema.clone());
1234        metadata.init_field(field.0, VectorIndexType::IvfTq);
1235        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1236
1237        // CoarseCentroids begins with num_clusters=1, dim=2, then the Vec
1238        // length. Bincode's standard varint marker 253 introduces a u64; this
1239        // tiny payload claims an impossible f32 vector and must hit the decode
1240        // limit before any large allocation is attempted.
1241        let mut bytes = vec![1, 2, 253];
1242        bytes.extend_from_slice(&u64::MAX.to_le_bytes());
1243        directory
1244            .write(Path::new("field_0_centroids.bin"), &bytes)
1245            .await
1246            .unwrap();
1247
1248        let error = IndexMetadata::try_load_trained_from_fields(
1249            &metadata.vector_fields,
1250            &schema,
1251            &directory,
1252        )
1253        .await
1254        .err()
1255        .expect("forged collection length must fail the bounded decoder")
1256        .to_string();
1257        assert!(error.contains("failed to deserialize"), "{error}");
1258    }
1259
1260    #[test]
1261    fn test_metadata_segments() {
1262        let mut meta = IndexMetadata::new(test_schema());
1263        meta.add_segment("abc123".to_string(), 50);
1264        meta.add_segment("def456".to_string(), 100);
1265        assert_eq!(meta.segment_metas.len(), 2);
1266        assert_eq!(meta.segment_doc_count("abc123"), Some(50));
1267        assert_eq!(meta.segment_doc_count("def456"), Some(100));
1268
1269        // Overwrites existing
1270        meta.add_segment("abc123".to_string(), 75);
1271        assert_eq!(meta.segment_metas.len(), 2);
1272        assert_eq!(meta.segment_doc_count("abc123"), Some(75));
1273
1274        meta.remove_segment("abc123");
1275        assert_eq!(meta.segment_metas.len(), 1);
1276        assert!(meta.has_segment("def456"));
1277        assert!(!meta.has_segment("abc123"));
1278    }
1279
1280    #[test]
1281    fn test_mark_field_built() {
1282        let mut meta = IndexMetadata::new(test_schema());
1283        meta.init_field(0, VectorIndexType::IvfTq);
1284        meta.total_vectors = 10000;
1285
1286        assert!(!meta.is_field_built(0));
1287
1288        meta.mark_field_built(0, 10000, 256, "field_0_centroids.bin".to_string(), None);
1289
1290        assert!(meta.is_field_built(0));
1291        let field = meta.get_field_meta(0).unwrap();
1292        assert_eq!(
1293            field.centroids_file.as_deref(),
1294            Some("field_0_centroids.bin")
1295        );
1296    }
1297
1298    #[test]
1299    fn scann_metadata_persists_generation_and_fingerprint_and_defaults_old_json() {
1300        let mut meta = IndexMetadata::new(test_schema());
1301        meta.init_field(3, VectorIndexType::Scann);
1302        meta.mark_scann_field_built(
1303            3,
1304            100_000,
1305            1_000,
1306            "field_3_scann_17.bin".to_string(),
1307            17,
1308            0xdecafbad,
1309        )
1310        .unwrap();
1311
1312        let bytes = meta.serialize_to_bytes().unwrap();
1313        let decoded: IndexMetadata = serde_json::from_slice(&bytes).unwrap();
1314        let field = decoded.get_field_meta(3).unwrap();
1315        assert_eq!(field.artifact_generation, Some(17));
1316        assert_eq!(field.artifact_id, Some(0xdecafbad));
1317
1318        let mut legacy_json = serde_json::to_value(&decoded).unwrap();
1319        legacy_json["vector_fields"]["3"]
1320            .as_object_mut()
1321            .unwrap()
1322            .remove("artifact_generation");
1323        legacy_json["vector_fields"]["3"]
1324            .as_object_mut()
1325            .unwrap()
1326            .remove("artifact_id");
1327        let legacy: IndexMetadata = serde_json::from_value(legacy_json).unwrap();
1328        let legacy_field = legacy.get_field_meta(3).unwrap();
1329        assert_eq!(legacy_field.artifact_generation, None);
1330        assert_eq!(legacy_field.artifact_id, None);
1331    }
1332
1333    #[test]
1334    fn scann_metadata_refuses_zero_or_non_scann_generation() {
1335        let mut meta = IndexMetadata::new(test_schema());
1336        meta.init_field(0, VectorIndexType::Scann);
1337        assert!(
1338            meta.mark_scann_field_built(0, 100_000, 1_000, "artifact.bin".into(), 0, 1)
1339                .is_err()
1340        );
1341        meta.init_field(1, VectorIndexType::IvfTq);
1342        assert!(
1343            meta.mark_scann_field_built(1, 100_000, 1_000, "artifact.bin".into(), 1, 2)
1344                .is_err()
1345        );
1346    }
1347
1348    #[test]
1349    fn scann_autopilot_accepts_resolved_billion_scale_three_level_geometry() {
1350        assert!(scann_explicit_geometry_matches(None, None, 10_000_000, 3));
1351        assert!(!scann_explicit_geometry_matches(
1352            Some(1_000_000),
1353            None,
1354            10_000_000,
1355            3
1356        ));
1357        assert!(!scann_explicit_geometry_matches(
1358            None,
1359            Some(1),
1360            10_000_000,
1361            3
1362        ));
1363    }
1364
1365    #[test]
1366    fn total_vectors_is_aggregate_of_built_field_counts() {
1367        let mut meta = IndexMetadata::new(test_schema());
1368        meta.init_field(7, VectorIndexType::IvfTq);
1369        meta.init_field(3, VectorIndexType::IvfTq);
1370
1371        // Build in reverse field-id order to ensure the result is not tied to
1372        // HashMap or training iteration order.
1373        meta.mark_field_built(7, 400, 20, "field_7_centroids.bin".to_string(), None);
1374        assert_eq!(meta.total_vectors, 400);
1375        meta.mark_field_built(3, 250, 15, "field_3_centroids.bin".to_string(), None);
1376        assert_eq!(meta.total_vectors, 650);
1377
1378        // Rebuilding a field replaces its contribution; it does not add a
1379        // duplicate training snapshot.
1380        meta.mark_field_built(7, 425, 20, "field_7_centroids.bin".to_string(), None);
1381        assert_eq!(meta.total_vectors, 675);
1382    }
1383
1384    #[test]
1385    fn test_should_build_field() {
1386        let mut meta = IndexMetadata::new(test_schema());
1387        meta.init_field(0, VectorIndexType::IvfTq);
1388
1389        // Below threshold
1390        meta.total_vectors = 500;
1391        assert!(!meta.should_build_field(0, 1000));
1392
1393        // Above threshold
1394        meta.total_vectors = 1500;
1395        assert!(meta.should_build_field(0, 1000));
1396
1397        // Already built - should not build again
1398        meta.mark_field_built(0, 1500, 256, "centroids.bin".to_string(), None);
1399        assert!(!meta.should_build_field(0, 1000));
1400    }
1401
1402    #[test]
1403    fn test_serialization() {
1404        let mut meta = IndexMetadata::new(test_schema());
1405        meta.add_segment("seg1".to_string(), 100);
1406        meta.init_field(0, VectorIndexType::IvfTq);
1407        meta.total_vectors = 5000;
1408
1409        let json = serde_json::to_string_pretty(&meta).unwrap();
1410        let loaded: IndexMetadata = serde_json::from_str(&json).unwrap();
1411
1412        assert_eq!(loaded.segment_ids().len(), meta.segment_ids().len());
1413        assert_eq!(loaded.segment_doc_count("seg1"), Some(100));
1414        assert_eq!(loaded.total_vectors, meta.total_vectors);
1415        assert!(loaded.vector_fields.contains_key(&0));
1416    }
1417
1418    #[test]
1419    fn old_metadata_defaults_the_bp_retry_counter() {
1420        let mut meta = IndexMetadata::new(test_schema());
1421        meta.add_segment("legacy".to_string(), 10);
1422        let mut json = serde_json::to_value(&meta).unwrap();
1423        json["segment_metas"]["legacy"]
1424            .as_object_mut()
1425            .unwrap()
1426            .remove("bp_unconverged_passes");
1427
1428        let loaded: IndexMetadata = serde_json::from_value(json).unwrap();
1429        assert_eq!(loaded.segment_metas["legacy"].bp_unconverged_passes, 0);
1430    }
1431
1432    #[test]
1433    fn test_merged_segment_lineage() {
1434        let mut meta = IndexMetadata::new(test_schema());
1435        meta.add_segment("a".to_string(), 50);
1436        meta.add_segment("b".to_string(), 75);
1437
1438        // Fresh segments: gen=0, no ancestors
1439        assert_eq!(meta.segment_metas["a"].generation, 0);
1440        assert!(meta.segment_metas["a"].ancestors.is_empty());
1441
1442        // Merge a+b → c
1443        meta.add_merged_segment(
1444            "c".to_string(),
1445            125,
1446            vec!["a".to_string(), "b".to_string()],
1447            1,
1448            false,
1449            true,
1450        );
1451        assert_eq!(meta.segment_metas["c"].generation, 1);
1452        assert_eq!(meta.segment_metas["c"].ancestors, vec!["a", "b"]);
1453        assert_eq!(meta.segment_doc_count("c"), Some(125));
1454
1455        // Merge c+d → e (gen should be 2)
1456        meta.add_segment("d".to_string(), 30);
1457        meta.add_merged_segment(
1458            "e".to_string(),
1459            155,
1460            vec!["c".to_string(), "d".to_string()],
1461            2,
1462            false,
1463            true,
1464        );
1465        assert_eq!(meta.segment_metas["e"].generation, 2);
1466    }
1467}