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 = 4;
34
35pub(crate) const MAX_TRAINED_ARTIFACT_BYTES: usize = 512 * 1024 * 1024;
40
41#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
43pub enum VectorIndexState {
44 #[default]
46 Flat,
47 Built {
49 vector_count: usize,
51 num_clusters: usize,
53 },
54}
55
56fn default_true() -> bool {
57 true
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SegmentMetaInfo {
64 pub num_docs: u32,
66 pub ancestors: Vec<String>,
68 pub generation: u32,
70 #[serde(default)]
74 pub reordered: bool,
75 #[serde(default = "default_true")]
80 pub bp_converged: bool,
81 #[serde(default)]
85 pub bp_unconverged_passes: u32,
86}
87
88#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
90#[serde(tag = "kind", content = "index", rename_all = "snake_case")]
91pub enum VectorFieldIndexType {
92 Float(VectorIndexType),
93 Binary(BinaryIndexType),
94}
95
96impl From<VectorIndexType> for VectorFieldIndexType {
97 fn from(value: VectorIndexType) -> Self {
98 Self::Float(value)
99 }
100}
101
102impl From<BinaryIndexType> for VectorFieldIndexType {
103 fn from(value: BinaryIndexType) -> Self {
104 Self::Binary(value)
105 }
106}
107
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct FieldVectorMeta {
110 pub field_id: u32,
112 pub index_type: VectorFieldIndexType,
114 pub state: VectorIndexState,
116 #[serde(skip_serializing_if = "Option::is_none")]
118 pub centroids_file: Option<String>,
119 #[serde(skip_serializing_if = "Option::is_none")]
123 pub codebook_file: Option<String>,
124 #[serde(default, skip_serializing_if = "Option::is_none")]
127 pub artifact_generation: Option<u64>,
128 #[serde(default, skip_serializing_if = "Option::is_none")]
131 pub artifact_id: Option<u64>,
132}
133
134#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct IndexMetadata {
137 pub version: u32,
139 #[serde(default)]
143 pub publication_generation: u64,
144 pub schema: Schema,
146 #[serde(default)]
149 pub segment_metas: HashMap<String, SegmentMetaInfo>,
150 #[serde(default)]
152 pub vector_fields: HashMap<u32, FieldVectorMeta>,
153 #[serde(default)]
160 pub total_vectors: usize,
161}
162
163impl IndexMetadata {
164 pub fn new(schema: Schema) -> Self {
166 Self {
167 version: INDEX_META_FORMAT_VERSION,
168 publication_generation: 0,
169 schema,
170 segment_metas: HashMap::new(),
171 vector_fields: HashMap::new(),
172 total_vectors: 0,
173 }
174 }
175
176 pub fn segment_ids(&self) -> Vec<String> {
178 let mut ids: Vec<String> = self.segment_metas.keys().cloned().collect();
179 ids.sort();
180 ids
181 }
182
183 pub fn add_segment(&mut self, segment_id: String, num_docs: u32) {
185 self.segment_metas.insert(
186 segment_id,
187 SegmentMetaInfo {
188 num_docs,
189 ancestors: Vec::new(),
190 generation: 0,
191 reordered: false,
192 bp_converged: true,
193 bp_unconverged_passes: 0,
194 },
195 );
196 }
197
198 pub fn add_merged_segment(
200 &mut self,
201 segment_id: String,
202 num_docs: u32,
203 ancestors: Vec<String>,
204 generation: u32,
205 reordered: bool,
206 bp_converged: bool,
207 ) {
208 self.add_segment_meta(
209 segment_id,
210 SegmentMetaInfo {
211 num_docs,
212 ancestors,
213 generation,
214 reordered,
215 bp_converged,
216 bp_unconverged_passes: 0,
217 },
218 );
219 }
220
221 pub(crate) fn add_segment_meta(&mut self, segment_id: String, info: SegmentMetaInfo) {
225 self.segment_metas.insert(segment_id, info);
226 }
227
228 pub fn remove_segment(&mut self, segment_id: &str) {
230 self.segment_metas.remove(segment_id);
231 }
232
233 pub fn has_segment(&self, segment_id: &str) -> bool {
235 self.segment_metas.contains_key(segment_id)
236 }
237
238 pub fn segment_doc_count(&self, segment_id: &str) -> Option<u32> {
240 self.segment_metas.get(segment_id).map(|m| m.num_docs)
241 }
242
243 pub fn is_field_built(&self, field_id: u32) -> bool {
245 self.vector_fields
246 .get(&field_id)
247 .map(|f| matches!(f.state, VectorIndexState::Built { .. }))
248 .unwrap_or(false)
249 }
250
251 pub fn get_field_meta(&self, field_id: u32) -> Option<&FieldVectorMeta> {
253 self.vector_fields.get(&field_id)
254 }
255
256 pub fn init_field(&mut self, field_id: u32, index_type: impl Into<VectorFieldIndexType>) {
258 let index_type = index_type.into();
259 self.vector_fields
260 .entry(field_id)
261 .or_insert(FieldVectorMeta {
262 field_id,
263 index_type,
264 state: VectorIndexState::Flat,
265 centroids_file: None,
266 codebook_file: None,
267 artifact_generation: None,
268 artifact_id: None,
269 });
270 }
271
272 pub fn mark_field_built(
274 &mut self,
275 field_id: u32,
276 vector_count: usize,
277 num_clusters: usize,
278 centroids_file: String,
279 codebook_file: Option<String>,
280 ) {
281 if let Some(field) = self.vector_fields.get_mut(&field_id) {
282 field.state = VectorIndexState::Built {
283 vector_count,
284 num_clusters,
285 };
286 field.centroids_file = Some(centroids_file);
287 field.codebook_file = codebook_file;
288 field.artifact_generation = None;
289 field.artifact_id = None;
290 self.refresh_total_vectors();
291 }
292 }
293
294 pub fn mark_scann_field_built(
298 &mut self,
299 field_id: u32,
300 vector_count: usize,
301 num_leaves: usize,
302 artifact_file: String,
303 artifact_generation: u64,
304 artifact_id: u64,
305 ) -> Result<()> {
306 if artifact_generation == 0 || artifact_id == 0 {
307 return Err(Error::Corruption(format!(
308 "ScaNN field {field_id} cannot publish a zero generation or artifact fingerprint"
309 )));
310 }
311 let field = self.vector_fields.get_mut(&field_id).ok_or_else(|| {
312 Error::Corruption(format!(
313 "ScaNN field {field_id} must be initialized before it is marked built"
314 ))
315 })?;
316 if !matches!(
317 field.index_type,
318 VectorFieldIndexType::Float(VectorIndexType::Scann)
319 | VectorFieldIndexType::Binary(BinaryIndexType::Scann)
320 ) {
321 return Err(Error::Corruption(format!(
322 "field {field_id} is not configured as ScaNN"
323 )));
324 }
325 field.state = VectorIndexState::Built {
326 vector_count,
327 num_clusters: num_leaves,
328 };
329 field.centroids_file = Some(artifact_file);
330 field.codebook_file = None;
331 field.artifact_generation = Some(artifact_generation);
332 field.artifact_id = Some(artifact_id);
333 self.refresh_total_vectors();
334 Ok(())
335 }
336
337 pub(crate) fn refresh_total_vectors(&mut self) {
343 self.total_vectors = self
344 .vector_fields
345 .values()
346 .filter_map(|field| match field.state {
347 VectorIndexState::Built { vector_count, .. } => Some(vector_count),
348 VectorIndexState::Flat => None,
349 })
350 .fold(0usize, usize::saturating_add);
351 }
352
353 pub fn should_build_field(&self, field_id: u32, threshold: usize) -> bool {
355 if self.is_field_built(field_id) {
357 return false;
358 }
359 self.total_vectors >= threshold
361 }
362
363 pub async fn load<D: crate::directories::Directory>(dir: &D) -> Result<Self> {
368 let path = Path::new(INDEX_META_FILENAME);
369 match dir.open_read(path).await {
370 Ok(slice) => {
371 let bytes = slice.read_bytes().await?;
372 Self::deserialize_versioned(bytes.as_slice())
373 }
374 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
375 let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
377 let slice = dir.open_read(tmp_path).await?;
378 let bytes = slice.read_bytes().await?;
379 let meta = Self::deserialize_versioned(bytes.as_slice())?;
380 log::warn!("Recovered metadata from temp file (previous crash during save)");
381 Ok(meta)
382 }
383 Err(e) => Err(Error::Io(e)),
384 }
385 }
386
387 fn deserialize_versioned(bytes: &[u8]) -> Result<Self> {
391 let meta: Self =
392 serde_json::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))?;
393 crate::dsl::reject_removed_vector_index_types(&meta.schema).map_err(Error::Schema)?;
394 if meta.version != INDEX_META_FORMAT_VERSION {
395 return Err(Error::Corruption(format!(
396 "metadata.json format version {} is incompatible with required version {}; \
397 rebuild and republish the index with this Hermes version",
398 meta.version, INDEX_META_FORMAT_VERSION
399 )));
400 }
401 Ok(meta)
402 }
403
404 pub async fn save<D: crate::directories::DirectoryWriter>(&self, dir: &D) -> Result<()> {
409 let bytes = self.serialize_to_bytes()?;
410 Self::save_bytes(dir, &bytes).await
411 }
412
413 pub fn serialize_to_bytes(&self) -> Result<Vec<u8>> {
416 serde_json::to_vec_pretty(self).map_err(|e| Error::Serialization(e.to_string()))
417 }
418
419 pub async fn save_bytes<D: crate::directories::DirectoryWriter>(
424 dir: &D,
425 bytes: &[u8],
426 ) -> Result<()> {
427 let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
428 let final_path = Path::new(INDEX_META_FILENAME);
429 let mut writer = dir.streaming_writer(tmp_path).await.map_err(Error::Io)?;
434 writer.write_all(bytes).map_err(Error::Io)?;
435 writer.finish().map_err(Error::Io)?;
436 dir.rename(tmp_path, final_path).await.map_err(Error::Io)?;
443 if let Err(error) = dir.sync().await {
444 log::error!(
445 "[metadata] directory fsync failed after committed rename: {}. \
446 Continuing with the renamed generation; crash durability is not guaranteed",
447 error,
448 );
449 }
450 Ok(())
451 }
452
453 #[cfg_attr(not(feature = "native"), allow(dead_code))]
455 pub(crate) async fn try_load_trained_from_fields<D: crate::directories::Directory>(
456 vector_fields: &HashMap<u32, FieldVectorMeta>,
457 schema: &Schema,
458 dir: &D,
459 ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
460 Self::load_trained_from_fields_impl(vector_fields, schema, dir).await
461 }
462
463 async fn load_trained_from_fields_impl<D: crate::directories::Directory>(
473 vector_fields: &HashMap<u32, FieldVectorMeta>,
474 schema: &Schema,
475 dir: &D,
476 ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
477 use std::sync::Arc;
478
479 let mut centroids = rustc_hash::FxHashMap::default();
480 let mut binary_quantizers = rustc_hash::FxHashMap::default();
481 let mut scann_artifacts = rustc_hash::FxHashMap::default();
482
483 let mut built_fields: Vec<_> = vector_fields
484 .iter()
485 .filter(|(_, meta)| matches!(meta.state, VectorIndexState::Built { .. }))
486 .collect();
487 built_fields.sort_unstable_by_key(|(field_id, _)| **field_id);
488
489 log::debug!(
490 "[trained] index={} loading trained structures, dense_vector_fields={:?}",
491 schema.index_label(),
492 vector_fields.keys().collect::<Vec<_>>()
493 );
494
495 for (field_id, field_meta) in built_fields {
496 log::debug!(
497 "[trained] index={} field {} state={:?} centroids_file={:?} codebook_file={:?}",
498 schema.index_label(),
499 field_id,
500 field_meta.state,
501 field_meta.centroids_file,
502 field_meta.codebook_file,
503 );
504 if field_meta.field_id != *field_id {
505 return Err(Error::Corruption(format!(
506 "trained vector metadata key {field_id} contains field_id {}",
507 field_meta.field_id
508 )));
509 }
510
511 let expected_clusters = match field_meta.state {
512 VectorIndexState::Built { num_clusters, .. } if num_clusters > 0 => num_clusters,
513 VectorIndexState::Built { .. } => {
514 return Err(Error::Corruption(format!(
515 "trained vector metadata field {field_id} has zero clusters"
516 )));
517 }
518 VectorIndexState::Flat => unreachable!("built_fields contains only Built entries"),
519 };
520
521 let centroids_file = field_meta.centroids_file.as_deref().ok_or_else(|| {
522 Error::Corruption(format!(
523 "trained vector metadata field {field_id} is Built but has no centroids_file"
524 ))
525 })?;
526 match field_meta.index_type {
527 VectorFieldIndexType::Float(VectorIndexType::IvfPq) => {
528 return Err(Error::Corruption(format!(
529 "field {field_id} was trained as IVF-PQ, which is no longer \
530 supported; recreate the index with `ivf_tq` and reindex \
531 (docs/turboquant-quantization.md)"
532 )));
533 }
534 VectorFieldIndexType::Float(index_type @ VectorIndexType::IvfTq) => {
535 let entry = schema
536 .get_field_entry(crate::dsl::Field(*field_id))
537 .ok_or_else(|| {
538 Error::Corruption(format!(
539 "trained vector metadata references missing field {field_id}"
540 ))
541 })?;
542 let schema_config = entry
543 .dense_vector_config
544 .as_ref()
545 .filter(|_| entry.field_type == crate::dsl::FieldType::DenseVector)
546 .ok_or_else(|| {
547 Error::Corruption(format!(
548 "trained vector metadata field {field_id} is not a float dense field"
549 ))
550 })?;
551 if schema_config.index_type != index_type {
552 return Err(Error::Corruption(format!(
553 "trained vector metadata field {field_id} uses {index_type:?}, schema requires {:?}",
554 schema_config.index_type
555 )));
556 }
557 let c: crate::structures::CoarseCentroids =
558 load_trained_artifact(dir, *field_id, "centroids", centroids_file).await?;
559 let expected_dim = schema_config.dim;
560 let actual_clusters = c.num_clusters as usize;
561 let expected_values =
562 actual_clusters.checked_mul(expected_dim).ok_or_else(|| {
563 Error::Corruption(format!(
564 "trained centroid dimensions overflow for field {field_id}"
565 ))
566 })?;
567 if actual_clusters == 0
568 || actual_clusters > expected_clusters
569 || c.dim == 0
570 || c.dim != expected_dim
571 || c.centroids.len() != expected_values
572 || c.centroids.iter().any(|value| !value.is_finite())
573 {
574 return Err(Error::Corruption(format!(
575 "trained centroids for field {field_id} do not match metadata/schema"
576 )));
577 }
578 if !crate::structures::is_ivf_tq_cosine_generation(c.version) {
579 return Err(Error::Corruption(format!(
580 "trained IVF-TQ centroids for field {field_id} use an \
581 unsupported legacy generation; rebuild the index"
582 )));
583 }
584 c.validate_routing(schema_config.ivf_routing)
585 .map_err(|error| {
586 Error::Corruption(format!(
587 "invalid trained centroid routing for field {field_id}: {error}"
588 ))
589 })?;
590 let _ = index_type;
593 if field_meta.codebook_file.is_some() {
594 return Err(Error::Corruption(format!(
595 "trained IVF-TQ field {field_id} unexpectedly references a codebook file"
596 )));
597 }
598 centroids.insert(*field_id, Arc::new(c));
599 }
600 VectorFieldIndexType::Binary(BinaryIndexType::Ivf) => {
601 let entry = schema
602 .get_field_entry(crate::dsl::Field(*field_id))
603 .ok_or_else(|| {
604 Error::Corruption(format!(
605 "trained vector metadata references missing field {field_id}"
606 ))
607 })?;
608 let schema_config = entry
609 .binary_dense_vector_config
610 .as_ref()
611 .filter(|config| {
612 entry.field_type == crate::dsl::FieldType::BinaryDenseVector
613 && config.index_type == BinaryIndexType::Ivf
614 })
615 .ok_or_else(|| {
616 Error::Corruption(format!(
617 "trained vector metadata field {field_id} is not a binary IVF field"
618 ))
619 })?;
620 let quantizer: crate::structures::BinaryCoarseQuantizer =
621 load_trained_artifact(dir, *field_id, "binary centroids", centroids_file)
622 .await?;
623 quantizer.validate().map_err(|error| {
624 Error::Corruption(format!(
625 "invalid binary coarse quantizer for field {field_id}: {error}"
626 ))
627 })?;
628 let actual_clusters = quantizer.num_clusters as usize;
629 if actual_clusters > expected_clusters
630 || schema_config.dim != quantizer.dim_bits
631 {
632 return Err(Error::Corruption(format!(
633 "binary coarse quantizer for field {field_id} does not match metadata/schema"
634 )));
635 }
636 quantizer
637 .validate_routing(schema_config.ivf_routing)
638 .map_err(|error| {
639 Error::Corruption(format!(
640 "invalid binary centroid routing for field {field_id}: {error}"
641 ))
642 })?;
643 binary_quantizers.insert(*field_id, Arc::new(quantizer));
644 }
645 VectorFieldIndexType::Float(VectorIndexType::Scann)
646 | VectorFieldIndexType::Binary(BinaryIndexType::Scann) => {
647 if field_meta.codebook_file.is_some() {
648 return Err(Error::Corruption(format!(
649 "trained ScaNN field {field_id} unexpectedly references a separate codebook file"
650 )));
651 }
652 let expected_generation = field_meta.artifact_generation.ok_or_else(|| {
653 Error::Corruption(format!(
654 "trained ScaNN field {field_id} has no artifact generation"
655 ))
656 })?;
657 let expected_artifact_id = field_meta.artifact_id.ok_or_else(|| {
658 Error::Corruption(format!(
659 "trained ScaNN field {field_id} has no artifact fingerprint"
660 ))
661 })?;
662 validate_trained_artifact_path(
663 field_id.to_owned(),
664 "ScaNN artifact",
665 centroids_file,
666 )?;
667 let path = Path::new(centroids_file);
668 let slice = dir.open_read(path).await.map_err(|error| {
669 Error::Corruption(format!(
670 "failed to open trained ScaNN artifact '{centroids_file}' for field {field_id}: {error}"
671 ))
672 })?;
673 let raw = slice.read_bytes().await.map_err(|error| {
674 Error::Corruption(format!(
675 "failed to map trained ScaNN artifact '{centroids_file}' for field {field_id}: {error}"
676 ))
677 })?;
678 let artifact =
679 crate::segment::ScannTrainedArtifactBytes::open(raw).map_err(|error| {
680 Error::Corruption(format!(
681 "invalid trained ScaNN artifact for field {field_id}: {error}"
682 ))
683 })?;
684 if artifact.generation() != expected_generation
685 || artifact.artifact_id() != expected_artifact_id
686 || artifact.config().num_leaves as usize != expected_clusters
687 {
688 return Err(Error::Corruption(format!(
689 "trained ScaNN artifact for field {field_id} does not match metadata"
690 )));
691 }
692 let entry = schema
693 .get_field_entry(crate::dsl::Field(*field_id))
694 .ok_or_else(|| {
695 Error::Corruption(format!(
696 "trained vector metadata references missing field {field_id}"
697 ))
698 })?;
699 let schema_matches = match field_meta.index_type {
700 VectorFieldIndexType::Float(VectorIndexType::Scann) => entry
701 .dense_vector_config
702 .as_ref()
703 .filter(|_| entry.field_type == crate::dsl::FieldType::DenseVector)
704 .is_some_and(|config| {
705 config.index_type == VectorIndexType::Scann
706 && config.dim == artifact.config().dimension as usize
707 && scann_explicit_geometry_matches(
708 config.num_clusters,
709 config.tree_levels,
710 artifact.config().num_leaves as usize,
711 artifact.config().tree_levels,
712 )
713 && matches!(
714 artifact.config().encoding,
715 crate::structures::vector::scann::ScannEncoding::AsymmetricHash { .. }
716 )
717 }),
718 VectorFieldIndexType::Binary(BinaryIndexType::Scann) => entry
719 .binary_dense_vector_config
720 .as_ref()
721 .filter(|_| {
722 entry.field_type == crate::dsl::FieldType::BinaryDenseVector
723 })
724 .is_some_and(|config| {
725 config.index_type == BinaryIndexType::Scann
726 && config.dim == artifact.config().dimension as usize
727 && scann_explicit_geometry_matches(
728 config.num_clusters,
729 config.tree_levels,
730 artifact.config().num_leaves as usize,
731 artifact.config().tree_levels,
732 )
733 && artifact.config().encoding
734 == crate::structures::vector::scann::ScannEncoding::BinaryHamming
735 }),
736 _ => false,
737 };
738 if !schema_matches {
739 return Err(Error::Corruption(format!(
740 "trained ScaNN artifact for field {field_id} does not match schema geometry/encoding"
741 )));
742 }
743 scann_artifacts.insert(*field_id, Arc::new(artifact));
744 }
745 unsupported => {
746 return Err(Error::Corruption(format!(
747 "field {field_id} is Built for {unsupported:?}, which has no global IVF artifacts"
748 )));
749 }
750 }
751 }
752
753 if centroids.is_empty() && binary_quantizers.is_empty() && scann_artifacts.is_empty() {
754 Ok(None)
755 } else {
756 let trained = crate::segment::TrainedVectorStructures {
757 #[cfg(feature = "native")]
758 _ann_pins: Default::default(),
759 centroids,
760 binary_quantizers,
761 scann_artifacts,
762 };
763 #[cfg(feature = "native")]
764 let trained = {
765 let mut trained = trained;
766 trained.pin_ann_structures(crate::segment::pin::pin_policy());
767 trained
768 };
769 Ok(Some(trained))
770 }
771 }
772}
773
774fn scann_explicit_geometry_matches(
778 configured_leaves: Option<usize>,
779 configured_levels: Option<u8>,
780 resolved_leaves: usize,
781 resolved_levels: u8,
782) -> bool {
783 configured_leaves.is_none_or(|leaves| leaves == resolved_leaves)
784 && configured_levels.is_none_or(|levels| levels == resolved_levels)
785}
786
787fn validate_trained_artifact_path(field_id: u32, kind: &str, filename: &str) -> Result<()> {
788 use std::path::Component;
789
790 let path = Path::new(filename);
791 if filename.is_empty()
792 || path.is_absolute()
793 || path.components().any(|component| {
794 matches!(
795 component,
796 Component::ParentDir | Component::RootDir | Component::Prefix(_)
797 )
798 })
799 {
800 return Err(Error::Corruption(format!(
801 "trained {kind} path for field {field_id} is not a safe relative path: '{filename}'"
802 )));
803 }
804 Ok(())
805}
806
807async fn load_trained_artifact<T, D>(
808 dir: &D,
809 field_id: u32,
810 kind: &str,
811 filename: &str,
812) -> Result<T>
813where
814 T: serde::de::DeserializeOwned,
815 D: crate::directories::Directory,
816{
817 validate_trained_artifact_path(field_id, kind, filename)?;
818 let path = Path::new(filename);
819 let file_size = dir.file_size(path).await.map_err(|error| {
820 Error::Corruption(format!(
821 "failed to stat trained {kind} '{filename}' for field {field_id}: {error}"
822 ))
823 })?;
824 validate_trained_artifact_size(field_id, kind, filename, file_size)?;
825 let slice = dir.open_read(path).await.map_err(|error| {
826 Error::Corruption(format!(
827 "failed to open trained {kind} '{filename}' for field {field_id}: {error}"
828 ))
829 })?;
830 validate_trained_artifact_size(field_id, kind, filename, slice.len())?;
831 let bytes = slice.read_bytes().await.map_err(|error| {
832 Error::Corruption(format!(
833 "failed to read trained {kind} '{filename}' for field {field_id}: {error}"
834 ))
835 })?;
836 let (artifact, consumed) = bincode::serde::decode_from_slice::<T, _>(
837 bytes.as_slice(),
838 bincode::config::standard().with_limit::<MAX_TRAINED_ARTIFACT_BYTES>(),
839 )
840 .map_err(|error| {
841 Error::Corruption(format!(
842 "failed to deserialize trained {kind} '{filename}' for field {field_id}: {error}"
843 ))
844 })?;
845 if consumed != bytes.len() {
846 return Err(Error::Corruption(format!(
847 "trained {kind} '{filename}' for field {field_id} has {} trailing bytes",
848 bytes.len() - consumed
849 )));
850 }
851 Ok(artifact)
852}
853
854fn validate_trained_artifact_size(
855 field_id: u32,
856 kind: &str,
857 filename: &str,
858 file_size: u64,
859) -> Result<()> {
860 if file_size > MAX_TRAINED_ARTIFACT_BYTES as u64 {
861 return Err(Error::Corruption(format!(
862 "trained {kind} '{filename}' for field {field_id} is {file_size} bytes, \
863 exceeding the {MAX_TRAINED_ARTIFACT_BYTES}-byte safety limit"
864 )));
865 }
866 Ok(())
867}
868
869#[cfg(test)]
870mod tests {
871 use super::*;
872 use crate::directories::DirectoryWriter;
873
874 #[derive(Clone, Default)]
875 struct SyncFailDirectory(crate::directories::RamDirectory);
876
877 #[async_trait::async_trait]
878 impl crate::directories::Directory for SyncFailDirectory {
879 async fn exists(&self, path: &Path) -> std::io::Result<bool> {
880 self.0.exists(path).await
881 }
882
883 async fn file_size(&self, path: &Path) -> std::io::Result<u64> {
884 self.0.file_size(path).await
885 }
886
887 async fn open_read(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
888 self.0.open_read(path).await
889 }
890
891 async fn read_range(
892 &self,
893 path: &Path,
894 range: std::ops::Range<u64>,
895 ) -> std::io::Result<crate::directories::OwnedBytes> {
896 self.0.read_range(path, range).await
897 }
898
899 async fn list_files(&self, prefix: &Path) -> std::io::Result<Vec<std::path::PathBuf>> {
900 self.0.list_files(prefix).await
901 }
902
903 async fn open_lazy(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
904 self.0.open_lazy(path).await
905 }
906 }
907
908 #[async_trait::async_trait]
909 impl crate::directories::DirectoryWriter for SyncFailDirectory {
910 async fn write(&self, path: &Path, data: &[u8]) -> std::io::Result<()> {
911 self.0.write(path, data).await
912 }
913
914 async fn delete(&self, path: &Path) -> std::io::Result<()> {
915 self.0.delete(path).await
916 }
917
918 async fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
919 self.0.rename(from, to).await
920 }
921
922 async fn sync(&self) -> std::io::Result<()> {
923 Err(std::io::Error::other("injected directory fsync failure"))
924 }
925
926 async fn streaming_writer(
927 &self,
928 path: &Path,
929 ) -> std::io::Result<Box<dyn crate::directories::StreamingWriter>> {
930 self.0.streaming_writer(path).await
931 }
932 }
933
934 fn test_schema() -> Schema {
935 Schema::default()
936 }
937
938 fn dense_schema(index_type: VectorIndexType) -> (Schema, crate::dsl::Field) {
939 let mut builder = crate::dsl::SchemaBuilder::default();
940 let config = match index_type {
941 VectorIndexType::IvfTq => crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1),
942 other => panic!("unsupported trained test index type: {other:?}"),
943 };
944 let field = builder.add_dense_vector_field_with_config("embedding", true, true, config);
945 (builder.build(), field)
946 }
947
948 fn test_centroids() -> crate::structures::CoarseCentroids {
949 crate::structures::CoarseCentroids {
950 num_clusters: 1,
951 dim: 2,
952 centroids: vec![0.25, 0.75],
953 version: crate::structures::mark_ivf_tq_cosine_generation(7),
954 soar_config: None,
955 routing_index: None,
956 }
957 }
958
959 async fn write_bincode(
960 directory: &crate::directories::RamDirectory,
961 filename: &str,
962 value: &impl serde::Serialize,
963 ) {
964 let bytes = bincode::serde::encode_to_vec(value, bincode::config::standard()).unwrap();
965 directory.write(Path::new(filename), &bytes).await.unwrap();
966 }
967
968 #[test]
969 fn test_metadata_init() {
970 let mut meta = IndexMetadata::new(test_schema());
971 assert_eq!(meta.total_vectors, 0);
972 assert!(meta.segment_metas.is_empty());
973 assert!(!meta.is_field_built(0));
974
975 meta.init_field(0, VectorIndexType::IvfTq);
976 assert!(!meta.is_field_built(0));
977 assert!(meta.vector_fields.contains_key(&0));
978 }
979
980 #[tokio::test]
981 async fn load_refuses_metadata_stamped_with_a_newer_format_version() {
982 let directory = crate::directories::RamDirectory::new();
983 let mut metadata = IndexMetadata::new(test_schema());
984 metadata.version = INDEX_META_FORMAT_VERSION + 1;
985 metadata.save(&directory).await.unwrap();
986
987 let error = IndexMetadata::load(&directory)
988 .await
989 .expect_err("metadata from a newer format version must be refused, not silently pruned")
990 .to_string();
991 assert!(error.contains("version 4"), "{error}");
992 assert!(error.contains("incompatible"), "{error}");
993 }
994
995 #[tokio::test]
996 async fn tmp_recovery_refuses_metadata_stamped_with_a_newer_format_version() {
997 let directory = crate::directories::RamDirectory::new();
998 let mut metadata = IndexMetadata::new(test_schema());
999 metadata.version = INDEX_META_FORMAT_VERSION + 1;
1000 let bytes = metadata.serialize_to_bytes().unwrap();
1001 directory
1003 .write(Path::new(INDEX_META_TMP_FILENAME), &bytes)
1004 .await
1005 .unwrap();
1006
1007 let error = IndexMetadata::load(&directory)
1008 .await
1009 .expect_err("temp-file recovery must apply the same version gate")
1010 .to_string();
1011 assert!(error.contains("version 4"), "{error}");
1012 }
1013
1014 #[tokio::test]
1015 async fn save_treats_post_rename_sync_failure_as_committed() {
1016 let directory = SyncFailDirectory::default();
1017 let mut metadata = IndexMetadata::new(test_schema());
1018 metadata.add_segment("committed".to_string(), 7);
1019
1020 metadata.save(&directory).await.unwrap();
1021
1022 let loaded = IndexMetadata::load(&directory).await.unwrap();
1023 assert_eq!(loaded.segment_doc_count("committed"), Some(7));
1024 }
1025
1026 #[tokio::test]
1027 async fn trained_artifacts_load_only_when_the_complete_built_set_is_valid() {
1028 let mut builder = crate::dsl::SchemaBuilder::default();
1029 let config = crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1);
1030 let first = builder.add_dense_vector_field_with_config(
1031 "first_embedding",
1032 true,
1033 true,
1034 config.clone(),
1035 );
1036 let second =
1037 builder.add_dense_vector_field_with_config("second_embedding", true, true, config);
1038 let schema = builder.build();
1039 let directory = crate::directories::RamDirectory::new();
1040 let mut metadata = IndexMetadata::new(schema.clone());
1041 metadata.init_field(first.0, VectorIndexType::IvfTq);
1042 metadata.init_field(second.0, VectorIndexType::IvfTq);
1043 metadata.mark_field_built(first.0, 10, 1, "field_0_centroids.bin".into(), None);
1044 metadata.mark_field_built(second.0, 10, 1, "field_1_centroids.bin".into(), None);
1045 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1046
1047 let error = IndexMetadata::try_load_trained_from_fields(
1048 &metadata.vector_fields,
1049 &schema,
1050 &directory,
1051 )
1052 .await
1053 .err()
1054 .expect("missing artifact must fail the complete load")
1055 .to_string();
1056 assert!(error.contains("field_1_centroids.bin"), "{error}");
1057 assert!(error.contains("field 1"), "{error}");
1058 }
1059
1060 #[tokio::test]
1061 async fn index_open_fails_closed_when_built_artifact_is_missing() {
1062 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1063 let directory = crate::directories::RamDirectory::new();
1064 let mut metadata = IndexMetadata::new(schema);
1065 metadata.init_field(field.0, VectorIndexType::IvfTq);
1066 metadata.mark_field_built(field.0, 10, 1, "missing_centroids.bin".into(), None);
1067 metadata.save(&directory).await.unwrap();
1068
1069 let error = match crate::index::Index::open(directory, crate::index::IndexConfig::default())
1070 .await
1071 {
1072 Ok(_) => panic!("Index::open accepted a Built field with no artifact"),
1073 Err(error) => error.to_string(),
1074 };
1075 assert!(error.contains("missing_centroids.bin"), "{error}");
1076 }
1077
1078 #[tokio::test]
1079 async fn ivf_tq_built_state_rejects_a_codebook_file() {
1080 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1081 let directory = crate::directories::RamDirectory::new();
1082 let mut metadata = IndexMetadata::new(schema.clone());
1083 metadata.init_field(field.0, VectorIndexType::IvfTq);
1084 metadata.mark_field_built(
1085 field.0,
1086 10,
1087 1,
1088 "field_0_centroids.bin".into(),
1089 Some("field_0_codebook.bin".into()),
1090 );
1091 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1092
1093 let error = IndexMetadata::try_load_trained_from_fields(
1094 &metadata.vector_fields,
1095 &schema,
1096 &directory,
1097 )
1098 .await
1099 .err()
1100 .expect("IVF-TQ Built state with a codebook file must fail")
1101 .to_string();
1102 assert!(error.contains("codebook"), "{error}");
1103 }
1104
1105 #[tokio::test]
1106 async fn legacy_ivf_tq_centroid_generation_is_rejected_while_loading() {
1107 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1108 let directory = crate::directories::RamDirectory::new();
1109 let mut metadata = IndexMetadata::new(schema.clone());
1110 metadata.init_field(field.0, VectorIndexType::IvfTq);
1111 metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1112 let mut legacy = test_centroids();
1113 legacy.version = 7;
1114 write_bincode(&directory, "field_0_centroids.bin", &legacy).await;
1115
1116 let error = IndexMetadata::try_load_trained_from_fields(
1117 &metadata.vector_fields,
1118 &schema,
1119 &directory,
1120 )
1121 .await
1122 .err()
1123 .expect("legacy IVF-TQ centroid state must fail while loading")
1124 .to_string();
1125 assert!(error.contains("unsupported legacy generation"), "{error}");
1126 assert!(error.contains("rebuild the index"), "{error}");
1127 }
1128
1129 #[tokio::test]
1130 async fn legacy_ivf_pq_trained_field_fails_with_actionable_error() {
1131 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1135 let directory = crate::directories::RamDirectory::new();
1136 let mut metadata = IndexMetadata::new(schema.clone());
1137 metadata.init_field(field.0, VectorIndexType::IvfTq);
1138 metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1139 metadata
1142 .vector_fields
1143 .get_mut(&field.0)
1144 .expect("field initialized")
1145 .index_type = VectorFieldIndexType::Float(VectorIndexType::IvfPq);
1146 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1147
1148 let error = IndexMetadata::try_load_trained_from_fields(
1149 &metadata.vector_fields,
1150 &schema,
1151 &directory,
1152 )
1153 .await
1154 .err()
1155 .expect("legacy IVF-PQ trained state must fail loudly")
1156 .to_string();
1157 assert!(error.contains("no longer"), "{error}");
1158 assert!(error.contains("ivf_tq"), "{error}");
1159 }
1160
1161 #[tokio::test]
1162 async fn requested_cluster_count_accepts_a_quality_clamped_artifact() {
1163 let mut builder = crate::dsl::SchemaBuilder::default();
1164 let field = builder.add_dense_vector_field_with_config(
1165 "embedding",
1166 true,
1167 true,
1168 crate::dsl::DenseVectorConfig::ivf_tq(2, Some(4), 1),
1169 );
1170 let schema = builder.build();
1171 let directory = crate::directories::RamDirectory::new();
1172 let mut metadata = IndexMetadata::new(schema.clone());
1173 metadata.init_field(field.0, VectorIndexType::IvfTq);
1174 metadata.mark_field_built(field.0, 1, 4, "field_0_centroids.bin".into(), None);
1175 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
1176
1177 let trained = IndexMetadata::try_load_trained_from_fields(
1178 &metadata.vector_fields,
1179 &schema,
1180 &directory,
1181 )
1182 .await
1183 .unwrap()
1184 .unwrap();
1185 assert_eq!(trained.centroids[&field.0].num_clusters, 1);
1186 }
1187
1188 #[tokio::test]
1189 async fn trained_artifact_loader_rejects_trailing_data() {
1190 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1191 let directory = crate::directories::RamDirectory::new();
1192 let mut metadata = IndexMetadata::new(schema.clone());
1193 metadata.init_field(field.0, VectorIndexType::IvfTq);
1194 metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1195 let mut bytes =
1196 bincode::serde::encode_to_vec(test_centroids(), bincode::config::standard()).unwrap();
1197 bytes.extend_from_slice(&[0xaa, 0xbb]);
1198 directory
1199 .write(Path::new("field_0_centroids.bin"), &bytes)
1200 .await
1201 .unwrap();
1202
1203 let error = IndexMetadata::try_load_trained_from_fields(
1204 &metadata.vector_fields,
1205 &schema,
1206 &directory,
1207 )
1208 .await
1209 .err()
1210 .expect("trailing artifact bytes must fail validation")
1211 .to_string();
1212 assert!(error.contains("trailing bytes"), "{error}");
1213 }
1214
1215 #[test]
1216 fn trained_artifact_size_limit_rejects_before_reading() {
1217 let error = validate_trained_artifact_size(
1218 3,
1219 "centroids",
1220 "field_3_centroids.bin",
1221 MAX_TRAINED_ARTIFACT_BYTES as u64 + 1,
1222 )
1223 .unwrap_err()
1224 .to_string();
1225 assert!(error.contains("exceeding"), "{error}");
1226 assert!(error.contains("field 3"), "{error}");
1227 }
1228
1229 #[tokio::test]
1230 async fn trained_artifact_decode_limit_rejects_forged_collection_length() {
1231 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1232 let directory = crate::directories::RamDirectory::new();
1233 let mut metadata = IndexMetadata::new(schema.clone());
1234 metadata.init_field(field.0, VectorIndexType::IvfTq);
1235 metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1236
1237 let mut bytes = vec![1, 2, 253];
1242 bytes.extend_from_slice(&u64::MAX.to_le_bytes());
1243 directory
1244 .write(Path::new("field_0_centroids.bin"), &bytes)
1245 .await
1246 .unwrap();
1247
1248 let error = IndexMetadata::try_load_trained_from_fields(
1249 &metadata.vector_fields,
1250 &schema,
1251 &directory,
1252 )
1253 .await
1254 .err()
1255 .expect("forged collection length must fail the bounded decoder")
1256 .to_string();
1257 assert!(error.contains("failed to deserialize"), "{error}");
1258 }
1259
1260 #[test]
1261 fn test_metadata_segments() {
1262 let mut meta = IndexMetadata::new(test_schema());
1263 meta.add_segment("abc123".to_string(), 50);
1264 meta.add_segment("def456".to_string(), 100);
1265 assert_eq!(meta.segment_metas.len(), 2);
1266 assert_eq!(meta.segment_doc_count("abc123"), Some(50));
1267 assert_eq!(meta.segment_doc_count("def456"), Some(100));
1268
1269 meta.add_segment("abc123".to_string(), 75);
1271 assert_eq!(meta.segment_metas.len(), 2);
1272 assert_eq!(meta.segment_doc_count("abc123"), Some(75));
1273
1274 meta.remove_segment("abc123");
1275 assert_eq!(meta.segment_metas.len(), 1);
1276 assert!(meta.has_segment("def456"));
1277 assert!(!meta.has_segment("abc123"));
1278 }
1279
1280 #[test]
1281 fn test_mark_field_built() {
1282 let mut meta = IndexMetadata::new(test_schema());
1283 meta.init_field(0, VectorIndexType::IvfTq);
1284 meta.total_vectors = 10000;
1285
1286 assert!(!meta.is_field_built(0));
1287
1288 meta.mark_field_built(0, 10000, 256, "field_0_centroids.bin".to_string(), None);
1289
1290 assert!(meta.is_field_built(0));
1291 let field = meta.get_field_meta(0).unwrap();
1292 assert_eq!(
1293 field.centroids_file.as_deref(),
1294 Some("field_0_centroids.bin")
1295 );
1296 }
1297
1298 #[test]
1299 fn scann_metadata_persists_generation_and_fingerprint_and_defaults_old_json() {
1300 let mut meta = IndexMetadata::new(test_schema());
1301 meta.init_field(3, VectorIndexType::Scann);
1302 meta.mark_scann_field_built(
1303 3,
1304 100_000,
1305 1_000,
1306 "field_3_scann_17.bin".to_string(),
1307 17,
1308 0xdecafbad,
1309 )
1310 .unwrap();
1311
1312 let bytes = meta.serialize_to_bytes().unwrap();
1313 let decoded: IndexMetadata = serde_json::from_slice(&bytes).unwrap();
1314 let field = decoded.get_field_meta(3).unwrap();
1315 assert_eq!(field.artifact_generation, Some(17));
1316 assert_eq!(field.artifact_id, Some(0xdecafbad));
1317
1318 let mut legacy_json = serde_json::to_value(&decoded).unwrap();
1319 legacy_json["vector_fields"]["3"]
1320 .as_object_mut()
1321 .unwrap()
1322 .remove("artifact_generation");
1323 legacy_json["vector_fields"]["3"]
1324 .as_object_mut()
1325 .unwrap()
1326 .remove("artifact_id");
1327 let legacy: IndexMetadata = serde_json::from_value(legacy_json).unwrap();
1328 let legacy_field = legacy.get_field_meta(3).unwrap();
1329 assert_eq!(legacy_field.artifact_generation, None);
1330 assert_eq!(legacy_field.artifact_id, None);
1331 }
1332
1333 #[test]
1334 fn scann_metadata_refuses_zero_or_non_scann_generation() {
1335 let mut meta = IndexMetadata::new(test_schema());
1336 meta.init_field(0, VectorIndexType::Scann);
1337 assert!(
1338 meta.mark_scann_field_built(0, 100_000, 1_000, "artifact.bin".into(), 0, 1)
1339 .is_err()
1340 );
1341 meta.init_field(1, VectorIndexType::IvfTq);
1342 assert!(
1343 meta.mark_scann_field_built(1, 100_000, 1_000, "artifact.bin".into(), 1, 2)
1344 .is_err()
1345 );
1346 }
1347
1348 #[test]
1349 fn scann_autopilot_accepts_resolved_billion_scale_three_level_geometry() {
1350 assert!(scann_explicit_geometry_matches(None, None, 10_000_000, 3));
1351 assert!(!scann_explicit_geometry_matches(
1352 Some(1_000_000),
1353 None,
1354 10_000_000,
1355 3
1356 ));
1357 assert!(!scann_explicit_geometry_matches(
1358 None,
1359 Some(1),
1360 10_000_000,
1361 3
1362 ));
1363 }
1364
1365 #[test]
1366 fn total_vectors_is_aggregate_of_built_field_counts() {
1367 let mut meta = IndexMetadata::new(test_schema());
1368 meta.init_field(7, VectorIndexType::IvfTq);
1369 meta.init_field(3, VectorIndexType::IvfTq);
1370
1371 meta.mark_field_built(7, 400, 20, "field_7_centroids.bin".to_string(), None);
1374 assert_eq!(meta.total_vectors, 400);
1375 meta.mark_field_built(3, 250, 15, "field_3_centroids.bin".to_string(), None);
1376 assert_eq!(meta.total_vectors, 650);
1377
1378 meta.mark_field_built(7, 425, 20, "field_7_centroids.bin".to_string(), None);
1381 assert_eq!(meta.total_vectors, 675);
1382 }
1383
1384 #[test]
1385 fn test_should_build_field() {
1386 let mut meta = IndexMetadata::new(test_schema());
1387 meta.init_field(0, VectorIndexType::IvfTq);
1388
1389 meta.total_vectors = 500;
1391 assert!(!meta.should_build_field(0, 1000));
1392
1393 meta.total_vectors = 1500;
1395 assert!(meta.should_build_field(0, 1000));
1396
1397 meta.mark_field_built(0, 1500, 256, "centroids.bin".to_string(), None);
1399 assert!(!meta.should_build_field(0, 1000));
1400 }
1401
1402 #[test]
1403 fn test_serialization() {
1404 let mut meta = IndexMetadata::new(test_schema());
1405 meta.add_segment("seg1".to_string(), 100);
1406 meta.init_field(0, VectorIndexType::IvfTq);
1407 meta.total_vectors = 5000;
1408
1409 let json = serde_json::to_string_pretty(&meta).unwrap();
1410 let loaded: IndexMetadata = serde_json::from_str(&json).unwrap();
1411
1412 assert_eq!(loaded.segment_ids().len(), meta.segment_ids().len());
1413 assert_eq!(loaded.segment_doc_count("seg1"), Some(100));
1414 assert_eq!(loaded.total_vectors, meta.total_vectors);
1415 assert!(loaded.vector_fields.contains_key(&0));
1416 }
1417
1418 #[test]
1419 fn old_metadata_defaults_the_bp_retry_counter() {
1420 let mut meta = IndexMetadata::new(test_schema());
1421 meta.add_segment("legacy".to_string(), 10);
1422 let mut json = serde_json::to_value(&meta).unwrap();
1423 json["segment_metas"]["legacy"]
1424 .as_object_mut()
1425 .unwrap()
1426 .remove("bp_unconverged_passes");
1427
1428 let loaded: IndexMetadata = serde_json::from_value(json).unwrap();
1429 assert_eq!(loaded.segment_metas["legacy"].bp_unconverged_passes, 0);
1430 }
1431
1432 #[test]
1433 fn test_merged_segment_lineage() {
1434 let mut meta = IndexMetadata::new(test_schema());
1435 meta.add_segment("a".to_string(), 50);
1436 meta.add_segment("b".to_string(), 75);
1437
1438 assert_eq!(meta.segment_metas["a"].generation, 0);
1440 assert!(meta.segment_metas["a"].ancestors.is_empty());
1441
1442 meta.add_merged_segment(
1444 "c".to_string(),
1445 125,
1446 vec!["a".to_string(), "b".to_string()],
1447 1,
1448 false,
1449 true,
1450 );
1451 assert_eq!(meta.segment_metas["c"].generation, 1);
1452 assert_eq!(meta.segment_metas["c"].ancestors, vec!["a", "b"]);
1453 assert_eq!(meta.segment_doc_count("c"), Some(125));
1454
1455 meta.add_segment("d".to_string(), 30);
1457 meta.add_merged_segment(
1458 "e".to_string(),
1459 155,
1460 vec!["c".to_string(), "d".to_string()],
1461 2,
1462 false,
1463 true,
1464 );
1465 assert_eq!(meta.segment_metas["e"].generation, 2);
1466 }
1467}