Use cached collection schema for query by ids (#1518)

Signed-off-by: yhmo <[email protected]>

Author: yhmo

docker-compose.yml

@@ -3,7 +3,7 @@ version: '3.5'
 services:
   etcd:
     container_name: milvus-javasdk-test-etcd
-    image: quay.io/coreos/etcd:v3.5.5
+    image: quay.io/coreos/etcd:v3.5.18
     volumes:
       - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
     command: etcd -listen-peer-urls=http://127.0.0.1:2380 -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379,http://0.0.0.0:4001 -initial-advertise-peer-urls=http://127.0.0.1:2380 --initial-cluster default=http://127.0.0.1:2380 --data-dir /etcd
@@ -14,7 +14,7 @@ services:
 
   minio:
     container_name: milvus-javasdk-test-minio
-    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
+    image: minio/minio:RELEASE.2024-05-28T17-19-04Z
     ports:
       - "9000:9000"
       - "9001:9001"
@@ -48,7 +48,7 @@ services:
 
   etcdslave:
     container_name: milvus-javasdk-test-etcd-slave
-    image: quay.io/coreos/etcd:v3.5.5
+    image: quay.io/coreos/etcd:v3.5.18
     volumes:
       - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd-slave:/etcd
     command: etcd -listen-peer-urls=http://127.0.0.1:2380 -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379,http://0.0.0.0:4001 -initial-advertise-peer-urls=http://127.0.0.1:2380 --initial-cluster default=http://127.0.0.1:2380 --data-dir /etcd
@@ -59,7 +59,7 @@ services:
 
   minioslave:
     container_name: milvus-javasdk-test-minio-slave
-    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
+    image: minio/minio:RELEASE.2024-05-28T17-19-04Z
     ports:
       - "19000:9000"
       - "19001:9001"

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

@@ -220,7 +220,9 @@ public UpsertResp upsert(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStu
     }
 
     public QueryResp query(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, QueryReq request) {
-        String title = String.format("QueryRequest collectionName:%s, databaseName:%s", request.getCollectionName(), request.getDatabaseName());
+        String dbName = request.getDatabaseName();
+        String collectionName = request.getCollectionName();
+        String title = String.format("QueryRequest collectionName:%s, databaseName:%s", collectionName, dbName);
         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) {
@@ -229,9 +231,19 @@ public QueryResp query(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
 
 
         if (request.getIds() != null && request.getFilter() == null) {
-            DescribeCollectionReq descReq = DescribeCollectionReq.builder().databaseName(request.getDatabaseName()).collectionName(request.getCollectionName()).build();
-            DescribeCollectionResp descResp = collectionService.describeCollection(blockingStub, descReq);
-            request.setFilter(vectorUtils.getExprById(descResp.getPrimaryFieldName(), request.getIds()));
+            DescribeCollectionResponse descResp = getCollectionInfo(blockingStub, dbName, collectionName, false);
+            String primaryKeyName = "";
+            List<FieldSchema> fields = descResp.getSchema().getFieldsList();
+            for (FieldSchema field : fields) {
+                if (field.getIsPrimaryKey()) {
+                    primaryKeyName = field.getName();
+                    break;
+                }
+            }
+            if (StringUtils.isEmpty(primaryKeyName)) {
+                throw new MilvusClientException(ErrorCode.SERVER_ERROR, "cannot find the primary key field in collection schema");
+            }
+            request.setFilter(vectorUtils.getExprById(primaryKeyName, request.getIds()));
         }
 
         // reset the db name so that the timestamp cache can set correct key for this collection
@@ -287,15 +299,17 @@ public SearchResp hybridSearch(MilvusServiceGrpc.MilvusServiceBlockingStub block
 
     public QueryIterator queryIterator(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
                                            QueryIteratorReq request) {
-        DescribeCollectionResponse descResp = describeCollection(blockingStub, request.getDatabaseName(), request.getCollectionName());
+        DescribeCollectionResponse descResp = getCollectionInfo(blockingStub, request.getDatabaseName(),
+                request.getCollectionName(), false);
         DescribeCollectionResp respR = convertUtils.convertDescCollectionResp(descResp);
         CreateCollectionReq.FieldSchema pkField = respR.getCollectionSchema().getField(respR.getPrimaryFieldName());
         return new QueryIterator(request, blockingStub, pkField);
     }
 
     public SearchIterator searchIterator(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
                                             SearchIteratorReq request) {
-        DescribeCollectionResponse descResp = describeCollection(blockingStub, request.getDatabaseName(), request.getCollectionName());
+        DescribeCollectionResponse descResp = getCollectionInfo(blockingStub, request.getDatabaseName(),
+                request.getCollectionName(), false);
         DescribeCollectionResp respR = convertUtils.convertDescCollectionResp(descResp);
         CreateCollectionReq.FieldSchema pkField = respR.getCollectionSchema().getField(respR.getPrimaryFieldName());
         return new SearchIterator(request, blockingStub, pkField);

sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java

@@ -1173,7 +1173,7 @@ void testDeleteUpsert() {
         // verify
         QueryResp queryResp = client.query(QueryReq.builder()
                 .collectionName(randomCollectionName)
-                .filter("pk == \"pk_2\" or pk == \"pk_5\"")
+                .ids(Arrays.asList("pk_2", "pk_5"))
                 .outputFields(Collections.singletonList("*"))
                 .build());
         List<QueryResp.QueryResult> queryResults = queryResp.getQueryResults();