1use 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
22pub const INDEX_META_FILENAME: &str = "metadata.json";
24const INDEX_META_TMP_FILENAME: &str = "metadata.json.tmp";
26
27pub const INDEX_META_FORMAT_VERSION: u32 = 6;
33
34pub(crate) const MAX_TRAINED_ARTIFACT_BYTES: usize = 512 * 1024 * 1024;
39
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
42pub enum VectorIndexState {
43 #[default]
45 Flat,
46 Built {
48 vector_count: usize,
50 num_clusters: usize,
52 },
53}
54
55fn default_true() -> bool {
56 true
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct SegmentMetaInfo {
63 pub num_docs: u32,
65 pub ancestors: Vec<String>,
67 pub generation: u32,
69 #[serde(default)]
73 pub reordered: bool,
74 #[serde(default = "default_true")]
79 pub bp_converged: bool,
80 #[serde(default)]
84 pub bp_unconverged_passes: u32,
85}
86
87#[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 pub field_id: u32,
111 pub index_type: VectorFieldIndexType,
113 pub state: VectorIndexState,
115 #[serde(skip_serializing_if = "Option::is_none")]
117 pub centroids_file: Option<String>,
118 #[serde(skip_serializing_if = "Option::is_none")]
122 pub codebook_file: Option<String>,
123 #[serde(default, skip_serializing_if = "Option::is_none")]
126 pub artifact_generation: Option<u64>,
127 #[serde(default, skip_serializing_if = "Option::is_none")]
130 pub artifact_id: Option<u64>,
131}
132
133#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct IndexMetadata {
136 pub version: u32,
138 #[serde(default)]
142 pub publication_generation: u64,
143 pub schema: Schema,
145 #[serde(default)]
148 pub segment_metas: HashMap<String, SegmentMetaInfo>,
149 #[serde(default)]
151 pub vector_fields: HashMap<u32, FieldVectorMeta>,
152 #[serde(default)]
159 pub total_vectors: usize,
160}
161
162impl IndexMetadata {
163 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 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 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 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 pub(crate) fn add_segment_meta(&mut self, segment_id: String, info: SegmentMetaInfo) {
224 self.segment_metas.insert(segment_id, info);
225 }
226
227 pub fn remove_segment(&mut self, segment_id: &str) {
229 self.segment_metas.remove(segment_id);
230 }
231
232 pub fn has_segment(&self, segment_id: &str) -> bool {
234 self.segment_metas.contains_key(segment_id)
235 }
236
237 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 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 pub fn get_field_meta(&self, field_id: u32) -> Option<&FieldVectorMeta> {
252 self.vector_fields.get(&field_id)
253 }
254
255 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 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 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 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 pub fn should_build_field(&self, field_id: u32, threshold: usize) -> bool {
354 if self.is_field_built(field_id) {
356 return false;
357 }
358 self.total_vectors >= threshold
360 }
361
362 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 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 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 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 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 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 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 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 #[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 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 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
773fn 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 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 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 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 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 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 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 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 meta.total_vectors = 500;
1390 assert!(!meta.should_build_field(0, 1000));
1391
1392 meta.total_vectors = 1500;
1394 assert!(meta.should_build_field(0, 1000));
1395
1396 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 assert_eq!(meta.segment_metas["a"].generation, 0);
1439 assert!(meta.segment_metas["a"].ancestors.is_empty());
1440
1441 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 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}