From 500b899589b648226261aaf62a06850b9e26e71d Mon Sep 17 00:00:00 2001
From: Yicai Yu <yuyicai@hotmail.com>
Date: Thu, 14 Sep 2023 15:52:39 +0800
Subject: [PATCH] feat: generate expr by pk type when delete entities by ids
 (#242)

Signed-off-by: Yicai Yu <yuyicai@hotmail.com>
---
 .gitignore                |  3 ++-
 milvus/grpc/Collection.ts | 48 +++++++++++++++++++++++++++++++++++++--
 milvus/grpc/Data.ts       | 18 +++++++++++----
 3 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 426c901c..187669ac 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,4 +4,5 @@ docs
 api
 coverage
 .DS_Store
-yarn-error.log
\ No newline at end of file
+yarn-error.log
+.idea
diff --git a/milvus/grpc/Collection.ts b/milvus/grpc/Collection.ts
index c58a7a06..d9182ea7 100644
--- a/milvus/grpc/Collection.ts
+++ b/milvus/grpc/Collection.ts
@@ -43,7 +43,7 @@ import {
   formatCollectionSchema,
   formatDescribedCol,
   validatePartitionNumbers,
-  METADATA,
+  METADATA, DataTypeMap, DataType,
 } from '../';
 
 /**
@@ -172,7 +172,7 @@ export class Collection extends Database {
  *  | :-- | :-- | :-- |
  *  | collection_name | String | Collection name |
  *  | timeout? | number | An optional duration of time in millisecond to allow for the RPC. If it is set to undefined, the client keeps waiting until the server responds or error occurs. Default is undefined |
- * 
+ *
  * @returns
  *  | Property | Description |
  *  | :-- | :-- |
@@ -985,4 +985,48 @@ export class Collection extends Database {
 
     return pkField;
   }
+
+  /**
+   * Get the primary key field type
+   *
+   * @param data
+   *  | Property | Type | Description |
+   *  | :-- | :-- | :-- |
+   *  | collection_name | string | the name of the collection |
+   *  | timeout? | number | An optional duration of time in milliseconds to allow for the RPC. If it is set to undefined, the client keeps waiting until the server responds or an error occurs. Default is undefined |
+   *
+   * @returns
+   *  | Property | Description |
+   *  | :-- | :-- |
+   *  | pkFieldType | the primary key field type |
+   *
+   * @throws {Error} if `collection_name` property is not present in `data`
+   *
+   * #### Example
+   *
+   * ```
+   *  new milvusClient(MILUVS_ADDRESS).getPkFieldType({
+   *    collection_name: 'my_collection',
+   *  });
+   * ```
+   */
+  async getPkFieldType(data: DescribeCollectionReq): Promise<keyof typeof DataType> {
+    // get collection info
+    const collectionInfo = await this.describeCollection(data);
+
+    // pk field type
+    let pkFieldType: keyof typeof DataType = 'None';
+    // extract key information
+    for (let i = 0; i < collectionInfo.schema.fields.length; i++) {
+      const f = collectionInfo.schema.fields[i];
+
+      // get pk field type info
+      if (f.is_primary_key) {
+        pkFieldType = f.data_type;
+        break;
+      }
+    }
+
+    return pkFieldType;
+  }
 }
diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts
index ffd27c0d..7e73fed8 100644
--- a/milvus/grpc/Data.ts
+++ b/milvus/grpc/Data.ts
@@ -103,7 +103,7 @@ export class Data extends Collection {
    *      scalar_field: 1
    *    }]
    *  });
-   * ``` 
+   * ```
    */
   private async _insert(
     data: InsertReq,
@@ -332,10 +332,11 @@ export class Data extends Collection {
     }
 
     const pkField = await this.getPkFieldName(data);
+    const pkFieldType = await this.getPkFieldType(data);
 
-    // generate expr by different type of ids
+    // generate expr by different type of pk
     const expr =
-      typeof data.ids[0] === 'string'
+      DataTypeMap[pkFieldType] === DataType.VarChar
         ? `${pkField} in ["${data.ids.join('","')}"]`
         : `${pkField} in [${data.ids.join(',')}]`;
     const req = { ...data, expr };
@@ -809,9 +810,16 @@ export class Data extends Collection {
       throw new Error(ERROR_REASONS.IDS_REQUIRED);
     }
 
-    // build query req
-    const req = { ...data, expr: `${pkField} in [${data.ids.join(',')}]` };
+    const pkFieldType = await this.getPkFieldType(data);
+
+    // generate expr by different type of pk
+    const expr =
+      DataTypeMap[pkFieldType] === DataType.VarChar
+        ? `${pkField} in ["${data.ids.join('","')}"]`
+        : `${pkField} in [${data.ids.join(',')}]`;
 
+    // build query req
+    const req = { ...data, expr };
     return this.query(req);
   }