From 45a0038757f73058a39ac8aa54d8232d1af5a71e Mon Sep 17 00:00:00 2001
From: drawnwren <drawnwren@gmail.com>
Date: Thu, 21 Nov 2024 15:10:04 -0800
Subject: [PATCH] add _async support to MilvusClient.search

---
 pymilvus/milvus_client/milvus_client.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py
index e7456aa8c..1d2646315 100644
--- a/pymilvus/milvus_client/milvus_client.py
+++ b/pymilvus/milvus_client/milvus_client.py
@@ -22,6 +22,7 @@
 from pymilvus.orm import utility
 from pymilvus.orm.collection import CollectionSchema
 from pymilvus.orm.connections import connections
+from pymilvus.orm.future import SearchFuture
 from pymilvus.orm.types import DataType
 
 from .index import IndexParams
@@ -366,7 +367,7 @@ def search(
         partition_names: Optional[List[str]] = None,
         anns_field: Optional[str] = None,
         **kwargs,
-    ) -> List[List[dict]]:
+    ) -> List[List[dict]] | SearchFuture:
         """Search for a query vector/vectors.
 
         In order for the search to process, a collection needs to have been either provided
@@ -388,6 +389,7 @@ def search(
         Returns:
             List[List[dict]]: A nested list of dicts containing the result data. Embeddings are
                 not included in the result data.
+            SearchFuture: If _async is True, return a SearchFuture object.
         """
         conn = self._get_connection()
         try:
@@ -408,6 +410,9 @@ def search(
             logger.error("Failed to search collection: %s", collection_name)
             raise ex from ex
 
+        if kwargs.get("_async", False):
+            return res
+
         ret = []
         for hits in res:
             query_result = []