Support partial load (#1091)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
milvus-io · Sep 30, 2024 · 6f37436 · 6f37436
1 parent 1b58ba3
commit 6f37436
Show file tree

Hide file tree

Showing 10 changed files with 134 additions and 7 deletions.
diff --git a/src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java b/src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java
@@ -665,7 +665,9 @@ public R<RpcStatus> loadCollection(@NonNull LoadCollectionParam requestParam) {
                     .setCollectionName(requestParam.getCollectionName())
                     .setReplicaNumber(requestParam.getReplicaNumber())
                     .addAllResourceGroups(requestParam.getResourceGroups())
-                    .setRefresh(requestParam.isRefresh());
+                    .setRefresh(requestParam.isRefresh())
+                    .addAllLoadFields(requestParam.getLoadFields())
+                    .setSkipLoadDynamicField(requestParam.isSkipLoadDynamicField());
             if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
                 builder.setDbName(requestParam.getDatabaseName());
             }
@@ -1066,7 +1068,9 @@ public R<RpcStatus> loadPartitions(@NonNull LoadPartitionsParam requestParam) {
                     .setReplicaNumber(requestParam.getReplicaNumber())
                     .addAllPartitionNames(requestParam.getPartitionNames())
                     .addAllResourceGroups(requestParam.getResourceGroups())
-                    .setRefresh(requestParam.isRefresh());
+                    .setRefresh(requestParam.isRefresh())
+                    .addAllLoadFields(requestParam.getLoadFields())
+                    .setSkipLoadDynamicField(requestParam.isSkipLoadDynamicField());
 
             if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
                 builder.setDbName(requestParam.getDatabaseName());

diff --git a/src/main/java/io/milvus/param/collection/FieldType.java b/src/main/java/io/milvus/param/collection/FieldType.java
@@ -265,8 +265,8 @@ public Builder withPartitionKey(boolean partitionKey) {
          * is generated to store the mapping relationship between segments and clustering key values. Once receiving
          * a search/query request that carries a clustering key value, it quickly finds out a search scope from
          * the PartitionStats which significantly improving search performance.
-         * Only scalar fields can be clustering key.
-         * Only one culstering key is allowed in one collection.
+         * Only scalar fields(except Array/JSON) can be clustering key.
+         * Only one clustering key is allowed in one collection.
          *
          * @param clusteringKey true is clustering key, false is not
          * @return <code>Builder</code>

diff --git a/src/main/java/io/milvus/param/collection/LoadCollectionParam.java b/src/main/java/io/milvus/param/collection/LoadCollectionParam.java
@@ -45,6 +45,8 @@ public class LoadCollectionParam {
     private final int replicaNumber;
     private final boolean refresh;
     private final List<String> resourceGroups;
+    private final List<String> loadFields;
+    private final boolean skipLoadDynamicField;
 
     public LoadCollectionParam(@NonNull Builder builder) {
         this.databaseName = builder.databaseName;
@@ -55,6 +57,8 @@ public LoadCollectionParam(@NonNull Builder builder) {
         this.replicaNumber = builder.replicaNumber;
         this.refresh = builder.refresh;
         this.resourceGroups = builder.resourceGroups;
+        this.loadFields = builder.loadFields;
+        this.skipLoadDynamicField = builder.skipLoadDynamicField;
     }
 
     public static Builder newBuilder() {
@@ -97,6 +101,15 @@ public static final class Builder {
         //   If not specified, the replicas will be loaded into the default resource group.
         private List<String> resourceGroups = new ArrayList<>();
 
+        // loadFields:
+        //   Specify load fields list needed during this load.
+        //   If not specified, all the fields will be loaded.
+        private List<String> loadFields = new ArrayList<>();
+
+        // skipLoadDynamicField:
+        //   Specify whether this load shall skip dynamic schema field.
+        private Boolean skipLoadDynamicField = Boolean.FALSE;
+
         private Builder() {
         }
 
@@ -200,6 +213,34 @@ public Builder withResourceGroups(@NonNull List<String> resourceGroups) {
             return this;
         }
 
+        /**
+         * Specify load fields list needed during this load.
+         * If not specified, all the fields will be loaded.
+         *
+         * @param loadFields a <code>List</code> of {@link String}
+         * @return <code>Builder</code>
+         */
+        public Builder withLoadFields(@NonNull List<String> loadFields) {
+            loadFields.forEach((field)->{
+                if (!this.loadFields.contains(field)) {
+                    this.loadFields.add(field);
+                }
+            });
+            return this;
+        }
+
+        /**
+         * Specify load fields list needed during this load. If not specified, all the fields will be loaded.
+         * Default is False.
+         *
+         * @param skip <code>Boolean.TRUE</code> skip dynamic field, <code>Boolean.FALSE</code> is not
+         * @return <code>Builder</code>
+         */
+        public Builder withSkipLoadDynamicField(@NonNull Boolean skip) {
+            this.skipLoadDynamicField = skip;
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link LoadCollectionParam} instance.
          *

diff --git a/src/main/java/io/milvus/param/partition/LoadPartitionsParam.java b/src/main/java/io/milvus/param/partition/LoadPartitionsParam.java
@@ -46,6 +46,8 @@ public class LoadPartitionsParam {
     private final int replicaNumber;
     private final boolean refresh;
     private final List<String> resourceGroups;
+    private final List<String> loadFields;
+    private final boolean skipLoadDynamicField;
 
     private LoadPartitionsParam(@NonNull Builder builder) {
         this.databaseName = builder.databaseName;
@@ -57,6 +59,8 @@ private LoadPartitionsParam(@NonNull Builder builder) {
         this.replicaNumber = builder.replicaNumber;
         this.refresh = builder.refresh;
         this.resourceGroups = builder.resourceGroups;
+        this.loadFields = builder.loadFields;
+        this.skipLoadDynamicField = builder.skipLoadDynamicField;
     }
 
     public static Builder newBuilder() {
@@ -100,6 +104,15 @@ public static final class Builder {
         //   If not specified, the replicas will be loaded into the default resource group.
         private List<String> resourceGroups = new ArrayList<>();
 
+        // loadFields:
+        //   Specify load fields list needed during this load.
+        //   If not specified, all the fields will be loaded.
+        private List<String> loadFields = new ArrayList<>();
+
+        // skipLoadDynamicField:
+        //   Specify whether this load shall skip dynamic schema field.
+        private Boolean skipLoadDynamicField = Boolean.FALSE;
+
         private Builder() {
         }
 
@@ -227,6 +240,34 @@ public Builder withResourceGroups(@NonNull List<String> resourceGroups) {
             return this;
         }
 
+        /**
+         * Specify load fields list needed during this load.
+         * If not specified, all the fields will be loaded.
+         *
+         * @param loadFields a <code>List</code> of {@link String}
+         * @return <code>Builder</code>
+         */
+        public Builder withLoadFields(@NonNull List<String> loadFields) {
+            loadFields.forEach((field)->{
+                if (!this.loadFields.contains(field)) {
+                    this.loadFields.add(field);
+                }
+            });
+            return this;
+        }
+
+        /**
+         * Specify load fields list needed during this load. If not specified, all the fields will be loaded.
+         * Default is False.
+         *
+         * @param skip <code>Boolean.TRUE</code> skip dynamic field, <code>Boolean.FALSE</code> is not
+         * @return <code>Builder</code>
+         */
+        public Builder withSkipLoadDynamicField(@NonNull Boolean skip) {
+            this.skipLoadDynamicField = skip;
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link LoadPartitionsParam} instance.
          *

diff --git a/src/main/java/io/milvus/v2/service/collection/CollectionService.java b/src/main/java/io/milvus/v2/service/collection/CollectionService.java
@@ -233,6 +233,9 @@ public Void loadCollection(MilvusServiceGrpc.MilvusServiceBlockingStub blockingS
         LoadCollectionRequest loadCollectionRequest = LoadCollectionRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
                 .setReplicaNumber(request.getNumReplicas())
+                .setRefresh(request.getRefresh())
+                .addAllLoadFields(request.getLoadFields())
+                .setSkipLoadDynamicField(request.getSkipLoadDynamicField())
                 .build();
         Status status = blockingStub.loadCollection(loadCollectionRequest);
         rpcUtils.handleResponse(title, status);

diff --git a/src/main/java/io/milvus/v2/service/collection/request/LoadCollectionReq.java b/src/main/java/io/milvus/v2/service/collection/request/LoadCollectionReq.java
@@ -23,6 +23,9 @@
 import lombok.Data;
 import lombok.experimental.SuperBuilder;
 
+import java.util.ArrayList;
+import java.util.List;
+
 @Data
 @SuperBuilder
 public class LoadCollectionReq {
@@ -33,4 +36,10 @@ public class LoadCollectionReq {
     private Boolean async = Boolean.TRUE;
     @Builder.Default
     private Long timeout = 60000L;
+    @Builder.Default
+    private Boolean refresh = Boolean.FALSE;
+    @Builder.Default
+    private List<String> loadFields = new ArrayList<>();
+    @Builder.Default
+    private Boolean skipLoadDynamicField = Boolean.FALSE;
 }
diff --git a/src/main/java/io/milvus/v2/service/partition/PartitionService.java b/src/main/java/io/milvus/v2/service/partition/PartitionService.java
@@ -84,7 +84,11 @@ public Void loadPartitions(MilvusServiceGrpc.MilvusServiceBlockingStub blockingS
 
         io.milvus.grpc.LoadPartitionsRequest loadPartitionsRequest = io.milvus.grpc.LoadPartitionsRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
-                .addAllPartitionNames(request.getPartitionNames()).build();
+                .addAllPartitionNames(request.getPartitionNames())
+                .setRefresh(request.getRefresh())
+                .addAllLoadFields(request.getLoadFields())
+                .setSkipLoadDynamicField(request.getSkipLoadDynamicField())
+                .build();
         Status status = blockingStub.loadPartitions(loadPartitionsRequest);
         rpcUtils.handleResponse(title, status);
 

diff --git a/src/main/java/io/milvus/v2/service/partition/request/LoadPartitionsReq.java b/src/main/java/io/milvus/v2/service/partition/request/LoadPartitionsReq.java
@@ -19,14 +19,23 @@
 
 package io.milvus.v2.service.partition.request;
 
+import lombok.Builder;
 import lombok.Data;
 import lombok.experimental.SuperBuilder;
 
+import java.util.ArrayList;
 import java.util.List;
 
 @Data
 @SuperBuilder
 public class LoadPartitionsReq {
     private String collectionName;
-    private List<String> partitionNames;
+    @Builder.Default
+    private List<String> partitionNames = new ArrayList<>();
+    @Builder.Default
+    private Boolean refresh = Boolean.FALSE;
+    @Builder.Default
+    private List<String> loadFields = new ArrayList<>();
+    @Builder.Default
+    private Boolean skipLoadDynamicField = Boolean.FALSE;
 }
diff --git a/src/test/java/io/milvus/client/MilvusClientDockerTest.java b/src/test/java/io/milvus/client/MilvusClientDockerTest.java
@@ -1222,9 +1222,14 @@ void testFloat16Vector() {
                 .build());
         Assertions.assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
 
-        // load collection
+        // load collection(partial load)
+        List<String> loadFields = new ArrayList<>();
+        loadFields.add("id");
+        loadFields.add(DataType.Float16Vector.name());
+        loadFields.add(DataType.BFloat16Vector.name());
         R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
                 .withCollectionName(randomCollectionName)
+                .withLoadFields(loadFields)
                 .build());
         Assertions.assertEquals(R.Status.Success.getCode(), loadR.getStatus().intValue());
 

diff --git a/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java b/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java
@@ -699,6 +699,17 @@ void testFloat16Vectors() {
                 .build();
         client.createCollection(requestCreate);
 
+        // partial load
+        List<String> loadFields = new ArrayList<>();
+        loadFields.add("id");
+        loadFields.add(float16Field);
+        loadFields.add(bfloat16Field);
+        client.releaseCollection(ReleaseCollectionReq.builder().collectionName(randomCollectionName).build());
+        client.loadCollection(LoadCollectionReq.builder()
+                .collectionName(randomCollectionName)
+                .loadFields(loadFields)
+                .build());
+
         // insert 10000 rows
         long count = 10000;
         List<JsonObject> data = generateRandomData(collectionSchema, count);