Fix intra-batch duplicate-ID orphaning (langchain, haystack)

turbovec-python/python/turbovec/haystack.py

@@ -177,28 +177,44 @@ def write_documents(
             policy = DuplicatePolicy.FAIL
 
         # First pass: validate and resolve duplicates according to policy.
+        # Duplicates are resolved against the batch-so-far as well as the
+        # existing store: InMemoryDocumentStore writes into its dict as it
+        # iterates, so a repeated id *within a single call* is resolved the
+        # same way a cross-call repeat would be. Without tracking the batch,
+        # every duplicate row still gets its own vector while _str_to_u64
+        # keeps only the last handle, orphaning the earlier vectors.
         to_write: List[Document] = []
+        batch_pos: Dict[str, int] = {}  # doc.id -> index into to_write
+        written = len(documents)
         for doc in documents:
             if doc.embedding is None:
                 raise ValueError(
                     f"Document {doc.id!r} has no embedding. "
                     "TurboQuantDocumentStore only stores documents with precomputed "
                     "embeddings — run an embedder component before writing."
                 )
-            if doc.id in self._str_to_u64:
+            present = doc.id in self._str_to_u64 or doc.id in batch_pos
+            if policy != DuplicatePolicy.OVERWRITE and present:
                 if policy == DuplicatePolicy.FAIL:
                     raise DuplicateDocumentError(
                         f"ID '{doc.id}' already exists in the document store."
                     )
                 if policy == DuplicatePolicy.SKIP:
+                    written -= 1
                     continue
-                if policy == DuplicatePolicy.OVERWRITE:
+            if policy == DuplicatePolicy.OVERWRITE:
+                if doc.id in self._str_to_u64:
                     self._remove_one(doc.id)
-                # fall through to add
+                if doc.id in batch_pos:
+                    # Last write wins: replace the earlier queued document
+                    # in place rather than appending a second vector.
+                    to_write[batch_pos[doc.id]] = doc
+                    continue
+            batch_pos[doc.id] = len(to_write)
             to_write.append(doc)
 
         if not to_write:
-            return 0
+            return written
 
         vectors = np.asarray(
             [doc.embedding for doc in to_write], dtype=np.float32
@@ -233,7 +249,7 @@ def write_documents(
                 "blob": doc.blob,
                 "sparse_embedding": doc.sparse_embedding,
             }
-        return len(to_write)
+        return written
 
     def delete_documents(self, document_ids: List[str]) -> None:
         # Haystack's protocol says silently ignore missing ids.

turbovec-python/python/turbovec/langchain.py

@@ -174,16 +174,30 @@ def _store_texts_and_vectors(
         metadatas: list[dict],
         ids: list[str],
     ) -> list[str]:
+        if vectors.ndim != 2:
+            raise ValueError(f"expected 2D embedding batch, got {vectors.ndim}D")
+
+        # Dedup intra-batch duplicate ids, keeping the last occurrence —
+        # matches InMemoryVectorStore, whose dict store silently overwrites
+        # on a repeated id. Without this every row is added to the index but
+        # _str_to_u64 keeps only the last handle per id, orphaning the
+        # earlier vectors. The returned id list still mirrors the input
+        # (one entry per input text), as the reference does.
+        result_ids = ids
+        if len(set(ids)) != len(ids):
+            keep = sorted({id_: i for i, id_ in enumerate(ids)}.values())
+            ids = [ids[i] for i in keep]
+            texts_list = [texts_list[i] for i in keep]
+            metadatas = [metadatas[i] for i in keep]
+            vectors = vectors[keep]
+
         # Upsert: any id that already exists is removed so the re-added
         # vector wins. Matches LangChain user expectation that `add_texts`
         # with an existing id updates in place.
         duplicates = [i for i in ids if i in self._str_to_u64]
         if duplicates:
             self.delete(duplicates)
 
-        if vectors.ndim != 2:
-            raise ValueError(f"expected 2D embedding batch, got {vectors.ndim}D")
-
         # IdMapIndex.add_with_ids handles both eager (dim must match) and
         # lazy (locks dim on first call) cases. Pre-check the eager case
         # so we surface a clean ValueError rather than a Rust panic.
@@ -205,7 +219,7 @@ def _store_texts_and_vectors(
             self._str_to_u64[id_] = h
             self._u64_to_str[h] = id_
             self._docs[id_] = (text, dict(meta))
-        return ids
+        return result_ids
 
     # ---- Read path (similarity search) --------------------------------
 

turbovec-python/tests/test_haystack.py

@@ -443,6 +443,55 @@ def test_duplicate_policy_overwrite_replaces():
     assert store.count_documents() == 3
 
 
+def test_intra_batch_duplicate_overwrite_keeps_last_no_orphan():
+    # Two docs sharing an id in a single call must not orphan a vector.
+    # InMemoryDocumentStore writes into a dict as it iterates, so the last
+    # write wins. count_documents and the id map must agree at one entry.
+    store = TurboQuantDocumentStore(dim=DIM, bit_width=4)
+    written = store.write_documents(
+        [
+            Document(id="dup", content="first", embedding=unit_vector(0)),
+            Document(id="dup", content="second", embedding=unit_vector(1)),
+        ],
+        policy=DuplicatePolicy.OVERWRITE,
+    )
+    # Reference counts every input row for OVERWRITE.
+    assert written == 2
+    # But only one survives — no orphaned vector.
+    assert store.count_documents() == 1
+    assert len(store._u64_to_doc) == 1
+    assert set(store._str_to_u64) == {"dup"}
+    assert store.filter_documents()[0].content == "second"
+
+
+def test_intra_batch_duplicate_fail_raises():
+    # FAIL must reject a repeat within the same call, not just across calls.
+    store = TurboQuantDocumentStore(dim=DIM, bit_width=4)
+    with pytest.raises(DuplicateDocumentError):
+        store.write_documents(
+            [
+                Document(id="dup", content="a", embedding=unit_vector(0)),
+                Document(id="dup", content="b", embedding=unit_vector(1)),
+            ],
+            policy=DuplicatePolicy.FAIL,
+        )
+
+
+def test_intra_batch_duplicate_skip_keeps_first():
+    # SKIP keeps the first occurrence and drops later in-batch repeats.
+    store = TurboQuantDocumentStore(dim=DIM, bit_width=4)
+    written = store.write_documents(
+        [
+            Document(id="dup", content="first", embedding=unit_vector(0)),
+            Document(id="dup", content="second", embedding=unit_vector(1)),
+        ],
+        policy=DuplicatePolicy.SKIP,
+    )
+    assert written == 1
+    assert store.count_documents() == 1
+    assert store.filter_documents()[0].content == "first"
+
+
 def test_write_document_without_embedding_raises():
     store = TurboQuantDocumentStore(dim=DIM, bit_width=4)
     with pytest.raises(ValueError, match="no embedding"):

turbovec-python/tests/test_langchain.py

@@ -176,6 +176,24 @@ def test_add_documents_with_ids_is_idempotent():
     assert set(store._docs.keys()) == {"a", "b"}
 
 
+def test_add_texts_intra_batch_duplicate_ids_keep_last():
+    # Two rows sharing an id in a single call must not orphan a vector.
+    # Reference InMemoryVectorStore overwrites on a repeated id (last wins);
+    # we dedup the batch the same way so the index holds one vector per id.
+    emb = StubEmbeddings(dim=64)
+    store = TurboQuantVectorStore.from_texts([], emb, bit_width=4)
+    ret = store.add_texts(["alpha", "beta"], ids=["dup", "dup"])
+
+    # Return value still mirrors the input (one entry per input text).
+    assert ret == ["dup", "dup"]
+    # No orphaned vector: index and id maps agree at one entry.
+    assert len(store._index) == 1
+    assert len(store._u64_to_str) == 1
+    assert set(store._str_to_u64) == {"dup"}
+    # Last occurrence wins.
+    assert store._docs["dup"][0] == "beta"
+
+
 def test_get_by_ids_empty_input_and_order_preserved():
     # Two contract points the reference makes that our existing tests
     # don't pin: (1) empty input returns [] without erroring; (2) output