Add checkHealth/getPersistentSegmentInfo/getQuerySegmentInfo/transferNode interfaces (#1373)

Signed-off-by: yhmo <yihua.mo@zilliz.com>

.github/workflows/maven.yml

@@ -12,15 +12,15 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
       - name: Setup JDK 11
-        uses: actions/setup-java@v2
+        uses: actions/setup-java@v4
         with:
           java-version: '11'
           distribution: 'adopt'
           cache: maven
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
       - name: Update submodule
         run: git submodule update --init
       - name: Build with Maven

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

@@ -812,6 +812,15 @@ public class MilvusClientV2 {
         return rpcUtils.retry(()->rgroupService.describeResourceGroup(this.getRpcStub(), request));
     }
 
+    /**
+     * Transfer query nodes from source resource group to target resource_group.
+     *
+     * @param request {@link TransferNodeReq}
+     */
+    public void transferNode(TransferNodeReq request) {
+        rpcUtils.retry(()->rgroupService.transferNode(this.getRpcStub(), request));
+    }
+
     /**
      * Transfer a replica from source resource group to target resource_group.
      *
@@ -886,6 +895,28 @@ public class MilvusClientV2 {
         utilityService.waitFlush(tempBlockingStub, response.getCollectionSegmentIDs(), response.getCollectionFlushTs());
     }
 
+    /**
+     * Gets the information of persistent segments from data node, including row count,
+     * persistence state(growing or flushed), etc.
+     *
+     * @param request get request
+     * @return GetPersistentSegmentInfoResp
+     */
+    GetPersistentSegmentInfoResp getPersistentSegmentInfo(GetPersistentSegmentInfoReq request) {
+        return rpcUtils.retry(()->utilityService.getPersistentSegmentInfo(this.getRpcStub(), request));
+    }
+
+    /**
+     * Gets the query information of segments in a collection from query node, including row count,
+     * memory usage size, index name, etc.
+     *
+     * @param request get request
+     * @return GetQuerySegmentInfoResp
+     */
+    GetQuerySegmentInfoResp getQuerySegmentInfo(GetQuerySegmentInfoReq request){
+        return rpcUtils.retry(()->utilityService.getQuerySegmentInfo(this.getRpcStub(), request));
+    }
+
     /**
      * trigger an asynchronous compaction in server side
      *
@@ -915,6 +946,15 @@ public class MilvusClientV2 {
         return rpcUtils.retry(()->clientUtils.getServerVersion(this.getRpcStub()));
     }
 
+    /**
+     * Check server health
+     *
+     * @return CheckHealthResp
+     */
+    CheckHealthResp checkHealth() {
+        return rpcUtils.retry(()->utilityService.checkHealth(this.getRpcStub()));
+    }
+
     /**
      * Disconnects from a Milvus server with configurable timeout
      *

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

@@ -155,6 +155,25 @@ public class ResourceGroupService extends BaseService {
                 .build();
     }
 
+    public Void transferNode(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, TransferNodeReq request) {
+        if (StringUtils.isEmpty(request.getSourceGroupName())) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid source group name");
+        }
+        if (StringUtils.isEmpty(request.getTargetGroupName())) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid target group name");
+        }
+
+        String title = String.format("TransferNode %d nodes from %s to %s", request.getNumOfNodes(),
+                request.getSourceGroupName(), request.getTargetGroupName());
+        Status response = blockingStub.transferNode(TransferNodeRequest.newBuilder()
+                .setSourceResourceGroup(request.getSourceGroupName())
+                .setTargetResourceGroup(request.getTargetGroupName())
+                .setNumNode(request.getNumOfNodes())
+                .build());
+        rpcUtils.handleResponse(title, response);
+        return null;
+    }
+
     public Void transferReplica(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
                                 TransferReplicaReq request) {
         if (StringUtils.isEmpty(request.getSourceGroupName())) {

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

@@ -0,0 +1,12 @@
+package io.milvus.v2.service.resourcegroup.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class TransferNodeReq {
+    String sourceGroupName;
+    String targetGroupName;
+    Integer numOfNodes;
+}

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

@@ -28,6 +28,7 @@ import io.milvus.v2.service.utility.request.*;
 import io.milvus.v2.service.utility.response.*;
 
 import java.util.*;
+import java.util.stream.Collectors;
 
 public class UtilityService extends BaseService {
     public FlushResp flush(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, FlushReq request) {
@@ -100,7 +101,7 @@ public class UtilityService extends BaseService {
 
     public GetCompactionStateResp getCompactionState(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
                                                      GetCompactionStateReq request) {
-        String title = "Get compaction state";
+        String title = "GetCompactionState";
         io.milvus.grpc.GetCompactionStateRequest getRequest = io.milvus.grpc.GetCompactionStateRequest.newBuilder()
                 .setCompactionID(request.getCompactionID())
                 .build();
@@ -116,7 +117,7 @@ public class UtilityService extends BaseService {
     }
 
     public Void createAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, CreateAliasReq request) {
-        String title = String.format("Create alias %s for collection %s", request.getAlias(), request.getCollectionName());
+        String title = String.format("CreateAlias %s for collection %s", request.getAlias(), request.getCollectionName());
         io.milvus.grpc.CreateAliasRequest createAliasRequest = io.milvus.grpc.CreateAliasRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
                 .setAlias(request.getAlias())
@@ -128,7 +129,7 @@ public class UtilityService extends BaseService {
     }
 
     public Void dropAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, DropAliasReq request) {
-        String title = String.format("Drop alias %s", request.getAlias());
+        String title = String.format("DropAlias %s", request.getAlias());
         io.milvus.grpc.DropAliasRequest dropAliasRequest = io.milvus.grpc.DropAliasRequest.newBuilder()
                 .setAlias(request.getAlias())
                 .build();
@@ -139,7 +140,7 @@ public class UtilityService extends BaseService {
     }
 
     public Void alterAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, AlterAliasReq request) {
-        String title = String.format("Alter alias %s for collection %s", request.getAlias(), request.getCollectionName());
+        String title = String.format("AlterAlias %s for collection %s", request.getAlias(), request.getCollectionName());
         io.milvus.grpc.AlterAliasRequest alterAliasRequest = io.milvus.grpc.AlterAliasRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
                 .setAlias(request.getAlias())
@@ -151,7 +152,7 @@ public class UtilityService extends BaseService {
     }
 
     public DescribeAliasResp describeAlias(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, DescribeAliasReq request) {
-        String title = String.format("Describe alias %s", request.getAlias());
+        String title = String.format("DescribeAlias %s", request.getAlias());
         io.milvus.grpc.DescribeAliasRequest describeAliasRequest = io.milvus.grpc.DescribeAliasRequest.newBuilder()
                 .setAlias(request.getAlias())
                 .build();
@@ -166,7 +167,7 @@ public class UtilityService extends BaseService {
     }
 
     public ListAliasResp listAliases(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, ListAliasesReq request) {
-        String title = "List aliases";
+        String title = "ListAliases";
         io.milvus.grpc.ListAliasesRequest listAliasesRequest = io.milvus.grpc.ListAliasesRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
                 .build();
@@ -179,4 +180,68 @@ public class UtilityService extends BaseService {
                 .alias(response.getAliasesList())
                 .build();
     }
+
+    public CheckHealthResp checkHealth(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub) {
+        String title = "CheckHealth";
+        CheckHealthResponse response = blockingStub.checkHealth(CheckHealthRequest.newBuilder().build());
+        rpcUtils.handleResponse(title, response.getStatus());
+
+        List<String> states = new ArrayList<>();
+        response.getQuotaStatesList().forEach(s->states.add(s.name()));
+        return CheckHealthResp.builder()
+                .isHealthy(response.getIsHealthy())
+                .reasons(response.getReasonsList().stream().collect(Collectors.toList()))
+                .quotaStates(states)
+                .build();
+    }
+
+    public GetPersistentSegmentInfoResp getPersistentSegmentInfo(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                    GetPersistentSegmentInfoReq request) {
+        String title = String.format("GetPersistentSegmentInfo collectionName %s", request.getCollectionName());
+        GetPersistentSegmentInfoResponse response = blockingStub.getPersistentSegmentInfo(GetPersistentSegmentInfoRequest.newBuilder()
+                .setCollectionName(request.getCollectionName())
+                .build());
+        rpcUtils.handleResponse(title, response.getStatus());
+
+        List<GetPersistentSegmentInfoResp.PersistentSegmentInfo> segmentInfos = new ArrayList<>();
+        response.getInfosList().forEach(info->{segmentInfos.add(GetPersistentSegmentInfoResp.PersistentSegmentInfo.builder()
+                .segmentID(info.getSegmentID())
+                .collectionID(info.getCollectionID())
+                .partitionID(info.getPartitionID())
+                .numOfRows(info.getNumRows())
+                .state(info.getState().name())
+                .level(info.getLevel().name())
+                .isSorted(info.getIsSorted())
+                .build());});
+        return GetPersistentSegmentInfoResp.builder()
+                .segmentInfos(segmentInfos)
+                .build();
+    }
+
+    public GetQuerySegmentInfoResp getQuerySegmentInfo(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                                            GetQuerySegmentInfoReq request) {
+        String title = String.format("GetQuerySegmentInfo collectionName %s", request.getCollectionName());
+        GetQuerySegmentInfoResponse response = blockingStub.getQuerySegmentInfo(GetQuerySegmentInfoRequest.newBuilder()
+                .setCollectionName(request.getCollectionName())
+                .build());
+        rpcUtils.handleResponse(title, response.getStatus());
+
+        List<GetQuerySegmentInfoResp.QuerySegmentInfo> segmentInfos = new ArrayList<>();
+        response.getInfosList().forEach(info->{segmentInfos.add(GetQuerySegmentInfoResp.QuerySegmentInfo.builder()
+                .segmentID(info.getSegmentID())
+                .collectionID(info.getCollectionID())
+                .partitionID(info.getPartitionID())
+                .memSize(info.getMemSize())
+                .numOfRows(info.getNumRows())
+                .indexName(info.getIndexName())
+                .indexID(info.getIndexID())
+                .state(info.getState().name())
+                .level(info.getLevel().name())
+                .nodeIDs(info.getNodeIdsList())
+                .isSorted(info.getIsSorted())
+                .build());});
+        return GetQuerySegmentInfoResp.builder()
+                .segmentInfos(segmentInfos)
+                .build();
+    }
 }

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

@@ -0,0 +1,10 @@
+package io.milvus.v2.service.utility.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class GetPersistentSegmentInfoReq {
+    private String collectionName;
+}

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

@@ -0,0 +1,10 @@
+package io.milvus.v2.service.utility.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class GetQuerySegmentInfoReq {
+    private String collectionName;
+}

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

@@ -0,0 +1,19 @@
+package io.milvus.v2.service.utility.response;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@Data
+@SuperBuilder
+public class CheckHealthResp {
+    @Builder.Default
+    Boolean isHealthy = false;
+    @Builder.Default
+    List<String> reasons = new ArrayList<>();
+    @Builder.Default
+    List<String> quotaStates = new ArrayList<>();
+}

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

@@ -0,0 +1,27 @@
+package io.milvus.v2.service.utility.response;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@Data
+@SuperBuilder
+public class GetPersistentSegmentInfoResp {
+    @Data
+    @SuperBuilder
+    public static class PersistentSegmentInfo {
+        private Long segmentID;
+        private Long collectionID;
+        private Long partitionID;
+        private Long numOfRows;
+        private String state;
+        private String level;
+        private Boolean isSorted;
+    }
+
+    @Builder.Default
+    private List<PersistentSegmentInfo> segmentInfos = new ArrayList<>();
+}

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

@@ -0,0 +1,32 @@
+package io.milvus.v2.service.utility.response;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@Data
+@SuperBuilder
+public class GetQuerySegmentInfoResp {
+    @Data
+    @SuperBuilder
+    public static class QuerySegmentInfo {
+        private Long segmentID;
+        private Long collectionID;
+        private Long partitionID;
+        private Long memSize;
+        private Long numOfRows;
+        private String indexName;
+        private Long indexID;
+        private String state;
+        private String level;
+        @Builder.Default
+        private List<Long> nodeIDs = new ArrayList<>();
+        private Boolean isSorted;
+    }
+
+    @Builder.Default
+    private List<QuerySegmentInfo> segmentInfos = new ArrayList<>();
+}

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

@@ -290,6 +290,9 @@ class MilvusClientV2DockerTest {
 
     @Test
     void testFloatVectors() {
+        CheckHealthResp healthy = client.checkHealth();
+        Assertions.assertTrue(healthy.getIsHealthy());
+
         String randomCollectionName = generator.generate(10);
 
         String vectorFieldName = "float_vector";
@@ -321,6 +324,20 @@ class MilvusClientV2DockerTest {
                 .collectionNames(Collections.singletonList(randomCollectionName))
                 .build());
 
+        // get persistent segment info
+        GetPersistentSegmentInfoResp pSegInfo = client.getPersistentSegmentInfo(GetPersistentSegmentInfoReq.builder()
+                .collectionName(randomCollectionName)
+                .build());
+        Assertions.assertEquals(1, pSegInfo.getSegmentInfos().size());
+        GetPersistentSegmentInfoResp.PersistentSegmentInfo pInfo = pSegInfo.getSegmentInfos().get(0);
+        Assertions.assertTrue(pInfo.getSegmentID() > 0L);
+        Assertions.assertTrue(pInfo.getCollectionID() > 0L);
+        Assertions.assertTrue(pInfo.getPartitionID() > 0L);
+        Assertions.assertEquals(count, pInfo.getNumOfRows());
+        Assertions.assertEquals("Flushed", pInfo.getState());
+        Assertions.assertEquals("L1", pInfo.getLevel());
+        Assertions.assertFalse(pInfo.getIsSorted());
+
         // compact
         CompactResp compactResp = client.compact(CompactReq.builder()
                 .collectionName(randomCollectionName)
@@ -347,6 +364,25 @@ class MilvusClientV2DockerTest {
                 .collectionName(randomCollectionName)
                 .build());
 
+        // get query segment info
+        GetQuerySegmentInfoResp qSegInfo = client.getQuerySegmentInfo(GetQuerySegmentInfoReq.builder()
+                .collectionName(randomCollectionName)
+                .build());
+        Assertions.assertEquals(1, qSegInfo.getSegmentInfos().size());
+        GetQuerySegmentInfoResp.QuerySegmentInfo qInfo = qSegInfo.getSegmentInfos().get(0);
+        Assertions.assertTrue(qInfo.getSegmentID() > 0L);
+        Assertions.assertTrue(qInfo.getCollectionID() > 0L);
+        Assertions.assertTrue(qInfo.getPartitionID() > 0L);
+        Assertions.assertTrue(qInfo.getMemSize() >= 0L);
+        Assertions.assertEquals(count, qInfo.getNumOfRows());
+        Assertions.assertEquals(vectorFieldName, qInfo.getIndexName());
+        Assertions.assertTrue(qInfo.getIndexID() > 0L);
+        Assertions.assertEquals("Sealed", qInfo.getState());
+        Assertions.assertEquals("L1", qInfo.getLevel());
+        Assertions.assertEquals(1, qInfo.getNodeIDs().size());
+        Assertions.assertTrue(qInfo.getNodeIDs().get(0) > 0L);
+        Assertions.assertTrue(qInfo.getIsSorted());
+
         // create partition, upsert one row to the partition
         String partitionName = "PPP";
         client.createPartition(CreatePartitionReq.builder()