Deletion optimisation

src/marqo/config.py

@@ -8,7 +8,8 @@ def __init__(
         url: str,
         timeout: Optional[int] = None,
         indexing_device: Optional[Union[enums.Device, str]] = None,
-        search_device: Optional[Union[enums.Device, str]] = None
+        search_device: Optional[Union[enums.Device, str]] = None,
+        backend: Optional[Union[enums.SearchDb, str]] = None,
     ) -> None:
         """
         Parameters
@@ -23,6 +24,7 @@ def __init__(
 
         self.indexing_device = indexing_device if indexing_device is not None else default_device
         self.search_device = search_device if search_device is not None else default_device
+        self.backend = backend if backend is not None else enums.SearchDb.opensearch
 
     def set_url(self, url):
         """Set the URL, and infers whether that url is remote"""

src/marqo/tensor_search/api.py

@@ -1,8 +1,9 @@
 """The API entrypoint for Tensor Search"""
 import typing
 from fastapi.responses import JSONResponse
-from fastapi import FastAPI, Request, Depends, HTTPException
-from fastapi.exceptions import RequestValidationError
+from fastapi import Request, Depends
+import marqo.tensor_search.delete_docs
+import marqo.tensor_search.tensor_search
 from marqo.errors import InvalidArgError, MarqoWebError, MarqoError
 from fastapi import FastAPI, Query
 import json
@@ -12,7 +13,6 @@
 import os
 from marqo.tensor_search.models.api_models import BulkSearchQuery, SearchQuery
 from marqo.tensor_search.web import api_validation, api_utils
-from marqo.tensor_search import utils
 from marqo.tensor_search.on_start_script import on_start
 from marqo import version
 from marqo.tensor_search.backend import get_index_info
@@ -237,7 +237,7 @@ def delete_index(index_name: str, marqo_config: config.Config = Depends(generate
 @app.post("/indexes/{index_name}/documents/delete-batch")
 def delete_docs(index_name: str, documentIds: List[str], refresh: bool = True,
                       marqo_config: config.Config = Depends(generate_config)):
-    return tensor_search.delete_documents(
+    return marqo.tensor_search.tensor_search.delete_documents(
         index_name=index_name, config=marqo_config, doc_ids=documentIds,
         auto_refresh=refresh
     )

src/marqo/tensor_search/delete_docs.py

@@ -0,0 +1,96 @@
+"""
+This module handles the delete documents endpoint
+"""
+import datetime
+from typing import List, NamedTuple, Literal
+import json
+from marqo import errors
+from marqo._httprequests import HttpRequests
+from marqo.config import Config
+from marqo.tensor_search import validation, utils, enums
+
+
+# -- Marqo delete endpoint interface: --
+
+
+class MqDeleteDocsResponse(NamedTuple):
+    """An object that holds the data we send back to users"""
+    index_name: str
+    status_string: Literal["succeeded"]
+    document_ids: List[str]
+    deleted_docments_count: int
+    deletion_start: datetime.datetime
+    deletion_end: datetime.datetime
+
+
+def format_delete_docs_response(marqo_response: MqDeleteDocsResponse) -> dict:
+    """This formats the delete response for users """
+    return {
+        "index_name": marqo_response.index_name, "status": marqo_response.status_string,
+        "type": "documentDeletion", "details": {
+            "receivedDocumentIds": len(marqo_response.document_ids),
+            "deletedDocuments": marqo_response.deleted_docments_count,
+        },
+        "duration": utils.create_duration_string(marqo_response.deletion_end - marqo_response.deletion_start),
+        "startedAt": utils.format_timestamp(marqo_response.deletion_start),
+        "finishedAt": utils.format_timestamp(marqo_response.deletion_end),
+    }
+
+
+class MqDeleteDocsRequest(NamedTuple):
+    """An object that holds the data from users for a delete request"""
+    index_name: str
+    document_ids: List[str]
+    auto_refresh: bool
+
+
+# -- Data-layer agnostic logic --
+
+
+def delete_documents(config: Config, del_request: MqDeleteDocsRequest) -> dict:
+    """entrypoint function for deleting documents"""
+    if not del_request.document_ids:
+        raise errors.InvalidDocumentIdError("doc_ids can't be empty!")
+
+    for _id in del_request.document_ids:
+        validation.validate_id(_id)
+
+    if config.backend == enums.SearchDb.opensearch:
+        del_response: MqDeleteDocsResponse = delete_documents_marqo_os(config=config, deletion_instruction=del_request)
+    else:
+        raise RuntimeError(f"Config set to use unknown backend `{config.backend}`. "
+                           f"See tensor_search.enums.SearchDB for allowed backends")
+
+    return format_delete_docs_response(del_response)
+
+
+# -- Marqo-OS-specific deletion implementation: --
+
+
+def delete_documents_marqo_os(config: Config, deletion_instruction: MqDeleteDocsRequest) -> MqDeleteDocsResponse:
+    """Deletes documents """
+
+    # Prepare bulk delete request body
+    bulk_request_body = ""
+    for doc_id in deletion_instruction.document_ids:
+        bulk_request_body += json.dumps({"delete": {"_index": deletion_instruction.index_name, "_id": doc_id}}) + "\n"
+
+    # Send bulk delete request
+    t0 = datetime.datetime.utcnow()
+    delete_res_backend = HttpRequests(config=config).post(
+        path="_bulk",
+        body=bulk_request_body,
+    )
+
+    if deletion_instruction.auto_refresh:
+        refresh_response = HttpRequests(config).post(path=f"{deletion_instruction.index_name}/_refresh")
+
+    t1 = datetime.datetime.utcnow()
+    deleted_documents_count = sum(1 for item in delete_res_backend["items"] if "delete" in item and item["delete"]["status"] == 200)
+
+    mq_delete_res = MqDeleteDocsResponse(
+        index_name=deletion_instruction.index_name, status_string='succeeded', document_ids=deletion_instruction.document_ids,
+        deleted_docments_count=deleted_documents_count, deletion_start=t0,
+        deletion_end=t1
+    )
+    return mq_delete_res

src/marqo/tensor_search/enums.py

@@ -115,6 +115,9 @@ class MappingsObjectType:
     multimodal_combination = "multimodal_combination"
 
 
+class SearchDb:
+    opensearch = 'opensearch'
+
 # Perhaps create a ThrottleType to differentiate thread_count and data_size throttling mechanisms
 
 

src/marqo/tensor_search/tensor_search.py

@@ -32,14 +32,13 @@
 """
 import copy
 import json
