fix: AsyncMilvusClient.search support EmbeddingList

Convert EmbeddingList to flat array in search() and hybrid_search()
before calling Prepare. Enhance is_legal_search_data() to recognize
EmbeddingList.

Add unit tests for search and hybrid_search with EmbeddingList.

Signed-off-by: silas.jiang <[email protected]>

Author: silas.jiang

pymilvus/client/async_grpc_handler.py

@@ -36,6 +36,7 @@
     is_legal_port,
 )
 from .constants import ITERATOR_SESSION_TS_FIELD
+from .embedding_list import EmbeddingList
 from .interceptor import _api_level_md
 from .prepare import Prepare
 from .search_result import SearchResult
@@ -844,6 +845,12 @@ async def search(
             guarantee_timestamp=kwargs.get("guarantee_timestamp"),
             timeout=timeout,
         )
+
+        # Convert EmbeddingList to flat array if present
+        if isinstance(data, list) and len(data) > 0 and isinstance(data[0], EmbeddingList):
+            data = [emb_list.to_flat_array() for emb_list in data]
+            kwargs["is_embedding_list"] = True
+
         request = Prepare.search_requests_with_expr(
             collection_name,
             data,
@@ -884,17 +891,24 @@ async def hybrid_search(
 
         requests = []
         for req in reqs:
+            data = req.data
+            req_kwargs = dict(kwargs)
+            # Convert EmbeddingList to flat array if present
+            if isinstance(data, list) and len(data) > 0 and isinstance(data[0], EmbeddingList):
+                data = [emb_list.to_flat_array() for emb_list in data]
+                req_kwargs["is_embedding_list"] = True
+
             search_request = Prepare.search_requests_with_expr(
                 collection_name,
-                req.data,
+                data,
                 req.anns_field,
                 req.param,
                 req.limit,
                 req.expr,
                 partition_names=partition_names,
                 round_decimal=round_decimal,
                 expr_params=req.expr_params,
-                **kwargs,
+                **req_kwargs,
             )
             requests.append(search_request)
 

pymilvus/client/check.py

@@ -205,6 +205,11 @@ def is_legal_search_data(data: Any) -> bool:
     if entity_helper.entity_is_sparse_matrix(data):
         return True
 
+    # Support EmbeddingList for array-of-vector searches
+    # Check for EmbeddingList by type name to avoid circular dependency
+    if isinstance(data, list) and len(data) > 0 and type(data[0]).__name__ == "EmbeddingList":
+        return True
+
     if not isinstance(data, (list, np.ndarray)):
         return False
 

tests/test_async_grpc_handler.py

@@ -1,8 +1,10 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import numpy as np
 import pytest
 from pymilvus.client.async_grpc_handler import AsyncGrpcHandler
 from pymilvus.exceptions import MilvusException
+from pymilvus.grpc_gen import schema_pb2
 
 
 class TestAsyncGrpcHandler:
@@ -242,3 +244,164 @@ async def test_load_partitions_with_resource_groups(self) -> None:
                 replica_number=2,
                 resource_groups=["rg1", "rg2"]
             )
+
+    @pytest.mark.asyncio
+    async def test_search_with_embedding_list(self) -> None:
+        """
+        Test that search works with EmbeddingList input data.
+        This test verifies the fix for issue where AsyncMilvusClient.search 
+        failed when using EmbeddingList for array-of-vector searches.
+        """
+        # Setup mock channel and stub
+        mock_channel = AsyncMock()
+        mock_channel.channel_ready = AsyncMock()
+        mock_channel._unary_unary_interceptors = []
+
+        handler = AsyncGrpcHandler(channel=mock_channel)
+        handler._is_channel_ready = True
+
+        mock_stub = AsyncMock()
+        handler._async_stub = mock_stub
+
+        # Mock Search response with proper SearchResultData structure
+        mock_search_result_data = schema_pb2.SearchResultData(
+            num_queries=2,
+            top_k=0,
+            scores=[],
+            ids=schema_pb2.IDs(int_id=schema_pb2.LongArray(data=[])),
+            topks=[],
+            primary_field_name="id"
+        )
+        mock_search_response = MagicMock()
+        mock_status = MagicMock()
+        mock_status.code = 0
+        mock_status.reason = ""
+        mock_search_response.status = mock_status
+        mock_search_response.results = mock_search_result_data
+        mock_search_response.session_ts = 0
+        mock_stub.Search = AsyncMock(return_value=mock_search_response)
+
+        # Create EmbeddingList data
+        from pymilvus.client.embedding_list import EmbeddingList
+        emb_list1 = EmbeddingList()
+        emb_list1.add([0.1, 0.2, 0.3, 0.4, 0.5])
+        emb_list2 = EmbeddingList()
+        emb_list2.add([0.5, 0.4, 0.3, 0.2, 0.1])
+        data = [emb_list1, emb_list2]
+
+        with patch('pymilvus.client.async_grpc_handler.Prepare') as mock_prepare, \
+             patch('pymilvus.client.async_grpc_handler.check_pass_param'), \
+             patch('pymilvus.client.async_grpc_handler.check_status'), \
+             patch('pymilvus.client.async_grpc_handler._api_level_md', return_value={}):
+            
+            # Mock search_requests_with_expr to return a request
+            mock_request = MagicMock()
+            mock_prepare.search_requests_with_expr.return_value = mock_request
+
+            await handler.search(
+                collection_name="test_collection",
+                data=data,
+                anns_field="vector",
+                param={"metric_type": "COSINE"},
+                limit=10
+            )
+
+            # Verify that Prepare.search_requests_with_expr was called
+            mock_prepare.search_requests_with_expr.assert_called_once()
+            call_args = mock_prepare.search_requests_with_expr.call_args
+            
+            # Verify that is_embedding_list was passed as True in kwargs
+            assert call_args.kwargs.get("is_embedding_list") is True
+            
+            # Verify data was converted (not EmbeddingList objects anymore)
+            passed_data = call_args[0][1]  # data is the second positional argument
+            assert isinstance(passed_data, list)
+            assert not isinstance(passed_data[0], EmbeddingList)
+            # The data should be converted to flat arrays
+            assert isinstance(passed_data[0], (list, np.ndarray))
+
+            # Verify Search was called
+            mock_stub.Search.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_hybrid_search_with_embedding_list(self) -> None:
+        """
+        Test that hybrid_search works with EmbeddingList input data.
+        """
+        # Setup mock channel and stub
+        mock_channel = AsyncMock()
+        mock_channel.channel_ready = AsyncMock()
+        mock_channel._unary_unary_interceptors = []
+
+        handler = AsyncGrpcHandler(channel=mock_channel)
+        handler._is_channel_ready = True
+
+        mock_stub = AsyncMock()
+        handler._async_stub = mock_stub
+
+        # Mock HybridSearch response with proper SearchResultData structure
+        mock_hybrid_result_data = schema_pb2.SearchResultData(
+            num_queries=1,
+            top_k=0,
+            scores=[],
+            ids=schema_pb2.IDs(int_id=schema_pb2.LongArray(data=[])),
+            topks=[],
+            primary_field_name="id"
+        )
+        mock_hybrid_response = MagicMock()
+        mock_status = MagicMock()
+        mock_status.code = 0
+        mock_status.reason = ""
+        mock_hybrid_response.status = mock_status
+        mock_hybrid_response.results = mock_hybrid_result_data
+        mock_stub.HybridSearch = AsyncMock(return_value=mock_hybrid_response)
+
+        # Create AnnSearchRequest with EmbeddingList
+        from pymilvus.client.embedding_list import EmbeddingList
+        from pymilvus.client.abstract import AnnSearchRequest
+        import numpy as np
+        
+        emb_list = EmbeddingList()
+        emb_list.add([0.1, 0.2, 0.3])
+        req = AnnSearchRequest(
+            data=[emb_list],
+            anns_field="vector",
+            param={"metric_type": "COSINE"},
+            limit=10
+        )
+
+        with patch('pymilvus.client.async_grpc_handler.Prepare') as mock_prepare, \
+             patch('pymilvus.client.async_grpc_handler.check_pass_param'), \
+             patch('pymilvus.client.async_grpc_handler.check_status'), \
+             patch('pymilvus.client.async_grpc_handler._api_level_md', return_value={}):
+            
+            # Mock search_requests_with_expr and hybrid_search_request_with_ranker
+            mock_search_request = MagicMock()
+            mock_hybrid_request = MagicMock()
+            mock_prepare.search_requests_with_expr.return_value = mock_search_request
+            mock_prepare.hybrid_search_request_with_ranker.return_value = mock_hybrid_request
+
+            # Mock rerank (BaseRanker)
+            mock_ranker = MagicMock()
+
+            await handler.hybrid_search(
+                collection_name="test_collection",
+                reqs=[req],
+                rerank=mock_ranker,
+                limit=10
+            )
+
+            # Verify that search_requests_with_expr was called with converted data
+            mock_prepare.search_requests_with_expr.assert_called_once()
+            call_args = mock_prepare.search_requests_with_expr.call_args
+            
+            # Verify is_embedding_list flag was set
+            assert call_args.kwargs.get("is_embedding_list") is True
+            
+            # Verify data was converted
+            passed_data = call_args[0][1]
+            assert isinstance(passed_data, list)
+            assert not isinstance(passed_data[0], EmbeddingList)
+
+            # Verify HybridSearch was called
+            mock_stub.HybridSearch.assert_called_once()