diff --git a/milvus/const/limits.ts b/milvus/const/limits.ts index 17595287..5a978dd2 100644 --- a/milvus/const/limits.ts +++ b/milvus/const/limits.ts @@ -1,3 +1,5 @@ // limits export const MAX_PARTITION_KEY_FIELD_COUNT = 1; -export const MAX_PARTITIONS_NUMBER = 4096; \ No newline at end of file +export const MAX_PARTITIONS_NUMBER = 4096; + +export const NO_LIMIT = -1; \ No newline at end of file diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts index 298a5b62..7a21cc74 100644 --- a/milvus/grpc/Data.ts +++ b/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 = diff --git a/milvus/types/Data.ts b/milvus/types/Data.ts index d2434515..b984c2b2 100644 --- a/milvus/types/Data.ts +++ b/milvus/types/Data.ts @@ -418,7 +418,7 @@ export interface QueryReq extends collectionNameReq { export interface QueryIteratorReq extends Omit { - limit: number; + limit?: number; batchSize: number; } diff --git a/test/grpc/Iterator.spec.ts b/test/grpc/Iterator.spec.ts index b0fd5db4..f70ccd64 100644 --- a/test/grpc/Iterator.spec.ts +++ b/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;