Skip to content

Commit b16e559

Browse files
yhmoyhmo
committed
Add checkHealth/getPersistentSegmentInfo/getQuerySegmentInfo/transferNode interfaces
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent c1a0c77 commit b16e559

10 files changed

Lines changed: 276 additions & 6 deletions

File tree

‎sdk-core/src/main/java/io/milvus/v2/client/MilvusClientV2.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -812,6 +812,15 @@ public DescribeResourceGroupResp describeResourceGroup(DescribeResourceGroupReq
812812
return rpcUtils.retry(()->rgroupService.describeResourceGroup(this.getRpcStub(), request));
813813
}
814814

815+
/**
816+
* Transfer query nodes from source resource group to target resource_group.
817+
*
818+
* @param request {@link TransferNodeReq}
819+
*/
820+
public void transferNode(TransferNodeReq request) {
821+
rpcUtils.retry(()->rgroupService.transferNode(this.getRpcStub(), request));
822+
}
823+
815824
/**
816825
* Transfer a replica from source resource group to target resource_group.
817826
*
@@ -886,6 +895,28 @@ public void flush(FlushReq request) {
886895
utilityService.waitFlush(tempBlockingStub, response.getCollectionSegmentIDs(), response.getCollectionFlushTs());
887896
}
888897

898+
/**
899+
* Gets the information of persistent segments from data node, including row count,
900+
* persistence state(growing or flushed), etc.
901+
*
902+
* @param request get request
903+
* @return GetPersistentSegmentInfoResp
904+
*/
905+
GetPersistentSegmentInfoResp getPersistentSegmentInfo(GetPersistentSegmentInfoReq request) {
906+
return rpcUtils.retry(()->utilityService.getPersistentSegmentInfo(this.getRpcStub(), request));
907+
}
908+
909+
/**
910+
* Gets the query information of segments in a collection from query node, including row count,
911+
* memory usage size, index name, etc.
912+
*
913+
* @param request get request
914+
* @return GetQuerySegmentInfoResp
915+
*/
916+
GetQuerySegmentInfoResp getQuerySegmentInfo(GetQuerySegmentInfoReq request){
917+
return rpcUtils.retry(()->utilityService.getQuerySegmentInfo(this.getRpcStub(), request));
918+
}
919+
889920
/**
890921
* trigger an asynchronous compaction in server side
891922
*
@@ -915,6 +946,15 @@ public String getServerVersion() {
915946
return rpcUtils.retry(()->clientUtils.getServerVersion(this.getRpcStub()));
916947
}
917948

949+
/**
950+
* Check server health
951+
*
952+
* @return CheckHealthResp
953+
*/
954+
CheckHealthResp checkHealth() {
955+
return rpcUtils.retry(()->utilityService.checkHealth(this.getRpcStub()));
956+
}
957+
918958
/**
919959
* Disconnects from a Milvus server with configurable timeout
920960
*

‎sdk-core/src/main/java/io/milvus/v2/service/resourcegroup/ResourceGroupService.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,25 @@ public DescribeResourceGroupResp describeResourceGroup(MilvusServiceGrpc.MilvusS
155155
.build();
156156
}
157157

158+
public Void transferNode(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, TransferNodeReq request) {
159+
if (StringUtils.isEmpty(request.getSourceGroupName())) {
160+
throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid source group name");
161+
}
162+
if (StringUtils.isEmpty(request.getTargetGroupName())) {
163+
throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid target group name");
164+
}
165+
166+
String title = String.format("TransferNode %d nodes from %s to %s", request.getNumOfNodes(),
167+
request.getSourceGroupName(), request.getTargetGroupName());
168+
Status response = blockingStub.transferNode(TransferNodeRequest.newBuilder()
169+
.setSourceResourceGroup(request.getSourceGroupName())
170+
.setTargetResourceGroup(request.getTargetGroupName())
171+
.setNumNode(request.getNumOfNodes())
172+
.build());
173+
rpcUtils.handleResponse(title, response);
174+
return null;
175+
}
176+
158177
public Void transferReplica(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
159178
TransferReplicaReq request) {
160179
if (StringUtils.isEmpty(request.getSourceGroupName())) {

‎sdk-core/src/main/java/io/milvus/v2/service/resourcegroup/request/TransferNodeReq.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
package io.milvus.v2.service.resourcegroup.request;
2+
3+
import lombok.Data;
4+
import lombok.experimental.SuperBuilder;
5+
6+
@Data
7+
@SuperBuilder
8+
public class TransferNodeReq {
9+
String sourceGroupName;
10+
String targetGroupName;
11+
Integer numOfNodes;
12+
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/UtilityService.java‎

Lines changed: 71 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import io.milvus.v2.service.utility.response.*;
2929

3030
import java.util.*;
31+
import java.util.stream.Collectors;
3132

3233
public class UtilityService extends BaseService {
3334
public FlushResp flush(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, FlushReq request) {
@@ -100,7 +101,7 @@ public CompactResp compact(MilvusServiceGrpc.MilvusServiceBlockingStub blockingS
100101

101102
public GetCompactionStateResp getCompactionState(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
102103
GetCompactionStateReq request) {
103-
String title = "Get compaction state";
104+
String title = "GetCompactionState";
104105
io.milvus.grpc.GetCompactionStateRequest getRequest = io.milvus.grpc.GetCompactionStateRequest.newBuilder()
105106
.setCompactionID(request.getCompactionID())
106107
.build();
@@ -116,7 +117,7 @@ public GetCompactionStateResp getCompactionState(MilvusServiceGrpc.MilvusService
116117
}
117118

118119
public Void createAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, CreateAliasReq request) {
119-
String title = String.format("Create alias %s for collection %s", request.getAlias(), request.getCollectionName());
120+
String title = String.format("CreateAlias %s for collection %s", request.getAlias(), request.getCollectionName());
120121
io.milvus.grpc.CreateAliasRequest createAliasRequest = io.milvus.grpc.CreateAliasRequest.newBuilder()
121122
.setCollectionName(request.getCollectionName())
122123
.setAlias(request.getAlias())
@@ -128,7 +129,7 @@ public Void createAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub
128129
}
129130

130131
public Void dropAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, DropAliasReq request) {
131-
String title = String.format("Drop alias %s", request.getAlias());
132+
String title = String.format("DropAlias %s", request.getAlias());
132133
io.milvus.grpc.DropAliasRequest dropAliasRequest = io.milvus.grpc.DropAliasRequest.newBuilder()
133134
.setAlias(request.getAlias())
134135
.build();
@@ -139,7 +140,7 @@ public Void dropAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
139140
}
140141

141142
public Void alterAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, AlterAliasReq request) {
142-
String title = String.format("Alter alias %s for collection %s", request.getAlias(), request.getCollectionName());
143+
String title = String.format("AlterAlias %s for collection %s", request.getAlias(), request.getCollectionName());
143144
io.milvus.grpc.AlterAliasRequest alterAliasRequest = io.milvus.grpc.AlterAliasRequest.newBuilder()
144145
.setCollectionName(request.getCollectionName())
145146
.setAlias(request.getAlias())
@@ -151,7 +152,7 @@ public Void alterAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
151152
}
152153

153154
public DescribeAliasResp describeAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, DescribeAliasReq request) {
154-
String title = String.format("Describe alias %s", request.getAlias());
155+
String title = String.format("DescribeAlias %s", request.getAlias());
155156
io.milvus.grpc.DescribeAliasRequest describeAliasRequest = io.milvus.grpc.DescribeAliasRequest.newBuilder()
156157
.setAlias(request.getAlias())
157158
.build();
@@ -166,7 +167,7 @@ public DescribeAliasResp describeAlias(MilvusServiceGrpc.MilvusServiceBlockingSt
166167
}
167168

168169
public ListAliasResp listAliases(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, ListAliasesReq request) {
169-
String title = "List aliases";
170+
String title = "ListAliases";
170171
io.milvus.grpc.ListAliasesRequest listAliasesRequest = io.milvus.grpc.ListAliasesRequest.newBuilder()
171172
.setCollectionName(request.getCollectionName())
172173
.build();
@@ -179,4 +180,68 @@ public ListAliasResp listAliases(MilvusServiceGrpc.MilvusServiceBlockingStub blo
179180
.alias(response.getAliasesList())
180181
.build();
181182
}
183+
184+
public CheckHealthResp checkHealth(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub) {
185+
String title = "CheckHealth";
186+
CheckHealthResponse response = blockingStub.checkHealth(CheckHealthRequest.newBuilder().build());
187+
rpcUtils.handleResponse(title, response.getStatus());
188+
189+
List<String> states = new ArrayList<>();
190+
response.getQuotaStatesList().forEach(s->states.add(s.name()));
191+
return CheckHealthResp.builder()
192+
.isHealthy(response.getIsHealthy())
193+
.reasons(response.getReasonsList().stream().collect(Collectors.toList()))
194+
.quotaStates(states)
195+
.build();
196+
}
197+
198+
public GetPersistentSegmentInfoResp getPersistentSegmentInfo(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
199+
GetPersistentSegmentInfoReq request) {
200+
String title = String.format("GetPersistentSegmentInfo collectionName %s", request.getCollectionName());
201+
GetPersistentSegmentInfoResponse response = blockingStub.getPersistentSegmentInfo(GetPersistentSegmentInfoRequest.newBuilder()
202+
.setCollectionName(request.getCollectionName())
203+
.build());
204+
rpcUtils.handleResponse(title, response.getStatus());
205+
206+
List<GetPersistentSegmentInfoResp.PersistentSegmentInfo> segmentInfos = new ArrayList<>();
207+
response.getInfosList().forEach(info->{segmentInfos.add(GetPersistentSegmentInfoResp.PersistentSegmentInfo.builder()
208+
.segmentID(info.getSegmentID())
209+
.collectionID(info.getCollectionID())
210+
.partitionID(info.getPartitionID())
211+
.numOfRows(info.getNumRows())
212+
.state(info.getState().name())
213+
.level(info.getLevel().name())
214+
.isSorted(info.getIsSorted())
215+
.build());});
216+
return GetPersistentSegmentInfoResp.builder()
217+
.segmentInfos(segmentInfos)
218+
.build();
219+
}
220+
221+
public GetQuerySegmentInfoResp getQuerySegmentInfo(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
222+
GetQuerySegmentInfoReq request) {
223+
String title = String.format("GetQuerySegmentInfo collectionName %s", request.getCollectionName());
224+
GetQuerySegmentInfoResponse response = blockingStub.getQuerySegmentInfo(GetQuerySegmentInfoRequest.newBuilder()
225+
.setCollectionName(request.getCollectionName())
226+
.build());
227+
rpcUtils.handleResponse(title, response.getStatus());
228+
229+
List<GetQuerySegmentInfoResp.QuerySegmentInfo> segmentInfos = new ArrayList<>();
230+
response.getInfosList().forEach(info->{segmentInfos.add(GetQuerySegmentInfoResp.QuerySegmentInfo.builder()
231+
.segmentID(info.getSegmentID())
232+
.collectionID(info.getCollectionID())
233+
.partitionID(info.getPartitionID())
234+
.memSize(info.getMemSize())
235+
.numOfRows(info.getNumRows())
236+
.indexName(info.getIndexName())
237+
.indexID(info.getIndexID())
238+
.state(info.getState().name())
239+
.level(info.getLevel().name())
240+
.nodeIDs(info.getNodeIdsList())
241+
.isSorted(info.getIsSorted())
242+
.build());});
243+
return GetQuerySegmentInfoResp.builder()
244+
.segmentInfos(segmentInfos)
245+
.build();
246+
}
182247
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/request/GetPersistentSegmentInfoReq.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
package io.milvus.v2.service.utility.request;
2+
3+
import lombok.Data;
4+
import lombok.experimental.SuperBuilder;
5+
6+
@Data
7+
@SuperBuilder
8+
public class GetPersistentSegmentInfoReq {
9+
private String collectionName;
10+
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/request/GetQuerySegmentInfoReq.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
package io.milvus.v2.service.utility.request;
2+
3+
import lombok.Data;
4+
import lombok.experimental.SuperBuilder;
5+
6+
@Data
7+
@SuperBuilder
8+
public class GetQuerySegmentInfoReq {
9+
private String collectionName;
10+
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/response/CheckHealthResp.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package io.milvus.v2.service.utility.response;
2+
3+
import lombok.Builder;
4+
import lombok.Data;
5+
import lombok.experimental.SuperBuilder;
6+
7+
import java.util.ArrayList;
8+
import java.util.List;
9+
10+
@Data
11+
@SuperBuilder
12+
public class CheckHealthResp {
13+
@Builder.Default
14+
Boolean isHealthy = false;
15+
@Builder.Default
16+
List<String> reasons = new ArrayList<>();
17+
@Builder.Default
18+
List<String> quotaStates = new ArrayList<>();
19+
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/response/GetPersistentSegmentInfoResp.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package io.milvus.v2.service.utility.response;
2+
3+
import lombok.Builder;
4+
import lombok.Data;
5+
import lombok.experimental.SuperBuilder;
6+
7+
import java.util.ArrayList;
8+
import java.util.List;
9+
10+
@Data
11+
@SuperBuilder
12+
public class GetPersistentSegmentInfoResp {
13+
@Data
14+
@SuperBuilder
15+
public static class PersistentSegmentInfo {
16+
private Long segmentID;
17+
private Long collectionID;
18+
private Long partitionID;
19+
private Long numOfRows;
20+
private String state;
21+
private String level;
22+
private Boolean isSorted;
23+
}
24+
25+
@Builder.Default
26+
private List<PersistentSegmentInfo> segmentInfos = new ArrayList<>();
27+
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/response/GetQuerySegmentInfoResp.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
package io.milvus.v2.service.utility.response;
2+
3+
import lombok.Builder;
4+
import lombok.Data;
5+
import lombok.experimental.SuperBuilder;
6+
7+
import java.util.ArrayList;
8+
import java.util.List;
9+
10+
@Data
11+
@SuperBuilder
12+
public class GetQuerySegmentInfoResp {
13+
@Data
14+
@SuperBuilder
15+
public static class QuerySegmentInfo {
16+
private Long segmentID;
17+
private Long collectionID;
18+
private Long partitionID;
19+
private Long memSize;
20+
private Long numOfRows;
21+
private String indexName;
22+
private Long indexID;
23+
private String state;
24+
private String level;
25+
@Builder.Default
26+
private List<Long> nodeIDs = new ArrayList<>();
27+
private Boolean isSorted;
28+
}
29+
30+
@Builder.Default
31+
private List<QuerySegmentInfo> segmentInfos = new ArrayList<>();
32+
}

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,9 @@ private long getRowCount(String collectionName) {
290290

291291
@Test
292292
void testFloatVectors() {
293+
CheckHealthResp healthy = client.checkHealth();
294+
Assertions.assertTrue(healthy.getIsHealthy());
295+
293296
String randomCollectionName = generator.generate(10);
294297

295298
String vectorFieldName = "float_vector";
@@ -321,6 +324,20 @@ void testFloatVectors() {
321324
.collectionNames(Collections.singletonList(randomCollectionName))
322325
.build());
323326

327+
// get persistent segment info
328+
GetPersistentSegmentInfoResp pSegInfo = client.getPersistentSegmentInfo(GetPersistentSegmentInfoReq.builder()
329+
.collectionName(randomCollectionName)
330+
.build());
331+
Assertions.assertEquals(1, pSegInfo.getSegmentInfos().size());
332+
GetPersistentSegmentInfoResp.PersistentSegmentInfo pInfo = pSegInfo.getSegmentInfos().get(0);
333+
Assertions.assertTrue(pInfo.getSegmentID() > 0L);
334+
Assertions.assertTrue(pInfo.getCollectionID() > 0L);
335+
Assertions.assertTrue(pInfo.getPartitionID() > 0L);
336+
Assertions.assertEquals(count, pInfo.getNumOfRows());
337+
Assertions.assertEquals("Flushed", pInfo.getState());
338+
Assertions.assertEquals("L1", pInfo.getLevel());
339+
Assertions.assertFalse(pInfo.getIsSorted());
340+
324341
// compact
325342
CompactResp compactResp = client.compact(CompactReq.builder()
326343
.collectionName(randomCollectionName)
@@ -347,6 +364,25 @@ void testFloatVectors() {
347364
.collectionName(randomCollectionName)
348365
.build());
349366

367+
// get query segment info
368+
GetQuerySegmentInfoResp qSegInfo = client.getQuerySegmentInfo(GetQuerySegmentInfoReq.builder()
369+
.collectionName(randomCollectionName)
370+
.build());
371+
Assertions.assertEquals(1, qSegInfo.getSegmentInfos().size());
372+
GetQuerySegmentInfoResp.QuerySegmentInfo qInfo = qSegInfo.getSegmentInfos().get(0);
373+
Assertions.assertTrue(qInfo.getSegmentID() > 0L);
374+
Assertions.assertTrue(qInfo.getCollectionID() > 0L);
375+
Assertions.assertTrue(qInfo.getPartitionID() > 0L);
376+
Assertions.assertTrue(qInfo.getMemSize() >= 0L);
377+
Assertions.assertEquals(count, qInfo.getNumOfRows());
378+
Assertions.assertEquals(vectorFieldName, qInfo.getIndexName());
379+
Assertions.assertTrue(qInfo.getIndexID() > 0L);
380+
Assertions.assertEquals("Sealed", qInfo.getState());
381+
Assertions.assertEquals("L1", qInfo.getLevel());
382+
Assertions.assertEquals(1, qInfo.getNodeIDs().size());
383+
Assertions.assertTrue(qInfo.getNodeIDs().get(0) > 0L);
384+
Assertions.assertTrue(qInfo.getIsSorted());
385+
350386
// create partition, upsert one row to the partition
351387
String partitionName = "PPP";
352388
client.createPartition(CreatePartitionReq.builder()

0 commit comments

Comments
 (0)