diff --git a/sdk-core/src/main/java/io/milvus/orm/iterator/QueryIterator.java b/sdk-core/src/main/java/io/milvus/orm/iterator/QueryIterator.java index 15733ef41..602659a80 100644 --- a/sdk-core/src/main/java/io/milvus/orm/iterator/QueryIterator.java +++ b/sdk-core/src/main/java/io/milvus/orm/iterator/QueryIterator.java @@ -36,6 +36,7 @@ import java.util.List; import static io.milvus.param.Constant.NO_CACHE_ID; +import static io.milvus.param.Constant.MAX_BATCH_SIZE; import static io.milvus.param.Constant.UNLIMITED; public class QueryIterator { @@ -114,11 +115,19 @@ private void seek() { return; } - QueryResults response = executeQuery(expr, 0L, offset, this.sessionTs); - QueryResultsWrapper queryWrapper = new QueryResultsWrapper(response); - List res = queryWrapper.getRowRecords(); - int resultIndex = Math.min(res.size(), (int) offset); - updateCursor(res.subList(0, resultIndex)); + long currentOffset = offset; + while (currentOffset > 0) { + long limit = Math.min(MAX_BATCH_SIZE, currentOffset); + String currentExpr = setupNextExpr(); + QueryResults response = executeQuery(currentExpr, 0L, limit, this.sessionTs); + QueryResultsWrapper queryWrapper = new QueryResultsWrapper(response); + List res = queryWrapper.getRowRecords(); + if (res.isEmpty()) { + break; + } + updateCursor(res); + currentOffset -= res.size(); + } offset = 0; } diff --git a/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java b/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java index eae58962f..c359799cc 100644 --- a/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java +++ b/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java @@ -1624,7 +1624,7 @@ public void testIterator() { client.createCollection(requestCreate); // insert rows - long count = 10000; + long count = 20000; List data = generateRandomData(collectionSchema, count); InsertResp insertResp = client.insert(InsertReq.builder() .collectionName(randomCollectionName) @@ -1710,13 +1710,15 @@ public void testIterator() { Assertions.assertTrue(counter > 0); // query iterator + long from = 17777; + long to = 18000; QueryIterator queryIterator = client.queryIterator(QueryIteratorReq.builder() .collectionName(randomCollectionName) - .expr("int64_field < 300") + .expr("int64_field < " + String.valueOf(to)) .outputFields(Lists.newArrayList("*")) .batchSize(50L) - .offset(5) - .limit(400) + .offset(from) + .limit(4000) .consistencyLevel(ConsistencyLevel.EVENTUALLY) .build()); @@ -1730,6 +1732,7 @@ public void testIterator() { } for (QueryResultsWrapper.RowRecord record : res) { + Assertions.assertInstanceOf(Long.class, record.get("id")); Assertions.assertInstanceOf(Boolean.class, record.get("bool_field")); Assertions.assertInstanceOf(Integer.class, record.get("int8_field")); Assertions.assertInstanceOf(Integer.class, record.get("int16_field")); @@ -1745,8 +1748,9 @@ public void testIterator() { Assertions.assertInstanceOf(ByteBuffer.class, record.get("bfloat16_vector")); Assertions.assertInstanceOf(SortedMap.class, record.get("sparse_vector")); - long int64Val = (long)record.get("int64_field"); - Assertions.assertTrue(int64Val < 300L); + long int64Val = (long)record.get("id"); + Assertions.assertTrue(int64Val >= from); + Assertions.assertTrue(int64Val < to); String varcharVal = (String)record.get("varchar_field"); Assertions.assertTrue(varcharVal.startsWith("varchar_")); @@ -1772,7 +1776,7 @@ public void testIterator() { counter++; } } - Assertions.assertEquals(295, counter); + Assertions.assertEquals(to - from, counter); // search iterator V2 SearchIteratorV2 searchIteratorV2 = client.searchIteratorV2(SearchIteratorReqV2.builder()