feat: implement delete_all_documents for weaviate integration (#2354)

* implement first test version

* weaviate: add delete_all_documents(recreate_index, batching)

Signed-off-by: phigep <philipp.geppner@gmail.com>

* fixed linter issues

* add requested tests and fix max query result issues with helper

* linter again

* forgot a type hint

* simplify and throw warning if batchsize too big

* refinements

* fix

---------

Signed-off-by: phigep <philipp.geppner@gmail.com>
Co-authored-by: anakin87 <stefanofiorucci@gmail.com>

Author: phigep

integrations/weaviate/src/haystack_integrations/document_stores/weaviate/document_store.py

@@ -505,6 +505,55 @@ def delete_documents(self, document_ids: List[str]) -> None:
         weaviate_ids = [generate_uuid5(doc_id) for doc_id in document_ids]
         self.collection.data.delete_many(where=weaviate.classes.query.Filter.by_id().contains_any(weaviate_ids))
 
+    def delete_all_documents(self, *, recreate_index: bool = False, batch_size: int = 1000) -> None:
+        """
+        Deletes all documents in a collection.
+
+        If recreate_index is False, it keeps the collection but deletes documents iteratively.
+        If recreate_index is True, the collection is dropped and faithfully recreated.
+        This is recommended for performance reasons.
+
+        :param recreate_index: Use drop and recreate strategy. (recommended for performance)
+        :param batch_size: Only relevant if recreate_index is false. Defines the deletion batch size.
+            Note that this parameter needs to be less or equal to the set `QUERY_MAXIMUM_RESULTS` variable
+            set for the weaviate deployment (default is 10000).
+            Reference: https://docs.weaviate.io/weaviate/manage-objects/delete#delete-all-objects
+        """
+
+        if recreate_index:
+            # get current up-to-date config from server, so we can recreate the collection faithfully
+            cfg = self.client.collections.get(self._collection_settings["class"]).config.get().to_dict()
+            class_name = cfg.get("class", self._collection_settings["class"])
+
+            self.client.collections.delete(class_name)
+            self.client.collections.create_from_dict(cfg)
+
+            self._collection_settings = cfg
+            self._collection = self.client.collections.get(class_name)
+            return
+
+        uuids = []
+        batch_size = max(1, int(batch_size))
+
+        for obj in self.collection.iterator(return_properties=[], include_vector=False):
+            uuids.append(obj.uuid)
+            if len(uuids) >= batch_size:
+                res = self.collection.data.delete_many(where=weaviate.classes.query.Filter.by_id().contains_any(uuids))
+                if res.successful < len(uuids):
+                    logger.warning(
+                        "Not all documents in the batch have been deleted. "
+                        "Make sure to specify a deletion `batch_size` which is less than `QUERY_MAXIMUM_RESULTS`.",
+                    )
+                uuids.clear()
+
+        if uuids:
+            res = self.collection.data.delete_many(where=weaviate.classes.query.Filter.by_id().contains_any(uuids))
+            if res.successful < len(uuids):
+                logger.warning(
+                    "Not all documents have been deleted. "
+                    "Make sure to specify a deletion `batch_size` which is less than `QUERY_MAXIMUM_RESULTS`.",
+                )
+
     def _bm25_retrieval(
         self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None
     ) -> List[Document]:

integrations/weaviate/tests/test_document_store.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import base64
+import logging
 import os
 from typing import List
 from unittest.mock import MagicMock, patch
@@ -794,3 +795,41 @@ def test_connect_to_local(self):
     def test_connect_to_embedded(self):
         document_store = WeaviateDocumentStore(embedded_options=EmbeddedOptions())
         assert document_store.client
+
+    def test_delete_all_documents(self, document_store):
+        docs = [Document(content="test doc 1"), Document(content="test doc 2")]
+        assert document_store.write_documents(docs) == 2
+        assert document_store.count_documents() == 2
+        document_store.delete_all_documents()
+        assert document_store.count_documents() == 0
+
+    def test_delete_all_documents_recreate(self, document_store):
+        docs = [Document(content="test doc 1"), Document(content="test doc 2")]
+        assert document_store.write_documents(docs) == 2
+        assert document_store.count_documents() == 2
+
+        cls = document_store._collection_settings["class"]
+        collection = document_store.client.collections.get(cls)
+        previous_config = collection.config.get().to_dict()
+
+        document_store.delete_all_documents(recreate_index=True)
+        assert document_store.count_documents() == 0
+
+        new_config = document_store.client.collections.get(cls).config.get().to_dict()
+        assert previous_config == new_config
+
+    def test_delete_all_documents_batch_size(self, document_store):
+        docs = [Document(content=str(i)) for i in range(0, 5)]
+        assert document_store.write_documents(docs) == 5
+        document_store.delete_all_documents(batch_size=2)
+        assert document_store.count_documents() == 0
+
+    def test_delete_all_documents_excessive_batch_size(self, document_store, caplog):
+        """Test that the deletion is not complete if the batch size exceeds the QUERY_MAXIMUM_RESULTS."""
+        # assume QUERY_MAXIMUM_RESULTS == 10000 with standard deployment
+        docs = [Document(content=str(i)) for i in range(0, 10005)]
+        assert document_store.write_documents(docs) == 10005
+        with caplog.at_level(logging.WARNING):
+            document_store.delete_all_documents(batch_size=20000)
+        assert document_store.count_documents() == 5
+        assert "Not all documents have been deleted." in caplog.text