-import datetime
 from collections import defaultdict
 from timeit import default_timer as timer
 import functools
 import pprint
 import typing
 import uuid
-from typing import List, Optional, Union, Iterable, Sequence, Dict, Any, Tuple, Set
+from typing import List, Optional, Union, Iterable, Sequence, Dict, Any, Tuple
 import numpy as np
 from PIL import Image
 import marqo.config as config
@@ -56,6 +55,7 @@
 from marqo.tensor_search.models.search import VectorisedJobs, VectorisedJobPointer, Qidx, JHash
 from marqo.tensor_search.models.index_info import IndexInfo
 from marqo.tensor_search.utils import add_timing
+from marqo.tensor_search import delete_docs
 from marqo.s2_inference.processing import text as text_processor
 from marqo.s2_inference.processing import image as image_processor
 from marqo.s2_inference.clip_utils import _is_image
@@ -959,41 +959,6 @@ def _get_documents_for_upsert(
     return res
 
 
-def delete_documents(config: Config, index_name: str, doc_ids: List[str], auto_refresh):
-    """Deletes documents """
-    if not doc_ids:
-        raise errors.InvalidDocumentIdError("doc_ids can't be empty!")
-
-    for _id in doc_ids:
-        validation.validate_id(_id)
-
-    # TODO: change to timer()
-    t0 = datetime.datetime.utcnow()
-    delete_res_backend = HttpRequests(config=config).post(
-        path=f"{index_name}/_delete_by_query", body={
-            "query": {
-                "terms": {
-                    "_id": doc_ids
-                }
-            }
-        }
-    )
-    if auto_refresh:
-        refresh_response = HttpRequests(config).post(path=F"{index_name}/_refresh")
-    t1 = datetime.datetime.utcnow()
-    delete_res = {
-        "index_name": index_name, "status": "succeeded",
-        "type": "documentDeletion", "details": {
-            "receivedDocumentIds": len(doc_ids),
-            "deletedDocuments": delete_res_backend["deleted"],
-        },
-        "duration": utils.create_duration_string(t1 - t0),
-        "startedAt": utils.format_timestamp(t0),
-        "finishedAt": utils.format_timestamp(t1),
-    }
-    return delete_res
-
-
 def refresh_index(config: Config, index_name: str):
     return HttpRequests(config).post(path=F"{index_name}/_refresh")
 
@@ -2593,4 +2558,15 @@ def _create_score_modifiers_tensor_search_query(result_count, offset, vector_fie
             "exclude": ["__chunks.__vector_*"]
         }
     }
-    return search_query
+    return search_query
+
+
+def delete_documents(config: Config, index_name: str, doc_ids: List[str], auto_refresh):
+    """Delete documents from the Marqo index with the given doc_ids """
+    return delete_docs.delete_documents(
+        config=config,
+        del_request=delete_docs.MqDeleteDocsRequest(
+            index_name=index_name,
+            document_ids=doc_ids,
+            auto_refresh=auto_refresh)
+    )

tests/tensor_search/test_config.py

@@ -76,3 +76,25 @@ def run():
 
     def test_device_for_clip(self):
         assert str(enums.Device.cpu) == "cpu"
+
+
+class TestConfigBackend(MarqoTestCase):
+
+    def setUp(self) -> None:
+        self.endpoint = self.authorized_url
+
+    class CustomSearchDb:
+        opensearch = "opensearch"
+        elasticsearch = "elasticsearch"
+
+    def test_init_default_backend(self):
+        c = config.Config(url=self.endpoint)
+        assert c.backend == enums.SearchDb.opensearch
+
+    def test_init_custom_backend(self):
+        c = config.Config(url=self.endpoint, backend=self.CustomSearchDb.elasticsearch)
+        assert c.backend == "elasticsearch"
+
+    def test_init_custom_backend_as_string(self):
+        c = config.Config(url=self.endpoint, backend="elasticsearch")
+        assert c.backend == "elasticsearch"