Skip to content

Commit c0ce51b

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

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
@@ -266,11 +266,12 @@ public SearchResp search(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStu
266266
SearchResults response = blockingStub.search(searchRequest);
267267
rpcUtils.handleResponse(title, response.getStatus());
268268

269-
return SearchResp.builder()
269+
SearchResp.SearchRespBuilder respBuilder = SearchResp.builder()
270270
.searchResults(convertUtils.getEntities(response))
271271
.sessionTs(response.getSessionTs())
272-
.recalls(response.getResults().getRecallsList())
273-
.build();
272+
.recalls(response.getResults().getRecallsList());
273+
fillSearchRespFromExtraInfo(respBuilder, response.getStatus().getExtraInfoMap());
274+
return respBuilder.build();
274275
}
275276

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

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

296314
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
@@ -139,6 +139,12 @@ public void setUp() {
139139
when(blockingStub.query(any())).thenReturn(QueryResults.newBuilder().build());
140140
when(blockingStub.delete(any())).thenReturn(MutationResult.newBuilder().setDeleteCnt(2L).build());
141141
SearchResults searchResults = SearchResults.newBuilder()
142+
.setStatus(Status.newBuilder().setCode(0)
143+
.putExtraInfo("report_value", "123")
144+
.putExtraInfo("scanned_remote_bytes", "456")
145+
.putExtraInfo("scanned_total_bytes", "789")
146+
.putExtraInfo("cache_hit_ratio", "0.5")
147+
.build())
142148
.setResults(SearchResultData.newBuilder().addScores(1L).addTopks(0L).build())
143149
.build();
144150
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
@@ -733,7 +733,9 @@ void testV2BuilderClasses() {
733733
VerifyClass(QueryResp.class.getName(), config);
734734
VerifyClass(QueryResp.QueryResult.class.getName(), config);
735735
VerifyClass(RunAnalyzerResp.class.getName(), config);
736+
config.setIgnoredMethods(Arrays.asList("getSearchResults", "setSearchResults", "getRecalls", "setRecalls"));
736737
VerifyClass(SearchResp.class.getName(), config);
738+
config.clearIgnoredMethods();
737739
VerifyClass(SearchResp.SearchResult.class.getName(), config);
738740
VerifyClass(UpsertResp.class.getName(), config);
739741
}

‎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
@@ -137,6 +137,10 @@ void testSearch() {
137137
.build();
138138
SearchResp statusR = client_v2.search(request);
139139
logger.info(statusR.toString());
140+
Assertions.assertEquals(123L, statusR.getCost());
141+
Assertions.assertEquals(456L, statusR.getScannedRemoteBytes());
142+
Assertions.assertEquals(789L, statusR.getScannedTotalBytes());
143+
Assertions.assertEquals(0.5f, statusR.getCacheHitRatio());
140144
}
141145

142146
@Test

0 commit comments

Comments
 (0)