From 367ed241dfa62b86fcc3b43a5a79a2c47a64cd9e Mon Sep 17 00:00:00 2001 From: mrorigo Date: Sat, 3 Oct 2026 23:22:06 +0200 Subject: [PATCH] Fix incremental updates and compact stale segments --- src/index.rs | 36 +++++++++++++++++- src/storage.rs | 30 +++++++-------- tests/index_tests.rs | 91 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 141 insertions(+), 16 deletions(-) diff --git a/src/index.rs b/src/index.rs index d7cfc0b..65994db 100644 --- a/src/index.rs +++ b/src/index.rs @@ -58,6 +58,7 @@ pub struct Index { } impl Index { + const COMPACT_SEGMENT_THRESHOLD: usize = 8; /// Opens an index from disk. /// /// # Arguments @@ -184,7 +185,40 @@ impl Index { /// Replaces existing documents with new versions. pub fn update_documents(&mut self, docs: Vec) -> Result<()> { - self.index_documents(docs) + if docs.is_empty() { + return Ok(()); + } + + let mut current_by_key = HashMap::new(); + for segment in &self.segments { + for doc in segment.document_views() { + let doc_id = doc.doc_id().unwrap_or_default(); + if self.manifest.tombstones.iter().any(|dead| dead == doc_id) { + continue; + } + if let Some(logical_key) = doc.logical_key() { + current_by_key.insert(logical_key.to_string(), doc.to_owned()); + } + } + } + + let changed = docs + .into_iter() + .filter(|incoming| { + current_by_key + .get(&incoming.logical_key) + .is_none_or(|current| current != incoming) + }) + .collect::>(); + if changed.is_empty() { + return Ok(()); + } + + self.index_documents(changed)?; + if self.manifest.segments.len() >= Self::COMPACT_SEGMENT_THRESHOLD { + self.compact()?; + } + Ok(()) } /// Searches the index. diff --git a/src/storage.rs b/src/storage.rs index f9fa13e..3a9cced 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -306,25 +306,25 @@ impl IndexStorage { /// Compacts live segments into a new segment. pub fn compact(&self, segments: &[SegmentView], manifest: &Manifest) -> Result { - let live_docs: Vec = segments - .iter() - .flat_map(|segment| { - segment - .document_views() - .into_iter() - .map(|doc| doc.to_owned()) - }) - .filter(|doc| !manifest.tombstones.iter().any(|dead| dead == &doc.doc_id)) - .collect(); + let mut latest_by_key = HashMap::new(); + for segment in segments { + let documents = segment.document_views(); + let embeddings = segment.embeddings(); + for (index, doc) in documents.into_iter().enumerate() { + let owned = doc.to_owned(); + if manifest.tombstones.iter().any(|dead| dead == &owned.doc_id) { + continue; + } + if let Some(embedding) = embeddings.get(index) { + latest_by_key.insert(owned.logical_key.clone(), (owned, embedding.clone())); + } + } + } + let (live_docs, embeddings): (Vec<_>, Vec<_>) = latest_by_key.into_values().unzip(); if live_docs.is_empty() { return Err(anyhow!("nothing to compact")); } let segment_id = format!("compact_{:016x}", manifest.generation + 1); - let embeddings = segments - .iter() - .flat_map(|segment| segment.embeddings().into_iter()) - .take(live_docs.len()) - .collect::>(); let bm25_pairs: Vec<(String, String)> = live_docs .iter() .map(|doc| (doc.doc_id.clone(), doc.searchable_text.clone())) diff --git a/tests/index_tests.rs b/tests/index_tests.rs index 361e78d..f2e4f10 100644 --- a/tests/index_tests.rs +++ b/tests/index_tests.rs @@ -244,6 +244,97 @@ fn add_is_idempotent_and_does_not_duplicate_scores() { assert!((results_first[0].fused_score - results_second[0].fused_score).abs() < f32::EPSILON); } +#[test] +fn update_skips_unchanged_documents_without_new_segment() { + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + let bundle = tmp.path().join("bundle"); + fs::create_dir_all(&bundle).unwrap(); + fs::create_dir_all(&index_dir).unwrap(); + + let mut index = open_index(&index_dir, Box::new(MockEmbeddingProvider::new(16))).unwrap(); + let doc = OkfDocumentBuilder::new(&bundle, bundle.join("a.md")) + .frontmatter_value("type", "Metric") + .frontmatter_value("title", "Orders") + .body("Orders completed by customers") + .build(); + index.update_documents(vec![doc.clone()]).unwrap(); + let generation = index.manifest().generation; + let segment_count = index.manifest().segments.len(); + + index.update_documents(vec![doc]).unwrap(); + + assert_eq!(index.manifest().generation, generation); + assert_eq!(index.manifest().segments.len(), segment_count); +} + +#[test] +fn update_replaces_changed_document_by_logical_key() { + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + let bundle = tmp.path().join("bundle"); + fs::create_dir_all(&bundle).unwrap(); + fs::create_dir_all(&index_dir).unwrap(); + + let mut index = open_index(&index_dir, Box::new(MockEmbeddingProvider::new(16))).unwrap(); + let old_doc = OkfDocumentBuilder::new(&bundle, bundle.join("a.md")) + .frontmatter_value("type", "Metric") + .frontmatter_value("title", "Orders") + .body("old searchable text") + .build(); + let new_doc = OkfDocumentBuilder::new(&bundle, bundle.join("a.md")) + .frontmatter_value("type", "Metric") + .frontmatter_value("title", "Orders") + .body("new searchable text") + .build(); + index.update_documents(vec![old_doc]).unwrap(); + + index.update_documents(vec![new_doc.clone()]).unwrap(); + + let (results, _) = index + .search("new searchable", SearchMode::Lexical, 10) + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].doc_id, new_doc.doc_id); + assert_eq!(index.manifest().tombstones.len(), 1); +} + +#[test] +fn update_compacts_after_segment_threshold() { + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + let bundle = tmp.path().join("bundle"); + fs::create_dir_all(&bundle).unwrap(); + fs::create_dir_all(&index_dir).unwrap(); + + let mut index = open_index(&index_dir, Box::new(MockEmbeddingProvider::new(16))).unwrap(); + for version in 0..10 { + let doc = OkfDocumentBuilder::new(&bundle, bundle.join("a.md")) + .frontmatter_value("title", "Orders") + .body(format!("version {version}")) + .build(); + index.update_documents(vec![doc]).unwrap(); + } + + assert!(index.manifest().segments.len() < 8); + assert_eq!(index.manifest().segments.len(), 3); + assert_eq!( + index + .manifest() + .segments + .iter() + .map(|s| s.doc_count) + .sum::(), + 3 + ); + assert_eq!(index.manifest().tombstones.len(), 2); + let results = index + .search("version 9", SearchMode::Lexical, 10) + .unwrap() + .0; + assert_eq!(results.len(), 1); +} + #[test] fn reindexing_deleted_doc_clears_tombstone() { let tmp = tempfile::tempdir().unwrap();