QueryIterator/SearchIterator support retry (#1270)

Signed-off-by: yhmo <[email protected]>
milvus-io · Dec 31, 2024 · 70c0802 · 70c0802
1 parent 417fbb8
commit 70c0802
Show file tree

Hide file tree

Showing 4 changed files with 198 additions and 192 deletions.
diff --git a/src/main/java/io/milvus/orm/iterator/QueryIterator.java b/src/main/java/io/milvus/orm/iterator/QueryIterator.java
@@ -190,8 +190,7 @@ private List<QueryResultsWrapper.RowRecord> getQueryResultsWrapper(String expr,
                 .build();
 
         QueryRequest queryRequest = ParamUtils.convertQueryParam(queryParam);
-        QueryResults response = blockingStub.query(queryRequest);
-
+        QueryResults response = rpcUtils.retry(()->blockingStub.query(queryRequest));
         String title = String.format("QueryRequest collectionName:%s", queryIteratorParam.getCollectionName());
         rpcUtils.handleResponse(title, response.getStatus());
 

diff --git a/src/main/java/io/milvus/orm/iterator/SearchIterator.java b/src/main/java/io/milvus/orm/iterator/SearchIterator.java
@@ -257,7 +257,7 @@ private SearchResultsWrapper executeNextSearch(Map<String, Object> params, Strin
         fillVectorsByPlType(searchParamBuilder);
 
         SearchRequest searchRequest = ParamUtils.convertSearchParam(searchParamBuilder.build());
-        SearchResults response = blockingStub.search(searchRequest);
+        SearchResults response = rpcUtils.retry(()->blockingStub.search(searchRequest));
 
         String title = String.format("SearchRequest collectionName:%s", searchIteratorParam.getCollectionName());
         rpcUtils.handleResponse(title, response.getStatus());