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