Fix incremental updates and compact stale segments

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<OkfDocument>) -> 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::<Vec<_>>();
+        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.

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<SegmentEntry> {
-        let live_docs: Vec<OkfDocument> = 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::<Vec<_>>();
         let bm25_pairs: Vec<(String, String)> = live_docs
             .iter()
             .map(|doc| (doc.doc_id.clone(), doc.searchable_text.clone()))

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::<usize>(),
+        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();