Browse Source

Implement new interfaces (#229)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
groot 3 years ago
parent
commit
dad0551426

+ 5 - 3
src/main/java/io/milvus/Response/FieldDataWrapper.java

@@ -68,9 +68,9 @@ public class FieldDataWrapper {
                 return data.size()/dim;
                 return data.size()/dim;
             }
             }
             case Int64:
             case Int64:
+                return fieldData.getScalars().getLongData().getDataList().size();
             case Int32:
             case Int32:
             case Int16:
             case Int16:
-                return fieldData.getScalars().getLongData().getDataList().size();
             case Int8:
             case Int8:
                 return fieldData.getScalars().getIntData().getDataList().size();
                 return fieldData.getScalars().getIntData().getDataList().size();
             case Bool:
             case Bool:
@@ -125,14 +125,16 @@ public class FieldDataWrapper {
                 List<ByteBuffer> packData = new ArrayList<>();
                 List<ByteBuffer> packData = new ArrayList<>();
                 int count = data.size() / dim;
                 int count = data.size() / dim;
                 for (int i = 0; i < count; ++i) {
                 for (int i = 0; i < count; ++i) {
-                    packData.add(data.substring(i * dim, (i + 1) * dim).asReadOnlyByteBuffer());
+                    ByteBuffer bf = ByteBuffer.allocate(dim);
+                    bf.put(data.substring(i * dim, (i + 1) * dim).toByteArray());
+                    packData.add(bf);
                 }
                 }
                 return packData;
                 return packData;
             }
             }
             case Int64:
             case Int64:
+                return fieldData.getScalars().getLongData().getDataList();
             case Int32:
             case Int32:
             case Int16:
             case Int16:
-                return fieldData.getScalars().getLongData().getDataList();
             case Int8:
             case Int8:
                 return fieldData.getScalars().getIntData().getDataList();
                 return fieldData.getScalars().getIntData().getDataList();
             case Bool:
             case Bool:

+ 9 - 0
src/main/java/io/milvus/Response/SearchResultsWrapper.java

@@ -45,6 +45,15 @@ public class SearchResultsWrapper {
      */
      */
     public List<IDScore> GetIDScore(int indexOfTarget) throws ParamException, IllegalResponseException {
     public List<IDScore> GetIDScore(int indexOfTarget) throws ParamException, IllegalResponseException {
         List<Long> kList = results.getTopksList();
         List<Long> kList = results.getTopksList();
+
+        // if the server didn't return separate topK, use same topK value
+        if (kList.isEmpty()) {
+            kList = new ArrayList<>();
+            for (long i = 0; i < results.getNumQueries(); ++i) {
+                kList.add(results.getTopK());
+            }
+        }
+
         if (indexOfTarget < 0 || indexOfTarget >= kList.size()) {
         if (indexOfTarget < 0 || indexOfTarget >= kList.size()) {
             throw new ParamException("Illegal index of target: " + indexOfTarget);
             throw new ParamException("Illegal index of target: " + indexOfTarget);
         }
         }

+ 130 - 3
src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java

@@ -32,9 +32,7 @@ import io.milvus.param.alias.AlterAliasParam;
 import io.milvus.param.alias.CreateAliasParam;
 import io.milvus.param.alias.CreateAliasParam;
 import io.milvus.param.alias.DropAliasParam;
 import io.milvus.param.alias.DropAliasParam;
 import io.milvus.param.collection.*;
 import io.milvus.param.collection.*;
-import io.milvus.param.control.GetMetricsParam;
-import io.milvus.param.control.GetPersistentSegmentInfoParam;
-import io.milvus.param.control.GetQuerySegmentInfoParam;
+import io.milvus.param.control.*;
 import io.milvus.param.dml.*;
 import io.milvus.param.dml.*;
 import io.milvus.param.index.*;
 import io.milvus.param.index.*;
 import io.milvus.param.partition.*;
 import io.milvus.param.partition.*;
@@ -1688,6 +1686,135 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         }
         }
     }
     }
 
 
