Update loadCollection() api (#288)

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

src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java

@@ -421,6 +421,7 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         try {
             LoadCollectionRequest loadCollectionRequest = LoadCollectionRequest.newBuilder()
                     .setCollectionName(requestParam.getCollectionName())
+                    .setReplicaNumber(requestParam.getReplicaNumber())
                     .build();
 
             Status response = blockingStub().loadCollection(loadCollectionRequest);
@@ -736,6 +737,7 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         try {
             LoadPartitionsRequest loadPartitionsRequest = LoadPartitionsRequest.newBuilder()
                     .setCollectionName(requestParam.getCollectionName())
+                    .setReplicaNumber(requestParam.getReplicaNumber())
                     .addAllPartitionNames(requestParam.getPartitionNames())
                     .build();
 
@@ -841,7 +843,7 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
             logError("GetPartitionStatisticsRequest RPC failed:\n{}", e.getStatus().toString());
             return R.failed(e);
         } catch (Exception e) {
-            logError("GetQuerySegmentInfoRequest failed:\n{}", e.getMessage());
+            logError("GetPartitionStatisticsRequest failed:\n{}", e.getMessage());
             return R.failed(e);
         }
     }

src/main/java/io/milvus/param/collection/LoadCollectionParam.java

@@ -35,12 +35,14 @@ public class LoadCollectionParam {
     private final boolean syncLoad;
     private final long syncLoadWaitingInterval;
     private final long syncLoadWaitingTimeout;
+    private final int replicaNumber;
 
     public LoadCollectionParam(@NonNull Builder builder) {
         this.collectionName = builder.collectionName;
         this.syncLoad = builder.syncLoad;
         this.syncLoadWaitingInterval = builder.syncLoadWaitingInterval;
         this.syncLoadWaitingTimeout = builder.syncLoadWaitingTimeout;
+        this.replicaNumber = builder.replicaNumber;
     }
 
     public static Builder newBuilder() {
@@ -67,6 +69,10 @@ public class LoadCollectionParam {
         //   this value control the waiting timeout. Unit: second. Default value: 60 seconds.
         private Long syncLoadWaitingTimeout = 60L;
 
+        // replicaNumber:
+        //   The replica number to load, default by 1
+        private Integer replicaNumber = 1;
+
         private Builder() {
         }
 
@@ -121,6 +127,17 @@ public class LoadCollectionParam {
             return this;
         }
 
+        /**
+         * Specify replica number to load
+         *
+         * @param replicaNumber replica number
+         * @return <code>Builder</code>
+         */
+        public Builder withReplicaNumber(@NonNull Integer replicaNumber) {
+            this.replicaNumber = replicaNumber;
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link LoadCollectionParam} instance.
          *

src/main/java/io/milvus/param/partition/LoadPartitionsParam.java

@@ -23,6 +23,7 @@ import io.milvus.exception.ParamException;
 import io.milvus.param.Constant;
 import io.milvus.param.ParamUtils;
 
+import io.milvus.param.collection.LoadCollectionParam;
 import lombok.Getter;
 import lombok.NonNull;
 import java.util.ArrayList;
@@ -38,6 +39,7 @@ public class LoadPartitionsParam {
     private final boolean syncLoad;
     private final long syncLoadWaitingInterval;
     private final long syncLoadWaitingTimeout;
+    private final int replicaNumber;
 
     private LoadPartitionsParam(@NonNull Builder builder) {
         this.collectionName = builder.collectionName;
@@ -45,6 +47,7 @@ public class LoadPartitionsParam {
         this.syncLoad = builder.syncLoad;
         this.syncLoadWaitingInterval = builder.syncLoadWaitingInterval;
         this.syncLoadWaitingTimeout = builder.syncLoadWaitingTimeout;
+        this.replicaNumber = builder.replicaNumber;
     }
 
     public static Builder newBuilder() {
@@ -72,6 +75,10 @@ public class LoadPartitionsParam {
         //   this value control the waiting timeout. Unit: second. Default value: 60 seconds.
         private Long syncLoadWaitingTimeout = 60L;
 
+        // replicaNumber:
+        //   The replica number to load, default by 1
+        private Integer replicaNumber = 1;
+
         private Builder() {
         }
 
@@ -150,6 +157,17 @@ public class LoadPartitionsParam {
             return this;
         }
 
+        /**
+         * Specify replica number to load
+         *
+         * @param replicaNumber replica number
+         * @return <code>Builder</code>
+         */
+        public Builder withReplicaNumber(@NonNull Integer replicaNumber) {
+            this.replicaNumber = replicaNumber;
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link LoadPartitionsParam} instance.
          *