support clustering key (#363)

* support clustering key Signed-off-by: ryjiang <[email protected]> * update test case Signed-off-by: ryjiang <[email protected]> --------- Signed-off-by: ryjiang <[email protected]>
milvus-io · Oct 8, 2024 · 88c2ee0 · 88c2ee0
1 parent 500ad2b
commit 88c2ee0
Show file tree

Hide file tree

Showing 6 changed files with 135 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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']);