support group size and strict_group_size (#371)

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

milvus/types/Data.ts

@@ -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
@@ -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
 }

milvus/utils/Format.ts

@@ -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;
 };

test/grpc/Data.spec.ts

@@ -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);

test/utils/Format.spec.ts

@@ -29,6 +29,8 @@ import {
   formatSearchData,
   buildSearchRequest,
   FieldSchema,
+  buildSearchParams,
+  SearchSimpleReq,
 } from '../../milvus';
 
 describe('utils/format', () => {
@@ -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,
+    });
+  });
 });