fix: queryIterator limit field should be optional (#361)

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

milvus/const/limits.ts

@@ -1,3 +1,5 @@
 // limits
 export const MAX_PARTITION_KEY_FIELD_COUNT = 1;
-export const MAX_PARTITIONS_NUMBER = 4096;
+export const MAX_PARTITIONS_NUMBER = 4096;
+
+export const NO_LIMIT = -1;

milvus/grpc/Data.ts

@@ -65,6 +65,7 @@ import {
   sparseRowsToBytes,
   getSparseDim,
   f32ArrayToBinaryBytes,
+  NO_LIMIT,
 } from '../';
 import { Collection } from './Collection';
 
@@ -704,6 +705,10 @@ export class Data extends Collection {
       collection_name: data.collection_name,
       expr: userExpr,
     });
+    // if limit not set, set it to count
+    if (!data.limit || data.limit === NO_LIMIT) {
+      data.limit = count.data;
+    }
     // total should be the minimum of total and count
     const total = data.limit > count.data ? count.data : data.limit;
     const batchSize =

milvus/types/Data.ts

@@ -418,7 +418,7 @@ export interface QueryReq extends collectionNameReq {
 
 export interface QueryIteratorReq
   extends Omit<QueryReq, 'ids' | 'offset' | 'limit'> {
-  limit: number;
+  limit?: number;
   batchSize: number;
 }
 

test/grpc/Iterator.spec.ts

@@ -1,4 +1,4 @@
-import { MilvusClient, DataType } from '../../milvus';
+import { MilvusClient, DataType, NO_LIMIT } from '../../milvus';
 import {
   IP,
   genCollectionParams,
@@ -299,6 +299,79 @@ describe(`Iterator API`, () => {
     });
   });
 
+  it(`query iterator with limit unset should success`, async () => {
+    // page size
+    const batchSize = 5000;
+    const iterator = await milvusClient.queryIterator({
+      collection_name: COLLECTION,
+      batchSize: batchSize,
+      expr: 'id > 0',
+      output_fields: ['id'],
+    });
+
+    const results: any = [];
+    let page = 0;
+    for await (const value of iterator) {
+      results.push(...value);
+      page += 1;
+    }
+
+    // page size should equal to page
+    expect(page).toEqual(Math.ceil(data.length / batchSize));
+    // results length should equal to data length
+    expect(results.length).toEqual(data.length);
+
+    // results id should be unique
+    const idSet = new Set();
+    results.forEach((result: any) => {
+      idSet.add(result.id);
+    });
+    expect(idSet.size).toEqual(data.length);
+
+    // every id in query result should be founded in the original data
+    results.forEach((result: any) => {
+      const item = dataMap.get(result.id.toString());
+      expect(typeof item !== 'undefined').toEqual(true);
+    });
+  });
+
+  it(`query iterator with limit = -1 should success`, async () => {
+    // page size
+    const batchSize = 5000;
+    const iterator = await milvusClient.queryIterator({
+      collection_name: COLLECTION,
+      batchSize: batchSize,
+      expr: 'id > 0',
+      output_fields: ['id'],
+      limit: NO_LIMIT,
+    });
+
+    const results: any = [];
+    let page = 0;
+    for await (const value of iterator) {
+      results.push(...value);
+      page += 1;
+    }
+
+    // page size should equal to page
+    expect(page).toEqual(Math.ceil(data.length / batchSize));
+    // results length should equal to data length
+    expect(results.length).toEqual(data.length);
+
+    // results id should be unique
+    const idSet = new Set();
+    results.forEach((result: any) => {
+      idSet.add(result.id);
+    });
+    expect(idSet.size).toEqual(data.length);
+
+    // every id in query result should be founded in the original data
+    results.forEach((result: any) => {
+      const item = dataMap.get(result.id.toString());
+      expect(typeof item !== 'undefined').toEqual(true);
+    });
+  });
+
   // it('search iterator with batch size = total should success', async () => {
   //   const batchSize = 100;
   //   const total = 100;