Merge pull request #189 from vespa-engine/tgm/delete-batch-async

Implement delete_batch with async mode as default

Author: Thiago G. Martins

vespa/application.py

@@ -270,14 +270,45 @@ def delete_data(self, schema: str, data_id: str) -> VespaResponse:
         with VespaSync(self) as sync_app:
             return sync_app.delete_data(schema=schema, data_id=data_id)
 
-    def delete_batch(self, batch: List):
+    def _delete_batch_sync(self, schema: str, batch: List[Dict]):
+        return [self.delete_data(schema, data_point["id"]) for data_point in batch]
+
+    async def _delete_batch_async(
+        self, schema: str, batch: List[Dict], connections, total_timeout
+    ):
+        async with VespaAsync(
+            app=self, connections=connections, total_timeout=total_timeout
+        ) as async_app:
+            return await async_app.delete_batch(schema=schema, batch=batch)
+
+    def delete_batch(
+        self,
+        schema: str,
+        batch: List[Dict],
+        asynchronous=True,
+        connections: Optional[int] = 100,
+        total_timeout: int = 100,
+    ):
         """
-        Async delete a batch of data from a Vespa app.
+        Delete a batch of data from a Vespa app.
 
-        :param batch: A list of tuples with 'schema' and 'id'
-        :return:
+        :param schema: The schema that we are deleting data from.
+        :param batch: A list of dict containing the key 'id'.
+        :param asynchronous: Set True to get data in async mode. Default to True.
+        :param connections: Number of allowed concurrent connections, valid only if `asynchronous=True`.
+        :param total_timeout: Total timeout in secs for each of the concurrent requests when using `asynchronous=True`.
+        :return: List of HTTP POST responses
         """
-        return [self.delete_data(schema, id) for schema, id in batch]
+        if asynchronous:
+            coro = self._delete_batch_async(
+                schema=schema,
+                batch=batch,
+                connections=connections,
+                total_timeout=total_timeout,
+            )
+            return self._check_for_running_loop_and_run_coroutine(coro=coro)
+        else:
+            return self._delete_batch_sync(schema=schema, batch=batch)
 
     def delete_all_docs(self, content_cluster_name: str, schema: str) -> Response:
         """
@@ -977,8 +1008,18 @@ async def delete_data(self, schema: str, data_id: str) -> VespaResponse:
             operation_type="delete",
         )
 
-    async def delete_batch(self, batch):
-        return await self._wait(self.delete_data, batch)
+    async def _delete_data_semaphore(
+        self, schema: str, data_id: str, semaphore: asyncio.Semaphore
+    ):
+        async with semaphore:
+            return await self.delete_data(schema=schema, data_id=data_id)
+
+    async def delete_batch(self, schema: str, batch: List[Dict]):
+        sem = asyncio.Semaphore(self.connections)
+        return await self._wait(
+            self._delete_data_semaphore,
+            [(schema, data_point["id"], sem) for data_point in batch],
+        )
 
     async def get_data(self, schema: str, data_id: str):
         end_point = "{}/document/v1/{}/{}/docid/{}".format(

vespa/test_integration_docker.py

@@ -648,6 +648,23 @@ def batch_operations_synchronous_mode(
             expected_updated_fields.update(fields_to_update[idx])
             self.assertDictEqual(response.json["fields"], expected_updated_fields)
 
+        #
+        # Delete data
+        #
+        result = app.delete_batch(schema=schema, batch=docs, asynchronous=False)
+        for idx, response in enumerate(result):
+            self.assertEqual(
+                response.json["id"],
+                "id:{}:{}::{}".format(schema, schema, docs[idx]["id"]),
+            )
+
+        #
+        # get batch deleted data
+        #
+        result = app.get_batch(schema=schema, batch=docs, asynchronous=False)
+        for idx, response in enumerate(result):
+            self.assertEqual(response.status_code, 404)
+
     def batch_operations_asynchronous_mode(
         self,
         app,
@@ -724,6 +741,23 @@ def batch_operations_asynchronous_mode(
             expected_updated_fields.update(fields_to_update[idx])
             self.assertDictEqual(response.json["fields"], expected_updated_fields)
 
+        #
+        # Delete data
+        #
+        result = app.delete_batch(schema=schema, batch=docs, asynchronous=True)
+        for idx, response in enumerate(result):
+            self.assertEqual(
+                response.json["id"],
+                "id:{}:{}::{}".format(schema, schema, docs[idx]["id"]),
+            )
+
+        #
+        # get batch deleted data
+        #
+        result = app.get_batch(schema=schema, batch=docs, asynchronous=True)
+        for idx, response in enumerate(result):
+            self.assertEqual(response.status_code, 404)
+
     @staticmethod
     def _parse_vespa_tensor(hit, feature):
         return [x["value"] for x in hit["fields"]["summaryfeatures"][feature]["cells"]]
@@ -911,7 +945,7 @@ def test_batch_operations_synchronous_mode(self):
             schema_name=self.app_package.name,
             fields_to_send=self.fields_to_send,
             expected_fields_from_get_operation=self.fields_to_send,
-            fields_to_update=self.fields_to_update
+            fields_to_update=self.fields_to_update,
         )
 
     def test_batch_operations_asynchronous_mode(self):
@@ -1004,7 +1038,7 @@ def test_batch_operations_synchronous_mode(self):
             schema_name=self.app_package.name,
             fields_to_send=self.fields_to_send,
             expected_fields_from_get_operation=self.expected_fields_from_get_operation,
-            fields_to_update=self.fields_to_update
+            fields_to_update=self.fields_to_update,
         )
 
     def test_batch_operations_asynchronous_mode(self):
@@ -1113,7 +1147,7 @@ def test_batch_operations_synchronous_mode(self):
             schema_name="sentence",
             fields_to_send=self.fields_to_send_sentence,
             expected_fields_from_get_operation=self.expected_fields_from_sentence_get_operation,
-            fields_to_update=self.fields_to_update
+            fields_to_update=self.fields_to_update,
         )
 
     def test_batch_operations_asynchronous_mode(self):