From 1a53e35b18a557b98666372f3862b7c91ca4f43e Mon Sep 17 00:00:00 2001 From: ryjiang Date: Thu, 18 Apr 2024 11:39:07 +0800 Subject: [PATCH] Support Milvus v2.4 (#306) * init 2.4 Signed-off-by: ryjiang * sync 2.4 proto Signed-off-by: ryjiang * update CI Signed-off-by: ryjiang * add 2.4 index (#272) * add INVERTED index Signed-off-by: ryjiang * add GPU_BRUTE_FORCE and GPU_CAGRA Signed-off-by: ryjiang * fix index test Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * support search group by (#273) * support group by Signed-off-by: ryjiang * remove console Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * update 2.4 version Signed-off-by: ryjiang * support mmap (#275) * support mmap Signed-off-by: ryjiang * comment out test for now Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * update proto Signed-off-by: ryjiang * support search with multiple vectors (#278) * update genCollectionParams to support multiple vectors field Signed-off-by: ryjiang * add test for generateInsertData about multiple vectors Signed-off-by: ryjiang * add compatible test for mutliple vectors Signed-off-by: ryjiang * update proto Signed-off-by: ruiyi.jiang * optimize search() Signed-off-by: ruiyi.jiang * part Signed-off-by: ruiyi.jiang * move buildSearchParams and formatSearchResult to utils Signed-off-by: ryjiang * refactor Signed-off-by: ryjiang * fix test Signed-off-by: ryjiang * refactor part2 Signed-off-by: ryjiang * improve code Signed-off-by: ryjiang * refactor part3 Signed-off-by: ryjiang * update part5 Signed-off-by: ryjiang * prefinish multivsearch Signed-off-by: ryjiang * add comments Signed-off-by: ryjiang * update test Signed-off-by: ryjiang * fix comments Signed-off-by: ryjiang * fix test Signed-off-by: ryjiang * fix build Signed-off-by: ryjiang --------- Signed-off-by: ryjiang Signed-off-by: ruiyi.jiang * sync 2.4 proto Signed-off-by: ruiyi.jiang * Add sparse vector support (#285) * add test data and const Signed-off-by: ruiyi.jiang * add sparse data generator Signed-off-by: ruiyi.jiang * fix create collection Signed-off-by: ruiyi.jiang * stash Signed-off-by: ruiyi.jiang * remove duplicate data sent to the sever Signed-off-by: ryjiang * refactor data part1 Signed-off-by: ryjiang * finish insert Signed-off-by: ryjiang * add query Signed-off-by: ryjiang * finish sparse vector Signed-off-by: ryjiang * fix alter collection should run before load Signed-off-by: ryjiang * fix http Signed-off-by: ryjiang * update test milvus verison Signed-off-by: ryjiang * fix test Signed-off-by: ryjiang * add http test Signed-off-by: ryjiang --------- Signed-off-by: ruiyi.jiang Signed-off-by: ryjiang * fix alterIndex Signed-off-by: ryjiang * Support f16 & bf16 (#287) * add f16 insert Signed-off-by: ryjiang * fp16 part2 Signed-off-by: ryjiang * stash Signed-off-by: ryjiang * f16 part3 Signed-off-by: ryjiang * finish query Signed-off-by: ryjiang * finish f16 Signed-off-by: ryjiang * fix type errors Signed-off-by: ryjiang * update parseFloat16VectorToBytes function Signed-off-by: ryjiang * add bf16 support Signed-off-by: ryjiang * update types Signed-off-by: ryjiang * fix test Signed-off-by: ryjiang * add test Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * Support more types of sparse vectors (#293) * generate different types of sparse vector Signed-off-by: ryjiang * fix sparse array in js Signed-off-by: ryjiang * add sparse array test Signed-off-by: ryjiang * add csr sparse vector test Signed-off-by: ryjiang * add coo support Signed-off-by: ryjiang * remove unused import Signed-off-by: ryjiang * refine comments Signed-off-by: ryjiang * refine comment Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * Sparse test update (#296) * update sparse test case Signed-off-by: ryjiang * log level Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * rename sparsevector test Signed-off-by: ryjiang * Add nq > 1 tests for sparse vectors and upgrade protos (#297) * add more test Signed-off-by: ryjiang * add nq > 1 tests for sparse vectors Signed-off-by: ryjiang * update test version Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * add nq>1 tests for hybridSearch (#298) Signed-off-by: ryjiang * Fix single search failed on mutliple vectors collection if the anns field is specified. (#300) * add more tests Signed-off-by: ryjiang * Fix multiple test Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * update README.md Signed-off-by: ryjiang * add more tests for client contstructor (#301) Signed-off-by: ryjiang * Add bfloat16 support (#302) * add test for bf16 Signed-off-by: ryjiang * add bf16 support Signed-off-by: ryjiang --------- Signed-off-by: ryjiang * add transformers (#303) Signed-off-by: ryjiang * rename types Signed-off-by: ryjiang * update readme (#304) Signed-off-by: ryjiang * Fix search transformers Signed-off-by: ryjiang * Revert "Fix search transformers" This reverts commit 2a14e0ec60b209e0d4548db8fcc9500523315d22. * make transformer optional (#305) * make transform optional Signed-off-by: shanghaikid * update readme Signed-off-by: shanghaikid * fix http test Signed-off-by: shanghaikid * fix format test Signed-off-by: shanghaikid * fix sparse array test Signed-off-by: ryjiang --------- Signed-off-by: shanghaikid Signed-off-by: ryjiang * feat: milvus api v2 (#295) Signed-off-by: Shuyou * bump version Signed-off-by: ryjiang * fix build Signed-off-by: ryjiang --------- Signed-off-by: ryjiang Signed-off-by: ruiyi.jiang Signed-off-by: shanghaikid Signed-off-by: Shuyou Co-authored-by: Shuyoou --- .github/workflows/check.yml | 2 + .gitmodules | 2 +- README.md | 187 +++++++-- milvus/HttpClient.ts | 38 +- milvus/const/defaults.ts | 2 +- milvus/const/error.ts | 2 +- milvus/const/milvus.ts | 24 ++ milvus/grpc/Data.ts | 382 +++++++----------- milvus/grpc/MilvusIndex.ts | 39 ++ milvus/http/Alias.ts | 56 +++ milvus/http/Collection.ts | 79 +++- milvus/http/Import.ts | 48 +++ milvus/http/MilvusIndex.ts | 48 +++ milvus/http/Partition.ts | 80 ++++ milvus/http/Role.ts | 65 +++ milvus/http/User.ts | 69 ++++ milvus/http/Vector.ts | 21 +- milvus/http/index.ts | 6 + milvus/types/Collection.ts | 4 +- milvus/types/Common.ts | 4 +- milvus/types/Data.ts | 233 +++++++---- milvus/types/Http.ts | 236 ++++++++++- milvus/types/MilvusIndex.ts | 5 + milvus/utils/Blob.ts | 15 - milvus/utils/Bytes.ts | 306 ++++++++++++++ milvus/utils/Format.ts | 573 ++++++++++++++++++++++----- milvus/utils/Function.ts | 23 +- milvus/utils/Validate.ts | 23 +- milvus/utils/index.ts | 2 +- package.json | 5 +- proto | 2 +- test/build/Collection.spec.ts | 6 +- test/grpc/Alias.spec.ts | 2 +- test/grpc/BFloat16Vector.spec.ts | 152 +++++++ test/grpc/Basic.spec.ts | 16 +- test/grpc/BinaryVector.spec.ts | 114 ++++++ test/grpc/Collection.spec.ts | 63 +-- test/grpc/Data.spec.ts | 53 ++- test/grpc/Database.spec.ts | 2 +- test/grpc/DynamicSchema.spec.ts | 8 +- test/grpc/Float16Vector.spec.ts | 146 +++++++ test/grpc/Import.spec.ts | 4 +- test/grpc/Index.spec.ts | 80 +++- test/grpc/Insert.spec.ts | 14 +- test/grpc/MilvusClient.spec.ts | 44 ++ test/grpc/MultipleVectors.spec.ts | 363 +++++++++++++++++ test/grpc/Partition.spec.ts | 2 +- test/grpc/PartitionKey.spec.ts | 24 +- test/grpc/Replica.spec.ts | 2 +- test/grpc/Resource.spec.ts | 2 +- test/grpc/SparseVector.array.spec.ts | 143 +++++++ test/grpc/SparseVector.coo.spec.ts | 143 +++++++ test/grpc/SparseVector.csr.spec.ts | 141 +++++++ test/grpc/SparseVector.dict.spec.ts | 135 +++++++ test/grpc/Upsert.spec.ts | 12 +- test/grpc/User.spec.ts | 4 +- test/http/client.spec.ts | 20 +- test/http/test.ts | 433 +++++++++++++++++++- test/tools/bench.ts | 2 +- test/tools/collection.ts | 30 +- test/tools/data.ts | 114 +++++- test/tools/utils.ts | 6 + test/utils/Bytes.spec.ts | 108 +++++ test/utils/Format.spec.ts | 152 ++++++- test/utils/Function.spec.ts | 76 +++- test/utils/Test.spec.ts | 159 +++++++- test/utils/Validate.spec.ts | 6 +- yarn.lock | 5 + 68 files changed, 4673 insertions(+), 694 deletions(-) create mode 100644 milvus/http/Alias.ts create mode 100644 milvus/http/Import.ts create mode 100644 milvus/http/MilvusIndex.ts create mode 100644 milvus/http/Partition.ts create mode 100644 milvus/http/Role.ts create mode 100644 milvus/http/User.ts delete mode 100644 milvus/utils/Blob.ts create mode 100644 milvus/utils/Bytes.ts create mode 100644 test/grpc/BFloat16Vector.spec.ts create mode 100644 test/grpc/BinaryVector.spec.ts create mode 100644 test/grpc/Float16Vector.spec.ts create mode 100644 test/grpc/MultipleVectors.spec.ts create mode 100644 test/grpc/SparseVector.array.spec.ts create mode 100644 test/grpc/SparseVector.coo.spec.ts create mode 100644 test/grpc/SparseVector.csr.spec.ts create mode 100644 test/grpc/SparseVector.dict.spec.ts create mode 100644 test/utils/Bytes.spec.ts diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 64bc900d..4426c54d 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -4,10 +4,12 @@ on: pull_request: branches: - main + - 2.4 types: [opened, synchronize] push: branches: - main + - 2.4 jobs: publish: runs-on: ubuntu-latest diff --git a/.gitmodules b/.gitmodules index b5e19c85..f1271435 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,4 +1,4 @@ [submodule "proto"] path = proto url = https://github.com/milvus-io/milvus-proto.git - branch = master + branch = 2.4 diff --git a/README.md b/README.md index 2574d230..07f3dbc2 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ # Milvus2-sdk-node [![typescript](https://badges.aleen42.com/src/typescript.svg)](https://badges.aleen42.com/src/typescript.svg) -[![version](https://img.shields.io/npm/v/@zilliz/milvus2-sdk-node?color=bright-green)](https://img.shields.io/npm/v/@zilliz/milvus2-sdk-node) -[![downloads](https://img.shields.io/npm/dw/@zilliz/milvus2-sdk-node?color=bright-green)](https://img.shields.io/npm/dw/@zilliz/milvus2-sdk-node) +[![version](https://img.shields.io/npm/v/@zilliz/milvus2-sdk-node?color=bright-green)](https://github.com/zilliztech/attu/releases) +[![downloads](https://img.shields.io/npm/dw/@zilliz/milvus2-sdk-node?color=bright-green)](https://www.npmjs.com/package/@zilliz/milvus2-sdk-node) [![codecov](https://codecov.io/gh/milvus-io/milvus-sdk-node/branch/main/graph/badge.svg?token=Zu5FwWstwI)](https://codecov.io/gh/milvus-io/milvus-sdk-node) The official [Milvus](https://github.com/milvus-io/milvus) client for Node.js. @@ -11,9 +11,11 @@ The official [Milvus](https://github.com/milvus-io/milvus) client for Node.js. The following table shows the recommended `@zilliz/milvus2-sdk-node` versions for different Milvus versions: -| Milvus version | Node sdk version | Installation | -| :------------: | :--------------: | :---------------------------------- | -| v2.2.0+ | **latest** | `yarn add @zilliz/milvus2-sdk-node` | +| Milvus version | Node sdk version | Installation | +| :------------: | :--------------: | :----------------------------------------- | +| v2.4.0+ | **latest** | `yarn add @zilliz/milvus2-sdk-node@latest` | +| v2.3.0+ | v2.3.5 | `yarn add @zilliz/milvus2-sdk-node@2.3.5` | +| v2.2.0+ | v2.3.5 | `yarn add @zilliz/milvus2-sdk-node@2.3.5` | ## Dependencies @@ -31,21 +33,137 @@ npm install @zilliz/milvus2-sdk-node yarn add @zilliz/milvus2-sdk-node ``` -This will download the Milvus Node.js client and add a dependency entry in your package.json file. +## What's new in v2.4.0 + +### New vector data types: float16 and bfloat16 + +Machine learning and neural networks often use half-precision data types, such as Float16 and BFloat16, [Milvus 2.4](https://milvus.io/docs/release_notes.md#Float16-and-BFloat16-Vector-DataType) supports inserting vectors in the BF16 and FP16 formats as bytes. + +> However, these data types are not natively available in the Node.js environment, To enable users to utilize these formats, the Node SDK provides support for transformers during insert, query, and search operations. +> +> There are four default transformers for performing a float32 to bytes transformation for BF16 and Float16 types: f32ArrayToF16Bytes, f16BytesToF32Array, f32ArrayToBf16Bytes, and bf16BytesToF32Array. If you wish to use your own transformers for Float16 and BFloat16, you can specify them. +> +> ```javascript +> import { +> f32ArrayToF16Bytes, +> f16BytesToF32Array, +> f32ArrayToBf16Bytes, +> bf16BytesToF32Array, +> } from '@zilliz/milvus2-sdk-node'; +> +> //Insert float32 array for the float16 field. Node SDK will transform it to bytes using `f32ArrayToF16Bytes`. You can use your own transformer. +> const insert = await milvusClient.insert({ +> collection_name: COLLECTION_NAME, +> data: data, +> // transformers: { +> // [DataType.BFloat16Vector]: f32ArrayToF16Bytes, // use your own transformer +> // }, +> }); +> // query: output float32 array other than bytes, +> const query = await milvusClient.query({ +> collection_name: COLLECTION_NAME, +> filter: 'id > 0', +> output_fields: ['vector', 'id'], +> // transformers: { +> // [DataType.BFloat16Vector]: bf16BytesToF32Array, // use your own transformer +> // }, +> }); +> // search: use bytes to search, output float32 array +> const search = await milvusClient.search({ +> vector: data[0].vector, // if you pass bytes, no transform will performed +> collection_name: COLLECTION_NAME, +> output_fields: ['id', 'vector'], +> limit: 5, +> // transformers: { +> // [DataType.BFloat16Vector]: bf16BytesToF32Array, // use your own transformer +> // }, +> }); +> ``` + +### New vector data types: sparse vector(beta) + +Sparse vectors in the Node SDK support four formats: `dict`, `coo`, `csr`, and `array`, However, query and search operations currently only output in the dict format. -## Code Examples +```javascript +// dict +const sparseObject = { + 3: 1.5, + 6: 2.0, + 9: -3.5, +}; +// coo +const sparseCOO = [ + { index: 2, value: 5 }, + { index: 5, value: 3 }, + { index: 8, value: 7 }, +]; +// csr +const sparseCSR = { + indices: [2, 5, 8], + values: [5, 3, 7], +}; +// array +const sparseArray = [undefined, 0.0, 0.5, 0.3, undefined, 0.2]; +``` -### Milvus examples +### Multi-vector and Hybrid Search -You can find code examples in the [examples/milvus](./examples/milvus) directory. These examples cover various aspects of working with Milvus, such as connecting to Milvus, vector search, data query, dynamic schema, partition key, and database operations. +Starting from Milvus 2.4, it supports [Multi-Vector Search](https://milvus.io/docs/multi-vector-search.md#API-overview), you can continue to utilize the search API with similar parameters to perform multi-vector searches, and the format of the results remains unchanged. -### Langchain.js example +```javascript +// single-vector search on a collection with multiple vector fields +const search = await milvusClient.search({ + collection_name: collection_name, + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', // required if you have multiple vector fields in the collection + params: { nprobe: 2 }, + filter: 'id > 100', + limit: 5, +}); -You can find a basic langchain.js example in the [examples/langchain](./examples/LangChain) directory. +// multi-vector search on a collection with multiple vector fields +const search = await milvusClient.search({ + collection_name: collection_name, + data: [ + { + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + { + data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + anns_field: 'vector1', + }, + ], + limit: 5, + filter: 'id > 100', +}); +``` + +### New Typescript client + +Starting from v2.4.0, we introduced a TypeScript client to provide better support for the [Milvus RESTful API V2](https://milvus.io/api-reference/restful/v2.3.x/About.md), take a look at our [test file](https://github.com/milvus-io/milvus-sdk-node/blob/main/test/http/test.ts). -### next.js example +```javascript +import { HttpClient } from '@zilliz/milvus2-sdk-node'; +const client = new HttpClient(config); +await client.createCollection(params); +await client.describeCollection(params); +await client.listCollections(params); +await client.insert(params); +await client.upsert(params); +await client.query(params); +await client.search(params); +``` -You can find nextjs app example in the [examples/nextjs](./examples/nextjs) directory. +## Code Examples + +This table organizes the examples by technology, providing a brief description and the directory where each example can be found. +| Technology | Example | Directory | +|------------------|--------------------------------------------|-----------------------------------| +| Next.js | Next.js app example | [examples/nextjs](./examples/nextjs) | +| Node.js | Basic Node.js examples for Milvus | [examples/milvus](./examples/milvus) | +| Langchain.js | Basic Langchain.js example | [examples/langchain](./examples/LangChain) | ## Basic usages @@ -54,16 +172,14 @@ This guide will show you how to set up a simple application using Node.js and Mi ### Start a Milvus server ```shell -# Download the milvus standalone yaml file -$ wget https://github.com/milvus-io/milvus/releases/latest/download/milvus-standalone-docker-compose.yml -O docker-compose.yml - -# start the milvus server -sudo docker-compose up -d +# Start Milvus with script +wget https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh +bash standalone_embed.sh start ``` ### Connect to Milvus -Create a new app.js file and add the following code to try out some basic vector operations using the Milvus node.js client. More details on the [API reference](https://milvus.io/api-reference/node/v2.2.x/Client/MilvusClient.md). +Create a new app.js file and add the following code to try out some basic vector operations using the Milvus node.js client. More details on the [API reference](https://milvus.io/api-reference/node/v2.3.x/Client/MilvusClient.md). ```javascript import { MilvusClient, DataType } from '@zilliz/milvus2-sdk-node'; @@ -82,7 +198,7 @@ In Milvus, the concept of the collection is like the table in traditional RDBMS, #### Define schema for collection -A schema defines the fields of a collection, such as the names and data types of the fields that make up the vectors. More details of how to define schema and advanced usage can be found in [API reference](https://milvus.io/api-reference/node/v2.2.x/Collection/createCollection.md). +A schema defines the fields of a collection, such as the names and data types of the fields that make up the vectors. More details of how to define schema and advanced usage can be found in [API reference](https://milvus.io/api-reference/node/v2.3.x/Collection/createCollection.md). ```javascript // define schema @@ -128,31 +244,31 @@ The data format utilized by the Milvus Node SDK comprises an array of objects. I ```javascript const fields_data = [ { + name: 'zlnmh', vector: [ 0.11878310581111173, 0.9694947902934701, 0.16443679307243175, 0.5484226189097237, 0.9839246709011924, 0.5178387104937776, 0.8716926129208069, 0.5616972243831446, ], height: 20405, - name: 'zlnmh', }, { + name: '5lr9y', vector: [ 0.9992090731236536, 0.8248790611809487, 0.8660083940881405, 0.09946359318481224, 0.6790698063908669, 0.5013786801063624, 0.795311915725105, 0.9183033261617566, ], height: 93773, - name: '5lr9y', }, { + name: 'nes0j', vector: [ 0.8761291569818763, 0.07127366044153227, 0.775648976160332, 0.5619757601304878, 0.6076543120476996, 0.8373907516027586, 0.8556140171597648, 0.4043893119391049, ], height: 85122, - name: 'nes0j', }, ]; ``` @@ -164,7 +280,7 @@ Once we have the data, you can insert data into the collection by calling the `i ```javascript await client.insert({ collection_name, - fields_data, + data, }); ``` @@ -175,8 +291,7 @@ By creating an index and loading the collection into memory, you can improve the ```javascript // create index await client.createIndex({ - // required - collection_name, + collection_name, // required field_name: 'vector', // optional if you are using milvus v2.2.9+ index_name: 'myindex', // optional index_type: 'HNSW', // optional if you are using milvus v2.2.9+ @@ -185,7 +300,7 @@ await client.createIndex({ }); ``` -Milvus supports [several different types of indexes](https://milvus.io/docs/index.md), each of which is optimized for different use cases and data distributions. Some of the most commonly used index types in Milvus include IVF_FLAT, IVF_SQ8, IVF_PQ, and HNSW. When creating an index in Milvus, you must choose an appropriate index type based on your specific use case and data distribution. +Milvus supports [several different types of indexes](https://milvus.io/docs/index.md), each of which is optimized for different use cases and data distributions. Some of the most commonly used index types in Milvus include HNSW, IVF_FLAT, IVF_SQ8, IVF_PQ. When creating an index in Milvus, you must choose an appropriate index type based on your specific use case and data distribution. ### Load collection @@ -210,22 +325,26 @@ const searchVector = fields_data[0].vector; const res = await client.search({ // required collection_name, // required, the collection name - vector: searchVector, // required, vector used to compare other vectors in milvus + data: searchVector, // required, vector used to compare other vectors in milvus // optionals - filter: 'height > 0', // optional, filter + filter: 'height > 0', // optional, filter expression params: { nprobe: 64 }, // optional, specify the search parameters limit: 10, // optional, specify the number of nearest neighbors to return - metric_type: 'L2', // optional, metric to calculate similarity of two vectors - output_fields: ['height', 'name'], // optional, specify the fields to return in the search results + output_fields: ['height', 'name'], // optional, specify the fields to return in the search results, }); ``` ## Next Steps +- [Attu, Using GUI to manage Milvus](https://github.com/zilliztech/attu) + ![attu home view +](https://github.com/zilliztech/attu/raw/main/.github/images/screenshot.png) + +## other useful links + - [What is Milvus](https://milvus.io/) -- [Milvus Node SDK API reference](https://milvus.io/api-reference/node/v2.2.x/About.md) -- [Attu, Milvus GUI tool, based on Milvus node.js SDK](https://github.com/zilliztech/attu) -- [Feder, anns index visuliazation tool](https://github.com/zilliztech/feder) +- [Milvus Node SDK API reference](https://milvus.io/api-reference/node/v2.3.x/About.md) +- [Feder, anns index visualization tool](https://github.com/zilliztech/feder) ## How to contribute diff --git a/milvus/HttpClient.ts b/milvus/HttpClient.ts index 68c9e43b..b74d5c03 100644 --- a/milvus/HttpClient.ts +++ b/milvus/HttpClient.ts @@ -1,5 +1,14 @@ import { HttpClientConfig, FetchOptions } from './types'; -import { Collection, Vector } from './http'; +import { + Collection, + Vector, + User, + Role, + Partition, + MilvusIndex, + Alias, + Import, +} from './http'; import { DEFAULT_DB, DEFAULT_HTTP_TIMEOUT, @@ -48,9 +57,7 @@ export class HttpBaseClient { get baseURL() { return ( this.config.baseURL || - `${this.config.endpoint}/${ - this.config.version || DEFAULT_HTTP_ENDPOINT_VERSION - }` + `${this.config.endpoint}/${DEFAULT_HTTP_ENDPOINT_VERSION}` ); } @@ -81,6 +88,10 @@ export class HttpBaseClient { Authorization: this.authorization, Accept: 'application/json', ContentType: 'application/json', + 'Accept-Type-Allow-Int64': + typeof this.config.acceptInt64 !== 'undefined' + ? this.config.acceptInt64.toString() + : 'false', }; } @@ -163,5 +174,20 @@ export class HttpBaseClient { } } -// The HttpClient class extends the functionality of the HttpBaseClient class by mixing in the Collection and Vector APIs. -export class HttpClient extends Collection(Vector(HttpBaseClient)) {} +/** + * The HttpClient class extends the functionality + * of the HttpBaseClient class by mixing in the + * - Collection + * - Vector + * - Alias + * - Partition + * - MilvusIndex + * - Import + * - Role + * - User APIs. + */ +export class HttpClient extends User( + Role( + MilvusIndex(Import(Alias(Partition(Collection(Vector(HttpBaseClient)))))) + ) +) {} diff --git a/milvus/const/defaults.ts b/milvus/const/defaults.ts index e991fe04..e3b1b949 100644 --- a/milvus/const/defaults.ts +++ b/milvus/const/defaults.ts @@ -13,7 +13,7 @@ export const DEFAULT_DB = 'default'; // default database name export const DEFAULT_DYNAMIC_FIELD = '$meta'; // default dynamic field name export const DEFAULT_COUNT_QUERY_STRING = 'count(*)'; // default count query string export const DEFAULT_HTTP_TIMEOUT = 60000; // default http timeout, 60s -export const DEFAULT_HTTP_ENDPOINT_VERSION = 'v1'; // api version, default v1 +export const DEFAULT_HTTP_ENDPOINT_VERSION = 'v2'; // api version, default v1 export const DEFAULT_POOL_MAX = 10; // default max pool client number export const DEFAULT_POOL_MIN = 2; // default min pool client number diff --git a/milvus/const/error.ts b/milvus/const/error.ts index be7da922..e8a95fc7 100644 --- a/milvus/const/error.ts +++ b/milvus/const/error.ts @@ -24,7 +24,7 @@ export const ERROR_REASONS = { COLLECTION_PARTITION_NAME_ARE_REQUIRED: 'The `collection_name` or the `partition_name` property is missing.', INSERT_CHECK_FIELD_DATA_IS_REQUIRED: - 'The type of the `fields_data` should be an array and length > 0.', + 'The type of the `data or field_data` should be an array and length > 0.', INSERT_CHECK_WRONG_FIELD: 'Insert fail: some field does not exist for this collection in line.', INSERT_CHECK_WRONG_DIM: diff --git a/milvus/const/milvus.ts b/milvus/const/milvus.ts index 6fc2b77e..0f39939a 100644 --- a/milvus/const/milvus.ts +++ b/milvus/const/milvus.ts @@ -124,17 +124,22 @@ export enum IndexType { DISKANN = 'DISKANN', AUTOINDEX = 'AUTOINDEX', ANNOY = 'ANNOY', + SPARSE_INVERTED_INDEX = 'SPARSE_INVERTED_INDEX', + SPARSE_WAND = 'SPARSE_WAND', // 2.3 GPU_FLAT = 'GPU_FLAT', GPU_IVF_FLAT = 'GPU_IVF_FLAT', GPU_IVF_PQ = 'GPU_IVF_PQ', GPU_IVF_SQ8 = 'GPU_IVF_SQ8', + GPU_BRUTE_FORCE = 'GPU_BRUTE_FORCE', + GPU_CAGRA = 'GPU_CAGRA', RAFT_IVF_FLAT = 'RAFT_IVF_FLAT', RAFT_IVF_PQ = 'RAFT_IVF_PQ', ScaNN = 'SCANN', // scalar STL_SORT = 'STL_SORT', TRIE = 'Trie', + INVERTED = 'INVERTED', } // MsgType @@ -261,8 +266,19 @@ export enum DataType { BinaryVector = 100, FloatVector = 101, + Float16Vector = 102, + BFloat16Vector = 103, + SparseFloatVector = 104, } +export const VectorDataTypes = [ + DataType.BinaryVector, + DataType.FloatVector, + DataType.Float16Vector, + DataType.BFloat16Vector, + DataType.SparseFloatVector, +]; + // data type map export const DataTypeMap: { [key in keyof typeof DataType]: number } = { None: 0, @@ -279,6 +295,9 @@ export const DataTypeMap: { [key in keyof typeof DataType]: number } = { JSON: 23, BinaryVector: 100, FloatVector: 101, + Float16Vector: 102, + BFloat16Vector: 103, + SparseFloatVector: 104, }; // RBAC: operate user role type @@ -385,3 +404,8 @@ export enum ShowCollectionsType { All, Loaded, } + +export enum RANKER_TYPE { + RRF = 'rrf', + WEIGHTED = 'weighted', +} diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts index 482b8c29..cdaa0e87 100644 --- a/milvus/grpc/Data.ts +++ b/milvus/grpc/Data.ts @@ -1,8 +1,8 @@ import { DataType, + VectorDataTypes, DataTypeMap, ERROR_REASONS, - DslType, DeleteEntitiesReq, FlushReq, GetFlushStateReq, @@ -12,22 +12,20 @@ import { LoadBalanceReq, ImportReq, ListImportTasksReq, - // ListIndexedSegmentReq, - // DescribeSegmentIndexDataReq, ErrorCode, FlushResult, GetFlushStateResponse, GetMetricsResponse, GetQuerySegmentInfoResponse, GePersistentSegmentInfoResponse, + buildSearchRequest, + formatSearchResult, MutationResult, QueryResults, ResStatus, SearchResults, ImportResponse, ListImportTasksResponse, - // ListIndexedSegmentResponse, - // DescribeSegmentIndexDataResponse, GetMetricsRequest, QueryReq, GetReq, @@ -38,35 +36,31 @@ import { SearchReq, SearchRes, SearchSimpleReq, - DEFAULT_TOPK, + HybridSearchReq, promisify, findKeyValue, sleep, - formatNumberPrecision, parseToKeyValue, checkCollectionName, - checkSearchParams, - parseBinaryVectorToBytes, - parseFloatVectorToBytes, DEFAULT_DYNAMIC_FIELD, buildDynamicRow, buildFieldDataMap, getDataKey, Field, buildFieldData, - Vectors, - BinaryVectors, + BinaryVector, RowData, CountReq, CountResult, DEFAULT_COUNT_QUERY_STRING, + SparseFloatVector, + sparseRowsToBytes, + getSparseDim, + f32ArrayToBinaryBytes, } from '../'; import { Collection } from './Collection'; export class Data extends Collection { - // vectorTypes - vectorTypes = [DataType.BinaryVector, DataType.FloatVector]; - /** * Upsert data into Milvus, view _insert for detail */ @@ -87,7 +81,8 @@ export class Data extends Collection { * @param {InsertReq} data - The request parameters. * @param {string} data.collection_name - The name of the collection. * @param {string} [data.partition_name] - The name of the partition (optional). - * @param {{ [x: string]: any }[]} data.fields_data - The data to be inserted. If the field type is binary, the vector data length needs to be dimension / 8. + * @param {{ [x: string]: any }[]} data.data - The data to be inserted. If the field type is binary, the vector data length needs to be dimension / 8. + * @param {InsertTransformers} data.transformers - The transformers for bf16 or f16 data, it accept an f32 array, it should output f16 or bf16 bytes (optional) * @param {number} [data.timeout] - 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 error occurs. Default is undefined. * * @returns {Promise} The result of the operation. @@ -179,7 +174,7 @@ export class Data extends Collection { } if ( DataTypeMap[field.type] === DataType.BinaryVector && - (rowData[name] as Vectors).length !== field.dim! / 8 + (rowData[name] as BinaryVector).length !== field.dim! / 8 ) { throw new Error(ERROR_REASONS.INSERT_CHECK_WRONG_DIM); } @@ -188,10 +183,16 @@ export class Data extends Collection { switch (DataTypeMap[field.type]) { case DataType.BinaryVector: case DataType.FloatVector: - field.data = field.data.concat(buildFieldData(rowData, field)); + field.data = (field.data as number[]).concat( + buildFieldData(rowData, field) as number[] + ); break; default: - field.data[rowIndex] = buildFieldData(rowData, field); + field.data[rowIndex] = buildFieldData( + rowData, + field, + data.transformers + ); break; } }); @@ -205,54 +206,83 @@ export class Data extends Collection { // milvus return string for field type, so we define the DataTypeMap to the value we need. // but if milvus change the string, may cause we cant find value. const type = DataTypeMap[field.type]; - const key = this.vectorTypes.includes(type) ? 'vectors' : 'scalars'; + const key = VectorDataTypes.includes(type) ? 'vectors' : 'scalars'; const dataKey = getDataKey(type); const elementType = DataTypeMap[field.elementType!]; const elementTypeKey = getDataKey(elementType); + // build key value + let keyValue; + switch (type) { + case DataType.FloatVector: + keyValue = { + dim: field.dim, + [dataKey]: { + data: field.data, + }, + }; + break; + case DataType.BFloat16Vector: + case DataType.Float16Vector: + keyValue = { + dim: field.dim, + [dataKey]: Buffer.concat(field.data as Buffer[]), + }; + break; + case DataType.BinaryVector: + keyValue = { + dim: field.dim, + [dataKey]: f32ArrayToBinaryBytes(field.data as BinaryVector), + }; + break; + case DataType.SparseFloatVector: + const dim = getSparseDim(field.data as SparseFloatVector[]); + keyValue = { + dim, + [dataKey]: { + dim, + contents: sparseRowsToBytes(field.data as SparseFloatVector[]), + }, + }; + break; + case DataType.Array: + keyValue = { + [dataKey]: { + data: field.data.map(d => { + return { + [elementTypeKey]: { + type: elementType, + data: d, + }, + }; + }), + element_type: elementType, + }, + }; + break; + default: + keyValue = { + [dataKey]: { + data: field.data, + }, + }; + break; + } + return { type, field_name: field.name, is_dynamic: field.name === DEFAULT_DYNAMIC_FIELD, - [key]: - type === DataType.FloatVector - ? { - dim: field.dim, - [dataKey]: { - data: field.data, - }, - } - : type === DataType.BinaryVector - ? { - dim: field.dim, - [dataKey]: parseBinaryVectorToBytes( - field.data as BinaryVectors - ), - } - : type === DataType.Array - ? { - [dataKey]: { - data: field.data.map(d => { - return { - [elementTypeKey]: { - type: elementType, - data: d, - }, - }; - }), - element_type: elementType, - }, - } - : { - [dataKey]: { - data: field.data, - }, - }, + [key]: keyValue, }; }); // if timeout is not defined, set timeout to 0 const timeout = typeof data.timeout === 'undefined' ? 0 : data.timeout; + // delete data + try { + delete params.data; + } catch (e) {} // execute Insert const promise = await promisify( this.channelPool, @@ -381,6 +411,9 @@ export class Data extends Collection { * @param {string} [data.filter] - Scalar field filter expression (optional). * @param {string[]} [data.output_fields] - Support scalar field (optional). * @param {object} [data.params] - Search params (optional). + * @param {OutputTransformers} data.transformers - The transformers for bf16 or f16 data, it accept bytes or sparse dic vector, it can ouput f32 array or other format(optional) + * @param {number} [data.timeout] - 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 error occurs. Default is undefined. + * * @returns {Promise} The result of the operation. * @returns {string} status.error_code - The error code of the operation. * @returns {string} status.reason - The reason for the error, if any. @@ -395,202 +428,59 @@ export class Data extends Collection { * }); * ``` */ - async search(data: SearchReq | SearchSimpleReq): Promise { - // params check - checkSearchParams(data); - - try { - // get collection info - const collectionInfo = await this.describeCollection({ - collection_name: data.collection_name, - cache: true, - }); - - // get information from collection info - let vectorType: DataType; - let defaultOutputFields = []; - let anns_field: string; - for (let i = 0; i < collectionInfo.schema.fields.length; i++) { - const f = collectionInfo.schema.fields[i]; - const type = DataTypeMap[f.data_type]; - - // filter vector field - if (type === DataType.FloatVector || type === DataType.BinaryVector) { - // anns field - anns_field = f.name; - // vector type - vectorType = type; - } else { - // save field name - defaultOutputFields.push(f.name); - } - } - - // create search params - const search_params = (data as SearchReq).search_params || { - anns_field: anns_field!, - topk: - (data as SearchSimpleReq).limit || - (data as SearchSimpleReq).topk || - DEFAULT_TOPK, - offset: (data as SearchSimpleReq).offset || 0, - metric_type: (data as SearchSimpleReq).metric_type || '', // leave it empty - params: JSON.stringify((data as SearchSimpleReq).params || {}), - ignore_growing: (data as SearchSimpleReq).ignore_growing || false, - }; - - // create search vectors - let searchVectors: number[] | number[][] = - (data as SearchReq).vectors || - (data as SearchSimpleReq).data || - (data as SearchSimpleReq).vector; - - // make sure the searchVectors format is correct - if (!Array.isArray(searchVectors[0])) { - searchVectors = [searchVectors as unknown] as number[][]; - } - - /** - * It will decide the score precision. - * If round_decimal is 3, need return like 3.142 - * And if Milvus return like 3.142, Node will add more number after this like 3.142000047683716. - * So the score need to slice by round_decimal - */ - const round_decimal = - (data as SearchReq).search_params?.round_decimal ?? - ((data as SearchSimpleReq).params?.round_decimal as number); - - // create placeholder_group - const PlaceholderGroup = this.milvusProto.lookupType( - 'milvus.proto.common.PlaceholderGroup' - ); - // tag $0 is hard code in milvus, when dsltype is expr - const placeholderGroupBytes = PlaceholderGroup.encode( - PlaceholderGroup.create({ - placeholders: [ - { - tag: '$0', - type: vectorType!, - values: searchVectors.map(v => - vectorType === DataType.BinaryVector - ? parseBinaryVectorToBytes(v) - : parseFloatVectorToBytes(v) - ), - }, - ], - }) - ).finish(); - - // get collection's consistency level - const collection_consistency_level = collectionInfo.consistency_level; - - const promise: SearchRes = await promisify( - this.channelPool, - 'Search', - { - collection_name: data.collection_name, - partition_names: data.partition_names, - output_fields: data.output_fields || defaultOutputFields, - nq: (data as SearchReq).nq || searchVectors.length, - dsl: - (data as SearchReq).expr || (data as SearchSimpleReq).filter || '', - dsl_type: DslType.BoolExprV1, - placeholder_group: placeholderGroupBytes, - search_params: parseToKeyValue(search_params), - consistency_level: - data.consistency_level || collection_consistency_level, - }, - data.timeout || this.timeout - ); + async search( + data: SearchReq | SearchSimpleReq | HybridSearchReq + ): Promise { + // get collection info + const collectionInfo = await this.describeCollection({ + collection_name: data.collection_name, + cache: true, + }); - // if search failed - // if nothing returned - // return empty with status - if ( - promise.status.error_code !== ErrorCode.SUCCESS || - promise.results.scores.length === 0 - ) { - return { - status: promise.status, - results: [], - }; - } + // build search params + const { request, nq, round_decimal, isHybridSearch } = buildSearchRequest( + data, + collectionInfo, + this.milvusProto + ); - // build final results array - const results: any[] = []; - const { topks, scores, fields_data, ids } = promise.results; - // build fields data map - const fieldsDataMap = buildFieldDataMap(fields_data); - // build output name array - const output_fields = [ - 'id', - ...(!!promise.results.output_fields?.length - ? promise.results.output_fields - : fields_data.map(f => f.field_name)), - ]; - - // vector id support int / str id. - const idData = ids ? ids[ids.id_field]!.data : {}; - // add id column - fieldsDataMap.set('id', idData as RowData[]); - // fieldsDataMap.set('score', scores); TODO: fieldDataMap to support formatter - - /** - * This code block formats the search results returned by Milvus into row data for easier use. - * Milvus supports multiple queries to search and returns all columns data, so we need to splice the data for each search result using the `topk` variable. - * The `topk` variable is the key we use to splice data for every search result. - * The `scores` array is spliced using the `topk` value, and the resulting scores are formatted to the specified precision using the `formatNumberPrecision` function. The resulting row data is then pushed to the `results` array. - */ - topks.forEach((v, index) => { - const topk = Number(v); - - scores.splice(0, topk).forEach((score, scoreIndex) => { - // get correct index - const i = index === 0 ? scoreIndex : scoreIndex + topk * index; - - // fix round_decimal - const fixedScore = - typeof round_decimal === 'undefined' || round_decimal === -1 - ? score - : formatNumberPrecision(score, round_decimal); - - // init result object - const result: any = { score: fixedScore }; - - // build result, - output_fields.forEach(field_name => { - // Check if the field_name exists in the fieldsDataMap - const isFixedSchema = fieldsDataMap.has(field_name); - - // Get the data for the field_name from the fieldsDataMap - // If the field_name is not in the fieldsDataMap, use the DEFAULT_DYNAMIC_FIELD - const data = fieldsDataMap.get( - isFixedSchema ? field_name : DEFAULT_DYNAMIC_FIELD - )!; - // make dynamic data[i] safe - data[i] = isFixedSchema ? data[i] : data[i] || {}; - // extract dynamic info from dynamic field if necessary - result[field_name] = isFixedSchema ? data[i] : data[i][field_name]; - }); - - // init result slot - results[index] = results[index] || []; - // push result data - results[index].push(result); - }); - }); + // execute search + const originSearchResult: SearchRes = await promisify( + this.channelPool, + isHybridSearch ? 'HybridSearch' : 'Search', + request, + data.timeout || this.timeout + ); + // if search failed + // if nothing returned + // return empty with status + if ( + originSearchResult.status.error_code !== ErrorCode.SUCCESS || + originSearchResult.results.scores.length === 0 + ) { return { - status: promise.status, - // if only searching 1 vector, return the first object of results array - results: searchVectors.length === 1 ? results[0] || [] : results, + status: originSearchResult.status, + results: [], }; - } catch (err) { - /* istanbul ignore next */ - throw new Error(err); } + + // build final results array + const results = formatSearchResult(originSearchResult, { + round_decimal, + transformers: data.transformers, + }); + + return { + status: originSearchResult.status, + // nq === 1, return the first object of results array + results: nq === 1 ? results[0] || [] : results, + }; } + // alias + hybridSearch = this.search; + /** * Flushes the newly inserted vectors that are temporarily buffered in the cache to the object storage. * This is an asynchronous function and may take some time to execute deponds on your data size. @@ -690,7 +580,8 @@ export class Data extends Collection { * @param {string[]} [data.partitions_names] - Array of partition names (optional). * @param {string[]} data.output_fields - Vector or scalar field to be returned. * @param {number} [data.timeout] - 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. - * @param {{key: value}[]} [data.params] - An optional key pair json array. + * @param {{key: value}[]} [data.params] - An optional key pair json array of search parameters. + * @param {OutputTransformers} data.transformers - The transformers for bf16 or f16 data, it accept bytes or sparse dic vector, it can ouput f32 array or other format(optional) * * @returns {Promise} The result of the operation. * @returns {string} status.error_code - The error code of the operation. @@ -753,7 +644,10 @@ export class Data extends Collection { // always get output_fields from fields_data const output_fields = promise.fields_data.map(f => f.field_name); - const fieldsDataMap = buildFieldDataMap(promise.fields_data); + const fieldsDataMap = buildFieldDataMap( + promise.fields_data, + data.transformers + ); // For each output field, check if it has a fixed schema or not const fieldDataContainer = output_fields.map(field_name => { diff --git a/milvus/grpc/MilvusIndex.ts b/milvus/grpc/MilvusIndex.ts index 4b53367e..29c55a23 100644 --- a/milvus/grpc/MilvusIndex.ts +++ b/milvus/grpc/MilvusIndex.ts @@ -7,6 +7,7 @@ import { CreateIndexRequest, GetIndexBuildProgressReq, GetIndexStateReq, + AlterIndexReq, ResStatus, DescribeIndexResponse, GetIndexStateResponse, @@ -311,4 +312,42 @@ export class Index extends Data { ); return promise; } + + /** + * Alters an existing index. + * + * @param {Object} data - The data for altering the index. + * @param {string} data.collection_name - The name of the collection. + * @param {Object} data.params - The new parameters for the index. For example, `{ nlist: number }`. + * @param {number} [data.timeout] - 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 {Promise} - A promise that resolves to a response status object. + * @returns {number} return.error_code - The error code number. + * @returns {string} return.reason - The cause of the error. + * + * @example + * ``` + * const milvusClient = new MilvusClient(MILUVS_ADDRESS); + * const alterIndexReq = { + * collection_name: 'my_collection', + * params: { nlist: 20 }, + * }; + * const res = await milvusClient.alterIndex(alterIndexReq); + * console.log(res); + * ``` + */ + async alterIndex(data: AlterIndexReq): Promise { + checkCollectionName(data); + const promise = await promisify( + this.channelPool, + 'AlterIndex', + { + collection_name: data.collection_name, + index_name: data.index_name, + extra_params: parseToKeyValue(data.params), + }, + data.timeout || this.timeout + ); + return promise; + } } diff --git a/milvus/http/Alias.ts b/milvus/http/Alias.ts new file mode 100644 index 00000000..e9196459 --- /dev/null +++ b/milvus/http/Alias.ts @@ -0,0 +1,56 @@ +import { HttpBaseClient } from '../HttpClient'; +import { + Constructor, + FetchOptions, + HttpAliasBaseReq, + HttpBaseResponse, + HttpAliasCreateReq, + HttpAliasAlterReq, + HttpAliasDescribeReq, + HttpAliasDropReq, + HttpAliasDescribeResponse, +} from '../types'; + +/** + * + * @param {Constructor} Base - The base class to be extended. + * @returns {class} - The extended class with additional methods for collection management. + * + * @method listAliases - Lists all aliases in a collection. + * @method createAlias - Creates a new alias in a collection. + * @method describeAlias - Describes an alias. + * @method dropAlias - Deletes an alias. + * @method alterAlias - Modifies an alias to another collection. + */ +export function Alias>(Base: T) { + return class extends Base { + get aliasPrefix() { + return '/vectordb/aliases'; + } + + async listAliases(params: HttpAliasBaseReq, options?: FetchOptions) { + const url = `${this.aliasPrefix}/list`; + return await this.POST>(url, params, options); + } + + async createAlias(params: HttpAliasCreateReq, options?: FetchOptions) { + const url = `${this.aliasPrefix}/create`; + return await this.POST(url, params, options); + } + + async describeAlias(params: HttpAliasDescribeReq, options?: FetchOptions) { + const url = `${this.aliasPrefix}/describe`; + return await this.POST(url, params, options); + } + + async dropAlias(params: HttpAliasDropReq, options?: FetchOptions) { + const url = `${this.aliasPrefix}/drop`; + return await this.POST(url, params, options); + } + + async alterAlias(params: HttpAliasAlterReq, options?: FetchOptions) { + const url = `${this.aliasPrefix}/alter`; + return await this.POST(url, params, options); + } + }; +} diff --git a/milvus/http/Collection.ts b/milvus/http/Collection.ts index 6e702a0d..624edf5b 100644 --- a/milvus/http/Collection.ts +++ b/milvus/http/Collection.ts @@ -8,11 +8,17 @@ import { HttpBaseResponse, HttpBaseReq, FetchOptions, + HttpCollectionRenameReq, + HttpCollectionHasResponse, + HttpCollectionStatisticsResponse, + HttpCollectionLoadStateReq, + HttpCollectionLoadStateResponse, } from '../types'; import { DEFAULT_PRIMARY_KEY_FIELD, DEFAULT_METRIC_TYPE, DEFAULT_VECTOR_FIELD, + DEFAULT_DB, } from '../const'; /** @@ -26,20 +32,31 @@ import { * @method describeCollection - Retrieves the description of a specific collection. * @method dropCollection - Deletes a specific collection from Milvus. * @method listCollections - Lists all collections in the Milvus cluster. + * @method hasCollection - Checks if a collection exists in the Milvus cluster. + * @method renameCollection - Renames a collection in the Milvus cluster. + * @method getCollectionStatistics - Retrieves statistics about a collection. + * @method loadCollection - Loads a collection into memory. + * @method releaseCollection - Releases a collection from memory. + * @method getCollectionLoadState - Retrieves the load state of a collection. */ export function Collection>(Base: T) { return class extends Base { + get collectionPrefix() { + return '/vectordb/collections'; + } + // POST create collection async createCollection( data: HttpCollectionCreateReq, options?: FetchOptions ): Promise { - const url = `/vector/collections/create`; + const url = `${this.collectionPrefix}/create`; // if some keys not provided, using default value data.metricType = data.metricType || DEFAULT_METRIC_TYPE; - data.primaryField = data.primaryField || DEFAULT_PRIMARY_KEY_FIELD; - data.vectorField = data.vectorField || DEFAULT_VECTOR_FIELD; + data.primaryFieldName = + data.primaryFieldName || DEFAULT_PRIMARY_KEY_FIELD; + data.vectorFieldName = data.vectorFieldName || DEFAULT_VECTOR_FIELD; return await this.POST(url, data, options); } @@ -49,8 +66,8 @@ export function Collection>(Base: T) { params: HttpBaseReq, options?: FetchOptions ): Promise { - const url = `/vector/collections/describe`; - return await this.GET( + const url = `${this.collectionPrefix}/describe`; + return await this.POST( url, params, options @@ -62,19 +79,63 @@ export function Collection>(Base: T) { data: HttpBaseReq, options?: FetchOptions ): Promise { - const url = `/vector/collections/drop`; + const url = `${this.collectionPrefix}/drop`; return await this.POST(url, data, options); } // GET list collections async listCollections( - params: HttpCollectionListReq = {}, + params: HttpCollectionListReq = { dbName: DEFAULT_DB }, options?: FetchOptions ): Promise { - const url = `/vector/collections`; + const url = `${this.collectionPrefix}/list`; + + return await this.POST(url, params, options); + } + + async hasCollection(params: Required, options?: FetchOptions) { + const url = `${this.collectionPrefix}/has`; + return await this.POST(url, params, options); + } + + async renameCollection( + params: HttpCollectionRenameReq, + options?: FetchOptions + ) { + const url = `${this.collectionPrefix}/rename`; + return await this.POST(url, params, options); + } + + async getCollectionStatistics(params: HttpBaseReq, options?: FetchOptions) { + const url = `${this.collectionPrefix}/get_stats`; + return await this.POST( + url, + params, + options + ); + } + + async loadCollection(params: HttpBaseReq, options?: FetchOptions) { + const url = `${this.collectionPrefix}/load`; + return await this.POST(url, params, options); + } + + async releaseCollection(params: HttpBaseReq, options?: FetchOptions) { + const url = `${this.collectionPrefix}/release`; + return await this.POST(url, params, options); + } - return await this.GET(url, params, options); + async getCollectionLoadState( + params: HttpCollectionLoadStateReq, + options?: FetchOptions + ) { + const url = `${this.collectionPrefix}/get_load_state`; + return await this.POST( + url, + params, + options + ); } }; } diff --git a/milvus/http/Import.ts b/milvus/http/Import.ts new file mode 100644 index 00000000..65a0f9cd --- /dev/null +++ b/milvus/http/Import.ts @@ -0,0 +1,48 @@ +import { HttpBaseClient } from '../HttpClient'; +import { + Constructor, + FetchOptions, + HttpBaseReq, + HttpImportListResponse, + HttpImportCreateReq, + HttpImportCreateResponse, + HttpImportProgressReq, +} from '../types'; + +/** + * + * @param {Constructor} Base - The base class to be extended. + * @returns {class} - The extended class with additional methods for collection management. + * + * @method listImportJobs - Lists all import jobs. + * @method createImportJobs - Creates new import jobs. + * @method getImportJobProgress - Retrieves the progress of an import job. + */ +export function Import>(Base: T) { + return class extends Base { + get importPrefix() { + return '/vectordb/jobs/import'; + } + + async listImportJobs(params: HttpBaseReq, options?: FetchOptions) { + const url = `${this.importPrefix}/list`; + return await this.POST(url, params, options); + } + + async createImportJobs( + params: HttpImportCreateReq, + options?: FetchOptions + ) { + const url = `${this.importPrefix}/create`; + return await this.POST(url, params, options); + } + + async getImportJobProgress( + params: HttpImportProgressReq, + options?: FetchOptions + ) { + const url = `${this.importPrefix}/get_progress`; + return await this.POST(url, params, options); + } + }; +} diff --git a/milvus/http/MilvusIndex.ts b/milvus/http/MilvusIndex.ts new file mode 100644 index 00000000..fc44b812 --- /dev/null +++ b/milvus/http/MilvusIndex.ts @@ -0,0 +1,48 @@ +import { HttpBaseClient } from '../HttpClient'; +import { + Constructor, + FetchOptions, + HttpBaseReq, + HttpBaseResponse, + HttpIndexCreateReq, + HttpIndexBaseReq, + HttpIndexDescribeResponse, +} from '../types'; + +/** + * + * @param {Constructor} Base - The base class to be extended. + * @returns {class} - The extended class with additional methods for collection management. + * + *@method createIndex - Creates an index. + *@method dropIndex - Deletes an index. + *@method describeIndex - Describes an index. + *@method listIndexes - Lists all indexes. + */ +export function MilvusIndex>(Base: T) { + return class extends Base { + get indexPrefix() { + return '/vectordb/indexes'; + } + + async createIndex(params: HttpIndexCreateReq, options?: FetchOptions) { + const url = `${this.indexPrefix}/create`; + return this.POST(url, params, options); + } + + async dropIndex(params: HttpIndexBaseReq, options?: FetchOptions) { + const url = `${this.indexPrefix}/drop`; + return this.POST(url, params, options); + } + + async describeIndex(params: HttpIndexBaseReq, options?: FetchOptions) { + const url = `${this.indexPrefix}/describe`; + return this.POST(url, params, options); + } + + async listIndexes(params: HttpBaseReq, options?: FetchOptions) { + const url = `${this.indexPrefix}/list`; + return this.POST>(url, params, options); + } + }; +} diff --git a/milvus/http/Partition.ts b/milvus/http/Partition.ts new file mode 100644 index 00000000..573c8cc9 --- /dev/null +++ b/milvus/http/Partition.ts @@ -0,0 +1,80 @@ +import { HttpBaseClient } from '../HttpClient'; +import { + Constructor, + FetchOptions, + HttpBaseReq, + HttpBaseResponse, + HttpPartitionBaseReq, + HttpPartitionListReq, + HttpPartitionHasResponse, + HttpPartitionStatisticsResponse, +} from '../types'; + +/** + * + * @param {Constructor} Base - The base class to be extended. + * @returns {class} - The extended class with additional methods for collection management. + * + * @method listPartitions - Lists all partitions in a collection. + * @method createPartition - Creates a new partition in a collection. + * @method dropPartition - Deletes a partition from a collection. + * @method loadPartitions - Loads partitions into memory. + * @method releasePartitions - Releases partitions from memory. + * @method hasPartition - Checks if a partition exists in a collection. + * @method getPartitionStatistics - Retrieves statistics about a partition. + */ +export function Partition>(Base: T) { + return class extends Base { + get partitionPrefix() { + return '/vectordb/partitions'; + } + + async listPartitions(params: HttpBaseReq, options?: FetchOptions) { + const url = `${this.partitionPrefix}/list`; + return await this.POST>(url, params, options); + } + + async createPartition( + params: HttpPartitionBaseReq, + options?: FetchOptions + ) { + const url = `${this.partitionPrefix}/create`; + return await this.POST(url, params, options); + } + + async dropPartition(params: HttpPartitionBaseReq, options?: FetchOptions) { + const url = `${this.partitionPrefix}/drop`; + return await this.POST(url, params, options); + } + + async loadPartitions(params: HttpPartitionListReq, options?: FetchOptions) { + const url = `${this.partitionPrefix}/load`; + return await this.POST(url, params, options); + } + + async releasePartitions( + params: HttpPartitionListReq, + options?: FetchOptions + ) { + const url = `${this.partitionPrefix}/release`; + return await this.POST(url, params, options); + } + + async hasPartition(params: HttpPartitionBaseReq, options?: FetchOptions) { + const url = `${this.partitionPrefix}/has`; + return await this.POST(url, params, options); + } + + async getPartitionStatistics( + params: HttpPartitionBaseReq, + options?: FetchOptions + ) { + const url = `${this.partitionPrefix}/get_stats`; + return await this.POST( + url, + params, + options + ); + } + }; +} diff --git a/milvus/http/Role.ts b/milvus/http/Role.ts new file mode 100644 index 00000000..77184c22 --- /dev/null +++ b/milvus/http/Role.ts @@ -0,0 +1,65 @@ +import { HttpBaseClient } from '../HttpClient'; +import { + Constructor, + FetchOptions, + HttpRolePrivilegeReq, + HttpRoleDescribeResponse, + HttpBaseResponse, + HttpRoleBaseReq, +} from '../types'; + +/** + * + * @param {Constructor} Base - The base class to be extended. + * @returns {class} - The extended class with additional methods for collection management. + * + * @method listRoles - Lists all roles in the system. + * @method describeRole - Describes a role. + * @method createRole - Creates a new role. + * @method dropRole - Deletes a role. + * @method grantPrivilegeToRole - Grants a privilege to a role. + * @method revokePrivilegeFromRole - Revokes a privilege from a role. + */ +export function Role>(Base: T) { + return class extends Base { + get rolePrefix() { + return '/vectordb/roles'; + } + + async listRoles(options?: FetchOptions) { + const url = `${this.rolePrefix}/list`; + return await this.POST>(url, {}, options); + } + + async describeRole(params: HttpRoleBaseReq, options?: FetchOptions) { + const url = `${this.rolePrefix}/describe`; + return await this.POST(url, params, options); + } + + async createRole(params: HttpRoleBaseReq, options?: FetchOptions) { + const url = `${this.rolePrefix}/create`; + return await this.POST(url, params, options); + } + + async dropRole(params: HttpRoleBaseReq, options?: FetchOptions) { + const url = `${this.rolePrefix}/drop`; + return await this.POST(url, params, options); + } + + async grantPrivilegeToRole( + params: HttpRolePrivilegeReq, + options?: FetchOptions + ) { + const url = `${this.rolePrefix}/grant_privilege`; + return await this.POST(url, params, options); + } + + async revokePrivilegeFromRole( + params: HttpRolePrivilegeReq, + options?: FetchOptions + ) { + const url = `${this.rolePrefix}/revoke_privilege`; + return await this.POST(url, params, options); + } + }; +} diff --git a/milvus/http/User.ts b/milvus/http/User.ts new file mode 100644 index 00000000..1bdad802 --- /dev/null +++ b/milvus/http/User.ts @@ -0,0 +1,69 @@ +import { HttpBaseClient } from '../HttpClient'; +import { + Constructor, + FetchOptions, + HttpUserBaseReq, + HttpUserCreateReq, + HttpUserRoleReq, + HttpUserUpdatePasswordReq, + HttpBaseResponse, +} from '../types'; + +/** + * + * @param {Constructor} Base - The base class to be extended. + * @returns {class} - The extended class with additional methods for collection management. + * + * @method createUser - Creates a new user in Milvus. + * @method updateUserPassword - Updates the password of a user. + * @method dropUser - Deletes a user from Milvus. + * @method describeUser - Retrieves the description of a specific user. + * @method listUsers - Lists all users in the Milvus cluster. + * @method grantRole - Grants a role to a user. + * @method revokeRole - Revokes a role from a user. + */ +export function User>(Base: T) { + return class extends Base { + get userPrefix() { + return '/vectordb/users'; + } + + async createUser(params: HttpUserCreateReq, options?: FetchOptions) { + const url = `${this.userPrefix}/create`; + return this.POST(url, params, options); + } + + async updateUserPassword( + params: HttpUserUpdatePasswordReq, + options?: FetchOptions + ) { + const url = `${this.userPrefix}/update_password`; + return this.POST(url, params, options); + } + + async dropUser(param: HttpUserBaseReq, options?: FetchOptions) { + const url = `${this.userPrefix}/drop`; + return this.POST(url, param, options); + } + + async describeUser(param: HttpUserBaseReq, options?: FetchOptions) { + const url = `${this.userPrefix}/describe`; + return this.POST>(url, param, options); + } + + async listUsers(options?: FetchOptions) { + const url = `${this.userPrefix}/list`; + return this.POST>(url, {}, options); + } + + async grantRoleToUser(params: HttpUserRoleReq, options?: FetchOptions) { + const url = `${this.userPrefix}/grant_role`; + return this.POST(url, params, options); + } + + async revokeRoleFromUser(params: HttpUserRoleReq, options?: FetchOptions) { + const url = `${this.userPrefix}/revoke_role`; + return this.POST(url, params, options); + } + }; +} diff --git a/milvus/http/Vector.ts b/milvus/http/Vector.ts index cd916424..1de986f7 100644 --- a/milvus/http/Vector.ts +++ b/milvus/http/Vector.ts @@ -11,6 +11,7 @@ import { HttpVectorSearchResponse, HttpBaseResponse, FetchOptions, + HttpVectorUpsertResponse, } from '../types'; /** @@ -29,12 +30,16 @@ import { */ export function Vector>(Base: T) { return class extends Base { + get vectorPrefix() { + return '/vectordb/entities'; + } + // GET get data async get( params: HttpVectorGetReq, options?: FetchOptions ): Promise { - const url = `/vector/get`; + const url = `${this.vectorPrefix}/get`; return await this.POST(url, params, options); } @@ -43,7 +48,7 @@ export function Vector>(Base: T) { data: HttpVectorInsertReq, options?: FetchOptions ): Promise { - const url = `/vector/insert`; + const url = `${this.vectorPrefix}/insert`; return await this.POST(url, data, options); } @@ -51,9 +56,9 @@ export function Vector>(Base: T) { async upsert( data: HttpVectorInsertReq, options?: FetchOptions - ): Promise { - const url = `/vector/upsert`; - return await this.POST(url, data, options); + ): Promise { + const url = `${this.vectorPrefix}/upsert`; + return await this.POST(url, data, options); } // POST query data @@ -61,7 +66,7 @@ export function Vector>(Base: T) { data: HttpVectorQueryReq, options?: FetchOptions ): Promise { - const url = `/vector/query`; + const url = `${this.vectorPrefix}/query`; return await this.POST(url, data, options); } @@ -70,7 +75,7 @@ export function Vector>(Base: T) { data: HttpVectorSearchReq, options?: FetchOptions ): Promise { - const url = `/vector/search`; + const url = `${this.vectorPrefix}/search`; return await this.POST(url, data, options); } @@ -79,7 +84,7 @@ export function Vector>(Base: T) { data: HttpVectorDeleteReq, options?: FetchOptions ): Promise { - const url = `/vector/delete`; + const url = `${this.vectorPrefix}/delete`; return await this.POST(url, data, options); } }; diff --git a/milvus/http/index.ts b/milvus/http/index.ts index 9297875e..c7ab6dcf 100644 --- a/milvus/http/index.ts +++ b/milvus/http/index.ts @@ -1,2 +1,8 @@ export * from './Collection'; export * from './Vector'; +export * from './User'; +export * from './Role'; +export * from './Partition'; +export * from './Alias'; +export * from './MilvusIndex'; +export * from './Import' diff --git a/milvus/types/Collection.ts b/milvus/types/Collection.ts index e0d92b4f..a59d173b 100644 --- a/milvus/types/Collection.ts +++ b/milvus/types/Collection.ts @@ -28,7 +28,7 @@ export interface FieldSchema { state: string; element_type?: keyof typeof DataType; default_value?: number | string; - dataType?: DataType; + dataType: DataType; is_partition_key?: boolean; is_dynamic?: boolean; is_clustering_key?: boolean; @@ -236,7 +236,7 @@ export interface GetLoadStateResponse extends resStatusResponse { } export interface AlterCollectionReq extends collectionNameReq { - properties: Record; + properties: Record; } export interface DescribeAliasResponse extends resStatusResponse { diff --git a/milvus/types/Common.ts b/milvus/types/Common.ts index dfb45670..2f994939 100644 --- a/milvus/types/Common.ts +++ b/milvus/types/Common.ts @@ -58,9 +58,7 @@ export interface TimeStampArray { created_utc_timestamps: string[]; } -export interface keyValueObj { - [key: string]: string | number; -} +export type keyValueObj = Record; export interface collectionNameReq extends GrpcTimeOut { collection_name: string; // required, collection name diff --git a/milvus/types/Data.ts b/milvus/types/Data.ts index 56e7278a..a59c5c0e 100644 --- a/milvus/types/Data.ts +++ b/milvus/types/Data.ts @@ -10,12 +10,35 @@ import { ConsistencyLevelEnum, collectionNameReq, resStatusResponse, + RANKER_TYPE, } from '../'; -// all types supported by milvus -export type FloatVectors = number[]; -export type BinaryVectors = number[]; -export type Vectors = FloatVectors | BinaryVectors; +// all value types supported by milvus +export type FloatVector = number[]; +export type Float16Vector = number[] | Uint8Array; +export type BFloat16Vector = number[] | Uint8Array; +export type BinaryVector = number[]; +export type SparseVectorArray = (number | undefined)[]; +export type SparseVectorDic = { [key: string]: number }; +export type SparseVectorCSR = { + indices: number[]; + values: number[]; +}; +export type SparseVectorCOO = { index: number; value: number }[]; + +export type SparseFloatVector = + | SparseVectorArray + | SparseVectorDic + | SparseVectorCSR + | SparseVectorCOO; + +// export type SparseFloatVector = { [key: string]: number }; +export type VectorTypes = + | FloatVector + | Float16Vector + | BinaryVector + | BFloat16Vector + | SparseFloatVector; export type Bool = boolean; export type Int8 = number; export type Int16 = number; @@ -48,9 +71,7 @@ export type FieldData = | VarChar | JSON | Array - | Vectors - | FloatVectors - | BinaryVectors; + | VectorTypes; // Represents a row of data in Milvus. export interface RowData { @@ -66,34 +87,43 @@ export interface Field { } export interface FlushReq extends GrpcTimeOut { - collection_names: string[]; + collection_names: string[]; // collection names } export interface CountReq extends collectionNameReq { - expr?: string; + expr?: string; // filter expression } +// because in javascript, there is no float16 and bfloat16 type +// we need to provide custom data transformer for these types +// milvus only accept bytes(buffer) for these types +export type InsertTransformers = { + [DataType.BFloat16Vector]?: (bf16: BFloat16Vector) => Buffer; + [DataType.Float16Vector]?: (f16: Float16Vector) => Buffer; +}; + export interface InsertReq extends collectionNameReq { - partition_name?: string; - fields_data?: RowData[]; - data?: RowData[]; + partition_name?: string; // partition name + data?: RowData[]; // data to insert + fields_data?: RowData[]; // alias for data hash_keys?: Number[]; // user can generate hash value depend on primarykey value + transformers?: InsertTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors } export interface DeleteEntitiesReq extends collectionNameReq { - expr?: string; - filter?: string; - partition_name?: string; + filter?: string; // filter expression + expr?: string; // alias for filter + partition_name?: string; // partition name } export interface DeleteByIdsReq extends collectionNameReq { - ids: string[] | number[]; - partition_name?: string; + ids: string[] | number[]; // primary key values + partition_name?: string; // partition name } export interface DeleteByFilterReq extends collectionNameReq { - filter: string; - partition_name?: string; + filter: string; // filter expression + partition_name?: string; // partition name } export type DeleteReq = DeleteByIdsReq | DeleteByFilterReq; @@ -105,24 +135,21 @@ export interface CalcDistanceReq extends GrpcTimeOut { } export interface GetFlushStateReq extends GrpcTimeOut { - segmentIDs: number[]; + segmentIDs: number[]; // segment id array } export interface LoadBalanceReq extends GrpcTimeOut { - // The source query node id to balance. - src_nodeID: number; - // The destination query node ids to balance. - dst_nodeIDs?: number[]; - // Sealed segment ids to balance. - sealed_segmentIDs?: number[]; + src_nodeID: number; // The source query node id to balance. + dst_nodeIDs?: number[]; // The destination query node ids to balance. + sealed_segmentIDs?: number[]; // Sealed segment ids to balance. } export interface GetQuerySegmentInfoReq extends GrpcTimeOut { - collectionName: string; + collectionName: string; // its collectioName, this is not colleciton_name :< } export interface GePersistentSegmentInfoReq extends GrpcTimeOut { - collectionName: string; + collectionName: string; // its collectioName, this is not colleciton_name:< } export interface ImportReq extends collectionNameReq { @@ -233,43 +260,81 @@ export interface GetMetricsRequest extends GrpcTimeOut { export interface SearchParam { anns_field: string; // your vector field name - topk: string; - metric_type: string; - params: string; - round_decimal?: number; - ignore_growing?: boolean; + topk: string | number; // how many results you want + metric_type: string; // distance metric type + params: string; // extra search parameters + offset?: number; // skip how many results + round_decimal?: number; // round decimal + ignore_growing?: boolean; // ignore growing + group_by_field?: string; // group by field } +// old search api parameter type +export interface SearchReq extends collectionNameReq { + anns_field?: string; // your vector field name + partition_names?: string[]; // partition names + expr?: string; // filter expression + search_params: SearchParam; // search parameters + vectors: VectorTypes[]; // vectors to search + output_fields?: string[]; // fields to return + travel_timestamp?: string; // time travel + vector_type: DataType.BinaryVector | DataType.FloatVector; // vector field type + nq?: number; // number of query vectors + consistency_level?: ConsistencyLevelEnum; // consistency level + transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors +} + +// simplified search api parameter type export interface SearchSimpleReq extends collectionNameReq { - vector?: number[]; - vectors?: number[][]; - data?: number[][] | number[]; + partition_names?: string[]; // partition names + anns_field?: string; // your vector field name + data?: VectorTypes[] | VectorTypes; // vector to search + vector?: VectorTypes; // alias for data + vectors?: VectorTypes[]; // alias for data output_fields?: string[]; - limit?: number; - topk?: number; // alias - offset?: number; - filter?: string; - expr?: string; // alias - partition_names?: string[]; - params?: keyValueObj; - metric_type?: string; - consistency_level?: ConsistencyLevelEnum; - ignore_growing?: boolean; -} + limit?: number; // how many results you want + topk?: number; // limit alias + offset?: number; // skip how many results + filter?: string; // filter expression + expr?: string; // alias for filter + params?: keyValueObj; // extra search parameters + metric_type?: string; // distance metric type + consistency_level?: ConsistencyLevelEnum; // consistency level + ignore_growing?: boolean; // ignore growing + group_by_field?: string; // group by field + round_decimal?: number; // round decimal + transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors +} + +export type HybridSearchSingleReq = Pick< + SearchParam, + 'anns_field' | 'ignore_growing' | 'group_by_field' +> & { + data: VectorTypes[] | VectorTypes; // vector to search + expr?: string; // filter expression + params?: keyValueObj; // extra search parameters + transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors +}; -export interface SearchReq extends collectionNameReq { - partition_names?: string[]; - expr?: string; - // dsl_type: DslType; - search_params: SearchParam; - vectors: number[][]; - output_fields?: string[]; - travel_timestamp?: string; - vector_type: DataType.BinaryVector | DataType.FloatVector; - nq?: number; - consistency_level?: ConsistencyLevelEnum; -} +// rerank strategy and parameters +export type RerankerObj = { + strategy: RANKER_TYPE | string; // rerank strategy + params: keyValueObj; // rerank parameters +}; + +// hybrid search api parameter type +export type HybridSearchReq = Omit< + SearchSimpleReq, + 'data' | 'vector' | 'vectors' | 'params' | 'anns_field' +> & { + // search requests + data: HybridSearchSingleReq[]; + // reranker + rerank?: RerankerObj; +}; + +// search api response type export interface SearchRes extends resStatusResponse { results: { top_k: number; @@ -304,27 +369,37 @@ export interface SearchRes extends resStatusResponse { num_queries: number; topks: number[]; output_fields: string[]; + group_by_field_value: string; }; } +// because in javascript, there is no float16 and bfloat16 type +// we need to provide custom data transformer for these types +export type OutputTransformers = { + [DataType.BFloat16Vector]?: (bf16bytes: Uint8Array) => BFloat16Vector; + [DataType.Float16Vector]?: (f16: Uint8Array) => Float16Vector; + [DataType.SparseFloatVector]?: (sparse: SparseVectorDic) => SparseFloatVector; +}; + export interface QueryReq extends collectionNameReq { - output_fields?: string[]; - partition_names?: string[]; - ids?: string[] | number[]; - expr?: string; - filter?: string; - offset?: number; - limit?: number; - consistency_level?: ConsistencyLevelEnum; + output_fields?: string[]; // fields to return + partition_names?: string[]; // partition names + ids?: string[] | number[]; // primary key values + expr?: string; // filter expression + filter?: string; // alias for expr + offset?: number; // skip how many results + limit?: number; // how many results you want + consistency_level?: ConsistencyLevelEnum; // consistency level + transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors } export interface GetReq extends collectionNameReq { - ids: string[] | number[]; - output_fields?: string[]; - partition_names?: string[]; - offset?: number; - limit?: number; - consistency_level?: ConsistencyLevelEnum; + ids: string[] | number[]; // primary key values + output_fields?: string[]; // fields to return + partition_names?: string[]; // partition names + offset?: number; // skip how many results + limit?: number; // how many results you want + consistency_level?: ConsistencyLevelEnum; // consistency level } export interface QueryRes extends resStatusResponse { @@ -356,19 +431,19 @@ export interface FlushResult extends resStatusResponse { } export interface ListIndexedSegmentReq extends collectionNameReq { - index_name: string; + index_name: string; // index name } export interface ListIndexedSegmentResponse extends resStatusResponse { - segmentIDs: number[]; + segmentIDs: number[]; // indexed segment id array } export interface DescribeSegmentIndexDataReq extends collectionNameReq { - index_name: string; - segmentsIDs: number[]; + index_name: string; // index name + segmentsIDs: number[]; // segment id array } export interface DescribeSegmentIndexDataResponse extends resStatusResponse { - index_params: any; - index_data: any; + index_params: any; // index parameters + index_data: any; // index data } diff --git a/milvus/types/Http.ts b/milvus/types/Http.ts index 81a7f819..6046c67c 100644 --- a/milvus/types/Http.ts +++ b/milvus/types/Http.ts @@ -1,4 +1,4 @@ -import { FloatVectors } from '..'; +import { FloatVector } from '..'; type Fetch = (input: any, init?: any) => Promise; // Class types @@ -6,13 +6,11 @@ export type Constructor = new (...args: any[]) => T; export type FetchOptions = { abortController: AbortController; timeout: number; -} +}; type HttpClientConfigBase = { // database name database?: string; - // version - version?: string; // token token?: string; // The username to use for authentication. @@ -23,6 +21,8 @@ type HttpClientConfigBase = { timeout?: number; // altenative fetch api fetch?: Fetch; + // accept int64 + acceptInt64?: boolean; }; type HttpClientConfigAddress = HttpClientConfigBase & { @@ -54,12 +54,56 @@ export interface HttpBaseResponse { } // collection operations +type CollectionIndexParam = { + metricType: string; + fieldName: string; + indexName: string; + params?: { + index_type: string; // The type of the index to create + nlist?: number; // The number of cluster units. This applies to IVF-related index types. + M?: string; // The maximum degree of the node and applies only when index_type is set to __HNSW__. + efConstruction?: string; // The search scope. This applies only when **index_type** is set to **HNSW** + }; +}; + +type CollectionCreateParams = { + max_length?: number; // The maximum number of characters in a VarChar field. This parameter is mandatory when the current field type is VarChar. + enableDynamicField?: boolean; // Whether to enable the reserved dynamic field. If set to true, non-schema-defined fields are saved in the reserved dynamic field as key-value pairs. + shardsNum?: number; // The number of shards to create along with the current collection. + consistencyLevel?: string; // The consistency level of the collection. Possible values are __STRONG__, __BOUNDED__, __SESSION__, and __EVENTUALLY__. + partitionsNum?: number; // The number of partitions to create along with the current collection. This parameter is mandatory if one field of the collection has been designated as the partition key. + ttlSeconds?: number; // The time-to-live (TTL) period of the collection. If set, the collection is to be dropped once the period ends. +}; + +type CollectionCreateField = { + fieldName: string; // The name of the field to create in the target collection + dataType: string; // The data type of the field values. + elementDataType?: string; // The data type of the elements in an array field. + isPrimary?: boolean; // Whether the current field is the primary field. Setting this to True makes the current field the primary field. + isPartitionKey?: boolean; // Whether the current field serves as the partition key. Setting this to True makes the current field serve as the partition key. In this case, MilvusZilliz Cloud manages all partitions in the current collection. + elementTypeParams?: { + max_length?: number; // An optional parameter for VarChar values that determines the maximum length of the value in the current field. + dim?: number; // An optional parameter for FloatVector or BinaryVector fields that determines the vector dimension. + max_capacity?: number; // An optional parameter for Array field values that determines the maximum number of elements in the current array field. + }; +}; + +type CollectionCreateSchema = { + autoID?: boolean; + enabledDynamicField?: boolean; + fields: CollectionCreateField[]; +}; + export interface HttpCollectionCreateReq extends HttpBaseReq { - dimension: number; - metricType?: string; - primaryField?: string; - vectorField?: string; - description?: string; + dimension?: number; // The number of dimensions a vector value should have.This is required if **dtype** of this field is set to **DataType.FLOAT_VECTOR**. + metricType?: string; // The metric type applied to this operation. Possible values are **L2**, **IP**, and **COSINE**. + idType?: string; // The data type of the primary field. This parameter is designed for the quick-setup of a collection and will be ignored if __schema__ is defined. + autoID?: boolean; // Whether the primary field automatically increments. This parameter is designed for the quick-setup of a collection and will be ignored if __schema__ is defined. + primaryFieldName?: string; // The name of the primary field. This parameter is designed for the quick-setup of a collection and will be ignored if __schema__ is defined. + vectorFieldName?: string; // The name of the vector field. This parameter is designed for the quick-setup of a collection and will be ignored if __schema__ is defined. + schema?: CollectionCreateSchema; // The schema is responsible for organizing data in the target collection. A valid schema should have multiple fields, which must include a primary key, a vector field, and several scalar fields. + indexParams?: CollectionIndexParam[]; // The parameters that apply to the index-building process. + params?: CollectionCreateParams; // Extra parameters for the collection. } // list collection request export interface HttpCollectionListReq @@ -94,6 +138,24 @@ export interface HttpCollectionDescribeResponse export interface HttpCollectionListResponse extends HttpBaseResponse {} +export interface HttpCollectionHasResponse + extends HttpBaseResponse<{ has: boolean }> {} + +export interface HttpCollectionRenameReq extends HttpBaseReq { + newCollectionName: string; + newDbName?: string; +} + +export interface HttpCollectionStatisticsResponse + extends HttpBaseResponse<{ rowCount: number }> {} + +export interface HttpCollectionLoadStateReq extends HttpBaseReq { + partitionNames?: string; +} + +export interface HttpCollectionLoadStateResponse + extends HttpBaseResponse<{ loadProgress: number; loadState: string }> {} + // vector operations // insert data request export interface HttpVectorInsertReq extends HttpBaseReq { @@ -107,6 +169,13 @@ export interface HttpVectorInsertResponse insertIds: number | string[]; }> {} +// upsert data response +export interface HttpVectorUpsertResponse + extends HttpBaseResponse<{ + upsertCount: number; + upsertIds: number | string[]; + }> {} + // get vector request export interface HttpVectorGetReq extends HttpBaseReq { id: number | number[] | string | string[]; @@ -114,8 +183,10 @@ export interface HttpVectorGetReq extends HttpBaseReq { } // delete vector request -export interface HttpVectorDeleteReq - extends Omit {} +export interface HttpVectorDeleteReq extends HttpBaseReq { + filter: string; + partitionName?: string; +} // query data request export interface HttpVectorQueryReq extends HttpBaseReq { @@ -135,10 +206,151 @@ export interface HttpVectorQueryResponse // search request export interface HttpVectorSearchReq extends Omit { - vector: FloatVectors; + data: FloatVector[]; filter?: string; } export interface HttpVectorSearchResponse extends HttpVectorQueryResponse { data: QueryResult & { distance: number | string }; } + +/* partition operation */ +export interface HttpPartitionBaseReq extends HttpBaseReq { + partitionName: string; +} + +export interface HttpPartitionListReq extends HttpBaseReq { + partitionNames: string[]; +} + +export interface HttpPartitionHasResponse + extends HttpBaseResponse<{ has: boolean }> {} + +export interface HttpPartitionStatisticsResponse + extends HttpBaseResponse<{ rowCount: number }> {} + +/* user operation */ +export interface HttpUserBaseReq { + userName: string; +} + +export interface HttpUserCreateReq extends HttpUserBaseReq { + password: string; +} + +export interface HttpUserUpdatePasswordReq extends HttpUserCreateReq { + newPassword: string; +} + +export interface HttpUserRoleReq extends HttpUserBaseReq { + roleName: string; +} + +/* role operation */ +export interface HttpRoleBaseReq { + roleName: string; +} + +export interface HttpRolePrivilegeReq extends HttpRoleBaseReq { + objectType: string; + objectName: string; + privilege: string; +} + +export interface HttpRoleDescribeResponse + extends HttpBaseResponse {} + +/* index operation */ +export interface HttpIndexCreateReq extends HttpBaseReq { + indexParams: CollectionIndexParam[]; +} + +export interface HttpIndexBaseReq extends HttpBaseReq { + indexName: string; +} + +type IndexDescribeType = { + failReason: string; + fieldName: string; + indexName: string; + indexState: string; + indexType: string; + indexedRows: number; + metricType: string; + pendingRows: number; + totalRows: number; +}; + +export interface HttpIndexDescribeResponse + extends HttpBaseResponse {} + +/* alias operation */ +export type HttpAliasBaseReq = Pick; + +export interface HttpAliasCreateReq extends HttpBaseReq { + aliasName: string; +} + +export type HttpAliasAlterReq = HttpAliasCreateReq; + +export interface HttpAliasDescribeReq extends HttpAliasBaseReq { + aliasName: string; +} + +export interface HttpAliasDescribeResponse + extends HttpBaseResponse<{ aliasName: string } & Required> {} + +export interface HttpAliasDropReq extends Partial { + aliasName: string; +} + +/* import operation */ +type ImportJobType = { + collectionName: string; + jobId: string; + progress: number; + state: string; +}; + +type ImportJobDetailType = { + completeTime: string; + fileName: string; + fileSize: number; + importedRows: number; + progress: number; + state: string; + totalRows: number; +}; + +export interface HttpImportListResponse + extends HttpBaseResponse<{ records: ImportJobType[] }> {} + +export interface HttpImportCreateReq extends HttpBaseReq { + files: string[][]; + options?: { + timeout: string; + }; +} + +export interface HttpImportCreateResponse + extends HttpBaseResponse<{ + jobId: string; + }> {} + +export interface HttpImportProgressReq extends Pick { + jobId: string; +} + +export interface HttpImportProgressResponse + extends HttpBaseResponse<{ + jobId: string; + progress: number; + state: string; + totalRows?: number; + importedRows?: number; + fileSize?: number; + completeTime?: string; + collectionName?: string; + details?: ImportJobDetailType[]; + reason?: string; + }> {} diff --git a/milvus/types/MilvusIndex.ts b/milvus/types/MilvusIndex.ts index d78637f0..5ea81769 100644 --- a/milvus/types/MilvusIndex.ts +++ b/milvus/types/MilvusIndex.ts @@ -66,3 +66,8 @@ export interface GetIndexBuildProgressResponse extends resStatusResponse { indexed_rows: number; total_rows: number; } + +export interface AlterIndexReq extends collectionNameReq { + index_name: string; + params: Record; +} diff --git a/milvus/utils/Blob.ts b/milvus/utils/Blob.ts deleted file mode 100644 index 635391d4..00000000 --- a/milvus/utils/Blob.ts +++ /dev/null @@ -1,15 +0,0 @@ -import { FloatVectors, BinaryVectors } from '../'; - -export const parseFloatVectorToBytes = (array: FloatVectors) => { - // create array buffer - const a = new Float32Array(array); - // need return bytes to milvus proto - return Buffer.from(a.buffer); -}; - -export const parseBinaryVectorToBytes = (array: BinaryVectors) => { - // create array buffer - const a = new Uint8Array(array); - // need return bytes to milvus proto - return Buffer.from(a.buffer); -}; diff --git a/milvus/utils/Bytes.ts b/milvus/utils/Bytes.ts new file mode 100644 index 00000000..66f354b6 --- /dev/null +++ b/milvus/utils/Bytes.ts @@ -0,0 +1,306 @@ +import { Root } from 'protobufjs'; +import { Float16Array } from '@petamoriken/float16'; +import { + FloatVector, + BinaryVector, + SparseFloatVector, + DataType, + VectorTypes, + Float16Vector, + SparseVectorCSR, + SparseVectorCOO, + BFloat16Vector, + SparseVectorArray, +} from '..'; + +/** + * Converts a float vector into bytes format. + * + * @param {FloatVector} array - The float vector to convert. + * @returns {Buffer} Bytes representing the float vector. + */ +export const f32ArrayToF32Bytes = (array: FloatVector) => { + // create array buffer + const a = new Float32Array(array); + // need return bytes to milvus proto + return Buffer.from(a.buffer); +}; + +/** + * Converts a binary vector into bytes format. + * + * @param {BinaryVector} array - The binary vector to convert. + * @returns {Buffer} Bytes representing the binary vector. + */ +export const f32ArrayToBinaryBytes = (array: BinaryVector) => { + const a = new Uint8Array(array); + // need return bytes to milvus proto + return Buffer.from(a.buffer); +}; + +/** + * Converts a float16 vector into bytes format. + * + * @param {Float16Vector} array - The float16 vector(f32 format) to convert. + * @returns {Buffer} Bytes representing the float16 vector. + */ +export const f32ArrayToF16Bytes = (array: Float16Vector) => { + const float16Bytes = new Float16Array(array); + return Buffer.from(float16Bytes.buffer); +}; + +/** + * Convert float16 bytes to float32 array. + * @param {Uint8Array} f16Bytes - The float16 bytes to convert. + * @returns {Array} The float32 array. + */ +export const f16BytesToF32Array = (f16Bytes: Uint8Array) => { + const buffer = new ArrayBuffer(f16Bytes.length); + const view = new Uint8Array(buffer); + view.set(f16Bytes); + + const f16Array = new Float16Array(buffer); + return Array.from(f16Array); +}; + +/** + * Convert float32 array to BFloat16 bytes, not a real conversion, just take the last 2 bytes of float32. + * @param {BFloat16Vector} array - The float32 array to convert. + * @returns {Buffer} The BFloat16 bytes. + */ +export const f32ArrayToBf16Bytes = (array: BFloat16Vector) => { + const totalBytesNeeded = array.length * 2; // 2 bytes per float32 + const buffer = new ArrayBuffer(totalBytesNeeded); + const bfloatView = new Uint8Array(buffer); + + let byteIndex = 0; + array.forEach(float32 => { + const floatBuffer = new ArrayBuffer(4); + const floatView = new Float32Array(floatBuffer); + const bfloatViewSingle = new Uint8Array(floatBuffer); + + floatView[0] = float32; + bfloatView.set(bfloatViewSingle.subarray(2, 4), byteIndex); + byteIndex += 2; + }); + + return Buffer.from(bfloatView); +}; + +/** + * Convert BFloat16 bytes to Float32 array. + * @param {Uint8Array} bf16Bytes - The BFloat16 bytes to convert. + * @returns {Array} The Float32 array. + */ +export const bf16BytesToF32Array = (bf16Bytes: Uint8Array) => { + const float32Array: number[] = []; + const totalFloats = bf16Bytes.length / 2; + + for (let i = 0; i < totalFloats; i++) { + const floatBuffer = new ArrayBuffer(4); + const floatView = new Float32Array(floatBuffer); + const bfloatView = new Uint8Array(floatBuffer); + + bfloatView.set(bf16Bytes.subarray(i * 2, i * 2 + 2), 2); + float32Array.push(floatView[0]); + } + + return float32Array; +}; + +/** + * Get SparseVector type. + * @param {SparseFloatVector} vector - The sparse float vector to convert. + * + * @returns string, 'array' | 'coo' | 'csr' | 'dict' + */ +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 if ( + (vector as SparseVectorCOO).every( + item => typeof item === 'object' && 'index' in item && 'value' in item + ) + ) { + return 'coo'; + } else { + return 'unknown'; + } + } else if ( + typeof vector === 'object' && + 'indices' in vector && + 'values' in vector + ) { + return 'csr'; + } else if ( + typeof vector === 'object' && + Object.keys(vector).every(key => typeof vector[key] === 'number') + ) { + return 'dict'; + } else { + return 'unknown'; + } +}; + +/** + * Converts a sparse float vector into bytes format. + * + * @param {SparseFloatVector} data - The sparse float vector to convert, support 'array' | 'coo' | 'csr' | 'dict'. + * + * @returns {Uint8Array} Bytes representing the sparse float vector. + * @throws {Error} If the length of indices and values is not the same, or if the index is not within the valid range, or if the value is NaN. + */ +export const sparseToBytes = (data: SparseFloatVector): Uint8Array => { + // detect the format of the sparse vector + const type = getSparseFloatVectorType(data); + + let indices: number[] = []; + let values: number[] = []; + + switch (type) { + case 'array': + for (let i = 0; i < (data as SparseVectorArray).length; i++) { + const element = (data as SparseVectorArray)[i]; + if (element !== undefined && !isNaN(element)) { + indices.push(i); + values.push(element); + } + } + break; + case 'coo': + indices = Object.values( + (data as SparseVectorCOO).map((item: any) => item.index) + ); + values = Object.values( + (data as SparseVectorCOO).map((item: any) => item.value) + ); + break; + case 'csr': + indices = (data as SparseVectorCSR).indices; + values = (data as SparseVectorCSR).values; + break; + case 'dict': + indices = Object.keys(data).map(Number); + values = Object.values(data); + break; + } + + // create a buffer to store the bytes + const bytes = new Uint8Array(8 * indices.length); + + // loop through the indices and values and add them to the buffer + for (let i = 0; i < indices.length; i++) { + const index = indices[i]; + const value = values[i]; + if (!(index >= 0 && index < Math.pow(2, 32) - 1)) { + throw new Error( + `Sparse vector index must be positive and less than 2^32-1: ${index}` + ); + } + + const indexBytes = new Uint32Array([index]); + const valueBytes = new Float32Array([value]); + bytes.set(new Uint8Array(indexBytes.buffer), i * 8); + bytes.set(new Uint8Array(valueBytes.buffer), i * 8 + 4); + } + return bytes; +}; + +/** + * Converts an array of sparse float vectors into an array of bytes format. + * + * @param {SparseFloatVector[]} data - The array of sparse float vectors to convert. + * + * @returns {Uint8Array[]} An array of bytes representing the sparse float vectors. + */ +export const sparseRowsToBytes = (data: SparseFloatVector[]): Uint8Array[] => { + const result: Uint8Array[] = []; + for (const row of data) { + result.push(sparseToBytes(row)); + } + return result; +}; + +/** + * Parses the provided buffer data into a sparse row representation. + * + * @param {Buffer} bufferData - The buffer data to parse. + * + * @returns {SparseFloatVector} The parsed sparse float vectors. + */ +export const bytesToSparseRow = (bufferData: Buffer): SparseFloatVector => { + const result: SparseFloatVector = {}; + for (let i = 0; i < bufferData.length; i += 8) { + const key: string = bufferData.readUInt32LE(i).toString(); + const value: number = bufferData.readFloatLE(i + 4); + if (value) { + result[key] = value; + } + } + return result; +}; + +/** + * This function builds a placeholder group in bytes format for Milvus. + * + * @param {Root} milvusProto - The root object of the Milvus protocol. + * @param {VectorTypes[]} vectors - An array of search vectors. + * @param {DataType} vectorDataType - The data type of the vectors. + * + * @returns {Uint8Array} The placeholder group in bytes format. + */ +export const buildPlaceholderGroupBytes = ( + milvusProto: Root, + vectors: VectorTypes[], + vectorDataType: DataType +) => { + // create placeholder_group value + let bytes; + // parse vectors to bytes + switch (vectorDataType) { + case DataType.FloatVector: + bytes = vectors.map(v => f32ArrayToF32Bytes(v as FloatVector)); + break; + case DataType.BinaryVector: + bytes = vectors.map(v => f32ArrayToBinaryBytes(v as BinaryVector)); + break; + case DataType.BFloat16Vector: + bytes = vectors.map(v => + Array.isArray(v) ? f32ArrayToBf16Bytes(v as BFloat16Vector) : v + ); + break; + case DataType.Float16Vector: + bytes = vectors.map(v => + Array.isArray(v) ? f32ArrayToF16Bytes(v as Float16Vector) : v + ); + break; + case DataType.SparseFloatVector: + bytes = vectors.map(v => sparseToBytes(v as SparseFloatVector)); + + break; + } + // create placeholder_group + const PlaceholderGroup = milvusProto.lookupType( + 'milvus.proto.common.PlaceholderGroup' + ); + // tag $0 is hard code in milvus, when dsltype is expr + const placeholderGroupBytes = PlaceholderGroup.encode( + PlaceholderGroup.create({ + placeholders: [ + { + tag: '$0', + type: vectorDataType, + values: bytes, + }, + ], + }) + ).finish(); + + return placeholderGroupBytes; +}; diff --git a/milvus/utils/Format.ts b/milvus/utils/Format.ts index 658f3a15..da08a988 100644 --- a/milvus/utils/Format.ts +++ b/milvus/utils/Format.ts @@ -1,4 +1,4 @@ -import { Type } from 'protobufjs'; +import { Type, Root } from 'protobufjs'; import { findKeyValue, ERROR_REASONS, @@ -15,13 +15,39 @@ import { FieldData, CreateCollectionWithFieldsReq, CreateCollectionWithSchemaReq, + SearchReq, + SearchSimpleReq, + VectorTypes, + SearchParam, + HybridSearchSingleReq, + HybridSearchReq, + DEFAULT_TOPK, + DslType, + SearchRes, + DEFAULT_DYNAMIC_FIELD, + ConsistencyLevelEnum, + isVectorType, + RANKER_TYPE, + RerankerObj, + bytesToSparseRow, + buildPlaceholderGroupBytes, + Float16Vector, + BFloat16Vector, + getSparseFloatVectorType, + InsertTransformers, + OutputTransformers, + SparseVectorArray, + f32ArrayToBf16Bytes, + f32ArrayToF16Bytes, + bf16BytesToF32Array, + f16BytesToF32Array, } from '../'; /** - * parse [{key:"row_count",value:4}] to {row_count:4} - * @param data key value pair array - * @param keys all keys in data - * @returns {key:value} + * Formats key-value data based on the provided keys. + * @param {KeyValuePair[]} data - The array of key-value pairs. + * @param {string[]} keys - The keys to include in the formatted result. + * @returns {Object} - The formatted key-value data as an object. */ export const formatKeyValueData = (data: KeyValuePair[], keys: string[]) => { const result: { [x: string]: any } = {}; @@ -73,6 +99,12 @@ export const formatNumberPrecision = (number: number, precision: number) => { const LOGICAL_BITS = BigInt(18); // const LOGICAL_BITS_MASK = (1 << LOGICAL_BITS) - 1; +/** + * Checks if the given time parameter is valid. + * + * @param ts - The time parameter to be checked. + * @returns A boolean value indicating whether the time parameter is valid or not. + */ export const checkTimeParam = (ts: any) => { switch (typeof ts) { case 'bigint': @@ -85,26 +117,10 @@ export const checkTimeParam = (ts: any) => { }; /** - * Convert a hybrid timestamp to UNIX Epoch time ignoring the logic part. - * - * @param data - * | Property | Type | Description | - * | :---------------- | :---- | :------------------------------- | - * | hybridts | String or BigInt | The known hybrid timestamp to convert to UNIX Epoch time. Non-negative interger range from 0 to 18446744073709551615. | - * - * - * - * @returns - * | Property | Description | - * | :-----------| :------------------------------- | - * | unixtime as string | The Unix Epoch time is the number of seconds that have elapsed since January 1, 1970 (midnight UTC/GMT). | - * - * - * #### Example - * - * ``` - * const res = hybridtsToUnixtime("429642767925248000"); - * ``` + * Converts a hybrid timestamp to Unix time. + * @param hybridts - The hybrid timestamp to convert. + * @returns The Unix time representation of the hybrid timestamp. + * @throws An error if the hybridts parameter fails the time parameter check. */ export const hybridtsToUnixtime = (hybridts: bigint | string) => { if (!checkTimeParam(hybridts)) { @@ -116,26 +132,10 @@ export const hybridtsToUnixtime = (hybridts: bigint | string) => { }; /** - * Generate a hybrid timestamp based on Unix Epoch time, timedelta and incremental time internval. - * - * @param data - * | Property | Type | Description | - * | :---------------- | :---- | :------------------------------- | - * | unixtime | string or bigint | The known Unix Epoch time used to generate a hybrid timestamp. The Unix Epoch time is the number of seconds that have elapsed since January 1, 1970 (midnight UTC/GMT). | - * - * - * - * @returns - * | Property | Type | Description | - * | :-----------| :--- | :------------------------------- | - * | Hybrid timetamp | String | Hybrid timetamp is a non-negative interger range from 0 to 18446744073709551615. | - * - * - * #### Example - * - * ``` - * const res = unixtimeToHybridts("429642767925248000"); - * ``` + * Converts a Unix timestamp to a hybrid timestamp. + * @param unixtime - The Unix timestamp to convert. + * @returns The hybrid timestamp as a string. + * @throws An error if the unixtime parameter fails the check. */ export const unixtimeToHybridts = (unixtime: bigint | string) => { if (!checkTimeParam(unixtime)) { @@ -148,26 +148,10 @@ export const unixtimeToHybridts = (unixtime: bigint | string) => { }; /** - * Generate a hybrid timestamp based on datetime。 - * - * @param data - * | Property | Type | Description | - * | :---------------- | :---- | :------------------------------- | - * | datetime | Date | The known datetime used to generate a hybrid timestamp. | - * - * - * - * @returns - * | Property | Type | Description | - * | :-----------| :--- | :------------------------------- | - * | Hybrid timetamp | String | Hybrid timetamp is a non-negative interger range from 0 to 18446744073709551615. | - * - * - * #### Example - * - * ``` - * const res = datetimeToHybrids("429642767925248000"); - * ``` + * Converts a JavaScript Date object to a hybridts timestamp. + * @param datetime - The JavaScript Date object to be converted. + * @returns The hybridts timestamp. + * @throws An error if the input is not a valid Date object. */ export const datetimeToHybrids = (datetime: Date) => { if (!(datetime instanceof Date)) { @@ -417,7 +401,10 @@ export const buildDynamicRow = ( * If the field is a vector, split the data into chunks of the appropriate size. * If the field is a scalar, decode the JSON/array data if necessary. */ -export const buildFieldDataMap = (fields_data: any[]) => { +export const buildFieldDataMap = ( + fields_data: any[], + transformers?: OutputTransformers +) => { const fieldsDataMap = new Map(); fields_data.forEach((item, i) => { @@ -426,42 +413,83 @@ export const buildFieldDataMap = (fields_data: any[]) => { // parse vector data if (item.field === 'vectors') { - const key = item.vectors!.data; - const vectorValue = - key === 'float_vector' - ? item.vectors![key]!.data - : item.vectors![key]!.toJSON().data; - - // if binary vector , need use dim / 8 to split vector data - const dim = - item.vectors?.data === 'float_vector' - ? Number(item.vectors!.dim) - : Number(item.vectors!.dim) / 8; - field_data = []; - - // parse number[] to number[][] by dim - vectorValue.forEach((v: any, i: number) => { - const index = Math.floor(i / dim); - if (!field_data[index]) { - field_data[index] = []; - } - field_data[index].push(v); - }); + const dataKey = item.vectors!.data; + + switch (dataKey) { + case 'float_vector': + case 'binary_vector': + const vectorValue = + dataKey === 'float_vector' + ? item.vectors![dataKey]!.data + : item.vectors![dataKey]!.toJSON().data; + + // if binary vector , need use dim / 8 to split vector data + const dim = + item.vectors?.data === 'float_vector' + ? Number(item.vectors!.dim) + : Number(item.vectors!.dim) / 8; + field_data = []; + + // parse number[] to number[][] by dim + vectorValue.forEach((v: any, i: number) => { + const index = Math.floor(i / dim); + if (!field_data[index]) { + field_data[index] = []; + } + field_data[index].push(v); + }); + break; + + case 'float16_vector': + case 'bfloat16_vector': + field_data = []; + const f16Dim = Number(item.vectors!.dim) * 2; // float16 is 2 bytes, so we need to multiply dim with 2 = one element length + const f16Bytes = item.vectors![dataKey]!; + + // split buffer data to float16 vector(bytes) + for (let i = 0; i < f16Bytes.byteLength; i += f16Dim) { + const slice = f16Bytes.slice(i, i + f16Dim); + const isFloat16 = dataKey === 'float16_vector'; + let dataType: DataType.BFloat16Vector | DataType.Float16Vector; + + dataType = isFloat16 + ? DataType.Float16Vector + : DataType.BFloat16Vector; + + const localTransformers = transformers || { + [DataType.BFloat16Vector]: bf16BytesToF32Array, + [DataType.Float16Vector]: f16BytesToF32Array, + }; + + field_data.push(localTransformers[dataType]!(slice)); + } + break; + case 'sparse_float_vector': + const sparseVectorValue = item.vectors![dataKey]!.contents; + field_data = []; + + sparseVectorValue.forEach((buffer: any, i: number) => { + field_data[i] = bytesToSparseRow(buffer); + }); + break; + default: + break; + } } else { // parse scalar data - const key = item.scalars!.data; - field_data = item.scalars![key]!.data; + const dataKey = item.scalars!.data; + field_data = item.scalars![dataKey]!.data; // we need to handle array element specifically here - if (key === 'array_data') { + if (dataKey === 'array_data') { field_data = field_data.map((f: any) => { - const key = f.data; - return key ? f[key].data : []; + const dataKey = f.data; + return dataKey ? f[dataKey].data : []; }); } // decode json - switch (key) { + switch (dataKey) { case 'json_data': field_data.forEach((buffer: any, i: number) => { // console.log(JSON.parse(buffer.toString())); @@ -512,18 +540,373 @@ export const getAuthString = (data: { * @param {Field} column - The column information. * @returns {FieldData} The field data for the row and column. */ -export const buildFieldData = (rowData: RowData, field: Field): FieldData => { +export const buildFieldData = ( + rowData: RowData, + field: Field, + transformers?: InsertTransformers +): FieldData => { const { type, elementType, name } = field; + const isFloat32 = Array.isArray(rowData[name]); + switch (DataTypeMap[type]) { case DataType.BinaryVector: case DataType.FloatVector: return rowData[name]; + case DataType.BFloat16Vector: + const bf16Transformer = + transformers?.[DataType.BFloat16Vector] || f32ArrayToBf16Bytes; + return isFloat32 + ? bf16Transformer(rowData[name] as BFloat16Vector) + : rowData[name]; + case DataType.Float16Vector: + const f16Transformer = + transformers?.[DataType.Float16Vector] || f32ArrayToF16Bytes; + return isFloat32 + ? f16Transformer(rowData[name] as Float16Vector) + : rowData[name]; case DataType.JSON: return Buffer.from(JSON.stringify(rowData[name] || {})); case DataType.Array: const elementField = { ...field, type: elementType! }; - return buildFieldData(rowData, elementField); + return buildFieldData(rowData, elementField, transformers); default: return rowData[name]; } }; + +/** + * Builds search parameters based on the provided data. + * @param data - The data object containing search parameters. + * @returns The search parameters in key-value format. + */ +export const buildSearchParams = ( + data: SearchSimpleReq | (HybridSearchSingleReq & HybridSearchReq), + anns_field: string +) => { + // create search params + const search_params: SearchParam = { + anns_field: data.anns_field || anns_field, + params: JSON.stringify(data.params ?? {}), + topk: data.limit ?? data.topk ?? DEFAULT_TOPK, + offset: data.offset ?? 0, + metric_type: data.metric_type ?? '', // leave it empty + ignore_growing: data.ignore_growing ?? false, + }; + + // if group_by_field is set, add it to the search params + if (data.group_by_field) { + search_params.group_by_field = data.group_by_field; + } + + return search_params; +}; + +/** + * Creates a RRFRanker object with the specified value of k. + * @param k - The value of k used in the RRFRanker strategy. + * @returns An object representing the RRFRanker strategy with the specified value of k. + */ +export const RRFRanker = (k: number = 60): RerankerObj => { + return { + strategy: RANKER_TYPE.RRF, + params: { + k, + }, + }; +}; + +/** + * Creates a weighted ranker object. + * @param weights - An array of numbers representing the weights. + * @returns The weighted ranker object. + */ +export const WeightedRanker = (weights: number[]): RerankerObj => { + return { + strategy: RANKER_TYPE.WEIGHTED, + params: { + weights, + }, + }; +}; + +/** + * Converts the rerank parameters object to a format suitable for API requests. + * @param rerank - The rerank parameters object. + * @returns The converted rerank parameters object. + */ +export const convertRerankParams = (rerank: RerankerObj) => { + const r = cloneObj(rerank) as any; + r.params = JSON.stringify(r.params); + return r; +}; + +/** + * This method is used to build search request for a given data. + * It first fetches the collection info and then constructs the search request based on the data type. + * It also creates search vectors and a placeholder group for the search. + * + * @param {SearchReq | SearchSimpleReq | HybridSearchReq} data - The data for which to build the search request. + * @param {DescribeCollectionResponse} collectionInfo - The collection information. + * @param {Root} milvusProto - The milvus protocol object. + * @returns {Object} An object containing the search requests and search vectors. + * @returns {Object} return.params - The search requests used in the operation. + * @returns {string} return.params.collection_name - The name of the collection. + * @returns {string[]} return.params.partition_names - The partition names. + * @returns {string[]} return.params.output_fields - The output fields. + * @returns {number} return.params.nq - The number of query vectors. + * @returns {string} return.params.dsl - The domain specific language. + * @returns {string} return.params.dsl_type - The type of the domain specific language. + * @returns {Uint8Array} return.params.placeholder_group - The placeholder group. + * @returns {Object} return.params.search_params - The search parameters. + * @returns {string} return.params.consistency_level - The consistency level. + * @returns {Number[][]} return.searchVectors - The search vectors used in the operation. + * @returns {number} return.round_decimal - The score precision. + */ +export const buildSearchRequest = ( + data: SearchReq | SearchSimpleReq | HybridSearchReq, + collectionInfo: DescribeCollectionResponse, + milvusProto: Root +) => { + // type cast + const searchReq = data as SearchReq; + const searchHybridReq = data as HybridSearchReq; + const searchSimpleReq = data as SearchSimpleReq; + + // Initialize requests array + const requests: { + collection_name: string; + partition_names: string[]; + output_fields: string[]; + nq: number; + dsl: string; + dsl_type: DslType; + placeholder_group: Uint8Array; + search_params: KeyValuePair[]; + consistency_level: ConsistencyLevelEnum; + }[] = []; + + // detect if the request is hybrid search request + const isHybridSearch = !!( + searchHybridReq.data && + searchHybridReq.data.length && + typeof searchHybridReq.data[0] === 'object' && + searchHybridReq.data[0].anns_field + ); + + // output fields(reference fields) + const default_output_fields: string[] = []; + + // Iterate through collection fields, create search request + for (let i = 0; i < collectionInfo.schema.fields.length; i++) { + const field = collectionInfo.schema.fields[i]; + const { name, dataType } = field; + + // if field type is vector, build the request + if (isVectorType(dataType)) { + let req: SearchSimpleReq | (HybridSearchReq & HybridSearchSingleReq) = + data as SearchSimpleReq; + + if (isHybridSearch) { + const singleReq = searchHybridReq.data.find(d => d.anns_field === name); + // if it is hybrid search and no request target is not found, skip + if (!singleReq) { + continue; + } + // merge single request with hybrid request + req = Object.assign(cloneObj(data), singleReq); + } else { + // if it is not hybrid search, and we have built one request, skip + const skip = + requests.length === 1 || + (typeof req.anns_field !== 'undefined' && req.anns_field !== name); + if (skip) { + continue; + } + } + + // get search vectors + let searchingVector: VectorTypes | VectorTypes[] = isHybridSearch + ? req.data! + : searchReq.vectors || + searchSimpleReq.vectors || + searchSimpleReq.vector || + searchSimpleReq.data; + + // format searching vector + searchingVector = formatSearchVector(searchingVector, field.dataType!); + + // create search request + requests.push({ + collection_name: req.collection_name, + partition_names: req.partition_names || [], + output_fields: req.output_fields || default_output_fields, + nq: searchReq.nq || searchingVector.length, + dsl: searchReq.expr || searchSimpleReq.filter || '', + dsl_type: DslType.BoolExprV1, + placeholder_group: buildPlaceholderGroupBytes( + milvusProto, + searchingVector as VectorTypes[], + field.dataType! + ), + search_params: parseToKeyValue( + searchReq.search_params || buildSearchParams(req, name) + ), + consistency_level: + req.consistency_level || (collectionInfo.consistency_level as any), + }); + } else { + // if field is not vector, add it to output fields + default_output_fields.push(name); + } + } + + /** + * It will decide the score precision. + * If round_decimal is 3, need return like 3.142 + * And if Milvus return like 3.142, Node will add more number after this like 3.142000047683716. + * So the score need to slice by round_decimal + */ + const round_decimal = + searchReq.search_params?.round_decimal ?? + (searchSimpleReq.params?.round_decimal as number) ?? + -1; + + return { + isHybridSearch, + request: isHybridSearch + ? { + collection_name: data.collection_name, + partition_names: data.partition_names, + requests: requests, + rank_params: [ + ...parseToKeyValue( + convertRerankParams(searchHybridReq.rerank || RRFRanker()) + ), + { key: 'round_decimal', value: round_decimal }, + { + key: 'limit', + value: + searchSimpleReq.limit ?? searchSimpleReq.topk ?? DEFAULT_TOPK, + }, + ], + output_fields: requests[0]?.output_fields, + consistency_level: requests[0]?.consistency_level, + } + : requests[0], + nq: requests[0].nq, + round_decimal, + }; +}; + +/** + * Formats the search results returned by Milvus into row data for easier use. + * + * @param {SearchRes} searchRes - The search results returned by Milvus. + * @param {Object} options - The options for formatting the search results. + * @param {number} options.round_decimal - The number of decimal places to which to round the scores. + * + * @returns {any[]} The formatted search results. + * + */ +export const formatSearchResult = ( + searchRes: SearchRes, + options: { + round_decimal: number; + transformers?: OutputTransformers; + } +) => { + const { round_decimal } = options; + // build final results array + const results: any[] = []; + const { topks, scores, fields_data, ids } = searchRes.results; + // build fields data map + const fieldsDataMap = buildFieldDataMap(fields_data, options.transformers); + // build output name array + const output_fields = [ + 'id', + ...(!!searchRes.results.output_fields?.length + ? searchRes.results.output_fields + : fields_data.map(f => f.field_name)), + ]; + + // vector id support int / str id. + const idData = ids ? ids[ids.id_field]!.data : {}; + // add id column + fieldsDataMap.set('id', idData as RowData[]); + // fieldsDataMap.set('score', scores); TODO: fieldDataMap to support formatter + + /** + * This code block formats the search results returned by Milvus into row data for easier use. + * Milvus supports multiple queries to search and returns all columns data, so we need to splice the data for each search result using the `topk` variable. + * The `topk` variable is the key we use to splice data for every search result. + * The `scores` array is spliced using the `topk` value, and the resulting scores are formatted to the specified precision using the `formatNumberPrecision` function. The resulting row data is then pushed to the `results` array. + */ + topks.forEach((v, index) => { + const topk = Number(v); + + scores.splice(0, topk).forEach((score, scoreIndex) => { + // get correct index + const i = index === 0 ? scoreIndex : scoreIndex + topk * index; + + // fix round_decimal + const fixedScore = + typeof round_decimal === 'undefined' || round_decimal === -1 + ? score + : formatNumberPrecision(score, round_decimal); + + // init result object + const result: any = { score: fixedScore }; + + // build result, + output_fields.forEach(field_name => { + // Check if the field_name exists in the fieldsDataMap + const isFixedSchema = fieldsDataMap.has(field_name); + + // Get the data for the field_name from the fieldsDataMap + // If the field_name is not in the fieldsDataMap, use the DEFAULT_DYNAMIC_FIELD + const data = fieldsDataMap.get( + isFixedSchema ? field_name : DEFAULT_DYNAMIC_FIELD + )!; + // make dynamic data[i] safe + data[i] = isFixedSchema ? data[i] : data[i] || {}; + // extract dynamic info from dynamic field if necessary + result[field_name] = isFixedSchema ? data[i] : data[i][field_name]; + }); + + // init result slot + results[index] = results[index] || []; + // push result data + results[index].push(result); + }); + }); + + 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 SparseVectorArray); + if (type !== 'unknown') { + return [searchVector] as VectorTypes[]; + } + default: + return searchVector as VectorTypes[]; + } +}; diff --git a/milvus/utils/Function.ts b/milvus/utils/Function.ts index cc797b63..61951d30 100644 --- a/milvus/utils/Function.ts +++ b/milvus/utils/Function.ts @@ -1,4 +1,4 @@ -import { KeyValuePair, DataType, ERROR_REASONS } from '../'; +import { KeyValuePair, DataType, ERROR_REASONS, SparseFloatVector } from '../'; import { Pool } from 'generic-pool'; /** @@ -89,9 +89,18 @@ export const getDataKey = (type: DataType, camelCase: boolean = false) => { case DataType.FloatVector: dataKey = 'float_vector'; break; + case DataType.Float16Vector: + dataKey = 'float16_vector'; + break; + case DataType.BFloat16Vector: + dataKey = 'bfloat16_vector'; + break; case DataType.BinaryVector: dataKey = 'binary_vector'; break; + case DataType.SparseFloatVector: + dataKey = 'sparse_float_vector'; + break; case DataType.Double: dataKey = 'double_data'; break; @@ -129,3 +138,15 @@ export const getDataKey = (type: DataType, camelCase: boolean = false) => { } return camelCase ? convertToCamelCase(dataKey) : dataKey; }; + +// get biggest size of sparse vector array +export const getSparseDim = (data: SparseFloatVector[]) => { + let dim = 0; + for (const row of data) { + const indices = Object.keys(row).map(Number); + if (indices.length > dim) { + dim = indices.length; + } + } + return dim; +}; diff --git a/milvus/utils/Validate.ts b/milvus/utils/Validate.ts index a7f3cb52..9a7c17cf 100644 --- a/milvus/utils/Validate.ts +++ b/milvus/utils/Validate.ts @@ -21,8 +21,6 @@ import { status as grpcStatus } from '@grpc/grpc-js'; * @param fields */ export const checkCollectionFields = (fields: FieldType[]) => { - // Define arrays of data types that are allowed for vector fields and primary keys, respectively - const vectorDataTypes = [DataType.BinaryVector, DataType.FloatVector]; const int64VarCharTypes = [DataType.Int64, DataType.VarChar]; let hasPrimaryKey = false; @@ -56,11 +54,11 @@ export const checkCollectionFields = (fields: FieldType[]) => { } // if this is the vector field, check dimension - const isVectorField = vectorDataTypes.includes(dataType!); + const isVectorField = isVectorType(dataType!); const typeParams = field.type_params; if (isVectorField) { const dim = Number(typeParams?.dim ?? field.dim); - if (!dim) { + if (!dim && dataType !== DataType.SparseFloatVector) { throw new Error(ERROR_REASONS.CREATE_COLLECTION_CHECK_MISS_DIM); } @@ -206,7 +204,7 @@ export const checkCreateCollectionCompatibility = ( if (hasDynamicSchemaEnabled) { throw new Error( - `Your milvus server doesn't support dynmaic schmea, please upgrade your server.` + `Your milvus server doesn't support dynamic schema, please upgrade your server.` ); } @@ -230,3 +228,18 @@ export const checkCreateCollectionCompatibility = ( ); } }; + +/** + * Checks if the given data type is a vector type. + * @param {DataType} type - The data type to check. + * @returns {Boolean} True if the data type is a vector type, false otherwise. + */ +export const isVectorType = (type: DataType) => { + return ( + type === DataType.BinaryVector || + type === DataType.FloatVector || + type === DataType.Float16Vector || + type === DataType.BFloat16Vector || + type === DataType.SparseFloatVector + ); +}; diff --git a/milvus/utils/index.ts b/milvus/utils/index.ts index 64804022..4703c7cb 100644 --- a/milvus/utils/index.ts +++ b/milvus/utils/index.ts @@ -1,5 +1,5 @@ export * from './Grpc'; -export * from './Blob'; +export * from './Bytes'; export * from './Format'; export * from './Validate'; export * from './Function'; diff --git a/package.json b/package.json index a6f3773d..8d143ddb 100644 --- a/package.json +++ b/package.json @@ -1,8 +1,8 @@ { "name": "@zilliz/milvus2-sdk-node", "author": "ued@zilliz.com", - "version": "2.3.6", - "milvusVersion": "v2.3.9", + "version": "2.4.0", + "milvusVersion": "v2.4.0", "main": "dist/milvus", "files": [ "dist" @@ -30,6 +30,7 @@ }, "devDependencies": { "@babel/plugin-transform-modules-commonjs": "^7.21.5", + "@petamoriken/float16": "^3.8.6", "@types/jest": "^29.5.1", "@types/node-fetch": "^2.6.8", "jest": "^29.5.0", diff --git a/proto b/proto index d367b5a5..e3012ae6 160000 --- a/proto +++ b/proto @@ -1 +1 @@ -Subproject commit d367b5a59df171628c736dfb8f05e5f799d13179 +Subproject commit e3012ae615bc1fdb0908cccf44c54c8a7da1a222 diff --git a/test/build/Collection.spec.ts b/test/build/Collection.spec.ts index 034de90e..0df9f9d0 100644 --- a/test/build/Collection.spec.ts +++ b/test/build/Collection.spec.ts @@ -18,7 +18,7 @@ const COLLECTION_NAME = GENERATE_NAME(); const LOAD_COLLECTION_NAME = GENERATE_NAME(); const COLLECTION_NAME_PARAMS = genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 128, + dim: [128], }); describe('Collection Api', () => { @@ -91,7 +91,7 @@ describe('Collection Api', () => { try { await milvusClient.createCollection( - genCollectionParams({ collectionName: 'any', dim: 8 }) + genCollectionParams({ collectionName: 'any', dim: [8] }) ); } catch (error) { expect(error.message).toEqual( @@ -102,7 +102,7 @@ describe('Collection Api', () => { it(`Create load Collection Successful`, async () => { const res = await milvusClient.createCollection( - genCollectionParams({ collectionName: LOAD_COLLECTION_NAME, dim: 128 }) + genCollectionParams({ collectionName: LOAD_COLLECTION_NAME, dim: [128] }) ); expect(res.error_code).toEqual(ErrorCode.SUCCESS); }); diff --git a/test/grpc/Alias.spec.ts b/test/grpc/Alias.spec.ts index e3e0c350..05310eb3 100644 --- a/test/grpc/Alias.spec.ts +++ b/test/grpc/Alias.spec.ts @@ -8,7 +8,7 @@ const COLLECTION_ALIAS = GENERATE_NAME('alias'); describe(`Alias API`, () => { beforeAll(async () => { await milvusClient.createCollection( - genCollectionParams({ collectionName: COLLECTION_NAME, dim: 8 }) + genCollectionParams({ collectionName: COLLECTION_NAME, dim: [8] }) ); }); diff --git a/test/grpc/BFloat16Vector.spec.ts b/test/grpc/BFloat16Vector.spec.ts new file mode 100644 index 00000000..8179f21e --- /dev/null +++ b/test/grpc/BFloat16Vector.spec.ts @@ -0,0 +1,152 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, + f32ArrayToBf16Bytes, + bf16BytesToF32Array, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'Bfloat_vector_16', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [DataType.BFloat16Vector], + dim: [8], +}; +const collectionParams = genCollectionParams(p); +const data = generateInsertData(collectionParams.fields, 10); + +// console.log( +// 'data to insert', +// data.map(d => d.vector) +// ); + +describe(`BFloat16 vector API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with Bfloat16 vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const BfloatVector16Fields = describe.schema.fields.filter( + (field: any) => field.data_type === 'BFloat16Vector' + ); + expect(BfloatVector16Fields.length).toBe(1); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert Bflaot16 vector data should be successful`, async () => { + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data: data, + }); + + // console.log(' insert', insert); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.L2, + index_type: IndexType.AUTOINDEX, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query Bfloat16 vector should be successful`, async () => { + const count = await milvusClient.count({ + collection_name: COLLECTION_NAME, + }); + + expect(count.data).toEqual(data.length); + + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'id'], + }); + + // console.dir(query, { depth: null }); + + // verify the query result + data.forEach((obj, index) => { + obj.vector.forEach((v: number, i: number) => { + expect(v).toBeCloseTo(query.data[index].vector[i], 2); + }); + }); + + expect(query.status.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`search with Bfloat16 vector should be successful`, async () => { + const search = await milvusClient.search({ + data: f32ArrayToBf16Bytes(data[0].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); + }); + + it(`search with Bfloat16 vector and nq > 0 should be successful`, async () => { + const search = await milvusClient.search({ + data: [ + f32ArrayToBf16Bytes(data[0].vector), + f32ArrayToBf16Bytes(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/Basic.spec.ts b/test/grpc/Basic.spec.ts index 05e2b13c..bd092fcf 100644 --- a/test/grpc/Basic.spec.ts +++ b/test/grpc/Basic.spec.ts @@ -75,11 +75,25 @@ describe(`Basic API without database`, () => { it(`search should be successful`, async () => { const search = await milvusClient.search({ collection_name: COLLECTION_NAME, - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], }); 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/BinaryVector.spec.ts b/test/grpc/BinaryVector.spec.ts new file mode 100644 index 00000000..cd27bfbd --- /dev/null +++ b/test/grpc/BinaryVector.spec.ts @@ -0,0 +1,114 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'binary_vector_test', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [DataType.BinaryVector], + dim: [16], +}; +const collectionParams = genCollectionParams(p); +const data = generateInsertData(collectionParams.fields, 10); + +describe(`Binary vectors API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with binary vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const binaryVectorFields = describe.schema.fields.filter( + (field: any) => field.data_type === 'BinaryVector' + ); + expect(binaryVectorFields.length).toBe(1); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert binary vector data should be successful`, async () => { + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data, + }); + + // console.log('data to insert', data); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.HAMMING, + index_type: IndexType.BIN_IVF_FLAT, + params: { + nlist: 10, + }, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query binary vector should be successful`, async () => { + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'id'], + }); + + expect(query.status.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`search with binary vector should be successful`, async () => { + const search = await milvusClient.search({ + data: data[0].vector, + collection_name: COLLECTION_NAME, + output_fields: ['id', 'vector'], + limit: 5, + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(search.results.length).toBeGreaterThan(0); + }); +}); diff --git a/test/grpc/Collection.spec.ts b/test/grpc/Collection.spec.ts index af925818..b6ae5986 100644 --- a/test/grpc/Collection.spec.ts +++ b/test/grpc/Collection.spec.ts @@ -31,7 +31,7 @@ const dbParam = { const COLLECTION_NAME_PARAMS = genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 128, + dim: [128], }); describe(`Collection API`, () => { @@ -62,7 +62,7 @@ describe(`Collection API`, () => { const res = await milvusClient.createCollection({ ...genCollectionParams({ collectionName: NUMBER_DIM_COLLECTION_NAME, - dim: 128, + dim: [128], }), consistency_level: 'Eventually', }); @@ -142,7 +142,7 @@ describe(`Collection API`, () => { try { const d = await milvusClient.createCollection( - genCollectionParams({ collectionName: 'any', dim: 10 }) + genCollectionParams({ collectionName: 'any', dim: [10] }) ); } catch (error) { expect(error.message).toEqual( @@ -193,7 +193,7 @@ describe(`Collection API`, () => { const res = await milvusClient.createCollection({ ...genCollectionParams({ collectionName: TEST_CONSISTENCY_LEVEL_COLLECTION_NAME, - dim: 128, + dim: [128], }), consistency_level: 'xxx' as any, }); @@ -203,12 +203,12 @@ describe(`Collection API`, () => { it(`Create load Collection Successful`, async () => { const res1 = await milvusClient.createCollection( - genCollectionParams({ collectionName: LOAD_COLLECTION_NAME, dim: 128 }) + genCollectionParams({ collectionName: LOAD_COLLECTION_NAME, dim: [128] }) ); const res2 = await milvusClient.createCollection( genCollectionParams({ collectionName: LOAD_COLLECTION_NAME_SYNC, - dim: 128, + dim: [128], }) ); // make sure load successful @@ -345,6 +345,38 @@ describe(`Collection API`, () => { expect(res.schema.fields[1].name).toEqual('id'); }); + it(`alter collection success`, async () => { + const key = 'collection.ttl.seconds'; + const value = 18000; + + const alter = await milvusClient.alterCollection({ + collection_name: LOAD_COLLECTION_NAME, + properties: { [key]: value }, + }); + expect(alter.error_code).toEqual(ErrorCode.SUCCESS); + + const key2 = 'mmap.enabled'; + const value2 = true; + + const alter2 = await milvusClient.alterCollection({ + collection_name: LOAD_COLLECTION_NAME, + properties: { [key2]: value2 }, + }); + + expect(alter2.error_code).toEqual(ErrorCode.SUCCESS); + const describe = await milvusClient.describeCollection({ + collection_name: LOAD_COLLECTION_NAME, + }); + + expect(Number(formatKeyValueData(describe.properties, [key])[key])).toEqual( + value + ); + + expect( + Boolean(formatKeyValueData(describe.properties, [key2])[key2]) + ).toEqual(value2); + }); + it(`Load Collection Sync throw COLLECTION_NAME_IS_REQUIRED`, async () => { try { await milvusClient.loadCollectionSync({} as any); @@ -506,25 +538,6 @@ describe(`Collection API`, () => { } }); - it(`alter collection success`, async () => { - const key = 'collection.ttl.seconds'; - const value = 18000; - const alter = await milvusClient.alterCollection({ - collection_name: LOAD_COLLECTION_NAME, - properties: { [key]: value }, - }); - - expect(alter.error_code).toEqual(ErrorCode.SUCCESS); - - const describe = await milvusClient.describeCollection({ - collection_name: LOAD_COLLECTION_NAME, - }); - - expect(Number(formatKeyValueData(describe.properties, [key])[key])).toEqual( - value - ); - }); - it(`Create alias success`, async () => { const res0 = await milvusClient.describeCollection({ collection_name: LOAD_COLLECTION_NAME, diff --git a/test/grpc/Data.spec.ts b/test/grpc/Data.spec.ts index 812f5541..74be5790 100644 --- a/test/grpc/Data.spec.ts +++ b/test/grpc/Data.spec.ts @@ -5,6 +5,7 @@ import { ERROR_REASONS, DEFAULT_TOPK, DEFAULT_COUNT_QUERY_STRING, + IndexType, } from '../../milvus'; import { IP, @@ -25,8 +26,8 @@ const dbParam = { }; const createCollectionParams = genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, }); const INDEX_NAME = 'collection_index'; @@ -51,14 +52,18 @@ describe(`Data.API`, () => { data: generateInsertData(createCollectionParams.fields, 1024), }); + await milvusClient.flush({ + collection_names: [COLLECTION_NAME], + }); + // create index await milvusClient.createIndex({ index_name: INDEX_NAME, collection_name: COLLECTION_NAME, field_name: VECTOR_FIELD_NAME, - index_type: 'IVF_FLAT', + index_type: IndexType.HNSW, metric_type: 'L2', - params: { nlist: 1024 }, + params: { M: 4, efConstruction: 8 }, }); // load await milvusClient.loadCollectionSync({ @@ -164,14 +169,6 @@ describe(`Data.API`, () => { } }); - it(`Exec search should throw SEARCH_PARAMS_IS_REQUIRED`, async () => { - try { - await milvusClient.search({ collection_name: 'asd' } as any); - } catch (error) { - expect(error.message).toEqual(ERROR_REASONS.VECTORS_OR_VECTOR_IS_MISSING); - } - }); - it(`Exec search should throw error`, async () => { try { await milvusClient.search({ @@ -187,7 +184,7 @@ describe(`Data.API`, () => { // const res = await milvusClient.search({ // collection_name: COLLECTION_NAME, // filter: '', - // vector: [1, 2, 3, 4], + // data: [1, 2, 3, 4], // limit: limit, // metric_type: 'IP', // }); @@ -205,10 +202,10 @@ describe(`Data.API`, () => { filter: '', data: [1, 2, 3, 4], limit: limit, + group_by_field: 'varChar', }); expect(searchWithData.status.error_code).toEqual(ErrorCode.SUCCESS); - expect(searchWithData.results.length).toEqual(limit); const searchWithData2 = await milvusClient.search({ collection_name: COLLECTION_NAME, @@ -220,12 +217,12 @@ describe(`Data.API`, () => { expect(searchWithData2.status.error_code).toEqual(ErrorCode.SUCCESS); expect(searchWithData2.results.length).toEqual(limit); - // parititon search + // partition search const partitionSearch = await milvusClient.search({ collection_name: COLLECTION_NAME, partition_names: [PARTITION_NAME], filter: '', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], topk: limit, }); @@ -237,7 +234,7 @@ describe(`Data.API`, () => { const res = await milvusClient.search({ collection_name: COLLECTION_NAME, filter: '', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], }); expect(res.status.error_code).toEqual(ErrorCode.SUCCESS); @@ -250,7 +247,7 @@ describe(`Data.API`, () => { const res = await milvusClient.search({ collection_name: COLLECTION_NAME, filter: '', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], limit: limit, params: { nprobe: 1024 }, }); @@ -258,11 +255,11 @@ describe(`Data.API`, () => { expect(res.status.error_code).toEqual(ErrorCode.SUCCESS); expect(res.results.length).toEqual(limit); - // mutitple vector search search + // multiple vector search const res2 = await milvusClient.search({ collection_name: COLLECTION_NAME, filter: '', - vectors: [[1, 2, 3, 4]], + data: [[1, 2, 3, 4]], limit: limit, offset: 2, params: { nprobe: 1024 }, @@ -277,7 +274,7 @@ describe(`Data.API`, () => { const res = await milvusClient.search({ collection_name: COLLECTION_NAME, filter: 'int64 < 10000', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], limit: limit, params: { nprobe: 1024 }, }); @@ -290,7 +287,7 @@ describe(`Data.API`, () => { const res2 = await milvusClient.search({ collection_name: COLLECTION_NAME, expr: 'int64 < 10000', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], limit: limit, params: { nprobe: 1024 }, }); @@ -305,7 +302,7 @@ describe(`Data.API`, () => { const res = await milvusClient.search({ collection_name: COLLECTION_NAME, filter: 'int64 < 10000', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], limit: limit, params: { nprobe: 1024, radius: 20, range_filter: 15 }, }); @@ -321,7 +318,7 @@ describe(`Data.API`, () => { collection_name: COLLECTION_NAME, // partition_names: [], filter: '', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], limit: 4, output_fields: ['id', 'json', VECTOR_FIELD_NAME], }; @@ -343,7 +340,7 @@ describe(`Data.API`, () => { collection_name: COLLECTION_NAME, // partition_names: [], filter: 'json["number"] >= 0', - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], limit: 4, output_fields: ['id', 'json'], }; @@ -364,7 +361,7 @@ describe(`Data.API`, () => { collection_name: COLLECTION_NAME, // partition_names: [], expr: '', - vectors: [[1, 2, 3, 4]], + data: [[1, 2, 3, 4]], search_params: { anns_field: VECTOR_FIELD_NAME, topk: '4', @@ -384,7 +381,7 @@ describe(`Data.API`, () => { collection_name: COLLECTION_NAME, // partition_names: [], expr: '', - vectors: [[1, 2, 3, 4]], + data: [[1, 2, 3, 4]], search_params: { anns_field: VECTOR_FIELD_NAME, topk: '4', @@ -405,7 +402,7 @@ describe(`Data.API`, () => { collection_name: COLLECTION_NAME, // partition_names: [], expr: '', - vectors: [[1, 2, 3]], + data: [[1, 2, 3]], search_params: { anns_field: VECTOR_FIELD_NAME, topk: '4', diff --git a/test/grpc/Database.spec.ts b/test/grpc/Database.spec.ts index d39d75cb..c62f16b1 100644 --- a/test/grpc/Database.spec.ts +++ b/test/grpc/Database.spec.ts @@ -23,7 +23,7 @@ describe(`Database API`, () => { // create collection on another db const create = await milvusClient.createCollection( - genCollectionParams({ collectionName: COLLECTION_NAME, dim: 4 }) + genCollectionParams({ collectionName: COLLECTION_NAME, dim: [4] }) ); expect(create.error_code).toEqual(ErrorCode.SUCCESS); diff --git a/test/grpc/DynamicSchema.spec.ts b/test/grpc/DynamicSchema.spec.ts index 9f05e0cc..64bebe98 100644 --- a/test/grpc/DynamicSchema.spec.ts +++ b/test/grpc/DynamicSchema.spec.ts @@ -22,8 +22,8 @@ const numPartitions = 3; // create const createCollectionParams = genCollectionParams({ collectionName: COLLECTION, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, partitionKeyEnabled: true, numPartitions, @@ -115,7 +115,7 @@ describe(`Dynamic schema API`, () => { const search = await milvusClient.search({ collection_name: COLLECTION, limit: 10, - vectors: [ + data: [ [1, 2, 3, 4], [1, 2, 3, 4], ], @@ -132,7 +132,7 @@ describe(`Dynamic schema API`, () => { const search2 = await milvusClient.search({ collection_name: COLLECTION, limit: 10, - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], expr: 'id > 0', output_fields: ['json', 'id', 'dynamic_int64', 'dynamic_varChar'], }); diff --git a/test/grpc/Float16Vector.spec.ts b/test/grpc/Float16Vector.spec.ts new file mode 100644 index 00000000..58aaaa54 --- /dev/null +++ b/test/grpc/Float16Vector.spec.ts @@ -0,0 +1,146 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, + f32ArrayToF16Bytes, + f16BytesToF32Array, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'float_vector_16', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [DataType.Float16Vector], + dim: [128], +}; +const collectionParams = genCollectionParams(p); +const data = generateInsertData(collectionParams.fields, 2); + +// console.log('data to insert', data); + +describe(`Float16 vector API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with float16 vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const floatVector16Fields = describe.schema.fields.filter( + (field: any) => field.data_type === 'Float16Vector' + ); + expect(floatVector16Fields.length).toBe(1); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert float16 vector data should be successful`, async () => { + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data, + }); + + // console.log(' insert', insert); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.L2, + index_type: IndexType.AUTOINDEX, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query float16 vector should be successful`, async () => { + const count = await milvusClient.count({ + collection_name: COLLECTION_NAME, + }); + + expect(count.data).toEqual(data.length); + + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'id'], + }); + + // verify the query result + data.forEach((obj, index) => { + obj.vector.forEach((v: number, i: number) => { + expect(v).toBeCloseTo(query.data[index].vector[i], 3); + }); + }); + + expect(query.status.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`search with float16 vector should be successful`, async () => { + const search = await milvusClient.search({ + data: data[0].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); + }); + + it(`search with float16 vector and nq > 0 should be successful`, async () => { + const search = await milvusClient.search({ + data: [ + f32ArrayToF16Bytes(data[0].vector), + f32ArrayToF16Bytes(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/Import.spec.ts b/test/grpc/Import.spec.ts index 1f2ea9fe..4936e614 100644 --- a/test/grpc/Import.spec.ts +++ b/test/grpc/Import.spec.ts @@ -22,8 +22,8 @@ describe(`Import API`, () => { await milvusClient.createCollection( genCollectionParams({ collectionName: COLLECTION_NAME, - dim: '4', - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, }) ); diff --git a/test/grpc/Index.spec.ts b/test/grpc/Index.spec.ts index 120fd7e2..44eb1bce 100644 --- a/test/grpc/Index.spec.ts +++ b/test/grpc/Index.spec.ts @@ -1,4 +1,10 @@ -import { MilvusClient, ErrorCode, MetricType, IndexType } from '../../milvus'; +import { + MilvusClient, + ErrorCode, + MetricType, + IndexType, + findKeyValue, +} from '../../milvus'; import { IP, genCollectionParams, @@ -38,18 +44,18 @@ describe(`Milvus Index API`, () => { await milvusClient.createDatabase(dbParam); await milvusClient.use(dbParam); await milvusClient.createCollection( - genCollectionParams({ collectionName: COLLECTION_NAME, dim: 8 }) + genCollectionParams({ collectionName: COLLECTION_NAME, dim: [8] }) ); await milvusClient.createCollection( genCollectionParams({ collectionName: COLLECTION_NAME_WITHOUT_INDEX_NAME, - dim: 8, + dim: [8], }) ); for (let i = 0; i < INDEX_COLLECTIONS.length; i++) { await milvusClient.createCollection( - genCollectionParams({ collectionName: INDEX_COLLECTIONS[i], dim: 32 }) + genCollectionParams({ collectionName: INDEX_COLLECTIONS[i], dim: [32] }) ); } }); @@ -235,6 +241,36 @@ describe(`Milvus Index API`, () => { expect(res.error_code).toEqual(ErrorCode.SUCCESS); }); + it(`Create STL_SORT index on int64 should success`, async () => { + const res = await milvusClient.createIndex({ + index_name: 'int64_index', + collection_name: COLLECTION_NAME, + field_name: 'int64', + index_type: IndexType.STL_SORT, + }); + expect(res.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`Create TRIE index on int64 varchar success`, async () => { + const res = await milvusClient.createIndex({ + index_name: 'varchar_index', + collection_name: COLLECTION_NAME, + field_name: 'varChar', + index_type: IndexType.TRIE, + }); + expect(res.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`Create INVERTED index on int64 varchar success`, async () => { + const res = await milvusClient.createIndex({ + index_name: 'float_index', + collection_name: COLLECTION_NAME, + field_name: 'float', + index_type: IndexType.INVERTED, + }); + expect(res.error_code).toEqual(ErrorCode.SUCCESS); + }); + it(`Create Index without name should success`, async () => { const res = await milvusClient.createIndex({ collection_name: COLLECTION_NAME_WITHOUT_INDEX_NAME, @@ -266,7 +302,8 @@ describe(`Milvus Index API`, () => { collection_name: COLLECTION_NAME, index_name: INDEX_NAME, }); - expect(res.index_descriptions[0].index_name).toEqual(INDEX_NAME); + const allIndexNames = res.index_descriptions.map(i => i.index_name); + expect(allIndexNames.includes(INDEX_NAME)).toEqual(true); expect(res.status.error_code).toEqual(ErrorCode.SUCCESS); }); @@ -276,7 +313,9 @@ describe(`Milvus Index API`, () => { field_name: VECTOR_FIELD_NAME, }); - expect(res.index_descriptions[0].field_name).toEqual(VECTOR_FIELD_NAME); + const field_names = res.index_descriptions.map(i => i.field_name); + expect(field_names.includes(VECTOR_FIELD_NAME)).toEqual(true); + expect(res.status.error_code).toEqual(ErrorCode.SUCCESS); }); @@ -332,14 +371,35 @@ describe(`Milvus Index API`, () => { expect(res.status.error_code).toEqual(ErrorCode.SUCCESS); }); - it(`Get Index progress with field name should be failed`, async () => { - const res = await milvusClient.getIndexBuildProgress({ + it(`Alter Index should be success`, async () => { + const alter = await milvusClient.alterIndex({ collection_name: COLLECTION_NAME, - field_name: VECTOR_FIELD_NAME, + index_name: INDEX_NAME, + params: { + 'mmap.enabled': true, + }, }); - expect(res.status.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeIndex({ + collection_name: COLLECTION_NAME, + index_name: INDEX_NAME, + }); + expect(alter.error_code).toEqual(ErrorCode.SUCCESS); + const params = describe.index_descriptions[0].params; + expect(findKeyValue(params, 'mmap.enabled')).toEqual("true"); + + // console.log('describe', describe.index_descriptions[0].params); }); + // @Deprecated + // it(`Get Index progress with field name should be failed`, async () => { + // const res = await milvusClient.getIndexBuildProgress({ + // collection_name: COLLECTION_NAME, + // field_name: VECTOR_FIELD_NAME, + // }); + // expect(res.status.error_code).toEqual(ErrorCode.SUCCESS); + // }); + it(`Drop Index with index name`, async () => { const res = await milvusClient.dropIndex({ collection_name: COLLECTION_NAME, diff --git a/test/grpc/Insert.spec.ts b/test/grpc/Insert.spec.ts index 2d6e11fd..433c73b1 100644 --- a/test/grpc/Insert.spec.ts +++ b/test/grpc/Insert.spec.ts @@ -26,14 +26,14 @@ const COLLECTION_NAME_AUTO_ID = GENERATE_NAME(); const MORE_SCALAR_COLLECTION_NAME = GENERATE_NAME(); const COLLECTION_NAME_PARAMS = genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, }); const COLLECTION_NAME_AUTO_ID_PARAMS = genCollectionParams({ collectionName: COLLECTION_NAME_AUTO_ID, - dim: 16, - vectorType: DataType.BinaryVector, + dim: [16], + vectorType: [DataType.BinaryVector], autoID: true, }); @@ -69,8 +69,8 @@ describe(`Insert API`, () => { await milvusClient.createCollection( genCollectionParams({ collectionName: BINARY_COLLECTION_NAME, - dim: 8, - vectorType: DataType.BinaryVector, + dim: [8], + vectorType: [DataType.BinaryVector], autoID: false, }) ); @@ -335,7 +335,7 @@ describe(`Insert API`, () => { const search = await milvusClient.search({ collection_name: COLLECTION_NAME, - vector: genFloatVector({ dim: 4 }) as number[], + data: genFloatVector({ dim: 4 }) as number[], output_fields: ['json', 'id', 'varChar_array'], }); // console.log('search', search.results); diff --git a/test/grpc/MilvusClient.spec.ts b/test/grpc/MilvusClient.spec.ts index 93842a60..024143d0 100644 --- a/test/grpc/MilvusClient.spec.ts +++ b/test/grpc/MilvusClient.spec.ts @@ -124,6 +124,50 @@ describe(`Milvus client`, () => { expect(m7.tlsMode).toEqual(TLS_MODE.ONE_WAY); }); + it(`should set tls to 2 if root cert and others provided`, async () => { + const m7s = new MilvusClient({ + address: IP, + ssl: true, + username: 'username', + password: 'password', + id: '1', + tls: { + rootCertPath: `test/cert/ca.pem`, + privateKeyPath: `test/cert/client.key`, + certChainPath: `test/cert/client.pem`, + }, + __SKIP_CONNECT__: true, + }); + + expect(m7s.tlsMode).toEqual(TLS_MODE.TWO_WAY); + }); + + it(`it should setup maxRetries and retryDelay successfully`, async () => { + const m8 = new MilvusClient({ + address: IP, + ssl: true, + username: 'username', + password: 'password', + id: '1', + retryDelay: 300, + maxRetries: 3, + __SKIP_CONNECT__: true, + }); + + expect(m8.config.maxRetries).toEqual(3); + expect(m8.config.retryDelay).toEqual(300); + }); + + it(`it should setup string timeout successfully`, async () => { + const m9 = new MilvusClient({ + address: IP, + timeout: '1h', + __SKIP_CONNECT__: true, + }); + + expect(m9.timeout).toEqual(3600000); + }); + it(`Should throw MILVUS_ADDRESS_IS_REQUIRED`, async () => { try { new MilvusClient(undefined as any); diff --git a/test/grpc/MultipleVectors.spec.ts b/test/grpc/MultipleVectors.spec.ts new file mode 100644 index 00000000..f8f24378 --- /dev/null +++ b/test/grpc/MultipleVectors.spec.ts @@ -0,0 +1,363 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, + RRFRanker, + WeightedRanker, + f32ArrayToF16Bytes, + f16BytesToF32Array, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'multiple_vectors', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [ + DataType.FloatVector, + DataType.FloatVector, + DataType.Float16Vector, + DataType.SparseFloatVector, + ], + dim: [8, 16, 4, 8], +}; +const collectionParams = genCollectionParams(p); + +describe(`Multiple vectors API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with multiple vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const floatVectorFields = describe.schema.fields.filter( + (field: any) => field.data_type === 'FloatVector' + ); + expect(floatVectorFields.length).toBe(2); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert multiple vector data should be successful`, async () => { + const data = generateInsertData(collectionParams.fields, 10); + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data, + transformers: { + [DataType.Float16Vector]: f32ArrayToF16Bytes, + }, + }); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create multiple index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.COSINE, + index_type: IndexType.HNSW, + params: { + M: 5, + efConstruction: 8, + }, + }, + { + collection_name: COLLECTION_NAME, + field_name: 'vector1', + metric_type: MetricType.COSINE, + index_type: IndexType.AUTOINDEX, + }, + { + collection_name: COLLECTION_NAME, + field_name: 'vector2', + metric_type: MetricType.COSINE, + index_type: IndexType.AUTOINDEX, + }, + { + collection_name: COLLECTION_NAME, + field_name: 'vector3', + metric_type: MetricType.IP, + index_type: IndexType.SPARSE_WAND, + params: { + drop_ratio_build: 0.2, + }, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load multiple vector collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query multiple vector collection should be successful`, async () => { + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'vector1', 'vector2', 'vector3'], + transformers: { + [DataType.Float16Vector]: f16BytesToF32Array, + }, + }); + + expect(query.status.error_code).toEqual(ErrorCode.SUCCESS); + + 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 single vector search should be successful`, async () => { + // search default first vector field + const search0 = await milvusClient.search({ + collection_name: COLLECTION_NAME, + 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, + 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, + 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 third vector field + const searchF16 = await milvusClient.search({ + collection_name: COLLECTION_NAME, + data: [1, 2, 3, 4], + anns_field: 'vector2', + }); + + expect(searchF16.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(searchF16.results.length).toEqual(search.results.length); + + // search the fourth vector field + const searchSparse = await milvusClient.search({ + collection_name: COLLECTION_NAME, + data: { 1: 2, 3: 4 }, + anns_field: 'vector3', + limit: 10, + }); + + expect(searchSparse.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(searchSparse.results.length).toBeGreaterThan(0); + }); + + it(`hybrid search with rrf ranker set should be successful`, async () => { + const search = await milvusClient.hybridSearch({ + collection_name: COLLECTION_NAME, + data: [ + { + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + { + data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + anns_field: 'vector1', + }, + { + data: f32ArrayToF16Bytes([1, 2, 3, 4]), + anns_field: 'vector2', + }, + { + data: { 1: 2, 3: 4 }, + anns_field: 'vector3', + }, + ], + rerank: RRFRanker(), + limit: 5, + output_fields: ['id', 'vector2', 'vector3'], + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(search.results.length).toEqual(5); + expect(Object.keys(search.results[0]).length).toEqual(4); + }); + + it(`hybrid search with weighted ranker set should be successful`, async () => { + const search = await milvusClient.hybridSearch({ + collection_name: COLLECTION_NAME, + data: [ + { + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + { + data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + anns_field: 'vector1', + }, + ], + rerank: WeightedRanker([0.9, 0.1]), + limit: 5, + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(search.results.length).toEqual(5); + }); + + it(`hybrid search without ranker set should be successful`, async () => { + const search = await milvusClient.hybridSearch({ + collection_name: COLLECTION_NAME, + data: [ + { + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + { + data: f32ArrayToF16Bytes([1, 2, 3, 4]), + anns_field: 'vector2', + }, + ], + limit: 5, + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(search.results.length).toEqual(5); + }); + + it(`hybrid search with nq > 1 should be successful`, async () => { + const search = await milvusClient.hybridSearch({ + collection_name: COLLECTION_NAME, + data: [ + { + data: [ + [1, 2, 3, 4, 5, 6, 7, 8], + [1, 2, 3, 4, 5, 6, 7, 8], + ], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + { + data: [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + ], + anns_field: 'vector1', + }, + ], + limit: 2, + output_fields: ['vector', 'vector1'], + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(search.results.length).toEqual(2); + expect(search.results[0].length).toEqual(2); + expect(search.results[1].length).toEqual(2); + }); + + it(`hybrid search with one vector should be successful`, async () => { + const search = await milvusClient.hybridSearch({ + collection_name: COLLECTION_NAME, + data: [ + { + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + ], + limit: 5, + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(search.results.length).toEqual(5); + }); + + it(`user can use search api for hybrid search`, async () => { + const search = await milvusClient.search({ + collection_name: COLLECTION_NAME, + data: [ + { + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + { + data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + anns_field: 'vector1', + }, + ], + limit: 5, + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(search.results.length).toEqual(5); + }); + + it(`hybrid search result should be equal to the original search result`, async () => { + const search = await milvusClient.search({ + collection_name: COLLECTION_NAME, + data: [ + { + data: [1, 2, 3, 4, 5, 6, 7, 8], + anns_field: 'vector', + params: { nprobe: 2 }, + }, + ], + limit: 5, + }); + + const originSearch = await milvusClient.search({ + collection_name: COLLECTION_NAME, + data: [1, 2, 3, 4, 5, 6, 7, 8], + limit: 5, + }); + + expect(search.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(originSearch.status.error_code).toEqual(search.status.error_code); + expect(originSearch.results.length).toEqual(search.results.length); + + expect(search.results.map(r => r.id)).toEqual( + originSearch.results.map(r => r.id) + ); + }); +}); diff --git a/test/grpc/Partition.spec.ts b/test/grpc/Partition.spec.ts index bb93b1c2..78907822 100644 --- a/test/grpc/Partition.spec.ts +++ b/test/grpc/Partition.spec.ts @@ -24,7 +24,7 @@ describe(`Partition API`, () => { await milvusClient.createDatabase(dbParam); await milvusClient.use(dbParam); await milvusClient.createCollection( - genCollectionParams({ collectionName: COLLECTION_NAME, dim: 128 }) + genCollectionParams({ collectionName: COLLECTION_NAME, dim: [128] }) ); await milvusClient.createIndex({ diff --git a/test/grpc/PartitionKey.spec.ts b/test/grpc/PartitionKey.spec.ts index 76b0de0f..f08fe5cd 100644 --- a/test/grpc/PartitionKey.spec.ts +++ b/test/grpc/PartitionKey.spec.ts @@ -29,8 +29,8 @@ describe(`Partition key API`, () => { // create const createCollectionParams = genCollectionParams({ collectionName: COLLECTION_DATA_NAME, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, partitionKeyEnabled: true, numPartitions, @@ -77,8 +77,8 @@ describe(`Partition key API`, () => { it(`Create Collection with 2 partition key fields should throw error`, async () => { const createCollectionParams = genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, partitionKeyEnabled: true, numPartitions, @@ -105,8 +105,8 @@ describe(`Partition key API`, () => { it(`Create Collection should be successful with numPartitions`, async () => { const createCollectionParams = genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, partitionKeyEnabled: true, numPartitions, @@ -119,8 +119,8 @@ describe(`Partition key API`, () => { it(`Create Collection should be successful without numPartitions`, async () => { const createCollectionParams = genCollectionParams({ collectionName: COLLECTION_NAME2, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, partitionKeyEnabled: true, }); @@ -133,8 +133,8 @@ describe(`Partition key API`, () => { it(`it should create collection successfully with partition_key_field set`, async () => { const createCollectionParams = genCollectionParams({ collectionName: COLLECTION_NAME2, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, partitionKeyEnabled: false, }); @@ -202,7 +202,7 @@ describe(`Partition key API`, () => { it(`Search Collection should be successful`, async () => { const res = await milvusClient.search({ collection_name: COLLECTION_DATA_NAME, - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], expr: 'varChar in ["apple"]', output_fields: ['varChar'], }); @@ -214,7 +214,7 @@ describe(`Partition key API`, () => { const search = await milvusClient.search({ collection_name: COLLECTION_DATA_NAME, partition_names: ['p'], - vector: [1, 2, 3, 4], + data: [1, 2, 3, 4], expr: 'varChar in ["apple"]', output_fields: ['varChar'], }); diff --git a/test/grpc/Replica.spec.ts b/test/grpc/Replica.spec.ts index e2d6fc95..06067df0 100644 --- a/test/grpc/Replica.spec.ts +++ b/test/grpc/Replica.spec.ts @@ -19,7 +19,7 @@ describe(`Replica API`, () => { await milvusClient.use(dbParam); await milvusClient.createCollection( - genCollectionParams({ collectionName: COLLECTION_NAME, dim: 8 }) + genCollectionParams({ collectionName: COLLECTION_NAME, dim: [8] }) ); await milvusClient.createIndex({ collection_name: COLLECTION_NAME, diff --git a/test/grpc/Resource.spec.ts b/test/grpc/Resource.spec.ts index 968e2cdc..b42c0e82 100644 --- a/test/grpc/Resource.spec.ts +++ b/test/grpc/Resource.spec.ts @@ -38,7 +38,7 @@ describe(`Resource API`, () => { // create collection await milvusClient.createCollection( - genCollectionParams({ collectionName: COLLECTION_NAME, dim: 128 }) + genCollectionParams({ collectionName: COLLECTION_NAME, dim: [128] }) ); await milvusClient.createIndex({ diff --git a/test/grpc/SparseVector.array.spec.ts b/test/grpc/SparseVector.array.spec.ts new file mode 100644 index 00000000..6349269d --- /dev/null +++ b/test/grpc/SparseVector.array.spec.ts @@ -0,0 +1,143 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'sparse_array_vector_DB', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [DataType.SparseFloatVector], + dim: [24], // useless +}; +const collectionParams = genCollectionParams(p); +const data = generateInsertData(collectionParams.fields, 10, { + sparseType: 'array', +}); + +describe(`Sparse vectors type:object API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with sparse vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const sparseFloatVectorFields = describe.schema.fields.filter( + (field: any) => field.data_type === 'SparseFloatVector' + ); + expect(sparseFloatVectorFields.length).toBe(1); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert sparse vector data should be successful`, async () => { + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data, + }); + + // console.log('insert', insert); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.IP, + index_type: IndexType.SPARSE_WAND, + params: { + drop_ratio_build: 0.2, + }, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query sparse vector should be successful`, async () => { + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'id'], + }); + + // console.dir(query, { depth: null }); + + const originKeys = Object.keys(query.data[0].vector); + const originValues = Object.values(query.data[0].vector); + + const outputKeys: string[] = Object.keys(query.data[0].vector); + const outputValues: number[] = Object.values(query.data[0].vector); + + expect(originKeys).toEqual(outputKeys); + + // filter undefined in originValues + originValues.forEach((value, index) => { + if (value) { + expect(value).toBeCloseTo(outputValues[index]); + } + }); + }); + + it(`search with sparse vector should be successful`, async () => { + const search = await milvusClient.search({ + data: data[0].vector, + collection_name: COLLECTION_NAME, + output_fields: ['id', 'vector'], + limit: 5, + }); + + 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({ + data: [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 new file mode 100644 index 00000000..bfe2a9de --- /dev/null +++ b/test/grpc/SparseVector.coo.spec.ts @@ -0,0 +1,143 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'sparse_coo_vector_DB', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [DataType.SparseFloatVector], + dim: [24], // useless + sparseType: 'coo', +}; +const collectionParams = genCollectionParams(p); +const data = generateInsertData(collectionParams.fields, 5, { + sparseType: 'coo', +}); + +// console.dir(data, { depth: null }); +describe(`Sparse vectors type:coo API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with sparse vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const sparseFloatVectorFields = describe.schema.fields.filter( + (field: any) => field.data_type === 'SparseFloatVector' + ); + expect(sparseFloatVectorFields.length).toBe(1); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert sparse vector data should be successful`, async () => { + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data, + }); + + // console.log('data to insert', data); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.IP, + index_type: IndexType.SPARSE_WAND, + params: { + drop_ratio_build: 0.2, + }, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query sparse vector should be successful`, async () => { + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'id'], + }); + + const originKeys = Object.values( + data[0].vector.map((item: any) => item.index) + ) as number[]; + const originValues = Object.values( + data[0].vector.map((item: any) => item.value) + ) as number[]; + + const outputKeys: number[] = Object.keys(query.data[0].vector).map(Number); + const outputValues: number[] = Object.values(query.data[0].vector); + + expect(originKeys).toEqual(outputKeys); + originValues.forEach((value: number, index: number) => { + expect(value).toBeCloseTo(outputValues[index]); + }); + }); + + it(`search with sparse vector should be successful`, async () => { + const search = await milvusClient.search({ + data: data[0].vector, + collection_name: COLLECTION_NAME, + output_fields: ['id', 'vector'], + limit: 5, + }); + + 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({ + data: [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 new file mode 100644 index 00000000..c3e63ff3 --- /dev/null +++ b/test/grpc/SparseVector.csr.spec.ts @@ -0,0 +1,141 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'sparse_csr_vector_DB', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [DataType.SparseFloatVector], + dim: [24], // useless + sparseType: 'csr', +}; +const collectionParams = genCollectionParams(p); +const data = generateInsertData(collectionParams.fields, 5, { + sparseType: 'csr', +}); + +// console.dir(data, { depth: null }); +describe(`Sparse vectors type:CSR API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with sparse vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const sparseFloatVectorFields = describe.schema.fields.filter( + (field: any) => field.data_type === 'SparseFloatVector' + ); + expect(sparseFloatVectorFields.length).toBe(1); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert sparse vector data should be successful`, async () => { + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data, + }); + + // console.log('data to insert', data); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.IP, + index_type: IndexType.SPARSE_WAND, + params: { + drop_ratio_build: 0.2, + }, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query sparse vector should be successful`, async () => { + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'id'], + }); + + 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); + const outputValues: number[] = Object.values(query.data[0].vector); + + expect(originKeys).toEqual(outputKeys); + originValues.forEach((value: number, index: number) => { + expect(value).toBeCloseTo(outputValues[index]); + }); + }); + + it(`search with sparse vector should be successful`, async () => { + const search = await milvusClient.search({ + data: data[0].vector, + collection_name: COLLECTION_NAME, + output_fields: ['id', 'vector'], + limit: 5, + }); + + 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({ + data: [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 new file mode 100644 index 00000000..62062844 --- /dev/null +++ b/test/grpc/SparseVector.dict.spec.ts @@ -0,0 +1,135 @@ +import { + MilvusClient, + ErrorCode, + DataType, + IndexType, + MetricType, +} from '../../milvus'; +import { + IP, + genCollectionParams, + GENERATE_NAME, + generateInsertData, +} from '../tools'; + +const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' }); +const COLLECTION_NAME = GENERATE_NAME(); + +const dbParam = { + db_name: 'sparse_dict_vector_DB', +}; + +const p = { + collectionName: COLLECTION_NAME, + vectorType: [DataType.SparseFloatVector], + dim: [24], // useless +}; +const collectionParams = genCollectionParams(p); +const data = generateInsertData(collectionParams.fields, 10); + +describe(`Sparse vectors type:dict API testing`, () => { + beforeAll(async () => { + await milvusClient.createDatabase(dbParam); + await milvusClient.use(dbParam); + }); + + afterAll(async () => { + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); + await milvusClient.dropDatabase(dbParam); + }); + + it(`Create collection with sparse vectors should be successful`, async () => { + const create = await milvusClient.createCollection(collectionParams); + expect(create.error_code).toEqual(ErrorCode.SUCCESS); + + const describe = await milvusClient.describeCollection({ + collection_name: COLLECTION_NAME, + }); + + const sparseFloatVectorFields = describe.schema.fields.filter( + (field: any) => field.data_type === 'SparseFloatVector' + ); + expect(sparseFloatVectorFields.length).toBe(1); + + // console.dir(describe.schema, { depth: null }); + }); + + it(`insert sparse vector data should be successful`, async () => { + const insert = await milvusClient.insert({ + collection_name: COLLECTION_NAME, + data, + }); + + // console.log('data to insert', data); + + expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS); + expect(insert.succ_index.length).toEqual(data.length); + }); + + it(`create index should be successful`, async () => { + const indexes = await milvusClient.createIndex([ + { + collection_name: COLLECTION_NAME, + field_name: 'vector', + metric_type: MetricType.IP, + index_type: IndexType.SPARSE_WAND, + params: { + drop_ratio_build: 0.2, + }, + }, + ]); + + expect(indexes.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`load collection should be successful`, async () => { + const load = await milvusClient.loadCollection({ + collection_name: COLLECTION_NAME, + }); + + expect(load.error_code).toEqual(ErrorCode.SUCCESS); + }); + + it(`query sparse vector should be successful`, async () => { + const query = await milvusClient.query({ + collection_name: COLLECTION_NAME, + filter: 'id > 0', + output_fields: ['vector', 'id'], + }); + + const originKeys = Object.keys(data[0].vector); + const originValues = Object.values(data[0].vector); + + const outputKeys: string[] = Object.keys(query.data[0].vector); + const outputValues: number[] = Object.values(query.data[0].vector); + + expect(originKeys).toEqual(outputKeys); + originValues.forEach((value, index) => { + expect(value).toBeCloseTo(outputValues[index]); + }); + }); + + it(`search with sparse vector should be successful`, async () => { + const search = await milvusClient.search({ + data: data[0].vector, + collection_name: COLLECTION_NAME, + output_fields: ['id', 'vector'], + limit: 5, + }); + + 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({ + data: [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/Upsert.spec.ts b/test/grpc/Upsert.spec.ts index 5b757c94..08fb450e 100644 --- a/test/grpc/Upsert.spec.ts +++ b/test/grpc/Upsert.spec.ts @@ -20,14 +20,14 @@ const COLLECTION_NAME_AUTO_ID = GENERATE_NAME(); const MORE_SCALAR_COLLECTION_NAME = GENERATE_NAME(); const COLLECTION_NAME_PARAMS = genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, }); const COLLECTION_NAME_AUTO_ID_PARAMS = genCollectionParams({ collectionName: COLLECTION_NAME_AUTO_ID, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: true, }); @@ -62,8 +62,8 @@ describe(`Upsert API`, () => { await milvusClient.createCollection( genCollectionParams({ collectionName: BINARY_COLLECTION_NAME, - dim: 8, - vectorType: DataType.BinaryVector, + dim: [8], + vectorType: [DataType.BinaryVector], autoID: false, }) ); diff --git a/test/grpc/User.spec.ts b/test/grpc/User.spec.ts index baf2ba52..f6c21f16 100644 --- a/test/grpc/User.spec.ts +++ b/test/grpc/User.spec.ts @@ -24,8 +24,8 @@ describe(`User Api`, () => { await authClient.createCollection( genCollectionParams({ collectionName: COLLECTION_NAME, - dim: 4, - vectorType: DataType.FloatVector, + dim: [4], + vectorType: [DataType.FloatVector], autoID: false, }) ); diff --git a/test/http/client.spec.ts b/test/http/client.spec.ts index 97c999e3..b41c5c4d 100644 --- a/test/http/client.spec.ts +++ b/test/http/client.spec.ts @@ -8,7 +8,6 @@ import { describe(`HTTP Client test`, () => { const baseURL = 'http://192.168.0.1:19530'; - const version = 'v3'; const username = 'user'; const password = 'pass'; const token = 'token'; @@ -39,30 +38,25 @@ describe(`HTTP Client test`, () => { ); }); - it('should return the correct baseURL if version is defined', () => { - // Mock configuration object - const config = { - endpoint: baseURL, - version, - }; - - // Create an instance of HttpBaseClient with the mock configuration - const client = new HttpClient(config); - expect(client.baseURL).toBe(`${config.endpoint}/${version}`); - }); - it('should return the correct authorization header', () => { // Mock configuration object const config = { baseURL, username, password, + acceptInt64: true, }; // Create an instance of HttpBaseClient with the mock configuration const client = new HttpClient(config); const expectedAuthorization = `Bearer ${config.username}:${config.password}`; expect(client.authorization).toBe(expectedAuthorization); + expect(client.headers).toEqual({ + Authorization: expectedAuthorization, + Accept: 'application/json', + ContentType: 'application/json', + 'Accept-Type-Allow-Int64': 'true', + }); const config2 = { baseURL, diff --git a/test/http/test.ts b/test/http/test.ts index fbebaae0..2e4935ad 100644 --- a/test/http/test.ts +++ b/test/http/test.ts @@ -4,6 +4,7 @@ import { DEFAULT_VECTOR_FIELD, HttpClientConfig, MilvusClient, + DEFAULT_DB, } from '../../milvus'; import { genCollectionParams, @@ -30,30 +31,83 @@ export function generateTests( // Mock configuration object const createParams = { dimension: 4, + dbName: config.database, collectionName: 'my_collection', metricType: 'L2', - primaryField: 'id', - vectorField: 'vector', - description: 'des', + primaryFieldName: 'id', + vectorFieldName: 'vector', }; const createDefaultParams = { + dbName: config.database, collectionName: 'default_collection_name', dimension: 128, }; + const createAliasParams = { + dbName: config.database, + collectionName: createParams.collectionName, + aliasName: 'my_alias', + }; + + const createUserParams = { + userName: `${config.database ?? DEFAULT_DB}_user`, + password: 'user1234', + }; + + const roleParams = { + roleName: `${config.database ?? DEFAULT_DB}_readOnly`, + objectType: 'Collection', + objectName: '*', + privilege: 'Search', + }; + + const createCustomSetupParams = { + collectionName: 'custom_setup_indexed', + schema: { + autoId: false, + enabledDynamicField: false, + fields: [ + { + fieldName: 'my_id', + dataType: 'Int64', + isPrimary: true, + }, + { + fieldName: 'my_vector', + dataType: 'FloatVector', + elementTypeParams: { + dim: 5, + }, + }, + ], + }, + }; + + const createIndexParams = { + metricType: 'L2', + fieldName: 'my_vector', + indexName: 'my_vector', + params: { + index_type: 'IVF_FLAT', + nlist: 128, + }, + }; + + const importFile = '/d1782fa1-6b65-4ff3-b05a-43a436342445/1.json'; + const count = 100; const data = generateInsertData( [ ...genCollectionParams({ collectionName: createParams.collectionName, - dim: createParams.dimension, + dim: [createParams.dimension], enableDynamic: true, }).fields, ...dynamicFields, ], count - ); + ).map((item, index) => ({ ...item, id: index + 1 })); // Create an instance of HttpBaseClient with the mock configuration const client = new HttpClient(config); @@ -72,17 +126,17 @@ export function generateTests( it('should describe collection successfully', async () => { const describe = await client.describeCollection({ + dbName: config.database, collectionName: createParams.collectionName, }); expect(describe.code).toEqual(200); expect(describe.data.collectionName).toEqual(createParams.collectionName); - expect(describe.data.description).toEqual(createParams.description); expect(describe.data.shardsNum).toEqual(1); expect(describe.data.enableDynamicField).toEqual(true); expect(describe.data.fields.length).toEqual(2); expect(describe.data.indexes[0].fieldName).toEqual( - createParams.vectorField + createParams.vectorFieldName ); expect(describe.data.indexes[0].metricType).toEqual( createParams.metricType @@ -91,6 +145,7 @@ export function generateTests( it('should describe default collection successfully', async () => { const describe = await client.describeCollection({ + dbName: createDefaultParams.dbName, collectionName: createDefaultParams.collectionName, }); @@ -106,7 +161,7 @@ export function generateTests( }); it('should list collections successfully', async () => { - const list = await client.listCollections(); + const list = await client.listCollections({ dbName: config.database }); expect(list.code).toEqual(200); expect(list.data.indexOf(createParams.collectionName) !== -1).toEqual( true @@ -123,16 +178,23 @@ export function generateTests( expect(insert.data.insertCount).toEqual(count); }); - // it('should upsert data successfully', async () => { - // const upsert = await client.upsert({ - // collectionName: createParams.collectionName, - // data: data, - // }); + it('should upsert data successfully', async () => { + const { data } = await client.query({ + collectionName: createParams.collectionName, + filter: 'id > 0', + limit: 1, + outputFields: ['*'], + }); + const target = data[0]; + const upsert = await client.upsert({ + collectionName: createParams.collectionName, + data: [{ ...target, int64: 0 }], + }); - // console.log(upsert); - // expect(upsert.code).toEqual(200); - // expect(upsert.data.insertCount).toEqual(count); - // }); + expect(upsert.code).toEqual(200); + expect(upsert.data.upsertCount).toEqual(1); + expect(upsert.data.upsertIds).toEqual([target.id]); + }); it('should query data and get data and delete successfully', async () => { const query = await client.query({ @@ -156,7 +218,7 @@ export function generateTests( const del = await client.delete({ collectionName: createParams.collectionName, - id: ids, + filter: `id in [${ids.join(',')}]`, }); expect(del.code).toEqual(200); }); @@ -165,7 +227,7 @@ export function generateTests( const search = await client.search({ collectionName: createParams.collectionName, outputFields: ['*'], - vector: [1, 2, 3, 4], + data: [[1, 2, 3, 4]], limit: 5, }); @@ -174,6 +236,263 @@ export function generateTests( expect(typeof search.data[0].distance).toEqual('number'); }); + it('should hasCollection successfully', async () => { + const has = await client.hasCollection({ + dbName: config.database ?? DEFAULT_DB, + collectionName: createParams.collectionName, + }); + + expect(has.code).toEqual(200); + expect(has.data.has).toEqual(true); + }); + + it('should rename collection successfully', async () => { + const newCollectionName = 'new_collection_name'; + const rename = await client.renameCollection({ + dbName: config.database, + collectionName: createParams.collectionName, + newCollectionName, + }); + + const describe = await client.describeCollection({ + dbName: config.database, + collectionName: newCollectionName, + }); + + expect(rename.code).toEqual(200); + expect(describe.code).toEqual(200); + expect(describe.data.collectionName).toEqual(newCollectionName); + + await client.renameCollection({ + dbName: config.database, + collectionName: newCollectionName, + newCollectionName: createParams.collectionName, + }); + }); + + it('should release collection successfully', async () => { + const release = await client.releaseCollection({ + dbName: config.database, + collectionName: createParams.collectionName, + }); + + expect(release.code).toEqual(200); + }); + + it('should load collection successfully', async () => { + const load = await client.loadCollection({ + dbName: config.database, + collectionName: createParams.collectionName, + }); + + expect(load.code).toEqual(200); + }); + + it('should getCollectionStatistics successfully', async () => { + const statistics = await client.getCollectionStatistics({ + dbName: config.database, + collectionName: createParams.collectionName, + }); + + expect(statistics.code).toEqual(200); + expect(statistics.data.rowCount).toEqual(0); + }); + + it('should getCollectionLoadState successfully', async () => { + const state = await client.getCollectionLoadState({ + dbName: config.database, + collectionName: createParams.collectionName, + }); + + expect(state.code).toEqual(200); + expect(state.data.loadState).toMatch('LoadState'); + expect(state.data.loadProgress).toBeLessThanOrEqual(100); + }); + + /* test index operations */ + it('should list indexes successfully', async () => { + const list = await client.listIndexes({ + dbName: config.database, + collectionName: createParams.collectionName, + }); + expect(list.code).toEqual(200); + expect(list.data.length).toEqual(1); + expect(list.data[0]).toEqual(createParams.vectorFieldName); + }); + + it('should create and describe index successfully', async () => { + await client.createCollection(createCustomSetupParams); + const create = await client.createIndex({ + indexParams: [createIndexParams], + collectionName: createCustomSetupParams.collectionName, + }); + const describe = await client.describeIndex({ + collectionName: createCustomSetupParams.collectionName, + indexName: createIndexParams.indexName, + }); + expect(create.code).toEqual(200); + expect(describe.code).toEqual(200); + expect(describe.data[0].indexName).toEqual(createIndexParams.indexName); + expect(describe.data[0].indexType).toEqual( + createIndexParams.params.index_type + ); + }); + + it('should drop index successfully', async () => { + const { collectionName } = createCustomSetupParams; + await client.releaseCollection({ collectionName }); + const drop = await client.dropIndex({ + dbName: config.database, + collectionName, + indexName: createIndexParams.indexName, + }); + await client.dropCollection({ collectionName }); + expect(drop.code).toEqual(200); + }); + + /* test alias operations */ + it('should create alias successfully', async () => { + const create = await client.createAlias(createAliasParams); + expect(create.code).toEqual(200); + }); + + it('should describe alias successfully', async () => { + /** + * https://github.com/milvus-io/milvus/issues/31978 + * TODO: Alias describe api has issue，temporarily comment + */ + + // const describe = await client.describeAlias({ + // dbName: config.database ?? DEFAULT_DB, + // aliasName: createAliasParams.aliasName, + // }); + + // expect(describe.code).toEqual(200); + // expect(describe.data.aliasName).toEqual(createAliasParams.aliasName); + // expect(describe.data.collectionName).toEqual( + // createAliasParams.collectionName + // ); + }); + + it('should list aliases successfully', async () => { + const list = await client.listAliases({ + dbName: config.database ?? DEFAULT_DB, + }); + expect(list.code).toEqual(200); + expect(list.data.length).toBeGreaterThanOrEqual(1); + expect(list.data[0]).toEqual(createAliasParams.aliasName); + }); + + it('should alter alias successfully', async () => { + const newCollectionName = 'new_collection'; + await client.createCollection({ + ...createParams, + collectionName: newCollectionName, + }); + const alter = await client.alterAlias({ + dbName: config.database ?? DEFAULT_DB, + collectionName: newCollectionName, + aliasName: createAliasParams.aliasName, + }); + await client.dropCollection({ collectionName: newCollectionName }); + expect(alter.code).toEqual(200); + }); + + it('should drop alias successfully', async () => { + const list = await client.listAliases({ + dbName: config.database ?? DEFAULT_DB, + }); + if (list.data.includes(createAliasParams.aliasName)) { + const drop = await client.dropAlias({ + dbName: config.database ?? DEFAULT_DB, + aliasName: createAliasParams.aliasName, + }); + expect(drop.code).toEqual(200); + } + }); + + /* test partition operations */ + it('should list partitions successfully', async () => { + const list = await client.listPartitions({ + collectionName: createParams.collectionName, + }); + expect(list.code).toEqual(200); + expect(list.data.length).toEqual(1); + expect(list.data[0]).toEqual('_default'); + }); + + it('should create partition successfully', async () => { + const create = await client.createPartition({ + collectionName: createParams.collectionName, + partitionName: 'my_partition', + }); + expect(create.code).toEqual(200); + }); + + it('should load partitions successfully', async () => { + const load = await client.loadPartitions({ + collectionName: createParams.collectionName, + partitionNames: ['my_partition'], + }); + expect(load.code).toEqual(200); + }); + + it('should release partitions successfully', async () => { + const release = await client.releasePartitions({ + collectionName: createParams.collectionName, + partitionNames: ['my_partition'], + }); + expect(release.code).toEqual(200); + }); + + it('should has partition successfully', async () => { + const has = await client.hasPartition({ + collectionName: createParams.collectionName, + partitionName: 'my_partition', + }); + expect(has.code).toEqual(200); + expect(has.data.has).toEqual(true); + }); + + it('should get partitions statistics successfully', async () => { + const statistics = await client.getPartitionStatistics({ + collectionName: createParams.collectionName, + partitionName: 'my_partition', + }); + expect(statistics.code).toEqual(200); + expect(statistics.data.rowCount).toEqual(0); + }); + + it('should drop partition successfully', async () => { + const drop = await client.dropPartition({ + collectionName: createParams.collectionName, + partitionName: 'my_partition', + }); + expect(drop.code).toEqual(200); + }); + + /* test import operations */ + it('should import jobs successfully', async () => { + const create = await client.createImportJobs({ + collectionName: createParams.collectionName, + files: [[importFile]], + }); + const jobId = create.data.jobId; + const list = await client.listImportJobs({ + collectionName: createParams.collectionName, + }); + const job = list.data.records.find(j => j.jobId === jobId); + const progress = await client.getImportJobProgress({ jobId }); + expect(create.code).toEqual(200); + expect(list.code).toEqual(200); + if (job) { + expect(job.collectionName).toEqual(createParams.collectionName); + expect(job.progress).toBeLessThanOrEqual(100); + } + expect(progress.code).toEqual(200); + expect(progress.data.jobId).toEqual(jobId); + }); + it('should drop collection successfully', async () => { const drop = await client.dropCollection({ collectionName: createParams.collectionName, @@ -186,5 +505,81 @@ export function generateTests( }); expect(dropDefault.code).toEqual(200); }); + + /* test role operations */ + it('should list roles successfully', async () => { + const list = await client.listRoles(); + expect(list.code).toEqual(200); + expect(list.data.length).toBeGreaterThanOrEqual(1); + }); + + it('should describe role successfully', async () => { + const describe = await client.describeRole({ roleName: 'public' }); + expect(describe.code).toEqual(200); + }); + + it('should create role successfully', async () => { + const create = await client.createRole({ + roleName: roleParams.roleName, + }); + expect(create.code).toEqual(200); + }); + + it('should drop role successfully', async () => { + const drop = await client.dropRole({ roleName: roleParams.roleName }); + expect(drop.code).toEqual(200); + }); + + /* test user operations */ + it('should create user successfully', async () => { + const create = await client.createUser(createUserParams); + expect(create.code).toEqual(200); + }); + + it('should list users successfully', async () => { + const list = await client.listUsers(); + expect(list.code).toEqual(200); + expect(list.data.length).toBeGreaterThanOrEqual(1); + }); + + it('should describe user successfully', async () => { + const describe = await client.describeUser({ + userName: createUserParams.userName, + }); + expect(describe.code).toEqual(200); + }); + + it('should update user password successfully', async () => { + const newPassword = 'test_new_password'; + const update = await client.updateUserPassword({ + userName: createUserParams.userName, + password: createUserParams.password, + newPassword, + }); + expect(update.code).toEqual(200); + }); + + it('should grant role to user successfully', async () => { + const grant = await client.grantRoleToUser({ + userName: createUserParams.userName, + roleName: 'public', + }); + expect(grant.code).toEqual(200); + }); + + it('should revoke role from user successfully', async () => { + const revoke = await client.revokeRoleFromUser({ + userName: createUserParams.userName, + roleName: 'public', + }); + expect(revoke.code).toEqual(200); + }); + + it('should drop user successfully', async () => { + const drop = await client.dropUser({ + userName: createUserParams.userName, + }); + expect(drop.code).toEqual(200); + }); }); } diff --git a/test/tools/bench.ts b/test/tools/bench.ts index 3329d64f..65856e14 100644 --- a/test/tools/bench.ts +++ b/test/tools/bench.ts @@ -70,7 +70,7 @@ const COLLECTION_NAME = 'bench_milvus'; console.time('Search time'); const search = await milvusClient.search({ collection_name: COLLECTION_NAME, - vector: vectorsData[i]['vector'], + data: vectorsData[i]['vector'], output_fields: ['id', 'int64', 'varChar'], limit: 5, }); diff --git a/test/tools/collection.ts b/test/tools/collection.ts index 51086f83..46b75bc3 100644 --- a/test/tools/collection.ts +++ b/test/tools/collection.ts @@ -1,5 +1,6 @@ import { DataType, ConsistencyLevelEnum } from '../../milvus'; import { VECTOR_FIELD_NAME, MAX_CAPACITY, MAX_LENGTH } from './const'; +import { GENERATE_VECTOR_NAME } from './'; export const dynamicFields = [ { @@ -31,8 +32,8 @@ export const dynamicFields = [ */ export const genCollectionParams = (data: { collectionName: string; - dim: number | string; - vectorType?: DataType; + dim: number[] | string[]; + vectorType?: DataType[]; autoID?: boolean; fields?: any[]; partitionKeyEnabled?: boolean; @@ -42,8 +43,8 @@ export const genCollectionParams = (data: { }) => { const { collectionName, - dim, - vectorType = DataType.FloatVector, + dim = [8], + vectorType = [DataType.FloatVector], autoID = true, fields = [], partitionKeyEnabled, @@ -52,16 +53,25 @@ export const genCollectionParams = (data: { maxCapacity, } = data; + const vectorFields = vectorType.map((type, i) => { + const res: any = { + name: GENERATE_VECTOR_NAME(i), + description: `vector type: ${type}`, + data_type: type, + }; + + if (type !== DataType.SparseFloatVector) { + res.dim = Number(dim[i]); + } + + return res; + }); + const params: any = { collection_name: collectionName, consistency_level: ConsistencyLevelEnum.Strong, fields: [ - { - name: VECTOR_FIELD_NAME, - description: 'Vector field', - data_type: vectorType, - dim: Number(dim), - }, + ...vectorFields, { name: 'id', description: 'ID field', diff --git a/test/tools/data.ts b/test/tools/data.ts index 0a9c725f..f1c4dfb2 100644 --- a/test/tools/data.ts +++ b/test/tools/data.ts @@ -3,6 +3,7 @@ import { FieldData, convertToDataType, FieldType, + SparseVectorCOO, } from '../../milvus'; import { MAX_LENGTH, P_KEY_VALUES } from './const'; import Long from 'long'; @@ -16,6 +17,7 @@ interface DataGenerator { max_capacity?: number; is_partition_key?: boolean; index?: number; + sparseType?: string; }): FieldData; } @@ -30,6 +32,13 @@ interface DataGenerator { export const genVarChar: DataGenerator = params => { const { max_length = MAX_LENGTH, is_partition_key = false, index } = params!; + const chance = Math.random(); + if (chance < 0.2) { + return 'apple'; + } else if (chance < 0.5) { + return 'orange'; + } + if (!is_partition_key) { let result = ''; const characters = @@ -141,6 +150,101 @@ export const genInt64: DataGenerator = () => { return Long.fromBits(low, high, true); // true for unsigned }; +// generate random sparse vector +// for example {2: 0.5, 3: 0.3, 4: 0.2} +export const genSparseVector: DataGenerator = params => { + const dim = params!.dim || 24; + const sparseType = params!.sparseType || 'object'; + const nonZeroCount = Math.floor(Math.random() * dim!) || 4; + + switch (sparseType) { + case 'array': + /* + const sparseArray = [ + undefined, + 0.0, + 0.5, + 0.3, + undefined, + 0.2] + */ + const sparseArray = Array.from({ length: dim! }, () => Math.random()); + for (let i = 0; i < nonZeroCount; i++) { + sparseArray[Math.floor(Math.random() * dim!)] = undefined as any; + } + return sparseArray; + + case 'csr': + /* + const sparseCSR = { + indices: [2, 5, 8], + values: [5, 3, 7] + }; + */ + const indicesSet = new Set(); + const csr = { + indices: Array.from({ length: nonZeroCount }, () => { + let index: number; + do { + index = Math.floor(Math.random() * dim!); + } while (indicesSet.has(index)); + indicesSet.add(index); + return index; + }).sort((a, b) => a - b), + values: Array.from({ length: nonZeroCount }, (_, i) => Math.random()), + }; + return csr; + + case 'coo': + /* + const sparseCOO = [ + { index: 2, value: 5 }, + { index: 5, value: 3 }, + { index: 8, value: 7 } + ]; + */ + const coo: SparseVectorCOO = []; + const indexSet = new Set(); + + while (coo.length < nonZeroCount) { + const index = Math.floor(Math.random() * dim!); + if (!indexSet.has(index)) { + coo.push({ + index: index, + value: Math.random(), + }); + indexSet.add(index); + } + } + + // sort by index + coo.sort((a, b) => a.index - b.index); + return coo; + + default: // object + /* + const sparseObject = { + 3: 1.5, + 6: 2.0, + 9: -3.5 + }; + */ + const sparseObject: { [key: number]: number } = {}; + for (let i = 0; i < nonZeroCount; i++) { + sparseObject[Math.floor(Math.random() * dim!)] = Math.random(); + } + return sparseObject; + } +}; + +export const genFloat16: DataGenerator = params => { + const float32Array = genFloatVector(params); + // console.log('origin float32array', float32Array); + // const float16Array = new Float16Array(float32Array as number[]); + // const float16Bytes = new Uint8Array(float16Array.buffer); + return float32Array; +}; + export const dataGenMap: { [key in DataType]: DataGenerator } = { [DataType.None]: genNone, [DataType.Bool]: genBool, @@ -155,6 +259,9 @@ export const dataGenMap: { [key in DataType]: DataGenerator } = { [DataType.JSON]: genJSON, [DataType.BinaryVector]: genBinaryVector, [DataType.FloatVector]: genFloatVector, + [DataType.Float16Vector]: genFloat16, + [DataType.BFloat16Vector]: genFloat16, + [DataType.SparseFloatVector]: genSparseVector, }; /** @@ -163,7 +270,11 @@ export const dataGenMap: { [key in DataType]: DataGenerator } = { * @param count The number of data points to generate * @returns An array of objects representing the generated data */ -export const generateInsertData = (fields: FieldType[], count: number = 10) => { +export const generateInsertData = ( + fields: FieldType[], + count: number = 10, + options?: { sparseType: string } +) => { const rows: { [x: string]: any }[] = []; // Initialize an empty array to store the generated data // Loop until we've generated the desired number of data points @@ -199,6 +310,7 @@ export const generateInsertData = (fields: FieldType[], count: number = 10) => { ), index: count, is_partition_key: field.is_partition_key, + sparseType: options && options.sparseType, }; // Generate data diff --git a/test/tools/utils.ts b/test/tools/utils.ts index 6ba7e2f6..18bf8495 100644 --- a/test/tools/utils.ts +++ b/test/tools/utils.ts @@ -1,3 +1,4 @@ +import {VECTOR_FIELD_NAME} from './' /** * Generates a random collection name with a prefix and a random string appended to it. * @param {string} [pre='collection'] - The prefix to use for the collection name. @@ -21,3 +22,8 @@ export const timeoutTest = (func: Function, args?: { [x: string]: any }) => { } }; }; + +export const GENERATE_VECTOR_NAME = (i = 0) => { + return i === 0 ? VECTOR_FIELD_NAME : `${VECTOR_FIELD_NAME}${i}` +} + \ No newline at end of file diff --git a/test/utils/Bytes.spec.ts b/test/utils/Bytes.spec.ts new file mode 100644 index 00000000..b3dec666 --- /dev/null +++ b/test/utils/Bytes.spec.ts @@ -0,0 +1,108 @@ +import { + bytesToSparseRow, + sparseRowsToBytes, + SparseFloatVector, + sparseToBytes, + getSparseFloatVectorType, + f32ArrayToF16Bytes, + f16BytesToF32Array, + f32ArrayToBf16Bytes, + bf16BytesToF32Array, +} from '../../milvus'; + +describe('Data <-> Bytes Test', () => { + it('should throw error if index is negative or exceeds 2^32-1', () => { + const invalidIndexData = { + 0: 1.5, + 4294967296: 2.7, // 2^32 + }; + expect(() => sparseToBytes(invalidIndexData)).toThrow(); + }); + + it('should return empty Uint8Array if data is empty', () => { + expect(sparseToBytes({})).toEqual(new Uint8Array(0)); + }); + + it('Conversion is reversible', () => { + const inputSparseRows = [ + { '12': 0.875, '17': 0.789, '19': 0.934 }, + ] as SparseFloatVector[]; + + const bytesArray = sparseRowsToBytes(inputSparseRows); + + const outputSparseRow = bytesToSparseRow(Buffer.concat(bytesArray)); + + const originKeys = Object.keys(inputSparseRows[0]); + const originValues = Object.values(inputSparseRows[0]); + const outputKeys = Object.keys(outputSparseRow); + const outputValues = Object.values(outputSparseRow); + + expect(originKeys).toEqual(outputKeys); + + originValues.forEach((value, index) => { + expect(value).toBeCloseTo(outputValues[index]); + }); + }); + + it('should return "array" if the input is an empty array', () => { + const data: any[] = []; + expect(getSparseFloatVectorType(data)).toEqual('array'); + }); + + it('should return "dict" if the input is an object', () => { + const data = { '12': 0.875, '17': 0.789, '19': 0.934 }; + expect(getSparseFloatVectorType(data)).toEqual('dict'); + }); + + it('should return "csr" if the input is an object with "indices" and "values"', () => { + const data = { indices: [12, 17, 19], values: [0.875, 0.789, 0.934] }; + expect(getSparseFloatVectorType(data)).toEqual('csr'); + }); + + it('should return "array" if the input is an array', () => { + const data = [0.875, 0.789, 0.934]; + expect(getSparseFloatVectorType(data)).toEqual('array'); + }); + + it('should return "coo" if the input is an array of objects with "index" and "value"', () => { + const data = [ + { index: 12, value: 0.875 }, + { index: 17, value: 0.789 }, + { index: 19, value: 0.934 }, + ]; + 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'); + }); + + it('should transform f16b -> f32 and f32 -> f16b successfully', () => { + const data = [0.123456789, -0.987654321, 3.14159265]; + const f16Bytes = f32ArrayToF16Bytes(data); + const f32Array = f16BytesToF32Array(f16Bytes); + + expect(f32Array.length).toEqual(data.length); + for (let i = 0; i < data.length; i++) { + expect(data[i]).toBeCloseTo(f32Array[i], 2); + } + }); + + it('should transform bf16b -> f32 and f32 -> bf16b successfully', () => { + const data = [0.123456789, -0.987654321, 3.14159265]; + const bf16Bytes = f32ArrayToBf16Bytes(data); + const f32Array = bf16BytesToF32Array(bf16Bytes); + + expect(f32Array.length).toEqual(data.length); + for (let i = 0; i < data.length; i++) { + expect(data[i]).toBeCloseTo(f32Array[i], 2); + } + }); +}); diff --git a/test/utils/Format.spec.ts b/test/utils/Format.spec.ts index e059d71d..c4a9eb7a 100644 --- a/test/utils/Format.spec.ts +++ b/test/utils/Format.spec.ts @@ -24,7 +24,9 @@ import { buildDynamicRow, getAuthString, buildFieldData, + formatSearchResult, Field, + formatSearchVector, } from '../../milvus'; describe('utils/format', () => { @@ -76,6 +78,10 @@ describe('utils/format', () => { const testValue = 3.1231241241234124124; const res = formatNumberPrecision(testValue, 3); expect(res).toBe(3.123); + + const testValue2 = -3.1231241241234124124; + const res2 = formatNumberPrecision(testValue2, 3); + expect(res2).toBe(-3.123); }); it(`hybridtsToUnixtime should success`, async () => { @@ -141,12 +147,12 @@ describe('utils/format', () => { it('does not throw an error if vectors or vector is defined', () => { const data1 = { collection_name: 'my_collection', - vectors: [[]], + data: [[]], }; const data2 = { collection_name: 'my_collection', - vector: [], + data: [], }; expect(() => checkSearchParams(data1)).not.toThrow(); @@ -367,6 +373,7 @@ describe('utils/format', () => { is_primary_key: false, description: 'vector field', data_type: 'FloatVector', + dataType: 101, autoID: false, state: 'created', }, @@ -378,6 +385,7 @@ describe('utils/format', () => { is_primary_key: true, description: '', data_type: 'Int64', + dataType: 5, autoID: true, state: 'created', }, @@ -535,4 +543,144 @@ describe('utils/format', () => { const field = { type: 'Int', name: 'age' }; expect(buildFieldData(row, field as Field)).toEqual(25); }); + + it('should format search results correctly', () => { + const searchPromise: any = { + results: { + fields_data: [ + { + type: 'Int64', + field_name: 'id', + field_id: '101', + is_dynamic: false, + scalars: { + long_data: { data: ['98286', '40057', '5878', '96232'] }, + data: 'long_data', + }, + field: 'scalars', + }, + ], + scores: [ + 14.632697105407715, 15.0767822265625, 15.287022590637207, + 15.357033729553223, + ], + topks: ['4'], + output_fields: ['id'], + num_queries: '1', + top_k: '4', + ids: { + int_id: { data: ['98286', '40057', '5878', '96232'] }, + id_field: 'int_id', + }, + group_by_field_value: null, + }, + }; + + const options = { round_decimal: 2 }; + + const expectedResults = [ + [ + { score: 14.63, id: '98286' }, + { score: 15.07, id: '40057' }, + { score: 15.28, id: '5878' }, + { score: 15.35, id: '96232' }, + ], + ]; + + const results = formatSearchResult(searchPromise, options); + + expect(results).toEqual(expectedResults); + }); + + it('should format search vector correctly', () => { + // float vector + const floatVector = [1, 2, 3]; + const formattedVector = formatSearchVector( + floatVector, + DataType.FloatVector + ); + expect(formattedVector).toEqual([floatVector]); + + const floatVectors = [ + [1, 2, 3], + [4, 5, 6], + ]; + expect(formatSearchVector(floatVectors, DataType.FloatVector)).toEqual( + floatVectors + ); + }); + + // sparse coo vector + const sparseCooVector = [ + { index: 1, value: 2 }, + { index: 3, value: 4 }, + ]; + const formattedSparseCooVector = formatSearchVector( + sparseCooVector, + DataType.SparseFloatVector + ); + expect(formattedSparseCooVector).toEqual([sparseCooVector]); + + // sparse csr vector + const sparseCsrVector = { + indices: [1, 3], + values: [2, 4], + }; + const formattedSparseCsrVector = formatSearchVector( + sparseCsrVector, + DataType.SparseFloatVector + ); + expect(formattedSparseCsrVector).toEqual([sparseCsrVector]); + + const sparseCsrVectors = [ + { + indices: [1, 3], + values: [2, 4], + }, + { + indices: [2, 4], + values: [3, 5], + }, + ]; + const formattedSparseCsrVectors = formatSearchVector( + sparseCsrVectors, + DataType.SparseFloatVector + ); + expect(formattedSparseCsrVectors).toEqual(sparseCsrVectors); + + // sparse array vector + const sparseArrayVector = [0.1, 0.2, 0.3]; + const formattedSparseArrayVector = formatSearchVector( + sparseArrayVector, + DataType.SparseFloatVector + ); + expect(formattedSparseArrayVector).toEqual([sparseArrayVector]); + + const sparseArrayVectors = [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + ]; + const formattedSparseArrayVectors = formatSearchVector( + sparseArrayVectors, + DataType.SparseFloatVector + ); + expect(formattedSparseArrayVectors).toEqual(sparseArrayVectors); + + // sparse dict vector + const sparseDictVector = { 1: 2, 3: 4 }; + const formattedSparseDictVector = formatSearchVector( + sparseDictVector, + DataType.SparseFloatVector + ); + expect(formattedSparseDictVector).toEqual([sparseDictVector]); + + const sparseDictVectors = [ + { 1: 2, 3: 4 }, + { 1: 2, 3: 4 }, + ]; + const formattedSparseDictVectors = formatSearchVector( + sparseDictVectors, + DataType.SparseFloatVector + ); + expect(formattedSparseDictVectors).toEqual(sparseDictVectors); }); diff --git a/test/utils/Function.spec.ts b/test/utils/Function.spec.ts index 7cb61466..793a6082 100644 --- a/test/utils/Function.spec.ts +++ b/test/utils/Function.spec.ts @@ -1,4 +1,10 @@ -import { promisify } from '../../milvus/utils'; +import { + promisify, + getSparseDim, + SparseFloatVector, + getDataKey, + DataType, +} from '../../milvus'; describe('promisify', () => { let pool: any; @@ -46,4 +52,72 @@ describe('promisify', () => { expect(pool.acquire).toHaveBeenCalled(); expect(pool.release).toHaveBeenCalled(); }); + + it('should return the correct dimension of the sparse vector', () => { + const data = [ + { '0': 1, '1': 2, '2': 3 }, + { '0': 1, '1': 2, '2': 3, '3': 4 }, + { '0': 1, '1': 2 }, + ] as SparseFloatVector[]; + const result = getSparseDim(data); + expect(result).toBe(4); + }); + + it('should return 0 for an empty array', () => { + const data = [] as SparseFloatVector[]; + const result = getSparseDim(data); + expect(result).toBe(0); + }); + + it('should return the correct dimension when the sparse vectors have different lengths', () => { + const data = [ + { '0': 1, '1': 2, '2': 3, '3': 4, '4': 5 }, + { '0': 1, '1': 2 }, + { '0': 1, '1': 2, '2': 3, '3': 4 }, + ] as SparseFloatVector[]; + const result = getSparseDim(data); + expect(result).toBe(5); + }); + + it('should return the correct data key for each data type without camel case conversion', () => { + expect(getDataKey(DataType.FloatVector)).toEqual('float_vector'); + expect(getDataKey(DataType.Float16Vector)).toEqual('float16_vector'); + expect(getDataKey(DataType.BFloat16Vector)).toEqual('bfloat16_vector'); + expect(getDataKey(DataType.BinaryVector)).toEqual('binary_vector'); + expect(getDataKey(DataType.SparseFloatVector)).toEqual( + 'sparse_float_vector' + ); + expect(getDataKey(DataType.Double)).toEqual('double_data'); + expect(getDataKey(DataType.Float)).toEqual('float_data'); + expect(getDataKey(DataType.Int64)).toEqual('long_data'); + expect(getDataKey(DataType.Int32)).toEqual('int_data'); + expect(getDataKey(DataType.Int16)).toEqual('int_data'); + expect(getDataKey(DataType.Int8)).toEqual('int_data'); + expect(getDataKey(DataType.Bool)).toEqual('bool_data'); + expect(getDataKey(DataType.VarChar)).toEqual('string_data'); + expect(getDataKey(DataType.Array)).toEqual('array_data'); + expect(getDataKey(DataType.JSON)).toEqual('json_data'); + expect(getDataKey(DataType.None)).toEqual('none'); + }); + + it('should return the correct data key for each data type with camel case conversion', () => { + expect(getDataKey(DataType.FloatVector, true)).toEqual('floatVector'); + expect(getDataKey(DataType.Float16Vector, true)).toEqual('float16Vector'); + expect(getDataKey(DataType.BFloat16Vector, true)).toEqual('bfloat16Vector'); + expect(getDataKey(DataType.BinaryVector, true)).toEqual('binaryVector'); + expect(getDataKey(DataType.SparseFloatVector, true)).toEqual( + 'sparseFloatVector' + ); + expect(getDataKey(DataType.Double, true)).toEqual('doubleData'); + expect(getDataKey(DataType.Float, true)).toEqual('floatData'); + expect(getDataKey(DataType.Int64, true)).toEqual('longData'); + expect(getDataKey(DataType.Int32, true)).toEqual('intData'); + expect(getDataKey(DataType.Int16, true)).toEqual('intData'); + expect(getDataKey(DataType.Int8, true)).toEqual('intData'); + expect(getDataKey(DataType.Bool, true)).toEqual('boolData'); + expect(getDataKey(DataType.VarChar, true)).toEqual('stringData'); + expect(getDataKey(DataType.Array, true)).toEqual('arrayData'); + expect(getDataKey(DataType.JSON, true)).toEqual('jsonData'); + expect(getDataKey(DataType.None, true)).toEqual('none'); + }); }); diff --git a/test/utils/Test.spec.ts b/test/utils/Test.spec.ts index 15c6f1f4..8d1e41a1 100644 --- a/test/utils/Test.spec.ts +++ b/test/utils/Test.spec.ts @@ -1,36 +1,57 @@ -import { generateInsertData, genCollectionParams } from '../tools'; -import { DataType } from '../../milvus'; +import { + generateInsertData, + genCollectionParams, + genSparseVector, +} from '../tools'; +import { + DataType, + SparseVectorArray, + SparseVectorDic, + SparseVectorCSR, + SparseVectorCOO, +} from '../../milvus'; describe(`utils/test`, () => { it('should generate data for schema created by genCollectionParams', () => { const param = genCollectionParams({ collectionName: 't', - dim: 10, + dim: [10], }); const data = generateInsertData(param.fields, 10); expect(data.length).toBe(10); expect(data[0].vector.length).toBe(10); }); + it('should generate multiple vector types for schema created by genCollectionParams', () => { + const param = genCollectionParams({ + collectionName: 't', + vectorType: [DataType.FloatVector, DataType.FloatVector], + dim: [10, 16], + }); + expect(param.fields); + const floatVectorFields = param.fields.filter( + (field: any) => field.data_type === DataType.FloatVector + ); + expect(floatVectorFields.length).toBe(2); + expect(floatVectorFields.some((field: any) => field.dim === 10)).toBe(true); + expect(floatVectorFields.some((field: any) => field.dim === 16)).toBe(true); + }); + it('should generate data for a collection with a vector field of type DataType.FloatVector', () => { - const fields = [ - { - name: 'vector', - description: 'vector field', - data_type: DataType.FloatVector, - dim: 10, - }, - { - name: 'id', - description: '', - data_type: DataType.Int64, - is_primary_key: true, - autoID: true, - }, - ]; - const data = generateInsertData(fields, 10); + const param = genCollectionParams({ + collectionName: 't', + vectorType: [ + DataType.FloatVector, + DataType.FloatVector, + DataType.BinaryVector, + ], + dim: [10, 10, 16], + }); + const data = generateInsertData(param.fields, 10); expect(data.length).toBe(10); expect(data[0].vector.length).toBe(10); + expect(data[0].vector1.length).toBe(10); + expect(data[0].vector2.length).toBe(2); }); it('should generate data for a collection with a vector field of type DataType.BinaryVector', () => { @@ -124,4 +145,104 @@ describe(`utils/test`, () => { expect(data.length).toBe(10); expect(typeof data[0].int_field).toBe('number'); }); + + it('should generate a sparse vector with default parameters', () => { + const fields = [ + { + name: 'sparse_vector', + description: '', + data_type: DataType.SparseFloatVector, + is_primary_key: true, + dim: 8, + }, + { + name: 'sparse_vector1', + description: '', + data_type: DataType.SparseFloatVector, + is_primary_key: true, + dim: 24, + }, + { + name: 'sparse_vector2', + description: '', + data_type: DataType.SparseFloatVector, + }, + ]; + + const data = generateInsertData(fields, 10); + expect(data.length).toBe(10); + data.forEach(d => { + expect( + Object.keys(d.sparse_vector).every(d => typeof d === 'string') + ).toBe(true); + expect( + Object.keys(d.sparse_vector1).every(d => typeof d === 'string') + ).toBe(true); + expect( + Object.keys(d.sparse_vector2).every(d => typeof d === 'string') + ).toBe(true); + }); + }); + + it('Generate sparse array vector', () => { + const params = { sparseType: 'array', dim: 24 } as any; + const sparseArray = genSparseVector(params) as SparseVectorArray; + expect(Array.isArray(sparseArray)).toBe(true); + expect(sparseArray.length).toBeLessThanOrEqual(24); + + // test some of items are zero + const nonZeroItems = sparseArray.filter(item => item !== undefined); + expect(nonZeroItems.length).toBeLessThanOrEqual(24); + // test some of items are undefined + const undefinedItems = sparseArray.filter(item => item === undefined); + expect(undefinedItems.length).toBeLessThanOrEqual(24); + }); + + it('Generate CSR sparse vector', () => { + const params = { sparseType: 'csr', dim: 24 } as any; + const csr = genSparseVector(params) as SparseVectorCSR; + expect(csr.hasOwnProperty('indices')).toBe(true); + expect(csr.hasOwnProperty('values')).toBe(true); + expect(Array.isArray(csr.indices)).toBe(true); + // test csr indices should be sorted + const sortedIndices = csr.indices.slice().sort((a, b) => a - b); + // test csr indices should be unique + const uniqueIndices = new Set(csr.indices); + expect(uniqueIndices.size).toBe(csr.indices.length); + expect(csr.indices).toEqual(sortedIndices); + expect(Array.isArray(csr.values)).toBe(true); + expect(csr.indices.length).toBeLessThanOrEqual(24); + expect(csr.values.length).toBeLessThanOrEqual(24); + expect(csr.indices.length).toEqual(csr.values.length); + }); + + it('Generate COO sparse vector', () => { + const params = { sparseType: 'coo', dim: 24 } as any; + const coo = genSparseVector(params) as SparseVectorCOO; + expect(Array.isArray(coo)).toBe(true); + expect(coo.length).toBeLessThanOrEqual(24); + // test every item should has index and value property, and value should be number + coo.forEach(item => { + expect(item.hasOwnProperty('index')).toBe(true); + expect(item.hasOwnProperty('value')).toBe(true); + expect(typeof item.index).toBe('number'); + expect(typeof item.value).toBe('number'); + }); + // test index should be unique + const indices = coo.map(item => item.index); + const uniqueIndices = new Set(indices); + expect(uniqueIndices.size).toBe(indices.length); + }); + + it('Generate dic sparse vector', () => { + const params = { sparseType: 'object', dim: 24 } as any; + const sparseObject = genSparseVector(params) as SparseVectorDic; + expect(typeof sparseObject).toBe('object'); + expect(Object.keys(sparseObject).length).toBeLessThanOrEqual(24); + for (const key in sparseObject) { + expect(parseInt(key, 10)).toBeGreaterThanOrEqual(0); + expect(parseInt(key, 10)).toBeLessThan(24); + expect(typeof sparseObject[key]).toBe('number'); + } + }); }); diff --git a/test/utils/Validate.spec.ts b/test/utils/Validate.spec.ts index 8091a253..f5467fab 100644 --- a/test/utils/Validate.spec.ts +++ b/test/utils/Validate.spec.ts @@ -26,12 +26,12 @@ describe('utils/validate', () => { it('does not throw an error if vectors or vector is defined', () => { const data1 = { collection_name: 'my_collection', - vectors: [[]], + data: [[]], }; const data2 = { collection_name: 'my_collection', - vector: [], + data: [], }; expect(() => checkSearchParams(data1)).not.toThrow(); @@ -330,7 +330,7 @@ describe('utils/validate', () => { }; expect(() => checkCreateCollectionCompatibility(data3)).toThrow( - `Your milvus server doesn't support dynmaic schmea, please upgrade your server.` + `Your milvus server doesn't support dynamic schema, please upgrade your server.` ); }); diff --git a/yarn.lock b/yarn.lock index 7e700986..f33c707a 100644 --- a/yarn.lock +++ b/yarn.lock @@ -698,6 +698,11 @@ "@jridgewell/resolve-uri" "3.1.0" "@jridgewell/sourcemap-codec" "1.4.14" +"@petamoriken/float16@^3.8.6": + version "3.8.6" + resolved "https://registry.yarnpkg.com/@petamoriken/float16/-/float16-3.8.6.tgz#580701cb97a510882342333d31c7cbfd9e14b4f4" + integrity sha512-GNJhABTtcmt9al/nqdJPycwFD46ww2+q2zwZzTjY0dFFwUAFRw9zszvEr9osyJRd9krRGy6hUDopWUg9fX7VVw== + "@protobufjs/aspromise@^1.1.1", "@protobufjs/aspromise@^1.1.2": version "1.1.2" resolved "https://registry.yarnpkg.com/@protobufjs/aspromise/-/aspromise-1.1.2.tgz#9b8b0cc663d669a7d8f6f5d0893a14d348f30fbf"