Browse Source

Enable query pagination (#410)

Signed-off-by: yhmo <yihua.mo@zilliz.com>

Signed-off-by: yhmo <yihua.mo@zilliz.com>
groot 2 years ago
parent
commit
979308b7a9

+ 2 - 2
docker-compose.yml

@@ -31,7 +31,7 @@ services:
 
   standalone:
     container_name: milvus-javasdk-test-standalone
-    image: milvusdb/milvus:master-20221103-94aafb50
+    image: milvusdb/milvus:master-20221122-f76ea292
     command: ["milvus", "run", "standalone"]
     environment:
       ETCD_ENDPOINTS: etcd:2379
@@ -75,7 +75,7 @@ services:
 
   standaloneslave:
     container_name: milvus-javasdk-test-slave-standalone
-    image: milvusdb/milvus:master-20221103-94aafb50
+    image: milvusdb/milvus:master-20221122-f76ea292
     command: ["milvus", "run", "standalone"]
     environment:
       ETCD_ENDPOINTS: etcdslave:2379

+ 22 - 3
src/main/java/io/milvus/param/ParamUtils.java

@@ -320,14 +320,33 @@ public class ParamUtils {
     public static QueryRequest convertQueryParam(@NonNull QueryParam requestParam) {
         long guaranteeTimestamp = getGuaranteeTimestamp(requestParam.getConsistencyLevel(),
                 requestParam.getGuaranteeTimestamp(), requestParam.getGracefulTime());
-        return QueryRequest.newBuilder()
+        QueryRequest.Builder builder = QueryRequest.newBuilder()
                 .setCollectionName(requestParam.getCollectionName())
                 .addAllPartitionNames(requestParam.getPartitionNames())
                 .addAllOutputFields(requestParam.getOutFields())
                 .setExpr(requestParam.getExpr())
                 .setTravelTimestamp(requestParam.getTravelTimestamp())
-                .setGuaranteeTimestamp(guaranteeTimestamp)
-                .build();
+                .setGuaranteeTimestamp(guaranteeTimestamp);
+
+        // set offset and limit value.
+        // directly pass the two values, the server will verify them.
+        long offset = requestParam.getOffset();
+        if (offset > 0) {
+            builder.addQueryParams(KeyValuePair.newBuilder()
+                    .setKey("offset")
+                    .setValue(String.valueOf(offset))
+                    .build());
+        }
+
+        long limit = requestParam.getLimit();
+        if (limit > 0) {
+            builder.addQueryParams(KeyValuePair.newBuilder()
+                    .setKey("limit")
+                    .setValue(String.valueOf(limit))
+                    .build());
+        }
+
+        return builder.build();
     }
 
     private static long getGuaranteeTimestamp(ConsistencyLevelEnum consistencyLevel,

+ 42 - 1
src/main/java/io/milvus/param/dml/QueryParam.java

@@ -41,8 +41,10 @@ public class QueryParam {
     private final String expr;
     private final long travelTimestamp;
     private final long guaranteeTimestamp;
-    private final Long gracefulTime;
+    private final long gracefulTime;
     private final ConsistencyLevelEnum consistencyLevel;
+    private final long offset;
+    private final long limit;
 
     private QueryParam(@NonNull Builder builder) {
         this.collectionName = builder.collectionName;
@@ -53,6 +55,8 @@ public class QueryParam {
         this.guaranteeTimestamp = builder.guaranteeTimestamp;
         this.consistencyLevel = builder.consistencyLevel;
         this.gracefulTime = builder.gracefulTime;
+        this.offset = builder.offset;
+        this.limit = builder.limit;
     }
 
     public static Builder newBuilder() {
@@ -71,6 +75,8 @@ public class QueryParam {
         private Long gracefulTime = 5000L;
         private Long guaranteeTimestamp = Constant.GUARANTEE_EVENTUALLY_TS;
         private ConsistencyLevelEnum consistencyLevel;
+        private Long offset = 0L;
+        private Long limit = 0L;
 
         private Builder() {
         }
@@ -199,6 +205,30 @@ public class QueryParam {
             return this;
         }
 
+        /**
+         * Specify a position to return results. Only take effect when the 'limit' value is specified.
+         * Default value is 0, start from begin.
+         *
+         * @param offset a value to define the position
+         * @return <code>Builder</code>
+         */
+        public Builder withOffset(@NonNull Long offset) {
+            this.offset = offset;
+            return this;
+        }
+
+        /**
+         * Specify a value to control the returned number of entities. Must be a positive value.
+         * Default value is 0, will return without limit.
+         *
+         * @param limit a value to define the limit of returned entities
+         * @return <code>Builder</code>
+         */
+        public Builder withLimit(@NonNull Long limit) {
+            this.limit = limit;
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link QueryParam} instance.
          *
@@ -216,6 +246,14 @@ public class QueryParam {
                 throw new ParamException("The guarantee timestamp must be greater than 0");
             }
 
+            if (offset < 0) {
+                throw new ParamException("The offset value cannot be less than 0");
+            }
+
+            if (limit < 0) {
+                throw new ParamException("The limit value cannot be less than 0");
+            }
+
             return new QueryParam(this);
         }
     }
@@ -230,8 +268,11 @@ public class QueryParam {
         return "QueryParam{" +
                 "collectionName='" + collectionName + '\'' +
                 ", partitionNames='" + partitionNames.toString() + '\'' +
+                ", outFields=" + outFields.toString() +
                 ", expr='" + expr + '\'' +
                 ", consistencyLevel='" + consistencyLevel + '\'' +
+                ", offset=" + offset +
+                ", limit=" + limit +
                 '}';
     }
 }

+ 16 - 0
src/test/java/io/milvus/client/MilvusClientDockerTest.java

@@ -423,6 +423,22 @@ class MilvusClientDockerTest {
             }
         }
 
+        // query with offset and limit
+        int queryLimit = 5;
+        expr = field5Name + " > 1"; // "age > 1"
+        queryParam = QueryParam.newBuilder()
+                .withCollectionName(randomCollectionName)
+                .withExpr(expr)
+                .withOffset(100L)
+                .withLimit((long) queryLimit)
+                .build();
+        queryR = client.query(queryParam);
+        assertEquals(R.Status.Success.getCode(), queryR.getStatus().intValue());
+
+        queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
+        // we didn't set the output fields, only primary key field is returned
+        List<?> out = queryResultsWrapper.getFieldWrapper(field1Name).getFieldData();
+        assertEquals(queryLimit, out.size());
 
         // pick some vectors to search
         List<Long> targetVectorIDs = new ArrayList<>();

+ 18 - 0
src/test/java/io/milvus/client/MilvusServiceClientTest.java

@@ -2073,6 +2073,24 @@ class MilvusServiceClientTest {
                 .withGuaranteeTimestamp(-1L)
                 .build()
         );
+
+        assertThrows(ParamException.class, () -> QueryParam.newBuilder()
+                .withCollectionName("collection1")
+                .withPartitionNames(partitions)
+                .withOutFields(outputFields)
+                .withExpr("dummy")
+                .withLimit(-1L)
+                .build()
+        );
+
+        assertThrows(ParamException.class, () -> QueryParam.newBuilder()
+                .withCollectionName("collection1")
+                .withPartitionNames(partitions)
+                .withOutFields(outputFields)
+                .withExpr("dummy")
+                .withOffset(-1L)
+                .build()
+        );
     }
 
     @Test