Skip to content

Commit

Permalink
support group size and strict_group_size (#371)
Browse files Browse the repository at this point in the history
Signed-off-by: ryjiang <jiangruiyi@gmail.com>
  • Loading branch information
shanghaikid authored Nov 6, 2024
1 parent 9f46da2 commit 17d86a7
Show file tree
Hide file tree
Showing 5 changed files with 62 additions and 1 deletion.
4 changes: 4 additions & 0 deletions milvus/types/Data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,8 @@ export interface SearchParam {
round_decimal?: number; // round decimal
ignore_growing?: boolean; // ignore growing
group_by_field?: string; // group by field
group_size?: number; // group size
strict_group_size?: boolean; // if strict group size
}

// old search api parameter type, deprecated
Expand Down Expand Up @@ -333,6 +335,8 @@ export interface SearchSimpleReq extends collectionNameReq {
consistency_level?: ConsistencyLevelEnum; // consistency level
ignore_growing?: boolean; // ignore growing
group_by_field?: string; // group by field
group_size?: number; // group size
strict_group_size?: boolean; // if strict group size
round_decimal?: number; // round decimal
transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors
}
Expand Down
6 changes: 6 additions & 0 deletions milvus/utils/Format.ts
Original file line number Diff line number Diff line change
Expand Up @@ -658,6 +658,12 @@ export const buildSearchParams = (
if (data.group_by_field) {
search_params.group_by_field = data.group_by_field;
}
if (data.strict_group_size) {
search_params.strict_group_size = data.strict_group_size;
}
if (data.group_size) {
search_params.group_size = data.group_size;
}

return search_params;
};
Expand Down
2 changes: 1 addition & 1 deletion proto
1 change: 1 addition & 0 deletions test/grpc/Data.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,7 @@ describe(`Data.API`, () => {
data: [1, 2, 3, 4],
limit: limit,
group_by_field: 'varChar2',
group_size: 2,
});

expect(searchWithData.status.error_code).toEqual(ErrorCode.SUCCESS);
Expand Down
50 changes: 50 additions & 0 deletions test/utils/Format.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ import {
formatSearchData,
buildSearchRequest,
FieldSchema,
buildSearchParams,
SearchSimpleReq,
} from '../../milvus';

describe('utils/format', () => {
Expand Down Expand Up @@ -937,4 +939,52 @@ describe('utils/format', () => {
}
);
});

it('should build search params correctly', () => {
const data: SearchSimpleReq = {
collection_name: 'test',
data: [1, 2, 3, 4, 5, 6, 7, 8],
params: { nprobe: 2 },
limit: 2,
output_fields: ['vector', 'vector1'],
};
const anns_field = 'anns_field2';

const newSearchParams = buildSearchParams(data, anns_field);

expect(newSearchParams).toEqual({
anns_field: 'anns_field2',
params: '{"nprobe":2}',
topk: 2,
offset: 0,
metric_type: '',
ignore_growing: false,
});

const data2: SearchSimpleReq = {
collection_name: 'test',
data: [1, 2, 3, 4, 5, 6, 7, 8],
anns_field: 'vector',
params: { nprobe: 2 },
limit: 2,
output_fields: ['vector', 'vector1'],
group_by_field: 'group_by_field_value',
group_size: 5,
strict_group_size: true,
};

const newSearchParams2 = buildSearchParams(data2, anns_field);

expect(newSearchParams2).toEqual({
anns_field: 'vector',
params: '{"nprobe":2}',
topk: 2,
offset: 0,
metric_type: '',
ignore_growing: false,
group_by_field: 'group_by_field_value',
group_size: 5,
strict_group_size: true,
});
});
});

0 comments on commit 17d86a7

Please sign in to comment.