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