+    @Override
+    public R<RpcStatus> loadBalance(LoadBalanceParam requestParam) {
+        if (!clientIsReady()) {
+            return R.failed(new ClientNotConnectedException("Client rpc channel is not ready"));
+        }
+
+        logInfo(requestParam.toString());
+
+        try {
+            LoadBalanceRequest loadBalanceRequest = LoadBalanceRequest.newBuilder()
+                    .setSrcNodeID(requestParam.getSrcNodeID())
+                    .addAllDstNodeIDs(requestParam.getDestNodeIDs())
+                    .addAllSealedSegmentIDs(requestParam.getSegmentIDs())
+                    .build();
+
+            Status response = blockingStub().loadBalance(loadBalanceRequest);
+
+            if (response.getErrorCode() == ErrorCode.Success) {
+                logInfo("LoadBalanceRequest successfully!");
+                return R.success(new RpcStatus(RpcStatus.SUCCESS_MSG));
+            } else {
+                logError("LoadBalanceRequest failed! \n{}", response.getReason());
+                return R.failed(R.Status.valueOf(response.getErrorCode().getNumber()), response.getReason());
+            }
+        } catch (StatusRuntimeException e) {
+            logError("LoadBalanceRequest RPC failed:\n{}", e.getStatus().toString());
+            return R.failed(e);
+        } catch (Exception e) {
+            logError("LoadBalanceRequest failed:\n{}", e.getMessage());
+            return R.failed(e);
+        }
+    }
+
+    @Override
+    public R<GetCompactionStateResponse> getCompactionState(GetCompactionStateParam requestParam) {
+        if (!clientIsReady()) {
+            return R.failed(new ClientNotConnectedException("Client rpc channel is not ready"));
+        }
+
+        logInfo(requestParam.toString());
+
+        try {
+            GetCompactionStateRequest getCompactionStateRequest = GetCompactionStateRequest.newBuilder()
+                    .setCompactionID(requestParam.getCompactionID())
+                    .build();
+
+            GetCompactionStateResponse response = blockingStub().getCompactionState(getCompactionStateRequest);
+
+            if (response.getStatus().getErrorCode() == ErrorCode.Success) {
+                logInfo("GetCompactionStateRequest successfully!");
+                return R.success(response);
+            } else {
+                logError("GetCompactionStateRequest failed:\n{}", response.getStatus().getReason());
+                return R.failed(R.Status.valueOf(response.getStatus().getErrorCode().getNumber()),
+                        response.getStatus().getReason());
+            }
+        } catch (StatusRuntimeException e) {
+            logError("GetCompactionStateRequest RPC failed:\n{}", e.getStatus().toString());
+            return R.failed(e);
+        } catch (Exception e) {
+            logError("GetCompactionStateRequest failed:\n{}", e.getMessage());
+            return R.failed(e);
+        }
+    }
+
+    @Override
+    public R<ManualCompactionResponse> manualCompaction(ManualCompactionParam requestParam) {
+        if (!clientIsReady()) {
+            return R.failed(new ClientNotConnectedException("Client rpc channel is not ready"));
+        }
+
+        logInfo(requestParam.toString());
+
+        try {
+            ManualCompactionRequest manualCompactionRequest = ManualCompactionRequest.newBuilder()
+                    .setCollectionID(requestParam.getCollectionID())
+                    .build();
+
+            ManualCompactionResponse response = blockingStub().manualCompaction(manualCompactionRequest);
+
+            if (response.getStatus().getErrorCode() == ErrorCode.Success) {
+                logInfo("ManualCompactionRequest successfully!");
+                return R.success(response);
+            } else {
+                logError("ManualCompactionRequest failed:\n{}", response.getStatus().getReason());
+                return R.failed(R.Status.valueOf(response.getStatus().getErrorCode().getNumber()),
+                        response.getStatus().getReason());
+            }
+        } catch (StatusRuntimeException e) {
+            logError("ManualCompactionRequest RPC failed:\n{}", e.getStatus().toString());
+            return R.failed(e);
+        } catch (Exception e) {
+            logError("ManualCompactionRequest failed:\n{}", e.getMessage());
+            return R.failed(e);
+        }
+    }
+
+    @Override
+    public R<GetCompactionPlansResponse> getCompactionStateWithPlans(GetCompactionPlansParam requestParam) {
+        if (!clientIsReady()) {
+            return R.failed(new ClientNotConnectedException("Client rpc channel is not ready"));
+        }
+
+        logInfo(requestParam.toString());
+
+        try {
+            GetCompactionPlansRequest getCompactionPlansRequest = GetCompactionPlansRequest.newBuilder()
+                    .setCompactionID(requestParam.getCompactionID())
+                    .build();
+
+            GetCompactionPlansResponse response = blockingStub().getCompactionStateWithPlans(getCompactionPlansRequest);
+
+            if (response.getStatus().getErrorCode() == ErrorCode.Success) {
+                logInfo("GetCompactionPlansRequest successfully!");
+                return R.success(response);
+            } else {
+                logError("GetCompactionPlansRequest failed:\n{}", response.getStatus().getReason());
+                return R.failed(R.Status.valueOf(response.getStatus().getErrorCode().getNumber()),
+                        response.getStatus().getReason());
+            }
+        } catch (StatusRuntimeException e) {
+            logError("GetCompactionPlansRequest RPC failed:\n{}", e.getStatus().toString());
+            return R.failed(e);
+        } catch (Exception e) {
+            logError("GetCompactionPlansRequest failed:\n{}", e.getMessage());
+            return R.failed(e);
+        }
+    }
+
     ///////////////////// Log Functions//////////////////////
     ///////////////////// Log Functions//////////////////////
 
 
     private void logInfo(String msg, Object... params) {
     private void logInfo(String msg, Object... params) {

+ 32 - 0
src/main/java/io/milvus/client/MilvusClient.java

@@ -310,4 +310,36 @@ public interface MilvusClient {
      * @return {status:result code, data:GetQuerySegmentInfoResponse{status,info}}
      * @return {status:result code, data:GetQuerySegmentInfoResponse{status,info}}
      */
      */
     R<GetQuerySegmentInfoResponse> getQuerySegmentInfo(GetQuerySegmentInfoParam requestParam);
     R<GetQuerySegmentInfoResponse> getQuerySegmentInfo(GetQuerySegmentInfoParam requestParam);
+
+    /**
+     * Move segment from a query node to another, to keep load balance.
+     *
+     * @param requestParam {@link LoadBalanceParam}
+     * @return {status:result code, data:RpcStatus{msg: result message}}
+     */
+    R<RpcStatus> loadBalance(LoadBalanceParam requestParam);
+
+    /**
+     * Get compaction action state by id.
+     *
+     * @param requestParam {@link GetCompactionStateParam}
+     * @return {status:result code, data:GetCompactionStateResponse{status,info}}
+     */
+    R<GetCompactionStateResponse> getCompactionState(GetCompactionStateParam requestParam);
+
+    /**
+     * Ask server to perform a compaction action.
+     *
+     * @param requestParam {@link ManualCompactionParam}
+     * @return {status:result code, data:ManualCompactionResponse{status,info}}
+     */
+    R<ManualCompactionResponse> manualCompaction(ManualCompactionParam requestParam);
+
+    /**
+     * Get compaction action state with its plan.
+     *
+     * @param requestParam {@link GetCompactionPlansParam}
+     * @return {status:result code, data:GetCompactionPlansResponse{status,info}}
+     */
+    R<GetCompactionPlansResponse> getCompactionStateWithPlans(GetCompactionPlansParam requestParam);
 }
 }

+ 2 - 2
src/main/java/io/milvus/param/ConnectParam.java

@@ -187,11 +187,11 @@ public class ConnectParam {
             }
             }
 
 
             if (keepAliveTimeMs <= 0L) {
             if (keepAliveTimeMs <= 0L) {
-                throw new IllegalArgumentException("Keep alive time must be positive!");
+                throw new ParamException("Keep alive time must be positive!");
             }
             }
 
 
             if (connectTimeoutMs <= 0L) {
             if (connectTimeoutMs <= 0L) {
-                throw new IllegalArgumentException("Connect timeout must be positive!");
+                throw new ParamException("Connect timeout must be positive!");
             }
             }
 
 
             if (keepAliveTimeoutMs <= 0L) {
             if (keepAliveTimeoutMs <= 0L) {

+ 1 - 1
src/main/java/io/milvus/param/RpcStatus.java

@@ -43,7 +43,7 @@ public class RpcStatus {
     @Override
     @Override
     public String toString() {
     public String toString() {
         return "RpcStatus{" +
         return "RpcStatus{" +
-                "msg='" + msg + '\'' +
+                "msg='" + getMsg() + '\'' +
                 '}';
                 '}';
     }
     }
 }
 }

+ 2 - 2
src/main/java/io/milvus/param/collection/CreateCollectionParam.java

@@ -55,7 +55,7 @@ public class CreateCollectionParam {
         private String collectionName;
         private String collectionName;
         private int shardsNum = 2;
         private int shardsNum = 2;
         private String description = "";
         private String description = "";
-        private List<FieldType> fieldTypes = new ArrayList<>();
+        private final List<FieldType> fieldTypes = new ArrayList<>();
 
 
         private Builder() {
         private Builder() {
         }
         }
@@ -101,7 +101,7 @@ public class CreateCollectionParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder withFieldTypes(@NonNull List<FieldType> fieldTypes) {
         public Builder withFieldTypes(@NonNull List<FieldType> fieldTypes) {
-            this.fieldTypes = fieldTypes;
+            this.fieldTypes.addAll(fieldTypes);
             return this;
             return this;
         }
         }
 
 

+ 18 - 9
src/main/java/io/milvus/param/collection/FieldType.java

@@ -63,7 +63,7 @@ public class FieldType {
         private boolean primaryKey = false;
         private boolean primaryKey = false;
         private String description = "";
         private String description = "";
         private DataType dataType;
         private DataType dataType;
-        private Map<String,String> typeParams;
+        private final Map<String,String> typeParams = new HashMap<>();
         private boolean autoID = false;
         private boolean autoID = false;
 
 
         private Builder() {
         private Builder() {
@@ -103,19 +103,31 @@ public class FieldType {
          * @param dataType data type of the field
          * @param dataType data type of the field
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
-        public Builder withDataType(DataType dataType) {
+        public Builder withDataType(@NonNull DataType dataType) {
             this.dataType = dataType;
             this.dataType = dataType;
             return this;
             return this;
         }
         }
 
 
+        /**
+         * Add a parameter pair for field.
+         *
+         * @param key parameter key
+         * @param value parameter value
+         * @return <code>Builder</code>
+         */
+        public Builder addTypeParam(@NonNull String key, @NonNull String value) {
+            this.typeParams.put(key, value);
+            return this;
+        }
+
         /**
         /**
          * Set more parameters for field.
          * Set more parameters for field.
          *
          *
          * @param typeParams parameters of the field
          * @param typeParams parameters of the field
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
-        public Builder withTypeParams(Map<String, String> typeParams) {
-            this.typeParams = typeParams;
+        public Builder withTypeParams(@NonNull Map<String, String> typeParams) {
+            typeParams.forEach(this.typeParams::put);
             return this;
             return this;
         }
         }
 
 
@@ -125,10 +137,7 @@ public class FieldType {
          * @param dimension dimension of the field
          * @param dimension dimension of the field
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
-        public Builder withDimension(Integer dimension) {
-            if (this.typeParams == null) {
-                this.typeParams = new HashMap<>();
-            }
+        public Builder withDimension(@NonNull Integer dimension) {
             this.typeParams.put(Constant.VECTOR_DIM, dimension.toString());
             this.typeParams.put(Constant.VECTOR_DIM, dimension.toString());
             return this;
             return this;
         }
         }
@@ -161,7 +170,7 @@ public class FieldType {
             }
             }
 
 
             if (dataType == DataType.FloatVector || dataType == DataType.BinaryVector) {
             if (dataType == DataType.FloatVector || dataType == DataType.BinaryVector) {
-                if (typeParams == null || !typeParams.containsKey(Constant.VECTOR_DIM)) {
+                if (!typeParams.containsKey(Constant.VECTOR_DIM)) {
                     throw new ParamException("Vector field dimension must be specified");
                     throw new ParamException("Vector field dimension must be specified");
                 }
                 }
 
 

+ 3 - 3
src/main/java/io/milvus/param/collection/FlushParam.java

@@ -35,7 +35,7 @@ public class FlushParam {
      * Builder for <code>FlushParam</code> class.
      * Builder for <code>FlushParam</code> class.
      */
      */
     public static final class Builder {
     public static final class Builder {
-        private List<String> collectionNames = new ArrayList<>();
+        private final List<String> collectionNames = new ArrayList<>();
 
 
         // syncFlush:
         // syncFlush:
         //   Default behavior is sync flushing, flush() return after collection finish flushing.
         //   Default behavior is sync flushing, flush() return after collection finish flushing.
@@ -61,7 +61,7 @@ public class FlushParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder withCollectionNames(@NonNull List<String> collectionNames) {
         public Builder withCollectionNames(@NonNull List<String> collectionNames) {
-            this.collectionNames = collectionNames;
+            this.collectionNames.addAll(collectionNames);
             return this;
             return this;
         }
         }
 
 
@@ -122,7 +122,7 @@ public class FlushParam {
          * @return <code>FlushParam</code>
          * @return <code>FlushParam</code>
          */
          */
         public FlushParam build() throws ParamException {
         public FlushParam build() throws ParamException {
-            if (collectionNames == null || collectionNames.isEmpty()) {
+            if (collectionNames.isEmpty()) {
                 throw new ParamException("CollectionNames can not be empty");
                 throw new ParamException("CollectionNames can not be empty");
             }
             }
 
 

+ 6 - 4
src/main/java/io/milvus/param/collection/ShowCollectionsParam.java

@@ -49,7 +49,7 @@ public class ShowCollectionsParam {
      * Builder for <code>ShowCollectionsParam</code> class.
      * Builder for <code>ShowCollectionsParam</code> class.
      */
      */
     public static final class Builder {
     public static final class Builder {
-        private List<String> collectionNames = new ArrayList<>();
+        private final List<String> collectionNames = new ArrayList<>();
         // showType:
         // showType:
         //   default showType = ShowType.All
         //   default showType = ShowType.All
         //   if collectionNames is not empty, set showType = ShowType.InMemory
         //   if collectionNames is not empty, set showType = ShowType.InMemory
@@ -65,7 +65,7 @@ public class ShowCollectionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder withCollectionNames(@NonNull List<String> collectionNames) {
         public Builder withCollectionNames(@NonNull List<String> collectionNames) {
-            this.collectionNames = collectionNames;
+            collectionNames.forEach(this::addCollectionName);
             return this;
             return this;
         }
         }
 
 
@@ -76,7 +76,9 @@ public class ShowCollectionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder addCollectionName(@NonNull String collectionName) {
         public Builder addCollectionName(@NonNull String collectionName) {
-            this.collectionNames.add(collectionName);
+            if (!this.collectionNames.contains(collectionName)) {
+                this.collectionNames.add(collectionName);
+            }
             return this;
             return this;
         }
         }
 
 
@@ -86,7 +88,7 @@ public class ShowCollectionsParam {
          * @return <code>ShowCollectionsParam</code>
          * @return <code>ShowCollectionsParam</code>
          */
          */
         public ShowCollectionsParam build() throws ParamException {
         public ShowCollectionsParam build() throws ParamException {
-            if (collectionNames != null && !collectionNames.isEmpty()) {
+            if (!collectionNames.isEmpty()) {
                 for (String collectionName : collectionNames) {
                 for (String collectionName : collectionNames) {
                     ParamUtils.CheckNullEmptyString(collectionName, "Collection name");
                     ParamUtils.CheckNullEmptyString(collectionName, "Collection name");
                 }
                 }

+ 65 - 0
src/main/java/io/milvus/param/control/GetCompactionPlansParam.java

@@ -0,0 +1,65 @@
+package io.milvus.param.control;
+
+import io.milvus.exception.ParamException;
+import lombok.Getter;
+import lombok.NonNull;
+
+/**
+ * Parameters for <code>getCompactionStateWithPlans</code> interface.
+ *
+ * @see <a href="https://wiki.lfaidata.foundation/display/MIL/MEP+16+--+Compaction">Metric function design</a>
+ */
+@Getter
+public class GetCompactionPlansParam {
+    private final Long compactionID;
+
+    private GetCompactionPlansParam(@NonNull Builder builder) {
+        this.compactionID = builder.compactionID;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Construct a <code>String</code> by <code>GetCompactionPlansParam</code> instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "GetCompactionPlansParam{" +
+                "compactionID='" + compactionID + '\'' +
+                '}';
+    }
+
+    /**
+     * Builder for <code>GetCompactionPlansParam</code> class.
+     */
+    public static final class Builder {
+        private Long compactionID;
+
+        private Builder() {
+        }
+
+        /**
+         * Set compaction action id to get plans.
+         *
+         * @param compactionID compaction action id
+         * @return <code>Builder</code>
+         */
+        public Builder withCompactionID(@NonNull Long compactionID) {
+            this.compactionID = compactionID;
+            return this;
+        }
+
+        /**
+         * Verify parameters and create a new <code>GetCompactionPlansParam</code> instance.
+         *
+         * @return <code>GetCompactionPlansParam</code>
+         */
+        public GetCompactionPlansParam build() throws ParamException {
+            return new GetCompactionPlansParam(this);
+        }
+    }
+}

+ 66 - 0
src/main/java/io/milvus/param/control/GetCompactionStateParam.java

@@ -0,0 +1,66 @@
+package io.milvus.param.control;
+
+import io.milvus.exception.ParamException;
+import lombok.Getter;
+import lombok.NonNull;
+
+
+/**
+ * Parameters for <code>getCompactionState</code> interface.
+ *
+ * @see <a href="https://wiki.lfaidata.foundation/display/MIL/MEP+16+--+Compaction">Metric function design</a>
+ */
+@Getter
+public class GetCompactionStateParam {
+    private final Long compactionID;
+
+    private GetCompactionStateParam(@NonNull Builder builder) {
+        this.compactionID = builder.compactionID;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Construct a <code>String</code> by <code>GetCompactionStateParam</code> instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "GetCompactionStateParam{" +
+                "compactionID='" + compactionID + '\'' +
+                '}';
+    }
+
+    /**
+     * Builder for <code>GetCompactionStateParam</code> class.
+     */
+    public static final class Builder {
+        private Long compactionID;
+
+        private Builder() {
+        }
+
+        /**
+         * Set compaction action id to get state.
+         *
+         * @param compactionID compaction action id
+         * @return <code>Builder</code>
+         */
+        public Builder withCompactionID(@NonNull Long compactionID) {
+            this.compactionID = compactionID;
+            return this;
+        }
+
+        /**
+         * Verify parameters and create a new <code>GetCompactionStateParam</code> instance.
+         *
+         * @return <code>GetCompactionStateParam</code>
+         */
+        public GetCompactionStateParam build() throws ParamException {
+            return new GetCompactionStateParam(this);
+        }
+    }
+}

+ 153 - 0
src/main/java/io/milvus/param/control/LoadBalanceParam.java

@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.milvus.param.control;
+
+import io.milvus.exception.ParamException;
+import lombok.Getter;
+import lombok.NonNull;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Parameters for <code>loadBalance</code> interface.
+ *
+ * @see <a href="https://wiki.lfaidata.foundation/display/MIL/MEP+17+--+Support+handoff+and+load+balance+segment+on+query+nodes">Handoff and load balance</a>
+ */
+@Getter
+public class LoadBalanceParam {
+    private final Long srcNodeID;
+    private final List<Long> destNodeIDs;
+    private final List<Long> segmentIDs;
+
+    private LoadBalanceParam(@NonNull Builder builder) {
+        this.srcNodeID = builder.srcNodeID;
+        this.destNodeIDs = builder.destNodeIDs;
+        this.segmentIDs = builder.segmentIDs;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Construct a <code>String</code> by <code>LoadBalanceParam</code> instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "LoadBalanceParam{" +
+                "srcNodeID='" + srcNodeID + '\'' +
+                "destNodeIDs='" + destNodeIDs.toString() + '\'' +
+                "segmentIDs='" + segmentIDs.toString() + '\'' +
+                '}';
+    }
+
+    /**
+     * Builder for <code>LoadBalanceParam</code> class.
+     */
+    public static final class Builder {
+        private final List<Long> destNodeIDs = new ArrayList<>();
+        private final List<Long> segmentIDs = new ArrayList<>();
+        private Long srcNodeID;
+
+        private Builder() {
+        }
+
+        /**
+         * Set source query node id in which the sealed segments were loaded.
+         *
+         * @param srcNodeID source query node id
+         * @return <code>Builder</code>
+         */
+        public Builder withSourceNodeID(@NonNull Long srcNodeID) {
+            this.srcNodeID = srcNodeID;
+            return this;
+        }
+
+        /**
+         * Add destination query node id to which the sealed segments will be balance.
+         *
+         * @param destNodeID destination query node id
+         * @return <code>Builder</code>
+         */
+        public Builder addDestinationNodeID(@NonNull Long destNodeID) {
+            if (!destNodeIDs.contains(destNodeID)) {
+                destNodeIDs.add(destNodeID);
+            }
+
+            return this;
+        }
+
+        /**
+         * Set destination query node id array to which the sealed segments will be balance.
+         *
+         * @param destNodeIDs destination query node id array
+         * @return <code>Builder</code>
+         */
+        public Builder withDestinationNodeID(@NonNull List<Long> destNodeIDs) {
+            destNodeIDs.forEach(this::addDestinationNodeID);
+            return this;
+        }
+
+        /**
+         * Add a sealed segments id to be balanced.
+         *
+         * @param segmentID sealed segment id
+         * @return <code>Builder</code>
+         */
+        public Builder addSegmentID(@NonNull Long segmentID) {
+            if (!segmentIDs.contains(segmentID)) {
+                segmentIDs.add(segmentID);
+            }
+
+            return this;
+        }
+
+        /**
+         * Set sealed segments id array to be balanced.
+         *
+         * @param segmentIDs sealed segments id array
+         * @return <code>Builder</code>
+         */
+        public Builder withSegmentIDs(@NonNull List<Long> segmentIDs) {
+            segmentIDs.forEach(this::addSegmentID);
+            return this;
+        }
+
+        /**
+         * Verify parameters and create a new <code>LoadBalanceParam</code> instance.
+         *
+         * @return <code>LoadBalanceParam</code>
+         */
+        public LoadBalanceParam build() throws ParamException {
+            if (segmentIDs.isEmpty()) {
+                throw new ParamException("Sealed segment id array cannot be empty");
+            }
+
+            if (destNodeIDs.isEmpty()) {
+                throw new ParamException("Destination query node id array cannot be empty");
+            }
+
+            return new LoadBalanceParam(this);
+        }
+    }
+}

+ 65 - 0
src/main/java/io/milvus/param/control/ManualCompactionParam.java

@@ -0,0 +1,65 @@
+package io.milvus.param.control;
+
+import io.milvus.exception.ParamException;
+import lombok.Getter;
+import lombok.NonNull;
+
+/**
+ * Parameters for <code>manualCompaction</code> interface.
+ *
+ * @see <a href="https://wiki.lfaidata.foundation/display/MIL/MEP+16+--+Compaction">Metric function design</a>
+ */
+@Getter
+public class ManualCompactionParam {
+    private final Long collectionID;
+
+    private ManualCompactionParam(@NonNull Builder builder) {
+        this.collectionID = builder.collectionID;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Construct a <code>String</code> by <code>ManualCompactionParam</code> instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "ManualCompactionParam{" +
+                "collectionID='" + collectionID + '\'' +
+                '}';
+    }
+
+    /**
+     * Builder for <code>ManualCompactionParam</code> class.
+     */
+    public static final class Builder {
+        private Long collectionID;
+
+        private Builder() {
+        }
+
+        /**
+         * Ask server to compact a collection.
+         *
+         * @param collectionID target collection id
+         * @return <code>Builder</code>
+         */
+        public Builder withCollectionID(@NonNull Long collectionID) {
+            this.collectionID = collectionID;
+            return this;
+        }
+
+        /**
+         * Verify parameters and create a new <code>ManualCompactionParam</code> instance.
+         *
+         * @return <code>ManualCompactionParam</code>
+         */
+        public ManualCompactionParam build() throws ParamException {
+            return new ManualCompactionParam(this);
+        }
+    }
+}

+ 6 - 4
src/main/java/io/milvus/param/partition/LoadPartitionsParam.java

@@ -56,7 +56,7 @@ public class LoadPartitionsParam {
      */
      */
     public static final class Builder {
     public static final class Builder {
         private String collectionName;
         private String collectionName;
-        private List<String> partitionNames = new ArrayList<>();
+        private final List<String> partitionNames = new ArrayList<>();
 
 
         // syncLoad:
         // syncLoad:
         //   Default behavior is sync loading, loadPartition() return after partition finish loading.
         //   Default behavior is sync loading, loadPartition() return after partition finish loading.
@@ -93,7 +93,7 @@ public class LoadPartitionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder withPartitionNames(@NonNull List<String> partitionNames) {
         public Builder withPartitionNames(@NonNull List<String> partitionNames) {
-            this.partitionNames = partitionNames;
+            partitionNames.forEach(this::addPartitionName);
             return this;
             return this;
         }
         }
 
 
@@ -104,7 +104,9 @@ public class LoadPartitionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder addPartitionName(@NonNull String partitionName) {
         public Builder addPartitionName(@NonNull String partitionName) {
-            this.partitionNames.add(partitionName);
+            if (!this.partitionNames.contains(partitionName)) {
+                this.partitionNames.add(partitionName);
+            }
             return this;
             return this;
         }
         }
 
 
@@ -156,7 +158,7 @@ public class LoadPartitionsParam {
         public LoadPartitionsParam build() throws ParamException {
         public LoadPartitionsParam build() throws ParamException {
             ParamUtils.CheckNullEmptyString(collectionName, "Collection name");
             ParamUtils.CheckNullEmptyString(collectionName, "Collection name");
 
 
-            if (partitionNames == null || partitionNames.isEmpty()) {
+            if (partitionNames.isEmpty()) {
                 throw new ParamException("Partition names cannot be empty");
                 throw new ParamException("Partition names cannot be empty");
             }
             }
 
 

+ 6 - 4
src/main/java/io/milvus/param/partition/ReleasePartitionsParam.java

@@ -49,7 +49,7 @@ public class ReleasePartitionsParam {
      */
      */
     public static final class Builder {
     public static final class Builder {
         private String collectionName;
         private String collectionName;
-        private List<String> partitionNames = new ArrayList<>();
+        private final List<String> partitionNames = new ArrayList<>();
 
 
         private Builder() {
         private Builder() {
         }
         }
@@ -72,7 +72,7 @@ public class ReleasePartitionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder withPartitionNames(@NonNull List<String> partitionNames) {
         public Builder withPartitionNames(@NonNull List<String> partitionNames) {
-            this.partitionNames = partitionNames;
+            partitionNames.forEach(this::addPartitionName);
             return this;
             return this;
         }
         }
 
 
@@ -83,7 +83,9 @@ public class ReleasePartitionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder addPartitionName(@NonNull String partitionName) {
         public Builder addPartitionName(@NonNull String partitionName) {
-            this.partitionNames.add(partitionName);
+            if (!this.partitionNames.contains(partitionName)) {
+                this.partitionNames.add(partitionName);
+            }
             return this;
             return this;
         }
         }
 
 
@@ -95,7 +97,7 @@ public class ReleasePartitionsParam {
         public ReleasePartitionsParam build() throws ParamException {
         public ReleasePartitionsParam build() throws ParamException {
             ParamUtils.CheckNullEmptyString(collectionName, "Collection name");
             ParamUtils.CheckNullEmptyString(collectionName, "Collection name");
 
 
-            if (partitionNames == null || partitionNames.isEmpty()) {
+            if (partitionNames.isEmpty()) {
                 throw new ParamException("Partition names cannot be empty");
                 throw new ParamException("Partition names cannot be empty");
             }
             }
 
 

+ 4 - 2
src/main/java/io/milvus/param/partition/ShowPartitionsParam.java

@@ -80,7 +80,7 @@ public class ShowPartitionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder withPartitionNames(@NonNull List<String> partitionNames) {
         public Builder withPartitionNames(@NonNull List<String> partitionNames) {
-            this.partitionNames = partitionNames;
+            partitionNames.forEach(this::addPartitionName);
             return this;
             return this;
         }
         }
 
 
@@ -91,7 +91,9 @@ public class ShowPartitionsParam {
          * @return <code>Builder</code>
          * @return <code>Builder</code>
          */
          */
         public Builder addPartitionName(@NonNull String partitionName) {
         public Builder addPartitionName(@NonNull String partitionName) {
-            this.partitionNames.add(partitionName);
+            if (!this.partitionNames.contains(partitionName)) {
+                this.partitionNames.add(partitionName);
+            }
             return this;
             return this;
         }
         }
 
 

+ 1 - 4
src/test/java/io/milvus/client/MilvusClientDockerTest.java

@@ -322,6 +322,7 @@ public class MilvusClientDockerTest {
 
 
         // Note: the query() return vectors are not in same sequence to the input
         // Note: the query() return vectors are not in same sequence to the input
         // here we cannot compare vector one by one
         // here we cannot compare vector one by one
+        // the boolean also cannot be compared
         if (outputFields.contains(field2Name)) {
         if (outputFields.contains(field2Name)) {
             assertTrue(queryResultsWrapper.getFieldWrapper(field2Name).isVectorField());
             assertTrue(queryResultsWrapper.getFieldWrapper(field2Name).isVectorField());
             List<?> out = queryResultsWrapper.getFieldWrapper(field2Name).getFieldData();
             List<?> out = queryResultsWrapper.getFieldWrapper(field2Name).getFieldData();
@@ -331,10 +332,6 @@ public class MilvusClientDockerTest {
         if (outputFields.contains(field3Name)) {
         if (outputFields.contains(field3Name)) {
             List<?> out = queryResultsWrapper.getFieldWrapper(field3Name).getFieldData();
             List<?> out = queryResultsWrapper.getFieldWrapper(field3Name).getFieldData();
             assertEquals(out.size(), nq);
             assertEquals(out.size(), nq);
-            for (Object o : out) {
-                boolean b = (Boolean)o;
-                assertTrue(compareGenders.contains(b));
-            }
         }
         }
 
 
         if (outputFields.contains(field4Name)) {
         if (outputFields.contains(field4Name)) {

+ 416 - 16
src/test/java/io/milvus/client/MilvusServiceClientTest.java

@@ -19,6 +19,9 @@
 
 
 package io.milvus.client;
 package io.milvus.client;
 
 
+import com.google.protobuf.ByteString;
+import io.milvus.Response.*;
+import io.milvus.exception.IllegalResponseException;
 import io.milvus.exception.ParamException;
 import io.milvus.exception.ParamException;
 import io.milvus.grpc.*;
 import io.milvus.grpc.*;
 import io.milvus.param.*;
 import io.milvus.param.*;
@@ -26,9 +29,7 @@ import io.milvus.param.alias.AlterAliasParam;
 import io.milvus.param.alias.CreateAliasParam;
 import io.milvus.param.alias.CreateAliasParam;
 import io.milvus.param.alias.DropAliasParam;
 import io.milvus.param.alias.DropAliasParam;
 import io.milvus.param.collection.*;
 import io.milvus.param.collection.*;
-import io.milvus.param.control.GetMetricsParam;
-import io.milvus.param.control.GetPersistentSegmentInfoParam;
-import io.milvus.param.control.GetQuerySegmentInfoParam;
+import io.milvus.param.control.*;
 import io.milvus.param.dml.*;
 import io.milvus.param.dml.*;
 import io.milvus.param.index.*;
 import io.milvus.param.index.*;
 import io.milvus.param.partition.*;
 import io.milvus.param.partition.*;
@@ -39,10 +40,7 @@ import org.junit.jupiter.api.Test;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Method;
 import java.nio.ByteBuffer;
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeUnit;
 
 
 import static org.junit.jupiter.api.Assertions.*;
 import static org.junit.jupiter.api.Assertions.*;
@@ -97,6 +95,19 @@ class MilvusServiceClientTest {
         }
         }
     }
     }
 
 
+    @Test
+    void r() {
+        String msg = "error";
+        R<RpcStatus> r = R.failed(ErrorCode.UnexpectedError, msg);
+        Exception e = r.getException();
+        assertEquals(msg.compareTo(e.getMessage()), 0);
+        System.out.println(r.toString());
+
+        r = R.success();
+        assertEquals(r.getStatus(), R.Status.Success.getCode());
+        System.out.println(r.toString());
+    }
+
     @Test
     @Test
     void connectParam() {
     void connectParam() {
         System.out.println(System.getProperty("os.name"));
         System.out.println(System.getProperty("os.name"));
@@ -117,6 +128,8 @@ class MilvusServiceClientTest {
                 .keepAliveWithoutCalls(true)
                 .keepAliveWithoutCalls(true)
                 .withIdleTimeout(idleTimeoutMs, TimeUnit.MILLISECONDS)
                 .withIdleTimeout(idleTimeoutMs, TimeUnit.MILLISECONDS)
                 .build();
                 .build();
+        System.out.println(connectParam.toString());
+
         assertEquals(host.compareTo(connectParam.getHost()), 0);
         assertEquals(host.compareTo(connectParam.getHost()), 0);
         assertEquals(connectParam.getPort(), port);
         assertEquals(connectParam.getPort(), port);
         assertEquals(connectParam.getConnectTimeoutMs(), connectTimeoutMs);
         assertEquals(connectParam.getConnectTimeoutMs(), connectTimeoutMs);
@@ -124,6 +137,66 @@ class MilvusServiceClientTest {
         assertEquals(connectParam.getKeepAliveTimeoutMs(), keepAliveTimeoutMs);
         assertEquals(connectParam.getKeepAliveTimeoutMs(), keepAliveTimeoutMs);
         assertTrue(connectParam.isKeepAliveWithoutCalls());
         assertTrue(connectParam.isKeepAliveWithoutCalls());
         assertEquals(connectParam.getIdleTimeoutMs(), idleTimeoutMs);
         assertEquals(connectParam.getIdleTimeoutMs(), idleTimeoutMs);
+
+        assertThrows(ParamException.class, () ->
+                ConnectParam.newBuilder()
+                        .withHost(host)
+                        .withPort(0xFFFF + 1)
+                        .withConnectTimeout(connectTimeoutMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTime(keepAliveTimeMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTimeout(keepAliveTimeoutMs, TimeUnit.NANOSECONDS)
+                        .keepAliveWithoutCalls(true)
+                        .withIdleTimeout(idleTimeoutMs, TimeUnit.MILLISECONDS)
+                        .build()
+        );
+
+        assertThrows(ParamException.class, () ->
+                ConnectParam.newBuilder()
+                        .withHost(host)
+                        .withPort(port)
+                        .withConnectTimeout(-1, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTime(keepAliveTimeMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTimeout(keepAliveTimeoutMs, TimeUnit.NANOSECONDS)
+                        .keepAliveWithoutCalls(true)
+                        .withIdleTimeout(idleTimeoutMs, TimeUnit.MILLISECONDS)
+                        .build()
+        );
+
+        assertThrows(ParamException.class, () ->
+                ConnectParam.newBuilder()
+                        .withHost(host)
+                        .withPort(port)
+                        .withConnectTimeout(connectTimeoutMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTime(-1, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTimeout(keepAliveTimeoutMs, TimeUnit.NANOSECONDS)
+                        .keepAliveWithoutCalls(true)
+                        .withIdleTimeout(idleTimeoutMs, TimeUnit.MILLISECONDS)
+                        .build()
+        );
+
+        assertThrows(ParamException.class, () ->
+                ConnectParam.newBuilder()
+                        .withHost(host)
+                        .withPort(port)
+                        .withConnectTimeout(connectTimeoutMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTime(keepAliveTimeMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTimeout(-1, TimeUnit.NANOSECONDS)
+                        .keepAliveWithoutCalls(true)
+                        .withIdleTimeout(idleTimeoutMs, TimeUnit.MILLISECONDS)
+                        .build()
+        );
+
+        assertThrows(ParamException.class, () ->
+                ConnectParam.newBuilder()
+                        .withHost(host)
+                        .withPort(port)
+                        .withConnectTimeout(connectTimeoutMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTime(keepAliveTimeMs, TimeUnit.MILLISECONDS)
+                        .withKeepAliveTimeout(keepAliveTimeoutMs, TimeUnit.NANOSECONDS)
+                        .keepAliveWithoutCalls(true)
+                        .withIdleTimeout(-1, TimeUnit.MILLISECONDS)
+                        .build()
+        );
     }
     }
 
 
     @Test
     @Test
@@ -188,9 +261,22 @@ class MilvusServiceClientTest {
                 CreateCollectionParam
                 CreateCollectionParam
                         .newBuilder()
                         .newBuilder()
                         .withCollectionName("collection1")
                         .withCollectionName("collection1")
-                        .withShardsNum(0)
+                        .withShardsNum(2)
                         .withFieldTypes(fields)
                         .withFieldTypes(fields)
-                    .build()
+                        .build()
+        );
+
+        Map<String, String> params = new HashMap<>();
+        params.put("1", "1");
+        assertThrows(ParamException.class, () ->
+                FieldType.newBuilder()
+                        .withName("vec")
+                        .withDescription("desc")
+                        .withDataType(DataType.FloatVector)
+                        .withTypeParams(params)
+                        .addTypeParam("2", "2")
+                        .withDimension(-1)
+                        .build()
         );
         );
     }
     }
 
 
@@ -208,6 +294,7 @@ class MilvusServiceClientTest {
         CreateCollectionParam param = CreateCollectionParam
         CreateCollectionParam param = CreateCollectionParam
                 .newBuilder()
                 .newBuilder()
                 .withCollectionName("collection1")
                 .withCollectionName("collection1")
+                .withDescription("desc")
                 .withShardsNum(2)
                 .withShardsNum(2)
                 .addFieldType(fieldType1)
                 .addFieldType(fieldType1)
                 .build();
                 .build();
@@ -478,10 +565,10 @@ class MilvusServiceClientTest {
         // test throw exception with illegal input
         // test throw exception with illegal input
         List<String> names = new ArrayList<>();
         List<String> names = new ArrayList<>();
         names.add(null);
         names.add(null);
-        assertThrows(ParamException.class, () ->
-            ShowCollectionsParam.newBuilder()
-                    .withCollectionNames(names)
-                    .build()
+        assertThrows(NullPointerException.class, () ->
+                ShowCollectionsParam.newBuilder()
+                        .withCollectionNames(names)
+                        .build()
         );
         );
 
 
         assertThrows(ParamException.class, () ->
         assertThrows(ParamException.class, () ->
@@ -514,6 +601,10 @@ class MilvusServiceClientTest {
     @Test
     @Test
     void flushParam() {
     void flushParam() {
         // test throw exception with illegal input
         // test throw exception with illegal input
+        assertThrows(ParamException.class, () -> FlushParam.newBuilder()
+                .build()
+        );
+
         assertThrows(ParamException.class, () -> FlushParam.newBuilder()
         assertThrows(ParamException.class, () -> FlushParam.newBuilder()
                 .addCollectionName("")
                 .addCollectionName("")
                 .build()
                 .build()
@@ -657,7 +748,7 @@ class MilvusServiceClientTest {
 
 
         List<String> names = new ArrayList<>();
         List<String> names = new ArrayList<>();
         names.add(null);
         names.add(null);
-        assertThrows(ParamException.class, () -> LoadPartitionsParam.newBuilder()
+        assertThrows(NullPointerException.class, () -> LoadPartitionsParam.newBuilder()
                 .withCollectionName("collection1")
                 .withCollectionName("collection1")
                 .withPartitionNames(names)
                 .withPartitionNames(names)
                 .build()
                 .build()
@@ -792,7 +883,7 @@ class MilvusServiceClientTest {
 
 
         List<String> names = new ArrayList<>();
         List<String> names = new ArrayList<>();
         names.add(null);
         names.add(null);
-        assertThrows(ParamException.class, () -> ReleasePartitionsParam.newBuilder()
+        assertThrows(NullPointerException.class, () -> ReleasePartitionsParam.newBuilder()
                 .withCollectionName("collection1")
                 .withCollectionName("collection1")
                 .withPartitionNames(names)
                 .withPartitionNames(names)
                 .build()
                 .build()
@@ -852,7 +943,7 @@ class MilvusServiceClientTest {
 
 
         List<String> names = new ArrayList<>();
         List<String> names = new ArrayList<>();
         names.add(null);
         names.add(null);
-        assertThrows(ParamException.class, () -> ShowPartitionsParam.newBuilder()
+        assertThrows(NullPointerException.class, () -> ShowPartitionsParam.newBuilder()
                 .withCollectionName("collection1`")
                 .withCollectionName("collection1`")
                 .withPartitionNames(names)
                 .withPartitionNames(names)
                 .build()
                 .build()
@@ -1692,4 +1783,313 @@ class MilvusServiceClientTest {
 
 
         testFuncByName("getQuerySegmentInfo", param);
         testFuncByName("getQuerySegmentInfo", param);
     }
     }
+
+    @Test
+    void loadBalanceParam() {
+        // test throw exception with illegal input
+        assertThrows(ParamException.class, () -> LoadBalanceParam
+                .newBuilder()
+                .withSourceNodeID(1L)
+                .withDestinationNodeID(Arrays.asList(2L, 3L))
+                .addDestinationNodeID(4L)
+                .build()
+        );
+
+        assertThrows(ParamException.class, () -> LoadBalanceParam
+                .newBuilder()
+                .withSourceNodeID(1L)
+                .withSegmentIDs(Arrays.asList(2L, 3L))
+                .addSegmentID(4L)
+                .build()
+        );
+    }
+
+    @Test
+    void loadBalance() {
+        LoadBalanceParam param = LoadBalanceParam.newBuilder()
+                .withSourceNodeID(1L)
+                .addDestinationNodeID(2L)
+                .addSegmentID(3L)
+                .build();
+
+        testFuncByName("loadBalance", param);
+    }
+
+    @Test
+    void getCompactionState() {
+        GetCompactionStateParam param = GetCompactionStateParam.newBuilder()
+                .withCompactionID(1L)
+                .build();
+
+        testFuncByName("getCompactionState", param);
+    }
+
+    @Test
+    void manualCompaction() {
+        ManualCompactionParam param = ManualCompactionParam.newBuilder()
+                .withCollectionID(1L)
+                .build();
+
+        testFuncByName("manualCompaction", param);
+    }
+
+    @Test
+    void getCompactionStateWithPlans() {
+        GetCompactionPlansParam param = GetCompactionPlansParam.newBuilder()
+                .withCompactionID(1L)
+                .build();
+
+        testFuncByName("getCompactionStateWithPlans", param);
+    }
+
+    ////////////////////////////////////////////////////////////////////////////////////
+    // Response wrapper test
+    private void testScalarField(ScalarField field, DataType type, long rowCount) {
+        FieldData fieldData = FieldData.newBuilder()
+                .setFieldName("scalar")
+                .setFieldId(1L)
+                .setType(type)
+                .setScalars(field)
+                .build();
+
+        FieldDataWrapper wrapper = new FieldDataWrapper(fieldData);
+        assertEquals(rowCount, wrapper.getRowCount());
+
+        List<?> data = wrapper.getFieldData();
+        assertEquals(rowCount, data.size());
+
+        assertThrows(IllegalResponseException.class, wrapper::getDim);
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    void fieldDataWrapper() {
+        // for float vector
+        long dim = 3;
+        List<Float> floatVectors = Arrays.asList(1F, 2F, 3F, 4F, 5F, 6F);
+        FieldData fieldData = FieldData.newBuilder()
+                .setFieldName("vec")
+                .setFieldId(1L)
+                .setType(DataType.FloatVector)
+                .setVectors(VectorField.newBuilder()
+                        .setDim(dim)
+                        .setFloatVector(FloatArray.newBuilder()
+                                .addAllData(floatVectors)
+                                .build())
+                        .build())
+                .build();
+
+        FieldDataWrapper wrapper = new FieldDataWrapper(fieldData);
+        assertEquals(dim, wrapper.getDim());
+        assertEquals(floatVectors.size() / dim, wrapper.getRowCount());
+
+        List<?> floatData = wrapper.getFieldData();
+        assertEquals(floatVectors.size() / dim, floatData.size());
+        for (Object obj : floatData) {
+            List<Float> vec = (List<Float>) obj;
+            assertEquals(dim, vec.size());
+        }
+
+        // for binary vector
+        dim = 16;
+        byte[] binary = new byte[(int) dim * 2];
+        for (int i = 0; i < binary.length; ++i) {
+            binary[i] = (byte) i;
+        }
+        fieldData = FieldData.newBuilder()
+                .setFieldName("vec")
+                .setFieldId(1L)
+                .setType(DataType.BinaryVector)
+                .setVectors(VectorField.newBuilder()
+                        .setDim(dim)
+                        .setBinaryVector(ByteString.copyFrom(binary))
+                        .build())
+                .build();
+
+        wrapper = new FieldDataWrapper(fieldData);
+        assertEquals(dim, wrapper.getDim());
+        assertEquals(binary.length / dim, wrapper.getRowCount());
+
+        List<?> binaryData = wrapper.getFieldData();
+        assertEquals(binary.length / dim, binaryData.size());
+        for (Object obj : binaryData) {
+            ByteBuffer vec = (ByteBuffer) obj;
+            assertEquals(dim, vec.position());
+        }
+
+        // for scalar field
+        LongArray.Builder int64Builder = LongArray.newBuilder();
+        for (long i = 0; i < dim; ++i) {
+            int64Builder.addData(i);
+        }
+        testScalarField(ScalarField.newBuilder().setLongData(int64Builder).build(),
+                DataType.Int64, dim);
+
+        IntArray.Builder intBuilder = IntArray.newBuilder();
+        for (int i = 0; i < dim; ++i) {
+            intBuilder.addData(i);
+        }
+        testScalarField(ScalarField.newBuilder().setIntData(intBuilder).build(),
+                DataType.Int32, dim);
+        testScalarField(ScalarField.newBuilder().setIntData(intBuilder).build(),
+                DataType.Int16, dim);
+        testScalarField(ScalarField.newBuilder().setIntData(intBuilder).build(),
+                DataType.Int8, dim);
+
+        BoolArray.Builder boolBuilder = BoolArray.newBuilder();
+        for (long i = 0; i < dim; ++i) {
+            boolBuilder.addData(i % 2 == 0);
+        }
+        testScalarField(ScalarField.newBuilder().setBoolData(boolBuilder).build(),
+                DataType.Bool, dim);
+
+        FloatArray.Builder floatBuilder = FloatArray.newBuilder();
+        for (long i = 0; i < dim; ++i) {
+            floatBuilder.addData((float) i);
+        }
+        testScalarField(ScalarField.newBuilder().setFloatData(floatBuilder).build(),
+                DataType.Float, dim);
+
+        DoubleArray.Builder doubleBuilder = DoubleArray.newBuilder();
+        for (long i = 0; i < dim; ++i) {
+            doubleBuilder.addData((double) i);
+        }
+        testScalarField(ScalarField.newBuilder().setDoubleData(doubleBuilder).build(),
+                DataType.Double, dim);
+
+        StringArray.Builder strBuilder = StringArray.newBuilder();
+        for (long i = 0; i < dim; ++i) {
+            strBuilder.addData(String.valueOf(i));
+        }
+        testScalarField(ScalarField.newBuilder().setStringData(strBuilder).build(),
+                DataType.String, dim);
+    }
+
+    @Test
+    void getCollStatResponseWrapper() {
+        GetCollectionStatisticsResponse response = GetCollectionStatisticsResponse.newBuilder()
+                .addStats(KeyValuePair.newBuilder().setKey("row_count").setValue("invalid").build())
+                .build();
+        GetCollStatResponseWrapper invalidWrapper = new GetCollStatResponseWrapper(response);
+        assertThrows(NumberFormatException.class, invalidWrapper::GetRowCount);
+
+        response = GetCollectionStatisticsResponse.newBuilder()
+                .addStats(KeyValuePair.newBuilder().setKey("row_count").setValue("10").build())
+                .build();
+        GetCollStatResponseWrapper wrapper = new GetCollStatResponseWrapper(response);
+        assertEquals(10, wrapper.GetRowCount());
+
+        response = GetCollectionStatisticsResponse.newBuilder().build();
+        wrapper = new GetCollStatResponseWrapper(response);
+        assertEquals(0, wrapper.GetRowCount());
+
+    }
+
+    @Test
+    void InsertResultWrapper() {
+        List<Long> nID = Arrays.asList(1L, 2L, 3L);
+        MutationResult results = MutationResult.newBuilder()
+                .setInsertCnt(nID.size())
+                .setIDs(IDs.newBuilder()
+                        .setIntId(LongArray.newBuilder()
+                                .addAllData(nID)
+                                .build()))
+                .build();
+        InsertResultWrapper longWrapper = new InsertResultWrapper(results);
+        assertEquals(nID.size(), longWrapper.getInsertCount());
+        assertThrows(ParamException.class, longWrapper::getStringIDs);
+
+        List<Long> longIDs = longWrapper.getLongIDs();
+        assertEquals(nID.size(), longIDs.size());
+        for (int i = 0; i < longIDs.size(); ++i) {
+            assertEquals(nID.get(i), longIDs.get(i));
+        }
+
+        List<String> sID = Arrays.asList("1", "2", "3");
+        results = MutationResult.newBuilder()
+                .setInsertCnt(sID.size())
+                .setIDs(IDs.newBuilder()
+                        .setStrId(StringArray.newBuilder()
+                                .addAllData(sID)
+                                .build()))
+                .build();
+        InsertResultWrapper strWrapper = new InsertResultWrapper(results);
+        assertEquals(sID.size(), strWrapper.getInsertCount());
+        assertThrows(ParamException.class, strWrapper::getLongIDs);
+
+        List<String> strIDs = strWrapper.getStringIDs();
+        assertEquals(sID.size(), strIDs.size());
+        for (int i = 0; i < strIDs.size(); ++i) {
+            assertEquals(sID.get(i), strIDs.get(i));
+        }
+    }
+
+    @Test
+    void QueryResultsWrapper() {
+        String fieldName = "test";
+        QueryResults results = QueryResults.newBuilder()
+                .addFieldsData(FieldData.newBuilder()
+                        .setFieldName(fieldName)
+                        .build())
+                .build();
+
+        QueryResultsWrapper wrapper = new QueryResultsWrapper(results);
+        assertThrows(ParamException.class, () -> wrapper.getFieldWrapper("invalid"));
+        assertNotNull(wrapper.getFieldWrapper(fieldName));
+    }
+
+    @Test
+    void SearchResultsWrapper() {
+        long topK = 5;
+        long numQueries = 2;
+        List<Long> longIDs = new ArrayList<>();
+        List<String> strIDs = new ArrayList<>();
+        List<Float> scores = new ArrayList<>();
+        for (long i = 0; i < topK * numQueries; ++i) {
+            longIDs.add(i);
+            strIDs.add(String.valueOf(i));
+            scores.add((float) i);
+        }
+
+        // for long id
+        String fieldName = "test";
+        SearchResultData results = SearchResultData.newBuilder()
+                .setTopK(topK)
+                .setNumQueries(numQueries)
+                .setIds(IDs.newBuilder()
+                        .setIntId(LongArray.newBuilder()
+                                .addAllData(longIDs)
+                                .build()))
+                .addAllScores(scores)
+                .addFieldsData(FieldData.newBuilder()
+                        .setFieldName(fieldName)
+                        .build())
+                .build();
+
+        SearchResultsWrapper intWrapper = new SearchResultsWrapper(results);
+        assertNotNull(intWrapper.GetFieldData(fieldName));
+        assertNull(intWrapper.GetFieldData("invalid"));
+
+        List<SearchResultsWrapper.IDScore> idScores = intWrapper.GetIDScore(1);
+        assertEquals(idScores.size(), topK);
+        assertThrows(ParamException.class, () -> intWrapper.GetIDScore((int) numQueries));
+
+        // for string id
+        results = SearchResultData.newBuilder()
+                .setTopK(topK)
+                .setNumQueries(numQueries)
+                .setIds(IDs.newBuilder()
+                        .setStrId(StringArray.newBuilder()
+                                .addAllData(strIDs)
+                                .build()))
+                .addAllScores(scores)
+                .addFieldsData(FieldData.newBuilder()
+                        .setFieldName(fieldName)
+                        .build())
+                .build();
+
+        SearchResultsWrapper strWrapper = new SearchResultsWrapper(results);
+        idScores = strWrapper.GetIDScore(0);
+        assertEquals(idScores.size(), topK);
+    }
 }
 }

+ 57 - 0
src/test/java/io/milvus/server/MockMilvusServerImpl.java

@@ -58,6 +58,11 @@ public class MockMilvusServerImpl extends MilvusServiceGrpc.MilvusServiceImplBas
     private io.milvus.grpc.GetQuerySegmentInfoResponse respGetQuerySegmentInfo;
     private io.milvus.grpc.GetQuerySegmentInfoResponse respGetQuerySegmentInfo;
     private io.milvus.grpc.GetMetricsResponse respGetMetrics;
     private io.milvus.grpc.GetMetricsResponse respGetMetrics;
 
 
+    private io.milvus.grpc.Status respLoadBalance;
+    private io.milvus.grpc.GetCompactionStateResponse respGetCompactionState;
+    private io.milvus.grpc.ManualCompactionResponse respManualCompaction;
+    private io.milvus.grpc.GetCompactionPlansResponse respGetCompactionPlans;
+
     public MockMilvusServerImpl() {
     public MockMilvusServerImpl() {
     }
     }
 
 
@@ -476,4 +481,56 @@ public class MockMilvusServerImpl extends MilvusServiceGrpc.MilvusServiceImplBas
     public void setGetMetricsResponse(io.milvus.grpc.GetMetricsResponse resp) {
     public void setGetMetricsResponse(io.milvus.grpc.GetMetricsResponse resp) {
         respGetMetrics = resp;
         respGetMetrics = resp;
     }
     }
+
+    @Override
+    public void loadBalance(io.milvus.grpc.LoadBalanceRequest request,
+                            io.grpc.stub.StreamObserver<io.milvus.grpc.Status> responseObserver) {
+        logger.info("getMetrics() call");
+
+        responseObserver.onNext(respLoadBalance);
+        responseObserver.onCompleted();
+    }
+
+    public void setLoadBalanceResponse(io.milvus.grpc.Status resp) {
+        respLoadBalance = resp;
+    }
+
+    @Override
+    public void getCompactionState(io.milvus.grpc.GetCompactionStateRequest request,
+                                   io.grpc.stub.StreamObserver<io.milvus.grpc.GetCompactionStateResponse> responseObserver) {
+        logger.info("getMetrics() call");
+
+        responseObserver.onNext(respGetCompactionState);
+        responseObserver.onCompleted();
+    }
+
+    public void setGetCompactionStateResponse(io.milvus.grpc.GetCompactionStateResponse resp) {
+        respGetCompactionState = resp;
+    }
+
+    @Override
+    public void manualCompaction(io.milvus.grpc.ManualCompactionRequest request,
+                                 io.grpc.stub.StreamObserver<io.milvus.grpc.ManualCompactionResponse> responseObserver) {
+        logger.info("getMetrics() call");
+
+        responseObserver.onNext(respManualCompaction);
+        responseObserver.onCompleted();
+    }
+
+    public void setManualCompactionResponse(io.milvus.grpc.ManualCompactionResponse resp) {
+        respManualCompaction = resp;
+    }
+
+    @Override
+    public void getCompactionStateWithPlans(io.milvus.grpc.GetCompactionPlansRequest request,
+                                            io.grpc.stub.StreamObserver<io.milvus.grpc.GetCompactionPlansResponse> responseObserver) {
+        logger.info("getMetrics() call");
+
+        responseObserver.onNext(respGetCompactionPlans);
+        responseObserver.onCompleted();
+    }
+
+    public void setGetCompactionPlansResponse(io.milvus.grpc.GetCompactionPlansResponse resp) {
+        respGetCompactionPlans = resp;
+    }
 }
 }