From e548a51ff7fe471b8d807d08bf85c7576fb20a22 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Mon, 1 Dec 2025 17:53:27 +0200
Subject: [PATCH] Optimized `rebuild --embed-only` to use batch updates via
 LanceDB `merge_insert` instead of individual chunk updates

---
 CHANGELOG.md                       |  4 ++++
 haiku_rag_slim/haiku/rag/client.py | 36 ++++++++++++++++++++++--------
 tests/test_rebuild.py              | 36 ++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8f8a5d96..162049e8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- **Rebuild Performance**: Optimized `rebuild --embed-only` to use batch updates via LanceDB's `merge_insert` instead of individual chunk updates, and skip chunks with unchanged embeddings
+
 ## [0.19.4] - 2025-11-28
 
 ### Added
diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py
index c09f2666..8e4902c4 100644
--- a/haiku_rag_slim/haiku/rag/client.py
+++ b/haiku_rag_slim/haiku/rag/client.py
@@ -748,21 +748,39 @@ class HaikuRAG:
         """Re-embed all chunks without changing chunk boundaries."""
         for doc in documents:
             assert doc.id is not None
-            chunks = await self.chunk_repository.get_by_document_id(doc.id)
-            if not chunks:
+
+            # Get raw chunk records directly from LanceDB
+            chunk_records = list(
+                self.store.chunks_table.search()
+                .where(f"document_id = '{doc.id}'")
+                .to_pydantic(self.store.ChunkRecord)
+            )
+            if not chunk_records:
                 continue
 
             # Batch embed all chunk contents
-            contents = [chunk.content for chunk in chunks]
+            contents = [rec.content for rec in chunk_records]
             embeddings = await self.chunk_repository.embedder.embed(contents)
 
-            # Update each chunk with new embedding
-            for chunk, embedding in zip(chunks, embeddings):
-                assert chunk.id is not None
-                self.store.chunks_table.update(
-                    where=f"id = '{chunk.id}'",
-                    values={"vector": embedding},
+            # Build updated records only for chunks with changed embeddings
+            updated_records = [
+                self.store.ChunkRecord(
+                    id=rec.id,
+                    document_id=rec.document_id,
+                    content=rec.content,
+                    metadata=rec.metadata,
+                    order=rec.order,
+                    vector=embedding,
                 )
+                for rec, embedding in zip(chunk_records, embeddings)
+                if rec.vector != embedding
+            ]
+
+            # Batch update chunks with changed embeddings
+            if updated_records:
+                self.store.chunks_table.merge_insert(
+                    "id"
+                ).when_matched_update_all().execute(updated_records)
 
             yield doc.id
 
diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py
index 60a04e81..d145fd4a 100644
--- a/tests/test_rebuild.py
+++ b/tests/test_rebuild.py
@@ -57,6 +57,42 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
             assert chunk.content == chunk_contents_before[chunk.id]
 
 
+@pytest.mark.asyncio
+async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_path):
+    """Test embed-only rebuild skips chunks with unchanged embeddings."""
+    async with HaikuRAG(temp_db_path) as client:
+        doc = await client.create_document(content=qa_corpus["document_extracted"][0])
+        assert doc.id is not None
+
+        # Get embeddings before rebuild
+        records_before = list(
+            client.store.chunks_table.search()
+            .where(f"document_id = '{doc.id}'")
+            .to_pydantic(client.store.ChunkRecord)
+        )
+        embeddings_before = {rec.id: rec.vector for rec in records_before}
+
+        # Run embed-only rebuild with same embedder - embeddings should be identical
+        processed_ids = [
+            doc_id
+            async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
+        ]
+        assert doc.id in processed_ids
+
+        # Get embeddings after rebuild
+        records_after = list(
+            client.store.chunks_table.search()
+            .where(f"document_id = '{doc.id}'")
+            .to_pydantic(client.store.ChunkRecord)
+        )
+        embeddings_after = {rec.id: rec.vector for rec in records_after}
+
+        # Embeddings should be identical (same content, same embedder)
+        assert embeddings_before.keys() == embeddings_after.keys()
+        for chunk_id in embeddings_before:
+            assert embeddings_before[chunk_id] == embeddings_after[chunk_id]
+
+
 @pytest.mark.asyncio
 async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
     """Test rechunk rebuild: re-chunks from content without accessing source files."""