Browse Source

optim clientV2 search/query to support db (#1324)

Signed-off-by: lentitude2tk <xushuang.hu@zilliz.com>
xushuang.hu 1 month ago
parent
commit
e1a42d4454

+ 9 - 0
sdk-core/src/main/java/io/milvus/v2/common/ConsistencyLevel.java

@@ -33,4 +33,13 @@ public enum ConsistencyLevel{
         this.name = name;
         this.code = code;
     }
+
+    public static ConsistencyLevel fromName(String name) {
+        for (ConsistencyLevel level : ConsistencyLevel.values()) {
+            if (level.getName().equals(name)) {
+                return level;
+            }
+        }
+        return null;
+    }
 }

+ 2 - 2
sdk-core/src/main/java/io/milvus/v2/service/vector/VectorService.java

@@ -177,7 +177,7 @@ public class VectorService extends BaseService {
     }
 
     public QueryResp query(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, QueryReq request) {
-        String title = String.format("QueryRequest collectionName:%s", request.getCollectionName());
+        String title = String.format("QueryRequest collectionName:%s, databaseName:%s", request.getCollectionName(), request.getDatabaseName());
         if (request.getFilter() == null && request.getIds() == null) {
             throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "filter and ids can't be null at the same time");
         } else if (request.getFilter() != null && request.getIds() != null) {
@@ -201,7 +201,7 @@ public class VectorService extends BaseService {
     }
 
     public SearchResp search(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, SearchReq request) {
-        String title = String.format("SearchRequest collectionName:%s", request.getCollectionName());
+        String title = String.format("SearchRequest collectionName:%s, databaseName:%s", request.getCollectionName(), request.getDatabaseName());
 
         //checkCollectionExist(blockingStub, request.getCollectionName());
 

+ 1 - 0
sdk-core/src/main/java/io/milvus/v2/service/vector/request/QueryReq.java

@@ -29,6 +29,7 @@ import java.util.*;
 @Data
 @SuperBuilder
 public class QueryReq {
+    private String databaseName;
     private String collectionName;
     @Builder.Default
     private List<String> partitionNames = new ArrayList<>();

+ 1 - 0
sdk-core/src/main/java/io/milvus/v2/service/vector/request/SearchReq.java

@@ -34,6 +34,7 @@ import java.util.Map;
 @Data
 @SuperBuilder
 public class SearchReq {
+    private String databaseName;
     private String collectionName;
     @Builder.Default
     private List<String> partitionNames = new ArrayList<>();

+ 8 - 0
sdk-core/src/main/java/io/milvus/v2/utils/VectorUtils.java

@@ -48,6 +48,10 @@ public class VectorUtils {
                 .addAllPartitionNames(request.getPartitionNames())
                 .addAllOutputFields(request.getOutputFields())
                 .setExpr(request.getFilter());
+        if (StringUtils.isNotEmpty(request.getDatabaseName())) {
+            builder.setDbName(request.getDatabaseName());
+        }
+
         if (request.getFilter() != null && !request.getFilter().isEmpty()) {
             Map<String, Object> filterTemplateValues = request.getFilterTemplateValues();
             filterTemplateValues.forEach((key, value)->{
@@ -140,6 +144,10 @@ public class VectorUtils {
             request.getPartitionNames().forEach(builder::addPartitionNames);
         }
 
+        if (StringUtils.isNotEmpty(request.getDatabaseName())) {
+            builder.setDbName(request.getDatabaseName());
+        }
+
         // prepare target, the input could be vectors or string list for doc-in-doc-out
         List<BaseVector> vectors = request.getData();
         if (vectors.isEmpty()) {