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}
125
126#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct IndexMetadata {
129 pub version: u32,
131 pub schema: Schema,
133 #[serde(default)]
136 pub segment_metas: HashMap<String, SegmentMetaInfo>,
137 #[serde(default)]
139 pub vector_fields: HashMap<u32, FieldVectorMeta>,
140 #[serde(default)]
147 pub total_vectors: usize,
148}
149
150impl IndexMetadata {
151 pub fn new(schema: Schema) -> Self {
153 Self {
154 version: INDEX_META_FORMAT_VERSION,
155 schema,
156 segment_metas: HashMap::new(),
157 vector_fields: HashMap::new(),
158 total_vectors: 0,
159 }
160 }
161
162 pub fn segment_ids(&self) -> Vec<String> {
164 let mut ids: Vec<String> = self.segment_metas.keys().cloned().collect();
165 ids.sort();
166 ids
167 }
168
169 pub fn add_segment(&mut self, segment_id: String, num_docs: u32) {
171 self.segment_metas.insert(
172 segment_id,
173 SegmentMetaInfo {
174 num_docs,
175 ancestors: Vec::new(),
176 generation: 0,
177 reordered: false,
178 bp_converged: true,
179 bp_unconverged_passes: 0,
180 },
181 );
182 }
183
184 pub fn add_merged_segment(
186 &mut self,
187 segment_id: String,
188 num_docs: u32,
189 ancestors: Vec<String>,
190 generation: u32,
191 reordered: bool,
192 bp_converged: bool,
193 ) {
194 self.add_segment_meta(
195 segment_id,
196 SegmentMetaInfo {
197 num_docs,
198 ancestors,
199 generation,
200 reordered,
201 bp_converged,
202 bp_unconverged_passes: 0,
203 },
204 );
205 }
206
207 pub(crate) fn add_segment_meta(&mut self, segment_id: String, info: SegmentMetaInfo) {
211 self.segment_metas.insert(segment_id, info);
212 }
213
214 pub fn remove_segment(&mut self, segment_id: &str) {
216 self.segment_metas.remove(segment_id);
217 }
218
219 pub fn has_segment(&self, segment_id: &str) -> bool {
221 self.segment_metas.contains_key(segment_id)
222 }
223
224 pub fn segment_doc_count(&self, segment_id: &str) -> Option<u32> {
226 self.segment_metas.get(segment_id).map(|m| m.num_docs)
227 }
228
229 pub fn is_field_built(&self, field_id: u32) -> bool {
231 self.vector_fields
232 .get(&field_id)
233 .map(|f| matches!(f.state, VectorIndexState::Built { .. }))
234 .unwrap_or(false)
235 }
236
237 pub fn get_field_meta(&self, field_id: u32) -> Option<&FieldVectorMeta> {
239 self.vector_fields.get(&field_id)
240 }
241
242 pub fn init_field(&mut self, field_id: u32, index_type: impl Into<VectorFieldIndexType>) {
244 let index_type = index_type.into();
245 self.vector_fields
246 .entry(field_id)
247 .or_insert(FieldVectorMeta {
248 field_id,
249 index_type,
250 state: VectorIndexState::Flat,
251 centroids_file: None,
252 codebook_file: None,
253 });
254 }
255
256 pub fn mark_field_built(
258 &mut self,
259 field_id: u32,
260 vector_count: usize,
261 num_clusters: usize,
262 centroids_file: String,
263 codebook_file: Option<String>,
264 ) {
265 if let Some(field) = self.vector_fields.get_mut(&field_id) {
266 field.state = VectorIndexState::Built {
267 vector_count,
268 num_clusters,
269 };
270 field.centroids_file = Some(centroids_file);
271 field.codebook_file = codebook_file;
272 self.refresh_total_vectors();
273 }
274 }
275
276 pub(crate) fn refresh_total_vectors(&mut self) {
282 self.total_vectors = self
283 .vector_fields
284 .values()
285 .filter_map(|field| match field.state {
286 VectorIndexState::Built { vector_count, .. } => Some(vector_count),
287 VectorIndexState::Flat => None,
288 })
289 .fold(0usize, usize::saturating_add);
290 }
291
292 pub fn should_build_field(&self, field_id: u32, threshold: usize) -> bool {
294 if self.is_field_built(field_id) {
296 return false;
297 }
298 self.total_vectors >= threshold
300 }
301
302 pub async fn load<D: crate::directories::Directory>(dir: &D) -> Result<Self> {
307 let path = Path::new(INDEX_META_FILENAME);
308 match dir.open_read(path).await {
309 Ok(slice) => {
310 let bytes = slice.read_bytes().await?;
311 Self::deserialize_versioned(bytes.as_slice())
312 }
313 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
314 let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
316 let slice = dir.open_read(tmp_path).await?;
317 let bytes = slice.read_bytes().await?;
318 let meta = Self::deserialize_versioned(bytes.as_slice())?;
319 log::warn!("Recovered metadata from temp file (previous crash during save)");
320 Ok(meta)
321 }
322 Err(e) => Err(Error::Io(e)),
323 }
324 }
325
326 fn deserialize_versioned(bytes: &[u8]) -> Result<Self> {
330 let meta: Self =
331 serde_json::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))?;
332 crate::dsl::reject_removed_vector_index_types(&meta.schema).map_err(Error::Schema)?;
333 if meta.version != INDEX_META_FORMAT_VERSION {
334 return Err(Error::Corruption(format!(
335 "metadata.json format version {} is incompatible with required version {}; \
336 rebuild and republish the index with this Hermes version",
337 meta.version, INDEX_META_FORMAT_VERSION
338 )));
339 }
340 Ok(meta)
341 }
342
343 pub async fn save<D: crate::directories::DirectoryWriter>(&self, dir: &D) -> Result<()> {
348 let bytes = self.serialize_to_bytes()?;
349 Self::save_bytes(dir, &bytes).await
350 }
351
352 pub fn serialize_to_bytes(&self) -> Result<Vec<u8>> {
355 serde_json::to_vec_pretty(self).map_err(|e| Error::Serialization(e.to_string()))
356 }
357
358 pub async fn save_bytes<D: crate::directories::DirectoryWriter>(
363 dir: &D,
364 bytes: &[u8],
365 ) -> Result<()> {
366 let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
367 let final_path = Path::new(INDEX_META_FILENAME);
368 let mut writer = dir.streaming_writer(tmp_path).await.map_err(Error::Io)?;
373 writer.write_all(bytes).map_err(Error::Io)?;
374 writer.finish().map_err(Error::Io)?;
375 dir.rename(tmp_path, final_path).await.map_err(Error::Io)?;
382 if let Err(error) = dir.sync().await {
383 log::error!(
384 "[metadata] directory fsync failed after committed rename: {}. \
385 Continuing with the renamed generation; crash durability is not guaranteed",
386 error,
387 );
388 }
389 Ok(())
390 }
391
392 #[cfg_attr(not(feature = "native"), allow(dead_code))]
394 pub(crate) async fn try_load_trained_from_fields<D: crate::directories::Directory>(
395 vector_fields: &HashMap<u32, FieldVectorMeta>,
396 schema: &Schema,
397 dir: &D,
398 ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
399 Self::load_trained_from_fields_impl(vector_fields, schema, dir).await
400 }
401
402 async fn load_trained_from_fields_impl<D: crate::directories::Directory>(
412 vector_fields: &HashMap<u32, FieldVectorMeta>,
413 schema: &Schema,
414 dir: &D,
415 ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
416 use std::sync::Arc;
417
418 let mut centroids = rustc_hash::FxHashMap::default();
419 let mut binary_quantizers = rustc_hash::FxHashMap::default();
420
421 let mut built_fields: Vec<_> = vector_fields
422 .iter()
423 .filter(|(_, meta)| matches!(meta.state, VectorIndexState::Built { .. }))
424 .collect();
425 built_fields.sort_unstable_by_key(|(field_id, _)| **field_id);
426
427 log::debug!(
428 "[trained] loading trained structures, dense_vector_fields={:?}",
429 vector_fields.keys().collect::<Vec<_>>()
430 );
431
432 for (field_id, field_meta) in built_fields {
433 log::debug!(
434 "[trained] field {} state={:?} centroids_file={:?} codebook_file={:?}",
435 field_id,
436 field_meta.state,
437 field_meta.centroids_file,
438 field_meta.codebook_file,
439 );
440 if field_meta.field_id != *field_id {
441 return Err(Error::Corruption(format!(
442 "trained vector metadata key {field_id} contains field_id {}",
443 field_meta.field_id
444 )));
445 }
446
447 let expected_clusters = match field_meta.state {
448 VectorIndexState::Built { num_clusters, .. } if num_clusters > 0 => num_clusters,
449 VectorIndexState::Built { .. } => {
450 return Err(Error::Corruption(format!(
451 "trained vector metadata field {field_id} has zero clusters"
452 )));
453 }
454 VectorIndexState::Flat => unreachable!("built_fields contains only Built entries"),
455 };
456
457 let centroids_file = field_meta.centroids_file.as_deref().ok_or_else(|| {
458 Error::Corruption(format!(
459 "trained vector metadata field {field_id} is Built but has no centroids_file"
460 ))
461 })?;
462 match field_meta.index_type {
463 VectorFieldIndexType::Float(VectorIndexType::IvfPq) => {
464 return Err(Error::Corruption(format!(
465 "field {field_id} was trained as IVF-PQ, which is no longer \
466 supported; recreate the index with `ivf_tq` and reindex \
467 (docs/turboquant-quantization.md)"
468 )));
469 }
470 VectorFieldIndexType::Float(index_type @ VectorIndexType::IvfTq) => {
471 let entry = schema
472 .get_field_entry(crate::dsl::Field(*field_id))
473 .ok_or_else(|| {
474 Error::Corruption(format!(
475 "trained vector metadata references missing field {field_id}"
476 ))
477 })?;
478 let schema_config = entry
479 .dense_vector_config
480 .as_ref()
481 .filter(|_| entry.field_type == crate::dsl::FieldType::DenseVector)
482 .ok_or_else(|| {
483 Error::Corruption(format!(
484 "trained vector metadata field {field_id} is not a float dense field"
485 ))
486 })?;
487 if schema_config.index_type != index_type {
488 return Err(Error::Corruption(format!(
489 "trained vector metadata field {field_id} uses {index_type:?}, schema requires {:?}",
490 schema_config.index_type
491 )));
492 }
493 let c: crate::structures::CoarseCentroids =
494 load_trained_artifact(dir, *field_id, "centroids", centroids_file).await?;
495 let expected_dim = schema_config.dim;
496 let actual_clusters = c.num_clusters as usize;
497 let expected_values =
498 actual_clusters.checked_mul(expected_dim).ok_or_else(|| {
499 Error::Corruption(format!(
500 "trained centroid dimensions overflow for field {field_id}"
501 ))
502 })?;
503 if actual_clusters == 0
504 || actual_clusters > expected_clusters
505 || c.dim == 0
506 || c.dim != expected_dim
507 || c.centroids.len() != expected_values
508 || c.centroids.iter().any(|value| !value.is_finite())
509 {
510 return Err(Error::Corruption(format!(
511 "trained centroids for field {field_id} do not match metadata/schema"
512 )));
513 }
514 c.validate_routing(schema_config.ivf_routing)
515 .map_err(|error| {
516 Error::Corruption(format!(
517 "invalid trained centroid routing for field {field_id}: {error}"
518 ))
519 })?;
520 let _ = index_type;
523 if field_meta.codebook_file.is_some() {
524 return Err(Error::Corruption(format!(
525 "trained IVF-TQ field {field_id} unexpectedly references a codebook file"
526 )));
527 }
528 centroids.insert(*field_id, Arc::new(c));
529 }
530 VectorFieldIndexType::Binary(BinaryIndexType::Ivf) => {
531 let entry = schema
532 .get_field_entry(crate::dsl::Field(*field_id))
533 .ok_or_else(|| {
534 Error::Corruption(format!(
535 "trained vector metadata references missing field {field_id}"
536 ))
537 })?;
538 let schema_config = entry
539 .binary_dense_vector_config
540 .as_ref()
541 .filter(|config| {
542 entry.field_type == crate::dsl::FieldType::BinaryDenseVector
543 && config.index_type == BinaryIndexType::Ivf
544 })
545 .ok_or_else(|| {
546 Error::Corruption(format!(
547 "trained vector metadata field {field_id} is not a binary IVF field"
548 ))
549 })?;
550 let quantizer: crate::structures::BinaryCoarseQuantizer =
551 load_trained_artifact(dir, *field_id, "binary centroids", centroids_file)
552 .await?;
553 quantizer.validate().map_err(|error| {
554 Error::Corruption(format!(
555 "invalid binary coarse quantizer for field {field_id}: {error}"
556 ))
557 })?;
558 let actual_clusters = quantizer.num_clusters as usize;
559 if actual_clusters > expected_clusters
560 || schema_config.dim != quantizer.dim_bits
561 {
562 return Err(Error::Corruption(format!(
563 "binary coarse quantizer for field {field_id} does not match metadata/schema"
564 )));
565 }
566 quantizer
567 .validate_routing(schema_config.ivf_routing)
568 .map_err(|error| {
569 Error::Corruption(format!(
570 "invalid binary centroid routing for field {field_id}: {error}"
571 ))
572 })?;
573 binary_quantizers.insert(*field_id, Arc::new(quantizer));
574 }
575 unsupported => {
576 return Err(Error::Corruption(format!(
577 "field {field_id} is Built for {unsupported:?}, which has no global IVF artifacts"
578 )));
579 }
580 }
581 }
582
583 if centroids.is_empty() && binary_quantizers.is_empty() {
584 Ok(None)
585 } else {
586 let trained = crate::segment::TrainedVectorStructures {
587 #[cfg(feature = "native")]
588 _ann_pins: Default::default(),
589 centroids,
590 binary_quantizers,
591 };
592 #[cfg(feature = "native")]
593 let trained = {
594 let mut trained = trained;
595 trained.pin_ann_structures(crate::segment::pin::pin_policy());
596 trained
597 };
598 Ok(Some(trained))
599 }
600 }
601}
602
603fn validate_trained_artifact_path(field_id: u32, kind: &str, filename: &str) -> Result<()> {
604 use std::path::Component;
605
606 let path = Path::new(filename);
607 if filename.is_empty()
608 || path.is_absolute()
609 || path.components().any(|component| {
610 matches!(
611 component,
612 Component::ParentDir | Component::RootDir | Component::Prefix(_)
613 )
614 })
615 {
616 return Err(Error::Corruption(format!(
617 "trained {kind} path for field {field_id} is not a safe relative path: '{filename}'"
618 )));
619 }
620 Ok(())
621}
622
623async fn load_trained_artifact<T, D>(
624 dir: &D,
625 field_id: u32,
626 kind: &str,
627 filename: &str,
628) -> Result<T>
629where
630 T: serde::de::DeserializeOwned,
631 D: crate::directories::Directory,
632{
633 validate_trained_artifact_path(field_id, kind, filename)?;
634 let path = Path::new(filename);
635 let file_size = dir.file_size(path).await.map_err(|error| {
636 Error::Corruption(format!(
637 "failed to stat trained {kind} '{filename}' for field {field_id}: {error}"
638 ))
639 })?;
640 validate_trained_artifact_size(field_id, kind, filename, file_size)?;
641 let slice = dir.open_read(path).await.map_err(|error| {
642 Error::Corruption(format!(
643 "failed to open trained {kind} '{filename}' for field {field_id}: {error}"
644 ))
645 })?;
646 validate_trained_artifact_size(field_id, kind, filename, slice.len())?;
647 let bytes = slice.read_bytes().await.map_err(|error| {
648 Error::Corruption(format!(
649 "failed to read trained {kind} '{filename}' for field {field_id}: {error}"
650 ))
651 })?;
652 let (artifact, consumed) = bincode::serde::decode_from_slice::<T, _>(
653 bytes.as_slice(),
654 bincode::config::standard().with_limit::<MAX_TRAINED_ARTIFACT_BYTES>(),
655 )
656 .map_err(|error| {
657 Error::Corruption(format!(
658 "failed to deserialize trained {kind} '{filename}' for field {field_id}: {error}"
659 ))
660 })?;
661 if consumed != bytes.len() {
662 return Err(Error::Corruption(format!(
663 "trained {kind} '{filename}' for field {field_id} has {} trailing bytes",
664 bytes.len() - consumed
665 )));
666 }
667 Ok(artifact)
668}
669
670fn validate_trained_artifact_size(
671 field_id: u32,
672 kind: &str,
673 filename: &str,
674 file_size: u64,
675) -> Result<()> {
676 if file_size > MAX_TRAINED_ARTIFACT_BYTES as u64 {
677 return Err(Error::Corruption(format!(
678 "trained {kind} '{filename}' for field {field_id} is {file_size} bytes, \
679 exceeding the {MAX_TRAINED_ARTIFACT_BYTES}-byte safety limit"
680 )));
681 }
682 Ok(())
683}
684
685#[cfg(test)]
686mod tests {
687 use super::*;
688 use crate::directories::DirectoryWriter;
689
690 #[derive(Clone, Default)]
691 struct SyncFailDirectory(crate::directories::RamDirectory);
692
693 #[async_trait::async_trait]
694 impl crate::directories::Directory for SyncFailDirectory {
695 async fn exists(&self, path: &Path) -> std::io::Result<bool> {
696 self.0.exists(path).await
697 }
698
699 async fn file_size(&self, path: &Path) -> std::io::Result<u64> {
700 self.0.file_size(path).await
701 }
702
703 async fn open_read(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
704 self.0.open_read(path).await
705 }
706
707 async fn read_range(
708 &self,
709 path: &Path,
710 range: std::ops::Range<u64>,
711 ) -> std::io::Result<crate::directories::OwnedBytes> {
712 self.0.read_range(path, range).await
713 }
714
715 async fn list_files(&self, prefix: &Path) -> std::io::Result<Vec<std::path::PathBuf>> {
716 self.0.list_files(prefix).await
717 }
718
719 async fn open_lazy(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
720 self.0.open_lazy(path).await
721 }
722 }
723
724 #[async_trait::async_trait]
725 impl crate::directories::DirectoryWriter for SyncFailDirectory {
726 async fn write(&self, path: &Path, data: &[u8]) -> std::io::Result<()> {
727 self.0.write(path, data).await
728 }
729
730 async fn delete(&self, path: &Path) -> std::io::Result<()> {
731 self.0.delete(path).await
732 }
733
734 async fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
735 self.0.rename(from, to).await
736 }
737
738 async fn sync(&self) -> std::io::Result<()> {
739 Err(std::io::Error::other("injected directory fsync failure"))
740 }
741
742 async fn streaming_writer(
743 &self,
744 path: &Path,
745 ) -> std::io::Result<Box<dyn crate::directories::StreamingWriter>> {
746 self.0.streaming_writer(path).await
747 }
748 }
749
750 fn test_schema() -> Schema {
751 Schema::default()
752 }
753
754 fn dense_schema(index_type: VectorIndexType) -> (Schema, crate::dsl::Field) {
755 let mut builder = crate::dsl::SchemaBuilder::default();
756 let config = match index_type {
757 VectorIndexType::IvfTq => crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1),
758 other => panic!("unsupported trained test index type: {other:?}"),
759 };
760 let field = builder.add_dense_vector_field_with_config("embedding", true, true, config);
761 (builder.build(), field)
762 }
763
764 fn test_centroids() -> crate::structures::CoarseCentroids {
765 crate::structures::CoarseCentroids {
766 num_clusters: 1,
767 dim: 2,
768 centroids: vec![0.25, 0.75],
769 version: 7,
770 soar_config: None,
771 routing_index: None,
772 }
773 }
774
775 async fn write_bincode(
776 directory: &crate::directories::RamDirectory,
777 filename: &str,
778 value: &impl serde::Serialize,
779 ) {
780 let bytes = bincode::serde::encode_to_vec(value, bincode::config::standard()).unwrap();
781 directory.write(Path::new(filename), &bytes).await.unwrap();
782 }
783
784 #[test]
785 fn test_metadata_init() {
786 let mut meta = IndexMetadata::new(test_schema());
787 assert_eq!(meta.total_vectors, 0);
788 assert!(meta.segment_metas.is_empty());
789 assert!(!meta.is_field_built(0));
790
791 meta.init_field(0, VectorIndexType::IvfTq);
792 assert!(!meta.is_field_built(0));
793 assert!(meta.vector_fields.contains_key(&0));
794 }
795
796 #[tokio::test]
797 async fn load_refuses_metadata_stamped_with_a_newer_format_version() {
798 let directory = crate::directories::RamDirectory::new();
799 let mut metadata = IndexMetadata::new(test_schema());
800 metadata.version = INDEX_META_FORMAT_VERSION + 1;
801 metadata.save(&directory).await.unwrap();
802
803 let error = IndexMetadata::load(&directory)
804 .await
805 .expect_err("metadata from a newer format version must be refused, not silently pruned")
806 .to_string();
807 assert!(error.contains("version 4"), "{error}");
808 assert!(error.contains("incompatible"), "{error}");
809 }
810
811 #[tokio::test]
812 async fn tmp_recovery_refuses_metadata_stamped_with_a_newer_format_version() {
813 let directory = crate::directories::RamDirectory::new();
814 let mut metadata = IndexMetadata::new(test_schema());
815 metadata.version = INDEX_META_FORMAT_VERSION + 1;
816 let bytes = metadata.serialize_to_bytes().unwrap();
817 directory
819 .write(Path::new(INDEX_META_TMP_FILENAME), &bytes)
820 .await
821 .unwrap();
822
823 let error = IndexMetadata::load(&directory)
824 .await
825 .expect_err("temp-file recovery must apply the same version gate")
826 .to_string();
827 assert!(error.contains("version 4"), "{error}");
828 }
829
830 #[tokio::test]
831 async fn save_treats_post_rename_sync_failure_as_committed() {
832 let directory = SyncFailDirectory::default();
833 let mut metadata = IndexMetadata::new(test_schema());
834 metadata.add_segment("committed".to_string(), 7);
835
836 metadata.save(&directory).await.unwrap();
837
838 let loaded = IndexMetadata::load(&directory).await.unwrap();
839 assert_eq!(loaded.segment_doc_count("committed"), Some(7));
840 }
841
842 #[tokio::test]
843 async fn trained_artifacts_load_only_when_the_complete_built_set_is_valid() {
844 let mut builder = crate::dsl::SchemaBuilder::default();
845 let config = crate::dsl::DenseVectorConfig::ivf_tq(2, Some(1), 1);
846 let first = builder.add_dense_vector_field_with_config(
847 "first_embedding",
848 true,
849 true,
850 config.clone(),
851 );
852 let second =
853 builder.add_dense_vector_field_with_config("second_embedding", true, true, config);
854 let schema = builder.build();
855 let directory = crate::directories::RamDirectory::new();
856 let mut metadata = IndexMetadata::new(schema.clone());
857 metadata.init_field(first.0, VectorIndexType::IvfTq);
858 metadata.init_field(second.0, VectorIndexType::IvfTq);
859 metadata.mark_field_built(first.0, 10, 1, "field_0_centroids.bin".into(), None);
860 metadata.mark_field_built(second.0, 10, 1, "field_1_centroids.bin".into(), None);
861 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
862
863 let error = IndexMetadata::try_load_trained_from_fields(
864 &metadata.vector_fields,
865 &schema,
866 &directory,
867 )
868 .await
869 .err()
870 .expect("missing artifact must fail the complete load")
871 .to_string();
872 assert!(error.contains("field_1_centroids.bin"), "{error}");
873 assert!(error.contains("field 1"), "{error}");
874 }
875
876 #[tokio::test]
877 async fn index_open_fails_closed_when_built_artifact_is_missing() {
878 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
879 let directory = crate::directories::RamDirectory::new();
880 let mut metadata = IndexMetadata::new(schema);
881 metadata.init_field(field.0, VectorIndexType::IvfTq);
882 metadata.mark_field_built(field.0, 10, 1, "missing_centroids.bin".into(), None);
883 metadata.save(&directory).await.unwrap();
884
885 let error = match crate::index::Index::open(directory, crate::index::IndexConfig::default())
886 .await
887 {
888 Ok(_) => panic!("Index::open accepted a Built field with no artifact"),
889 Err(error) => error.to_string(),
890 };
891 assert!(error.contains("missing_centroids.bin"), "{error}");
892 }
893
894 #[tokio::test]
895 async fn ivf_tq_built_state_rejects_a_codebook_file() {
896 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
897 let directory = crate::directories::RamDirectory::new();
898 let mut metadata = IndexMetadata::new(schema.clone());
899 metadata.init_field(field.0, VectorIndexType::IvfTq);
900 metadata.mark_field_built(
901 field.0,
902 10,
903 1,
904 "field_0_centroids.bin".into(),
905 Some("field_0_codebook.bin".into()),
906 );
907 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
908
909 let error = IndexMetadata::try_load_trained_from_fields(
910 &metadata.vector_fields,
911 &schema,
912 &directory,
913 )
914 .await
915 .err()
916 .expect("IVF-TQ Built state with a codebook file must fail")
917 .to_string();
918 assert!(error.contains("codebook"), "{error}");
919 }
920
921 #[tokio::test]
922 async fn legacy_ivf_pq_trained_field_fails_with_actionable_error() {
923 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
927 let directory = crate::directories::RamDirectory::new();
928 let mut metadata = IndexMetadata::new(schema.clone());
929 metadata.init_field(field.0, VectorIndexType::IvfTq);
930 metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
931 metadata
934 .vector_fields
935 .get_mut(&field.0)
936 .expect("field initialized")
937 .index_type = VectorFieldIndexType::Float(VectorIndexType::IvfPq);
938 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
939
940 let error = IndexMetadata::try_load_trained_from_fields(
941 &metadata.vector_fields,
942 &schema,
943 &directory,
944 )
945 .await
946 .err()
947 .expect("legacy IVF-PQ trained state must fail loudly")
948 .to_string();
949 assert!(error.contains("no longer"), "{error}");
950 assert!(error.contains("ivf_tq"), "{error}");
951 }
952
953 #[tokio::test]
954 async fn requested_cluster_count_accepts_a_quality_clamped_artifact() {
955 let mut builder = crate::dsl::SchemaBuilder::default();
956 let field = builder.add_dense_vector_field_with_config(
957 "embedding",
958 true,
959 true,
960 crate::dsl::DenseVectorConfig::ivf_tq(2, Some(4), 1),
961 );
962 let schema = builder.build();
963 let directory = crate::directories::RamDirectory::new();
964 let mut metadata = IndexMetadata::new(schema.clone());
965 metadata.init_field(field.0, VectorIndexType::IvfTq);
966 metadata.mark_field_built(field.0, 1, 4, "field_0_centroids.bin".into(), None);
967 write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
968
969 let trained = IndexMetadata::try_load_trained_from_fields(
970 &metadata.vector_fields,
971 &schema,
972 &directory,
973 )
974 .await
975 .unwrap()
976 .unwrap();
977 assert_eq!(trained.centroids[&field.0].num_clusters, 1);
978 }
979
980 #[tokio::test]
981 async fn trained_artifact_loader_rejects_trailing_data() {
982 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
983 let directory = crate::directories::RamDirectory::new();
984 let mut metadata = IndexMetadata::new(schema.clone());
985 metadata.init_field(field.0, VectorIndexType::IvfTq);
986 metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
987 let mut bytes =
988 bincode::serde::encode_to_vec(test_centroids(), bincode::config::standard()).unwrap();
989 bytes.extend_from_slice(&[0xaa, 0xbb]);
990 directory
991 .write(Path::new("field_0_centroids.bin"), &bytes)
992 .await
993 .unwrap();
994
995 let error = IndexMetadata::try_load_trained_from_fields(
996 &metadata.vector_fields,
997 &schema,
998 &directory,
999 )
1000 .await
1001 .err()
1002 .expect("trailing artifact bytes must fail validation")
1003 .to_string();
1004 assert!(error.contains("trailing bytes"), "{error}");
1005 }
1006
1007 #[test]
1008 fn trained_artifact_size_limit_rejects_before_reading() {
1009 let error = validate_trained_artifact_size(
1010 3,
1011 "centroids",
1012 "field_3_centroids.bin",
1013 MAX_TRAINED_ARTIFACT_BYTES as u64 + 1,
1014 )
1015 .unwrap_err()
1016 .to_string();
1017 assert!(error.contains("exceeding"), "{error}");
1018 assert!(error.contains("field 3"), "{error}");
1019 }
1020
1021 #[tokio::test]
1022 async fn trained_artifact_decode_limit_rejects_forged_collection_length() {
1023 let (schema, field) = dense_schema(VectorIndexType::IvfTq);
1024 let directory = crate::directories::RamDirectory::new();
1025 let mut metadata = IndexMetadata::new(schema.clone());
1026 metadata.init_field(field.0, VectorIndexType::IvfTq);
1027 metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
1028
1029 let mut bytes = vec![1, 2, 253];
1034 bytes.extend_from_slice(&u64::MAX.to_le_bytes());
1035 directory
1036 .write(Path::new("field_0_centroids.bin"), &bytes)
1037 .await
1038 .unwrap();
1039
1040 let error = IndexMetadata::try_load_trained_from_fields(
1041 &metadata.vector_fields,
1042 &schema,
1043 &directory,
1044 )
1045 .await
1046 .err()
1047 .expect("forged collection length must fail the bounded decoder")
1048 .to_string();
1049 assert!(error.contains("failed to deserialize"), "{error}");
1050 }
1051
1052 #[test]
1053 fn test_metadata_segments() {
1054 let mut meta = IndexMetadata::new(test_schema());
1055 meta.add_segment("abc123".to_string(), 50);
1056 meta.add_segment("def456".to_string(), 100);
1057 assert_eq!(meta.segment_metas.len(), 2);
1058 assert_eq!(meta.segment_doc_count("abc123"), Some(50));
1059 assert_eq!(meta.segment_doc_count("def456"), Some(100));
1060
1061 meta.add_segment("abc123".to_string(), 75);
1063 assert_eq!(meta.segment_metas.len(), 2);
1064 assert_eq!(meta.segment_doc_count("abc123"), Some(75));
1065
1066 meta.remove_segment("abc123");
1067 assert_eq!(meta.segment_metas.len(), 1);
1068 assert!(meta.has_segment("def456"));
1069 assert!(!meta.has_segment("abc123"));
1070 }
1071
1072 #[test]
1073 fn test_mark_field_built() {
1074 let mut meta = IndexMetadata::new(test_schema());
1075 meta.init_field(0, VectorIndexType::IvfTq);
1076 meta.total_vectors = 10000;
1077
1078 assert!(!meta.is_field_built(0));
1079
1080 meta.mark_field_built(0, 10000, 256, "field_0_centroids.bin".to_string(), None);
1081
1082 assert!(meta.is_field_built(0));
1083 let field = meta.get_field_meta(0).unwrap();
1084 assert_eq!(
1085 field.centroids_file.as_deref(),
1086 Some("field_0_centroids.bin")
1087 );
1088 }
1089
1090 #[test]
1091 fn total_vectors_is_aggregate_of_built_field_counts() {
1092 let mut meta = IndexMetadata::new(test_schema());
1093 meta.init_field(7, VectorIndexType::IvfTq);
1094 meta.init_field(3, VectorIndexType::IvfTq);
1095
1096 meta.mark_field_built(7, 400, 20, "field_7_centroids.bin".to_string(), None);
1099 assert_eq!(meta.total_vectors, 400);
1100 meta.mark_field_built(3, 250, 15, "field_3_centroids.bin".to_string(), None);
1101 assert_eq!(meta.total_vectors, 650);
1102
1103 meta.mark_field_built(7, 425, 20, "field_7_centroids.bin".to_string(), None);
1106 assert_eq!(meta.total_vectors, 675);
1107 }
1108
1109 #[test]
1110 fn test_should_build_field() {
1111 let mut meta = IndexMetadata::new(test_schema());
1112 meta.init_field(0, VectorIndexType::IvfTq);
1113
1114 meta.total_vectors = 500;
1116 assert!(!meta.should_build_field(0, 1000));
1117
1118 meta.total_vectors = 1500;
1120 assert!(meta.should_build_field(0, 1000));
1121
1122 meta.mark_field_built(0, 1500, 256, "centroids.bin".to_string(), None);
1124 assert!(!meta.should_build_field(0, 1000));
1125 }
1126
1127 #[test]
1128 fn test_serialization() {
1129 let mut meta = IndexMetadata::new(test_schema());
1130 meta.add_segment("seg1".to_string(), 100);
1131 meta.init_field(0, VectorIndexType::IvfTq);
1132 meta.total_vectors = 5000;
1133
1134 let json = serde_json::to_string_pretty(&meta).unwrap();
1135 let loaded: IndexMetadata = serde_json::from_str(&json).unwrap();
1136
1137 assert_eq!(loaded.segment_ids().len(), meta.segment_ids().len());
1138 assert_eq!(loaded.segment_doc_count("seg1"), Some(100));
1139 assert_eq!(loaded.total_vectors, meta.total_vectors);
1140 assert!(loaded.vector_fields.contains_key(&0));
1141 }
1142
1143 #[test]
1144 fn old_metadata_defaults_the_bp_retry_counter() {
1145 let mut meta = IndexMetadata::new(test_schema());
1146 meta.add_segment("legacy".to_string(), 10);
1147 let mut json = serde_json::to_value(&meta).unwrap();
1148 json["segment_metas"]["legacy"]
1149 .as_object_mut()
1150 .unwrap()
1151 .remove("bp_unconverged_passes");
1152
1153 let loaded: IndexMetadata = serde_json::from_value(json).unwrap();
1154 assert_eq!(loaded.segment_metas["legacy"].bp_unconverged_passes, 0);
1155 }
1156
1157 #[test]
1158 fn test_merged_segment_lineage() {
1159 let mut meta = IndexMetadata::new(test_schema());
1160 meta.add_segment("a".to_string(), 50);
1161 meta.add_segment("b".to_string(), 75);
1162
1163 assert_eq!(meta.segment_metas["a"].generation, 0);
1165 assert!(meta.segment_metas["a"].ancestors.is_empty());
1166
1167 meta.add_merged_segment(
1169 "c".to_string(),
1170 125,
1171 vec!["a".to_string(), "b".to_string()],
1172 1,
1173 false,
1174 true,
1175 );
1176 assert_eq!(meta.segment_metas["c"].generation, 1);
1177 assert_eq!(meta.segment_metas["c"].ancestors, vec!["a", "b"]);
1178 assert_eq!(meta.segment_doc_count("c"), Some(125));
1179
1180 meta.add_segment("d".to_string(), 30);
1182 meta.add_merged_segment(
1183 "e".to_string(),
1184 155,
1185 vec!["c".to_string(), "d".to_string()],
1186 2,
1187 false,
1188 true,
1189 );
1190 assert_eq!(meta.segment_metas["e"].generation, 2);
1191 }
1192}