Add nq > 1 tests for sparse vectors and upgrade protos (#297)

* add more test Signed-off-by: ryjiang <jiangruiyi@gmail.com> * add nq > 1 tests for sparse vectors Signed-off-by: ryjiang <jiangruiyi@gmail.com> * update test version Signed-off-by: ryjiang <jiangruiyi@gmail.com> --------- Signed-off-by: ryjiang <jiangruiyi@gmail.com>
milvus-io · Apr 11, 2024 · 1183230 · 1183230
1 parent 60523c5
commit 1183230
Show file tree

Hide file tree

Showing 13 changed files with 267 additions and 27 deletions.
diff --git a/milvus/utils/Bytes.ts b/milvus/utils/Bytes.ts
@@ -58,17 +58,37 @@ export const parseBytesToFloat16Vector = (float16Bytes: Uint8Array) => {
  *
  * @returns string, 'array' | 'coo' | 'csr' | 'dict'
  */
-export const getSparseFloatVectorType = (vector: SparseFloatVector) => {
+export const getSparseFloatVectorType = (
+  vector: SparseFloatVector
+): 'array' | 'coo' | 'csr' | 'dict' | 'unknown' => {
   if (Array.isArray(vector)) {
+    if (vector.length === 0) {
+      return 'array';
+    }
     if (typeof vector[0] === 'number' || typeof vector[0] === 'undefined') {
       return 'array';
-    } else {
+    } else if (
+      (vector as SparseVectorCOO).every(
+        item => typeof item === 'object' && 'index' in item && 'value' in item
+      )
+    ) {
       return 'coo';
+    } else {
+      return 'unknown';
     }
-  } else if ('indices' in vector && 'values' in vector) {
+  } else if (
+    typeof vector === 'object' &&
+    'indices' in vector &&
+    'values' in vector
+  ) {
     return 'csr';
-  } else {
+  } else if (
+    typeof vector === 'object' &&
+    Object.keys(vector).every(key => typeof vector[key] === 'number')
+  ) {
     return 'dict';
+  } else {
+    return 'unknown';
   }
 };
 
@@ -86,8 +106,8 @@ export const parseSparseVectorToBytes = (
   // detect the format of the sparse vector
   const type = getSparseFloatVectorType(data);
 
-  let indices: number[];
-  let values: number[];
+  let indices: number[] = [];
+  let values: number[] = [];
 
   switch (type) {
     case 'array':

diff --git a/milvus/utils/Format.ts b/milvus/utils/Format.ts
@@ -34,6 +34,7 @@ import {
   parseBytesToFloat16Vector,
   parseFloat16VectorToBytes,
   Float16Vector,
+  getSparseFloatVectorType,
 } from '../';
 
 /**
@@ -691,10 +692,8 @@ export const buildSearchRequest = (
           searchSimpleReq.vector ||
           searchSimpleReq.data;
 
-      // make sure the vector format
-      if (!Array.isArray(searchingVector[0])) {
-        searchingVector = [searchingVector as unknown] as VectorTypes[];
-      }
+      // format saerching vector
+      searchingVector = formatSearchVector(searchingVector, field.dataType!);
 
       // create search request
       requests.push({
@@ -842,3 +841,31 @@ export const formatSearchResult = (
 
   return results;
 };
+
+/**
+ * Formats the search vector to match a specific data type.
+ * @param {VectorTypes | VectorTypes[]} searchVector - The search vector or array of vectors to be formatted.
+ * @param {DataType} dataType - The specified data type.
+ * @returns {VectorTypes[]} The formatted search vector or array of vectors.
+ */
+export const formatSearchVector = (
+  searchVector: VectorTypes | VectorTypes[],
+  dataType: DataType
+): VectorTypes[] => {
+  switch (dataType) {
+    case DataType.FloatVector:
+    case DataType.BinaryVector:
+    case DataType.Float16Vector:
+    case DataType.BFloat16Vector:
+      if (!Array.isArray(searchVector)) {
+        return [searchVector] as VectorTypes[];
+      }
+    case DataType.SparseFloatVector:
+      const type = getSparseFloatVectorType(searchVector as VectorTypes);
+      if (type !== 'unknown') {
+        return [searchVector] as VectorTypes[];
+      }
+    default:
+      return searchVector as VectorTypes[];
+  }
+};
diff --git a/package.json b/package.json
@@ -2,7 +2,7 @@
   "name": "@zilliz/milvus2-sdk-node",
   "author": "ued@zilliz.com",
   "version": "2.3.6",
-  "milvusVersion": "2.4-20240327-96cec787-amd64",
+  "milvusVersion": "2.4-20240411-246ef454-amd64",
   "main": "dist/milvus",
   "files": [
     "dist"

diff --git a/proto b/proto
diff --git a/test/grpc/Basic.spec.ts b/test/grpc/Basic.spec.ts
@@ -80,6 +80,20 @@ describe(`Basic API without database`, () => {
     expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
   });
 
+  it(`search nq > 1 should be successful`, async () => {
+    const search = await milvusClient.search({
+      collection_name: COLLECTION_NAME,
+      data: [
+        [1, 2, 3, 4],
+        [5, 6, 7, 8],
+      ],
+    });
+    expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
+    expect(search.results.length).toEqual(2);
+    expect(search.results[0].length).toEqual(10);
+    expect(search.results[1].length).toEqual(10);
+  });
+
   it(`release and drop should be successful`, async () => {
     // releases
     const release = await milvusClient.releaseCollection({

diff --git a/test/grpc/Float16Vector.spec.ts b/test/grpc/Float16Vector.spec.ts
@@ -4,7 +4,6 @@ import {
   DataType,
   IndexType,
   MetricType,
-  parseBytesToFloat16Vector,
 } from '../../milvus';
 import {
   IP,
@@ -126,4 +125,18 @@ describe(`Float16 vector API testing`, () => {
     expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
     expect(search.results.length).toBeGreaterThan(0);
   });
+
+  it(`search with float16 vector and nq > 0 should be successful`, async () => {
+    const search = await milvusClient.search({
+      vector: [data[0].vector, data[1].vector],
+      collection_name: COLLECTION_NAME,
+      output_fields: ['id', 'vector'],
+      limit: 5,
+    });
+
+    // console.log('search', search);
+
+    expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
+    expect(search.results.length).toBeGreaterThan(0);
+  });
 });
diff --git a/test/grpc/MultipleVectors.spec.ts b/test/grpc/MultipleVectors.spec.ts
@@ -14,7 +14,7 @@ import {
   generateInsertData,
 } from '../tools';
 
-const milvusClient = new MilvusClient({ address: IP });
+const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' });
 const COLLECTION_NAME = GENERATE_NAME();
 
 const dbParam = {
@@ -23,7 +23,11 @@ const dbParam = {
 
 const p = {
   collectionName: COLLECTION_NAME,
-  vectorType: [DataType.FloatVector, DataType.FloatVector],
+  vectorType: [
+    DataType.FloatVector,
+    DataType.FloatVector,
+    // DataType.Float16Vector,
+  ],
   dim: [8, 16],
 };
 const collectionParams = genCollectionParams(p);
@@ -85,6 +89,12 @@ describe(`Multiple vectors API testing`, () => {
         metric_type: MetricType.COSINE,
         index_type: IndexType.AUTOINDEX,
       },
+      // {
+      //   collection_name: COLLECTION_NAME,
+      //   field_name: 'vector2',
+      //   metric_type: MetricType.COSINE,
+      //   index_type: IndexType.AUTOINDEX,
+      // },
     ]);
 
     expect(indexes.error_code).toEqual(ErrorCode.SUCCESS);
@@ -110,44 +120,46 @@ describe(`Multiple vectors API testing`, () => {
     const item = query.data[0];
     expect(item.vector.length).toEqual(p.dim[0]);
     expect(item.vector1.length).toEqual(p.dim[1]);
+    // expect(item.vector2.length).toEqual(p.dim[2]);
   });
 
   it(`search multiple vector collection with old search api should be successful`, async () => {
     // search default first vector field
     const search0 = await milvusClient.search({
       collection_name: COLLECTION_NAME,
-      vector: [1, 2, 3, 4, 5, 6, 7, 8],
+      data: [1, 2, 3, 4, 5, 6, 7, 8],
     });
     expect(search0.status.error_code).toEqual(ErrorCode.SUCCESS);
 
     // search specific vector field
     const search = await milvusClient.search({
       collection_name: COLLECTION_NAME,
-      vector: [1, 2, 3, 4, 5, 6, 7, 8],
+      data: [1, 2, 3, 4, 5, 6, 7, 8],
       anns_field: 'vector',
     });
     expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
 
     // search second vector field
     const search2 = await milvusClient.search({
       collection_name: COLLECTION_NAME,
-      vector: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
+      data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
       anns_field: 'vector1',
       limit: 5,
     });
 
     expect(search2.status.error_code).toEqual(ErrorCode.SUCCESS);
     expect(search2.results.length).toEqual(5);
 
-    // search first vector field
-    const search3 = await milvusClient.search({
-      collection_name: COLLECTION_NAME,
-      vector: [1, 2, 3, 4, 5, 6, 7, 8],
-    });
+    // // search third vector field
+    // const search3 = await milvusClient.search({
+    //   collection_name: COLLECTION_NAME,
+    //   data: [1, 2, 3, 4, 5, 6, 7, 8],
+    //   anns_field: 'vector2',
+    // });
 
-    expect(search3.status.error_code).toEqual(ErrorCode.SUCCESS);
-    expect(search3.results.length).toEqual(search.results.length);
-    expect(search3.results).toEqual(search.results);
+    // expect(search3.status.error_code).toEqual(ErrorCode.SUCCESS);
+    // expect(search3.results.length).toEqual(search.results.length);
+    // expect(search3.results).toEqual(search.results);
   });
 
   it(`hybrid search with rrf ranker set should be successful`, async () => {

diff --git a/test/grpc/SparseVector.array.spec.ts b/test/grpc/SparseVector.array.spec.ts
@@ -128,4 +128,16 @@ describe(`Sparse vectors type:object API testing`, () => {
     expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
     expect(search.results.length).toBeGreaterThan(0);
   });
+
+  it(`search with sparse vector with nq > 1 should be successful`, async () => {
+    const search = await milvusClient.search({
+      vectors: [data[0].vector, data[1].vector],
+      collection_name: COLLECTION_NAME,
+      output_fields: ['id', 'vector'],
+      limit: 5,
+    });
+
+    expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
+    expect(search.results.length).toEqual(2);
+  });
 });
diff --git a/test/grpc/SparseVector.coo.spec.ts b/test/grpc/SparseVector.coo.spec.ts
@@ -128,4 +128,16 @@ describe(`Sparse vectors type:coo API testing`, () => {
     expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
     expect(search.results.length).toBeGreaterThan(0);
   });
+
+  it(`search with sparse vector with nq > 1 should be successful`, async () => {
+    const search = await milvusClient.search({
+      vectors: [data[0].vector, data[1].vector],
+      collection_name: COLLECTION_NAME,
+      output_fields: ['id', 'vector'],
+      limit: 5,
+    });
+
+    expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
+    expect(search.results.length).toEqual(2);
+  });
 });
diff --git a/test/grpc/SparseVector.csr.spec.ts b/test/grpc/SparseVector.csr.spec.ts
@@ -101,7 +101,9 @@ describe(`Sparse vectors type:CSR API testing`, () => {
       output_fields: ['vector', 'id'],
     });
 
-    const originKeys = data[0].vector.indices.map((index: number) => index.toString());
+    const originKeys = data[0].vector.indices.map((index: number) =>
+      index.toString()
+    );
     const originValues = data[0].vector.values;
 
     const outputKeys: string[] = Object.keys(query.data[0].vector);
@@ -124,4 +126,16 @@ describe(`Sparse vectors type:CSR API testing`, () => {
     expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
     expect(search.results.length).toBeGreaterThan(0);
   });
+
+  it(`search with sparse vector with nq > 1 should be successful`, async () => {
+    const search = await milvusClient.search({
+      vectors: [data[0].vector, data[1].vector],
+      collection_name: COLLECTION_NAME,
+      output_fields: ['id', 'vector'],
+      limit: 5,
+    });
+
+    expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
+    expect(search.results.length).toEqual(2);
+  });
 });
diff --git a/test/grpc/SparseVector.dict.spec.ts b/test/grpc/SparseVector.dict.spec.ts
@@ -120,4 +120,16 @@ describe(`Sparse vectors type:dict API testing`, () => {
     expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
     expect(search.results.length).toBeGreaterThan(0);
   });
+
+  it(`search with sparse vector with nq > 1 should be successful`, async () => {
+    const search = await milvusClient.search({
+      vectors: [data[0].vector, data[1].vector],
+      collection_name: COLLECTION_NAME,
+      output_fields: ['id', 'vector'],
+      limit: 5,
+    });
+
+    expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
+    expect(search.results.length).toEqual(2);
+  });
 });
diff --git a/test/utils/Bytes.spec.ts b/test/utils/Bytes.spec.ts
@@ -63,4 +63,15 @@ describe('Sparse rows <-> Bytes conversion', () => {
     ];
     expect(getSparseFloatVectorType(data)).toEqual('coo');
   });
+
+  it('should return "unknown" if the input is not recognized', () => {
+    const data: any = 'invalid';
+    expect(getSparseFloatVectorType(data)).toEqual('unknown');
+
+    const data2: any = [
+      [1, 2, 3],
+      [4, 5, 6],
+    ];
+    expect(getSparseFloatVectorType(data2)).toEqual('unknown');
+  });
 });
+7 −0		.github/mergify.yml
+1 −0		.github/workflows/check.yaml
+229 −228		go-api/commonpb/common.pb.go
+27 −27		go-api/rgpb/rg.pb.go
+1 −1		proto/common.proto
+3 −4		proto/rg.proto