fix: rebuild FAISS index on vector deletion (#4178)
This commit is contained in:
@@ -309,8 +309,23 @@ class FAISS(VectorStoreBase):
|
|||||||
break
|
break
|
||||||
|
|
||||||
if index_to_delete is not None:
|
if index_to_delete is not None:
|
||||||
|
# Reconstruct remaining vectors and rebuild the FAISS index
|
||||||
|
remaining_vectors = []
|
||||||
|
new_index_to_id = {}
|
||||||
|
new_idx = 0
|
||||||
|
for old_idx in sorted(self.index_to_id.keys()):
|
||||||
|
if old_idx == index_to_delete:
|
||||||
|
continue
|
||||||
|
remaining_vectors.append(self.index.reconstruct(int(old_idx)))
|
||||||
|
new_index_to_id[new_idx] = self.index_to_id[old_idx]
|
||||||
|
new_idx += 1
|
||||||
|
|
||||||
|
self.index.reset()
|
||||||
|
if remaining_vectors:
|
||||||
|
self.index.add(np.array(remaining_vectors, dtype=np.float32))
|
||||||
|
|
||||||
self.docstore.pop(vector_id, None)
|
self.docstore.pop(vector_id, None)
|
||||||
self.index_to_id.pop(index_to_delete, None)
|
self.index_to_id = new_index_to_id
|
||||||
|
|
||||||
self._save()
|
self._save()
|
||||||
|
|
||||||
|
|||||||
@@ -165,19 +165,29 @@ def test_search_with_filters(faiss_instance, mock_faiss_index):
|
|||||||
assert results[0].payload == {"name": "vector1", "category": "A"}
|
assert results[0].payload == {"name": "vector1", "category": "A"}
|
||||||
|
|
||||||
|
|
||||||
def test_delete(faiss_instance):
|
def test_delete(faiss_instance, mock_faiss_index):
|
||||||
# Setup the docstore and index_to_id mapping
|
# Setup the docstore and index_to_id mapping
|
||||||
faiss_instance.docstore = {"id1": {"name": "vector1"}, "id2": {"name": "vector2"}}
|
faiss_instance.docstore = {"id1": {"name": "vector1"}, "id2": {"name": "vector2"}}
|
||||||
faiss_instance.index_to_id = {0: "id1", 1: "id2"}
|
faiss_instance.index_to_id = {0: "id1", 1: "id2"}
|
||||||
|
|
||||||
|
# Mock reconstruct to return vectors for remaining entries
|
||||||
|
mock_faiss_index.reconstruct.side_effect = lambda idx: np.array(
|
||||||
|
[0.1, 0.2, 0.3] if idx == 0 else [0.4, 0.5, 0.6], dtype=np.float32
|
||||||
|
)
|
||||||
|
|
||||||
# Call delete
|
# Call delete
|
||||||
faiss_instance.delete(vector_id="id1")
|
faiss_instance.delete(vector_id="id1")
|
||||||
|
|
||||||
# Verify the vector was removed from docstore and index_to_id
|
# Verify the vector was removed from docstore
|
||||||
assert "id1" not in faiss_instance.docstore
|
assert "id1" not in faiss_instance.docstore
|
||||||
assert 0 not in faiss_instance.index_to_id
|
|
||||||
assert "id2" in faiss_instance.docstore
|
assert "id2" in faiss_instance.docstore
|
||||||
assert 1 in faiss_instance.index_to_id
|
|
||||||
|
# Verify the FAISS index was rebuilt
|
||||||
|
mock_faiss_index.reset.assert_called_once()
|
||||||
|
mock_faiss_index.add.assert_called_once()
|
||||||
|
|
||||||
|
# Verify index_to_id was remapped contiguously
|
||||||
|
assert faiss_instance.index_to_id == {0: "id2"}
|
||||||
|
|
||||||
|
|
||||||
def test_update(faiss_instance, mock_faiss_index):
|
def test_update(faiss_instance, mock_faiss_index):
|
||||||
|
|||||||
Reference in New Issue
Block a user