Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: queryIterator limit field should be optional #361

Merged
merged 1 commit into from
Sep 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion milvus/const/limits.ts
Original file line number Diff line number Diff line change
@@ -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;
5 changes: 5 additions & 0 deletions milvus/grpc/Data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ import {
sparseRowsToBytes,
getSparseDim,
f32ArrayToBinaryBytes,
NO_LIMIT,
} from '../';
import { Collection } from './Collection';

Expand Down Expand Up @@ -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 =
Expand Down
2 changes: 1 addition & 1 deletion milvus/types/Data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,7 @@ export interface QueryReq extends collectionNameReq {

export interface QueryIteratorReq
extends Omit<QueryReq, 'ids' | 'offset' | 'limit'> {
limit: number;
limit?: number;
batchSize: number;
}

Expand Down
75 changes: 74 additions & 1 deletion test/grpc/Iterator.spec.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { MilvusClient, DataType } from '../../milvus';
import { MilvusClient, DataType, NO_LIMIT } from '../../milvus';
import {
IP,
genCollectionParams,
Expand Down Expand Up @@ -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;
Expand Down
Loading