diff --git a/docarray/index/backends/hnswlib.py b/docarray/index/backends/hnswlib.py index 5d7f951fe9f..44033a7bace 100644 --- a/docarray/index/backends/hnswlib.py +++ b/docarray/index/backends/hnswlib.py @@ -439,7 +439,7 @@ def _create_docs_table(self): def _send_docs_to_sqlite(self, docs: Sequence[BaseDoc]): ids = (self._to_hashed_id(doc.id) for doc in docs) self._sqlite_cursor.executemany( - 'INSERT INTO docs VALUES (?, ?)', + 'INSERT OR REPLACE INTO docs VALUES (?, ?)', ((id_, self._doc_to_bytes(doc)) for id_, doc in zip(ids, docs)), ) diff --git a/docs/user_guide/storing/index_hnswlib.md b/docs/user_guide/storing/index_hnswlib.md index e05a1ff0c9f..b85d903852d 100644 --- a/docs/user_guide/storing/index_hnswlib.md +++ b/docs/user_guide/storing/index_hnswlib.md @@ -208,4 +208,56 @@ To delete nested data, you need to specify the `id`. del doc_index[index_docs[6].id] ``` -Check [here](docindex#nested-data-with-subindex) for nested data with subindex. \ No newline at end of file +Check [here](../docindex#nested-data-with-subindex) for nested data with subindex. + +### Update elements +In order to update a Document inside the index, you only need to reindex it with the updated attributes. + +First lets create a schema for our Index +```python +import numpy as np +from docarray import BaseDoc, DocList +from docarray.typing import NdArray +from docarray.index import HnswDocumentIndex +class MyDoc(BaseDoc): + text: str + embedding: NdArray[128] +``` +Now we can instantiate our Index and index some data. + +```python +docs = DocList[MyDoc]( + [MyDoc(embedding=np.random.rand(10), text=f'I am the first version of Document {i}') for i in range(100)] +) +index = HnswDocumentIndex[MyDoc]() +index.index(docs) +assert index.num_docs() == 100 +``` + +Now we can find relevant documents + +```python +res = index.find(query=docs[0], search_field='tens', limit=100) +assert len(res.documents) == 100 +for doc in res.documents: + assert 'I am the first version' in doc.text +``` + +and update all of the text of this documents and reindex them + +```python +for i, doc in enumerate(docs): + doc.text = f'I am the second version of Document {i}' + +index.index(docs) +assert index.num_docs() == 100 +``` + +When we retrieve them again we can see that their text attribute has been updated accordingly + +```python +res = index.find(query=docs[0], search_field='tens', limit=100) +assert len(res.documents) == 100 +for doc in res.documents: + assert 'I am the second version' in doc.text +``` diff --git a/tests/index/hnswlib/test_index_get_del.py b/tests/index/hnswlib/test_index_get_del.py index 4ecb60e465a..ebf06ce4b14 100644 --- a/tests/index/hnswlib/test_index_get_del.py +++ b/tests/index/hnswlib/test_index_get_del.py @@ -362,3 +362,53 @@ def test_num_docs(ten_simple_docs, tmp_path): del index[more_docs[2].id, ten_simple_docs[7].id] assert index.num_docs() == 10 + + +def test_update_payload(tmp_path): + class TextSimpleDoc(SimpleDoc): + text: str = 'hey' + + docs = DocList[TextSimpleDoc]( + [TextSimpleDoc(tens=np.random.rand(10), text=f'hey {i}') for i in range(100)] + ) + index = HnswDocumentIndex[TextSimpleDoc](work_dir=str(tmp_path)) + index.index(docs) + assert index.num_docs() == 100 + + for doc in docs: + doc.text += '_changed' + + index.index(docs) + assert index.num_docs() == 100 + + res = index.find(query=docs[0], search_field='tens', limit=100) + assert len(res.documents) == 100 + for doc in res.documents: + assert '_changed' in doc.text + + +def test_update_embedding(tmp_path): + class TextSimpleDoc(SimpleDoc): + text: str = 'hey' + + docs = DocList[TextSimpleDoc]( + [TextSimpleDoc(tens=np.random.rand(10), text=f'hey {i}') for i in range(100)] + ) + index = HnswDocumentIndex[TextSimpleDoc](work_dir=str(tmp_path)) + index.index(docs) + assert index.num_docs() == 100 + + new_tensor = np.random.rand(10) + docs[0].tens = new_tensor + + index.index(docs[0]) + assert index.num_docs() == 100 + + res = index.find(query=docs[0], search_field='tens', limit=100) + assert len(res.documents) == 100 + found = False + for doc in res.documents: + if doc.id == docs[0].id: + found = True + assert (doc.tens == new_tensor).all() + assert found