Fix a bug of partition key (#522)

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

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

@@ -69,21 +69,6 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
 
     protected abstract boolean clientIsReady();
 
-    ///////////////////// Internal Functions//////////////////////
-    private List<KeyValuePair> assembleKvPair(Map<String, String> sourceMap) {
-        List<KeyValuePair> result = new ArrayList<>();
-
-        if (MapUtils.isNotEmpty(sourceMap)) {
-            sourceMap.forEach((key, value) -> {
-                KeyValuePair kv = KeyValuePair.newBuilder()
-                        .setKey(key)
-                        .setValue(value).build();
-                result.add(kv);
-            });
-        }
-        return result;
-    }
-
     private void waitForLoadingCollection(String databaseName, String collectionName, List<String> partitionNames,
                                           long waitingInterval, long timeout) throws IllegalResponseException {
         long tsBegin = System.currentTimeMillis();
@@ -497,23 +482,7 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
 
             long fieldID = 0;
             for (FieldType fieldType : requestParam.getFieldTypes()) {
-                FieldSchema.Builder fieldSchemaBuilder = FieldSchema.newBuilder()
-                        .setFieldID(fieldID)
-                        .setName(fieldType.getName())
-                        .setIsPrimaryKey(fieldType.isPrimaryKey())
-                        .setIsPartitionKey(fieldType.isPartitionKey())
-                        .setDescription(fieldType.getDescription())
-                        .setDataType(fieldType.getDataType())
-                        .setAutoID(fieldType.isAutoID())
-                        .setIsDynamic(fieldType.isDynamic());
-
-                // assemble typeParams for CollectionSchema
-                List<KeyValuePair> typeParamsList = assembleKvPair(fieldType.getTypeParams());
-                if (CollectionUtils.isNotEmpty(typeParamsList)) {
-                    typeParamsList.forEach(fieldSchemaBuilder::addTypeParams);
-                }
-
-                collectionSchemaBuilder.addFields(fieldSchemaBuilder.build());
+                collectionSchemaBuilder.addFields(ParamUtils.ConvertField(fieldType));
                 fieldID++;
             }
 
@@ -526,6 +495,9 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
             if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
                 builder.setDbName(requestParam.getDatabaseName());
             }
+            if (requestParam.getPartitionsNum() > 0) {
+                builder.setNumPartitions(requestParam.getPartitionsNum());
+            }
 
             CreateCollectionRequest createCollectionRequest = builder.build();
 
@@ -824,7 +796,7 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
 
         try {
             AlterCollectionRequest.Builder alterCollRequestBuilder = AlterCollectionRequest.newBuilder();
-            List<KeyValuePair> propertiesList = assembleKvPair(requestParam.getProperties());
+            List<KeyValuePair> propertiesList = ParamUtils.AssembleKvPair(requestParam.getProperties());
             if (CollectionUtils.isNotEmpty(propertiesList)) {
                 propertiesList.forEach(alterCollRequestBuilder::addProperties);
             }
@@ -1287,7 +1259,7 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
 
         try {
             CreateIndexRequest.Builder createIndexRequestBuilder = CreateIndexRequest.newBuilder();
-            List<KeyValuePair> extraParamList = assembleKvPair(requestParam.getExtraParam());
+            List<KeyValuePair> extraParamList = ParamUtils.AssembleKvPair(requestParam.getExtraParam());
             if (CollectionUtils.isNotEmpty(extraParamList)) {
                 extraParamList.forEach(createIndexRequestBuilder::addExtraParams);
             }
@@ -2554,7 +2526,7 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
                 importRequest.setPartitionName(requestParam.getPartitionName());
             }
 
-            List<KeyValuePair> options = assembleKvPair(requestParam.getOptions());
+            List<KeyValuePair> options = ParamUtils.AssembleKvPair(requestParam.getOptions());
             if (CollectionUtils.isNotEmpty(options)) {
                 options.forEach(importRequest::addOptions);
             }

src/main/java/io/milvus/param/ParamUtils.java

@@ -18,6 +18,7 @@ import lombok.Builder;
 import lombok.Getter;
 import lombok.NonNull;
 import org.apache.commons.collections4.CollectionUtils;
+import org.apache.commons.collections4.MapUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 
@@ -666,9 +667,10 @@ public class ParamUtils {
                 .withName(field.getName())
                 .withDescription(field.getDescription())
                 .withPrimaryKey(field.getIsPrimaryKey())
+                .withPartitionKey(field.getIsPartitionKey())
                 .withAutoID(field.getAutoID())
                 .withDataType(field.getDataType())
-                .withPartitionKey(field.getIsPartitionKey());
+                .withIsDynamic(field.getIsDynamic());
 
         if (field.getIsDynamic()) {
             builder.withIsDynamic(true);
@@ -688,18 +690,37 @@ public class ParamUtils {
      */
     public static FieldSchema ConvertField(@NonNull FieldType field) {
         FieldSchema.Builder builder = FieldSchema.newBuilder()
-                .setIsPrimaryKey(field.isPrimaryKey())
-                .setAutoID(field.isAutoID())
                 .setName(field.getName())
                 .setDescription(field.getDescription())
-                .setDataType(field.getDataType());
-        Map<String, String> params = field.getTypeParams();
-        params.forEach((key, value) -> builder.addTypeParams(KeyValuePair.newBuilder()
-                .setKey(key).setValue(value).build()));
+                .setIsPrimaryKey(field.isPrimaryKey())
+                .setIsPartitionKey(field.isPartitionKey())
+                .setAutoID(field.isAutoID())
+                .setDataType(field.getDataType())
+                .setIsDynamic(field.isDynamic());
+
+        // assemble typeParams for CollectionSchema
+        List<KeyValuePair> typeParamsList = AssembleKvPair(field.getTypeParams());
+        if (CollectionUtils.isNotEmpty(typeParamsList)) {
+            typeParamsList.forEach(builder::addTypeParams);
+        }
 
         return builder.build();
     }
 
+    public static List<KeyValuePair> AssembleKvPair(Map<String, String> sourceMap) {
+        List<KeyValuePair> result = new ArrayList<>();
+
+        if (MapUtils.isNotEmpty(sourceMap)) {
+            sourceMap.forEach((key, value) -> {
+                KeyValuePair kv = KeyValuePair.newBuilder()
+                        .setKey(key)
+                        .setValue(value).build();
+                result.add(kv);
+            });
+        }
+        return result;
+    }
+
     @Builder
     @Getter
     public static class InsertDataInfo {

src/test/java/io/milvus/client/MilvusClientDockerTest.java

@@ -256,6 +256,16 @@ class MilvusClientDockerTest {
         DescCollResponseWrapper desc = new DescCollResponseWrapper(response.getData());
         System.out.println(desc.toString());
 
+        R<ShowPartitionsResponse> spResp = client.showPartitions(ShowPartitionsParam.newBuilder()
+                .withCollectionName(randomCollectionName)
+                .build());
+        System.out.println(spResp);
+
+        ShowPartResponseWrapper wra = new ShowPartResponseWrapper(spResp.getData());
+        List<ShowPartResponseWrapper.PartitionInfo> parts = wra.getPartitionsInfo();
+        System.out.println("Partition num: "+parts.size());
+
+
         // insert data
         int rowCount = 10000;
         List<Long> ids = new ArrayList<>();