From 88c2ee009fbc5ff41d6a9b6a069758647ea80321 Mon Sep 17 00:00:00 2001 From: ryjiang Date: Tue, 8 Oct 2024 14:52:18 +0800 Subject: [PATCH] support clustering key (#363) * support clustering key Signed-off-by: ryjiang * update test case Signed-off-by: ryjiang --------- 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 | 111 ++++++++++++++++++++++++++++++++ test/tools/collection.ts | 6 ++ test/utils/Format.spec.ts | 5 +- 6 files changed, 135 insertions(+), 2 deletions(-) create mode 100644 test/grpc/ClusteringKey.spec.ts diff --git a/milvus/types/Collection.ts b/milvus/types/Collection.ts index eeba992..6493d12 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 b5e8aa4..b4fb324 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 49762eb..e88f55f 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 0000000..cb473b1 --- /dev/null +++ b/test/grpc/ClusteringKey.spec.ts @@ -0,0 +1,111 @@ +import { MilvusClient, DataType, ErrorCode } from '../../milvus'; +import { IP, GENERATE_NAME, genCollectionParams } 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 res = await milvusClient.createCollection( + genCollectionParams({ + collectionName: COLLECTION_NAME, + dim: [4], + vectorType: [DataType.FloatVector], + autoID: true, + clusterKeyEnabled: true, + }) + ); + 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(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 7740c71..e1f8385 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 e98d34d..8317f2b 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']);