Skip to content

Commit 6bae5bf

Browse files
yhmoyhmo
committed
Align search response with pymilvus
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent ccbee12 commit 6bae5bf

5 files changed

Lines changed: 104 additions & 6 deletions

File tree

‎sdk-core/src/main/java/io/milvus/v2/service/vector/VectorService.java‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -265,11 +265,12 @@ public SearchResp search(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStu
265265
SearchResults response = blockingStub.search(searchRequest);
266266
rpcUtils.handleResponse(title, response.getStatus());
267267

268-
return SearchResp.builder()
268+
SearchResp.SearchRespBuilder respBuilder = SearchResp.builder()
269269
.searchResults(convertUtils.getEntities(response))
270270
.sessionTs(response.getSessionTs())
271-
.recalls(response.getResults().getRecallsList())
272-
.build();
271+
.recalls(response.getResults().getRecallsList());
272+
fillSearchRespFromExtraInfo(respBuilder, response.getStatus().getExtraInfoMap());
273+
return respBuilder.build();
273274
}
274275

275276
public SearchResp hybridSearch(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, HybridSearchReq request) {
@@ -286,10 +287,27 @@ public SearchResp hybridSearch(MilvusServiceGrpc.MilvusServiceBlockingStub block
286287
SearchResults response = blockingStub.hybridSearch(searchRequest);
287288
rpcUtils.handleResponse(title, response.getStatus());
288289

289-
return SearchResp.builder()
290+
SearchResp.SearchRespBuilder respBuilder = SearchResp.builder()
290291
.searchResults(convertUtils.getEntities(response))
291-
.recalls(response.getResults().getRecallsList())
292-
.build();
292+
.sessionTs(response.getSessionTs())
293+
.recalls(response.getResults().getRecallsList());
294+
fillSearchRespFromExtraInfo(respBuilder, response.getStatus().getExtraInfoMap());
295+
return respBuilder.build();
296+
}
297+
298+
private void fillSearchRespFromExtraInfo(SearchResp.SearchRespBuilder respBuilder, java.util.Map<String, String> extraInfo) {
299+
if (extraInfo.containsKey("report_value")) {
300+
respBuilder.cost(Long.parseLong(extraInfo.get("report_value")));
301+
}
302+
if (extraInfo.containsKey("scanned_remote_bytes")) {
303+
respBuilder.scannedRemoteBytes(Long.parseLong(extraInfo.get("scanned_remote_bytes")));
304+
}
305+
if (extraInfo.containsKey("scanned_total_bytes")) {
306+
respBuilder.scannedTotalBytes(Long.parseLong(extraInfo.get("scanned_total_bytes")));
307+
}
308+
if (extraInfo.containsKey("cache_hit_ratio")) {
309+
respBuilder.cacheHitRatio(Float.parseFloat(extraInfo.get("cache_hit_ratio")));
310+
}
293311
}
294312

295313
public QueryIterator queryIterator(RpcStubWrapper blockingStub,

‎sdk-core/src/main/java/io/milvus/v2/service/vector/response/SearchResp.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,19 @@ public class SearchResp {
2828
private List<List<SearchResult>> searchResults;
2929
private long sessionTs; // default eventually ts
3030
private List<Float> recalls;
31+
private Long cost;
32+
private Long scannedRemoteBytes;
33+
private Long scannedTotalBytes;
34+
private Float cacheHitRatio;
3135

3236
private SearchResp(SearchRespBuilder builder) {
3337
this.searchResults = builder.searchResults;
3438
this.sessionTs = builder.sessionTs;
3539
this.recalls = builder.recalls;
40+
this.cost = builder.cost;
41+
this.scannedRemoteBytes = builder.scannedRemoteBytes;
42+
this.scannedTotalBytes = builder.scannedTotalBytes;
43+
this.cacheHitRatio = builder.cacheHitRatio;
3644
}
3745

3846
public static SearchRespBuilder builder() {
@@ -63,19 +71,59 @@ public void setRecalls(List<Float> recalls) {
6371
this.recalls = recalls;
6472
}
6573

74+
public Long getCost() {
75+
return cost;
76+
}
77+
78+
public void setCost(Long cost) {
79+
this.cost = cost;
80+
}
81+
82+
public Long getScannedRemoteBytes() {
83+
return scannedRemoteBytes;
84+
}
85+
86+
public void setScannedRemoteBytes(Long scannedRemoteBytes) {
87+
this.scannedRemoteBytes = scannedRemoteBytes;
88+
}
89+
90+
public Long getScannedTotalBytes() {
91+
return scannedTotalBytes;
92+
}
93+
94+
public void setScannedTotalBytes(Long scannedTotalBytes) {
95+
this.scannedTotalBytes = scannedTotalBytes;
96+
}
97+
98+
public Float getCacheHitRatio() {
99+
return cacheHitRatio;
100+
}
101+
102+
public void setCacheHitRatio(Float cacheHitRatio) {
103+
this.cacheHitRatio = cacheHitRatio;
104+
}
105+
66106
@Override
67107
public String toString() {
68108
return "SearchResp{" +
69109
"searchResults=" + searchResults +
70110
", sessionTs=" + sessionTs +
71111
", recalls=" + recalls +
112+
", cost=" + cost +
113+
", scannedRemoteBytes=" + scannedRemoteBytes +
114+
", scannedTotalBytes=" + scannedTotalBytes +
115+
", cacheHitRatio=" + cacheHitRatio +
72116
'}';
73117
}
74118

75119
public static class SearchRespBuilder {
76120
private List<List<SearchResult>> searchResults = new ArrayList<>();
77121
private long sessionTs = 1L; // default eventually ts
78122
private List<Float> recalls = new ArrayList<>();
123+
private Long cost;
124+
private Long scannedRemoteBytes;
125+
private Long scannedTotalBytes;
126+
private Float cacheHitRatio;
79127

80128
public SearchRespBuilder searchResults(List<List<SearchResult>> searchResults) {
81129
this.searchResults = searchResults;
@@ -92,6 +140,26 @@ public SearchRespBuilder recalls(List<Float> recalls) {
92140
return this;
93141
}
94142

143+
public SearchRespBuilder cost(Long cost) {
144+
this.cost = cost;
145+
return this;
146+
}
147+
148+
public SearchRespBuilder scannedRemoteBytes(Long scannedRemoteBytes) {
149+
this.scannedRemoteBytes = scannedRemoteBytes;
150+
return this;
151+
}
152+
153+
public SearchRespBuilder scannedTotalBytes(Long scannedTotalBytes) {
154+
this.scannedTotalBytes = scannedTotalBytes;
155+
return this;
156+
}
157+
158+
public SearchRespBuilder cacheHitRatio(Float cacheHitRatio) {
159+
this.cacheHitRatio = cacheHitRatio;
160+
return this;
161+
}
162+
95163
public SearchResp build() {
96164
return new SearchResp(this);
97165
}

‎sdk-core/src/test/java/io/milvus/v2/BaseTest.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,12 @@ public void setUp() {
138138
when(blockingStub.query(any())).thenReturn(QueryResults.newBuilder().build());
139139
when(blockingStub.delete(any())).thenReturn(MutationResult.newBuilder().setDeleteCnt(2L).build());
140140
SearchResults searchResults = SearchResults.newBuilder()
141+
.setStatus(Status.newBuilder().setCode(0)
142+
.putExtraInfo("report_value", "123")
143+
.putExtraInfo("scanned_remote_bytes", "456")
144+
.putExtraInfo("scanned_total_bytes", "789")
145+
.putExtraInfo("cache_hit_ratio", "0.5")
146+
.build())
141147
.setResults(SearchResultData.newBuilder().addScores(1L).addTopks(0L).build())
142148
.build();
143149
when(blockingStub.search(any())).thenReturn(searchResults);

‎sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2Test.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -707,7 +707,9 @@ void testV2BuilderClasses() {
707707
VerifyClass(QueryResp.class.getName(), config);
708708
VerifyClass(QueryResp.QueryResult.class.getName(), config);
709709
VerifyClass(RunAnalyzerResp.class.getName(), config);
710+
config.setIgnoredMethods(Arrays.asList("getSearchResults", "setSearchResults", "getRecalls", "setRecalls"));
710711
VerifyClass(SearchResp.class.getName(), config);
712+
config.clearIgnoredMethods();
711713
VerifyClass(SearchResp.SearchResult.class.getName(), config);
712714
VerifyClass(UpsertResp.class.getName(), config);
713715
}

‎sdk-core/src/test/java/io/milvus/v2/service/vector/VectorTest.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,10 @@ void testSearch() {
134134
.build();
135135
SearchResp statusR = client_v2.search(request);
136136
logger.info(statusR.toString());
137+
Assertions.assertEquals(123L, statusR.getCost());
138+
Assertions.assertEquals(456L, statusR.getScannedRemoteBytes());
139+
Assertions.assertEquals(789L, statusR.getScannedTotalBytes());
140+
Assertions.assertEquals(0.5f, statusR.getCacheHitRatio());
137141
}
138142

139143
@Test

0 commit comments

Comments
 (0)