From 28d664a0874e04f0e1896293cf074a69add48a2d Mon Sep 17 00:00:00 2001 From: ryjiang Date: Tue, 8 Oct 2024 14:41:17 +0800 Subject: [PATCH 1/2] support clustering key Signed-off-by: ryjiang --- milvus/types/Collection.ts | 2 + milvus/utils/Format.ts | 3 + test/grpc/Basic.spec.ts | 10 ++- test/grpc/ClusteringKey.spec.ts | 139 ++++++++++++++++++++++++++++++++ test/tools/collection.ts | 6 ++ test/utils/Format.spec.ts | 5 +- 6 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 test/grpc/ClusteringKey.spec.ts diff --git a/milvus/types/Collection.ts b/milvus/types/Collection.ts index eeba9927..6493d129 100644 --- a/milvus/types/Collection.ts +++ b/milvus/types/Collection.ts @@ -67,6 +67,7 @@ export interface FieldType { element_type?: DataType | keyof typeof DataTypeMap; is_primary_key?: boolean; is_partition_key?: boolean; + is_clustering_key?: boolean; type_params?: { [key: string]: TypeParam; }; @@ -98,6 +99,7 @@ export interface BaseCreateCollectionReq extends GrpcTimeOut { | 'Customized'; // optional,consistency level, default is 'Bounded' num_partitions?: number; // optional, partitions number, default is 1 partition_key_field?: string; // optional, partition key field + clustring_key_field?: string; // optional, clustring key field enable_dynamic_field?: boolean; // optional, enable dynamic field, default is false enableDynamicField?: boolean; // optional, alias of enable_dynamic_field properties?: Properties; diff --git a/milvus/utils/Format.ts b/milvus/utils/Format.ts index b5e8aa47..b4fb3247 100644 --- a/milvus/utils/Format.ts +++ b/milvus/utils/Format.ts @@ -295,6 +295,7 @@ export const formatCollectionSchema = ( enable_dynamic_field, enableDynamicField, partition_key_field, + clustring_key_field, } = data; let fields = (data as CreateCollectionWithFieldsReq).fields; @@ -318,6 +319,8 @@ export const formatCollectionSchema = ( isPrimaryKey: !!field.is_primary_key, isPartitionKey: !!field.is_partition_key || field.name === partition_key_field, + isClusteringKey: + !!field.is_clustering_key || field.name === clustring_key_field, }; // if element type exist and diff --git a/test/grpc/Basic.spec.ts b/test/grpc/Basic.spec.ts index 49762eb9..e88f55f2 100644 --- a/test/grpc/Basic.spec.ts +++ b/test/grpc/Basic.spec.ts @@ -1,7 +1,7 @@ import { MilvusClient, ErrorCode, DataType } from '../../milvus'; import { IP, GENERATE_NAME, generateInsertData } from '../tools'; -const milvusClient = new MilvusClient({ address: IP }); +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); const COLLECTION_NAME = GENERATE_NAME(); const schema = [ { @@ -49,6 +49,14 @@ describe(`Basic API without database`, () => { collection_name: COLLECTION_NAME, }); expect(desc.schema.fields.length).toEqual(schema.length); + expect(desc.schema.fields[0].name).toEqual('vector'); + expect(desc.schema.fields[1].name).toEqual('id'); + expect(desc.schema.fields[2].name).toEqual('varChar'); + expect(desc.schema.fields[3].name).toEqual('array'); + // test primary key + expect(desc.schema.fields[1].is_primary_key).toEqual(true); + // test partition key + expect(desc.schema.fields[2].is_partition_key).toEqual(false); }); it(`Create index should be successful`, async () => { diff --git a/test/grpc/ClusteringKey.spec.ts b/test/grpc/ClusteringKey.spec.ts new file mode 100644 index 00000000..06c6982f --- /dev/null +++ b/test/grpc/ClusteringKey.spec.ts @@ -0,0 +1,139 @@ +import { MilvusClient, DataType, ErrorCode } from '../../milvus'; +import { IP, GENERATE_NAME } from '../tools'; + +const milvusClient = new MilvusClient({ address: IP }); +const COLLECTION_NAME = GENERATE_NAME(); +const COLLECTION_NAME2 = GENERATE_NAME(); +const dbParam = { + db_name: 'ClusteringKey', +}; + +describe(`Clustering key API`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ + collection_name: COLLECTION_NAME, + }); + await milvusClient.dropCollection({ + collection_name: COLLECTION_NAME2, + }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create Collection with schema should successfully`, async () => { + const schema = [ + { + name: 'vector', + description: 'Vector field', + data_type: DataType.FloatVector, + dim: Number(4), + }, + { + name: 'id', + description: 'ID field', + data_type: DataType.Int64, + is_primary_key: true, + autoID: true, + }, + { + name: 'varChar', + description: 'VarChar field', + data_type: DataType.VarChar, + max_length: 128, + is_partition_key: false, + is_clustering_key: true, + }, + { + name: 'array', + description: 'array field', + data_type: DataType.Array, + element_type: DataType.VarChar, + max_capacity: 128, + max_length: 128, + is_partition_key: false, + }, + ]; + const res = await milvusClient.createCollection({ + collection_name: COLLECTION_NAME, + fields: schema, + }); + expect(res.error_code).toEqual(ErrorCode.SUCCESS); + + // describe + const desc = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + // test clustering key + let found = 0; + desc.schema.fields.forEach(field => { + if (field.is_clustering_key) { + found++; + expect(field.name).toEqual('varChar'); + } + }); + + expect(found).toEqual(1); + }); + + it(`Create Collection should be successful with clusteringkey`, async () => { + const schema = [ + { + name: 'vector', + description: 'Vector field', + data_type: DataType.FloatVector, + dim: Number(4), + }, + { + name: 'id', + description: 'ID field', + data_type: DataType.Int64, + is_primary_key: true, + autoID: true, + }, + { + name: 'varChar', + description: 'VarChar field', + data_type: DataType.VarChar, + max_length: 128, + is_partition_key: false, + is_clustering_key: false, + }, + { + name: 'array', + description: 'array field', + data_type: DataType.Array, + element_type: DataType.VarChar, + max_capacity: 128, + max_length: 128, + is_partition_key: false, + }, + ]; + const res = await milvusClient.createCollection({ + collection_name: COLLECTION_NAME2, + fields: schema, + clustring_key_field: 'varChar', + }); + expect(res.error_code).toEqual(ErrorCode.SUCCESS); + + // describe + const desc = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME2, + }); + + // test clustering key + let found = 0; + desc.schema.fields.forEach(field => { + if (field.is_clustering_key) { + found++; + expect(field.name).toEqual('varChar'); + } + }); + + expect(found).toEqual(1); + }); +}); diff --git a/test/tools/collection.ts b/test/tools/collection.ts index 7740c71b..e1f83851 100644 --- a/test/tools/collection.ts +++ b/test/tools/collection.ts @@ -37,6 +37,7 @@ export const genCollectionParams = (data: { autoID?: boolean; fields?: any[]; partitionKeyEnabled?: boolean; + clusterKeyEnabled?: boolean; numPartitions?: number; enableDynamic?: boolean; maxCapacity?: number; @@ -53,6 +54,7 @@ export const genCollectionParams = (data: { enableDynamic = false, maxCapacity, idType = DataType.Int64, + clusterKeyEnabled = false, } = data; const vectorFields = vectorType.map((type, i) => { @@ -152,5 +154,9 @@ export const genCollectionParams = (data: { params.num_partitions = numPartitions; } + if (clusterKeyEnabled) { + params.clustring_key_field = 'int64'; + } + return params; }; diff --git a/test/utils/Format.spec.ts b/test/utils/Format.spec.ts index e98d34d0..8317f2b7 100644 --- a/test/utils/Format.spec.ts +++ b/test/utils/Format.spec.ts @@ -301,6 +301,7 @@ describe('utils/format', () => { dataType: 5, isPrimaryKey: true, isPartitionKey: false, + isClusteringKey: false, }, { typeParams: [ @@ -317,6 +318,7 @@ describe('utils/format', () => { dataType: 101, isPrimaryKey: false, isPartitionKey: false, + isClusteringKey: false, }, { typeParams: [ @@ -334,6 +336,7 @@ describe('utils/format', () => { isPrimaryKey: false, isPartitionKey: false, elementType: 5, + isClusteringKey: false, }, ], }; @@ -866,7 +869,7 @@ describe('utils/format', () => { describeCollectionResponse, milvusProto ); - console.dir(searchRequest, { depth: null }); + // console.dir(searchRequest, { depth: null }); expect(searchRequest.isHybridSearch).toEqual(true); expect(searchRequest.request.collection_name).toEqual('test'); expect(searchRequest.request.output_fields).toEqual(['vector', 'vector1']); From 3af64063d67ccd4f81946ed1dd2bbc88c320dbb8 Mon Sep 17 00:00:00 2001 From: ryjiang Date: Tue, 8 Oct 2024 14:48:00 +0800 Subject: [PATCH 2/2] update test case Signed-off-by: ryjiang --- test/grpc/ClusteringKey.spec.ts | 46 +++++++-------------------------- 1 file changed, 9 insertions(+), 37 deletions(-) diff --git a/test/grpc/ClusteringKey.spec.ts b/test/grpc/ClusteringKey.spec.ts index 06c6982f..cb473b1e 100644 --- a/test/grpc/ClusteringKey.spec.ts +++ b/test/grpc/ClusteringKey.spec.ts @@ -1,5 +1,5 @@ import { MilvusClient, DataType, ErrorCode } from '../../milvus'; -import { IP, GENERATE_NAME } from '../tools'; +import { IP, GENERATE_NAME, genCollectionParams } from '../tools'; const milvusClient = new MilvusClient({ address: IP }); const COLLECTION_NAME = GENERATE_NAME(); @@ -25,42 +25,15 @@ describe(`Clustering key API`, () => { }); it(`Create Collection with schema should successfully`, async () => { - const schema = [ - { - name: 'vector', - description: 'Vector field', - data_type: DataType.FloatVector, - dim: Number(4), - }, - { - name: 'id', - description: 'ID field', - data_type: DataType.Int64, - is_primary_key: true, + const res = await milvusClient.createCollection( + genCollectionParams({ + collectionName: COLLECTION_NAME, + dim: [4], + vectorType: [DataType.FloatVector], autoID: true, - }, - { - name: 'varChar', - description: 'VarChar field', - data_type: DataType.VarChar, - max_length: 128, - is_partition_key: false, - is_clustering_key: true, - }, - { - name: 'array', - description: 'array field', - data_type: DataType.Array, - element_type: DataType.VarChar, - max_capacity: 128, - max_length: 128, - is_partition_key: false, - }, - ]; - const res = await milvusClient.createCollection({ - collection_name: COLLECTION_NAME, - fields: schema, - }); + clusterKeyEnabled: true, + }) + ); expect(res.error_code).toEqual(ErrorCode.SUCCESS); // describe @@ -73,7 +46,6 @@ describe(`Clustering key API`, () => { desc.schema.fields.forEach(field => { if (field.is_clustering_key) { found++; - expect(field.name).toEqual('varChar'); } });