Browse Source

add fluent api for CreateCollectionReq (#1054)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
groot 10 months ago
parent
commit
33cc20cd2d

+ 1 - 0
.gitignore

@@ -26,6 +26,7 @@ hs_err_pid*
 .DS_Store
 
 .idea/
+.vscode/
 target/
 volumes/
 *.iml

+ 17 - 2
src/main/java/io/milvus/v2/service/collection/request/CreateCollectionReq.java

@@ -62,18 +62,32 @@ public class CreateCollectionReq {
     // create collections with schema
     private CollectionSchema collectionSchema;
 
-    private List<IndexParam> indexParams;
+    @Builder.Default
+    private List<IndexParam> indexParams = new ArrayList<>();
 
     //private String partitionKeyField;
     private Integer numPartitions;
 
+    public static abstract class CreateCollectionReqBuilder<C extends CreateCollectionReq, B extends CreateCollectionReq.CreateCollectionReqBuilder<C, B>> {
+        public B indexParam(IndexParam indexParam) {
+            try {
+                this.indexParams$value.add(indexParam);
+            }catch (UnsupportedOperationException _e){
+                this.indexParams$value = new ArrayList<>(this.indexParams$value);
+                this.indexParams$value.add(indexParam);
+            }
+            this.indexParams$set = true;
+            return self();
+        }
+    }
+
     @Data
     @SuperBuilder
     public static class CollectionSchema {
         @Builder.Default
         private List<CreateCollectionReq.FieldSchema> fieldSchemaList = new ArrayList<>();
 
-        public void addField(AddFieldReq addFieldReq) {
+        public CollectionSchema addField(AddFieldReq addFieldReq) {
             CreateCollectionReq.FieldSchema fieldSchema = FieldSchema.builder()
                     .name(addFieldReq.getFieldName())
                     .dataType(addFieldReq.getDataType())
@@ -97,6 +111,7 @@ public class CreateCollectionReq {
                 fieldSchema.setDimension(addFieldReq.getDimension());
             }
             fieldSchemaList.add(fieldSchema);
+            return this;
         }
 
         public CreateCollectionReq.FieldSchema getField(String fieldName) {

+ 4 - 3
src/test/java/io/milvus/v2/service/collection/CollectionTest.java

@@ -54,9 +54,10 @@ class CollectionTest extends BaseTest {
 
         CreateCollectionReq.CollectionSchema collectionSchema = CreateCollectionReq.CollectionSchema.builder()
                 .build();
-        collectionSchema.addField(AddFieldReq.builder().fieldName("id").dataType(DataType.Int64).build());
-        collectionSchema.addField(AddFieldReq.builder().fieldName("vector").dataType(DataType.FloatVector).dimension(2).build());
-
+        collectionSchema
+                .addField(AddFieldReq.builder().fieldName("id").dataType(DataType.Int64).build())
+                .addField(AddFieldReq.builder().fieldName("vector").dataType(DataType.FloatVector).dimension(2).build())
+                .addField(AddFieldReq.builder().fieldName("description").dataType(DataType.VarChar).maxLength(64).build());
         IndexParam indexParam = IndexParam.builder()
                 .fieldName("vector")
                 .metricType(IndexParam.MetricType.L2)