Support setting properties for CreateCollectionRequest (#1128)

Signed-off-by: yhmo <[email protected]>
milvus-io · Oct 21, 2024 · f5f26aa · f5f26aa
1 parent 23be2a6
commit f5f26aa
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 7 deletions.
diff --git a/src/main/java/io/milvus/v2/service/collection/CollectionService.java b/src/main/java/io/milvus/v2/service/collection/CollectionService.java
@@ -120,15 +120,18 @@ public Void createCollectionWithSchema(MilvusServiceGrpc.MilvusServiceBlockingSt
         }
 
         //create collection
-        CreateCollectionRequest createCollectionRequest = CreateCollectionRequest.newBuilder()
+        CreateCollectionRequest.Builder builder = CreateCollectionRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
                 .setSchema(grpcSchema.toByteString())
-                .setShardsNum(request.getNumShards())
-                .build();
+                .setShardsNum(request.getNumShards());
+        List<KeyValuePair> propertiesList = ParamUtils.AssembleKvPair(request.getProperties());
+        if (CollectionUtils.isNotEmpty(propertiesList)) {
+            propertiesList.forEach(builder::addProperties);
+        }
         if (request.getNumPartitions() != null) {
-            createCollectionRequest = createCollectionRequest.toBuilder().setNumPartitions(request.getNumPartitions()).build();
+            builder.setNumPartitions(request.getNumPartitions());
         }
-        Status createCollectionResponse = blockingStub.createCollection(createCollectionRequest);
+        Status createCollectionResponse = blockingStub.createCollection(builder.build());
         rpcUtils.handleResponse(title, createCollectionResponse);
 
         //create index

diff --git a/src/main/java/io/milvus/v2/service/collection/request/CreateCollectionReq.java b/src/main/java/io/milvus/v2/service/collection/request/CreateCollectionReq.java
@@ -30,7 +30,9 @@
 import lombok.experimental.SuperBuilder;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 @Data
 @SuperBuilder
@@ -74,6 +76,9 @@ public class CreateCollectionReq {
     @Builder.Default
     private ConsistencyLevel consistencyLevel = ConsistencyLevel.BOUNDED;
 
+    @Builder.Default
+    private final Map<String, String> properties = new HashMap<>();
+
     public static abstract class CreateCollectionReqBuilder<C extends CreateCollectionReq, B extends CreateCollectionReq.CreateCollectionReqBuilder<C, B>> {
         public B indexParam(IndexParam indexParam) {
             try {
@@ -106,6 +111,15 @@ public B collectionSchema(CollectionSchema collectionSchema) {
             this.enableDynamicField$set = true;
             return self();
         }
+
+        public B property(String key, String value) {
+            if(null == this.properties$value ){
+                this.properties$value = new HashMap<>();
+            }
+            this.properties$value.put(key, value);
+            this.properties$set = true;
+            return self();
+        }
     }
 
     @Data

diff --git a/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java b/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java
@@ -1225,8 +1225,19 @@ void testIndex() {
                 .collectionName(randomCollectionName)
                 .collectionSchema(collectionSchema)
                 .indexParams(indexes)
+                .property(Constant.TTL_SECONDS, "5")
+                .property(Constant.MMAP_ENABLED, "false")
                 .build());
 
+        DescribeCollectionResp descCollResp = client.describeCollection(DescribeCollectionReq.builder()
+                .collectionName(randomCollectionName)
+                .build());
+        Map<String, String> collProps = descCollResp.getProperties();
+        Assertions.assertTrue(collProps.containsKey(Constant.TTL_SECONDS));
+        Assertions.assertTrue(collProps.containsKey(Constant.MMAP_ENABLED));
+        Assertions.assertEquals("5", collProps.get(Constant.TTL_SECONDS));
+        Assertions.assertEquals("false", collProps.get(Constant.MMAP_ENABLED));
+
         client.releaseCollection(ReleaseCollectionReq.builder()
                 .collectionName(randomCollectionName)
                 .build());
@@ -1239,10 +1250,10 @@ void testIndex() {
                 .properties(properties)
                 .property("prop", "val")
                 .build());
-        DescribeCollectionResp descCollResp = client.describeCollection(DescribeCollectionReq.builder()
+        descCollResp = client.describeCollection(DescribeCollectionReq.builder()
                 .collectionName(randomCollectionName)
                 .build());
-        Map<String, String> collProps = descCollResp.getProperties();
+        collProps = descCollResp.getProperties();
         Assertions.assertTrue(collProps.containsKey(Constant.TTL_SECONDS));
         Assertions.assertTrue(collProps.containsKey(Constant.MMAP_ENABLED));
         Assertions.assertTrue(collProps.containsKey("prop"));