add: COSINE metric and text sparse search

.gitignore

@@ -17,3 +17,7 @@ Cargo.lock
 /target
 #Cargo.lock
 /.idea/
+
+/volumes
+
+/.vscode/

Cargo.toml

@@ -27,6 +27,7 @@ dashmap = "5.5.3"
 tonic-build = { version = "0.8.2", default-features = false, features = [
     "prost",
 ] }
+prost-build = "0.11.0"
 
 [dev-dependencies]
 rand = "0.8.5"

README.md

@@ -1,15 +1,19 @@
 # Milvus Rust SDK
+
 Rust SDK for Milvus.
 
 **This is still in progress, be careful to use it in your production, and we are looking for active maintianers of this repo**
 
 ## Get Started
+
 Add the SDK into your project:
+
 ```
 cargo add milvus-sdk-rust
 ```
 
 Connect to milvus service and create collection:
+
 ```rust
 #[tokio::main]
 async fn main() -> Result<(), Error> {
@@ -19,16 +23,18 @@ async fn main() -> Result<(), Error> {
 
     let schema =
         CollectionSchemaBuilder::new("hello_milvus", "a guide example for milvus rust SDK")
-            .add_field(FieldSchema::new_primary_int64(
-                "id",
-                "primary key field",
-                true,
-            ))
-            .add_field(FieldSchema::new_float_vector(
-                DEFAULT_VEC_FIELD,
-                "feature field",
-                256,
-            ))
+            .add_field(FieldSchemaBuilder::new()
+                .with_name("id")
+                .with_primary(true)
+                .with_dtype(DataType::Int64)
+                .with_description("primary key field")
+                .build())
+            .add_field(FieldSchemaBuilder::new()
+                .with_name(DEFAULT_VEC_FIELD)
+                .with_dtype(DataType::FloatVector)
+                .with_dim(256)
+                .with_description("feature field")
+                .build())
             .build()?;
     let collection = client.create_collection(schema.clone(), None).await?;
     Ok(())
@@ -38,23 +44,29 @@ async fn main() -> Result<(), Error> {
 ## Development
 
 Pre-requisites:
+
 - cargo
 - protocol-compiler
 - docker (for testing)
 
 ### How to test
+
 Many tests require the Milvus server, the project provide a docker-compose file to setup a Milvus cluster:
+
 ```
 docker-compose -f ./docker-compose.yml up -d
 ```
+
 You may need to wait for seconds until the system ready
 
 Run all tests:
+
 ```
 cargo test
 ```
 
 Enable the full backtrace for debugging:
+
 ```
 RUST_BACKTRACE=1 cargo test
 ```

build.rs

@@ -1,8 +1,12 @@
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let mut config = prost_build::Config::new();
+    config.protoc_arg("--experimental_allow_proto3_optional");
+
     tonic_build::configure()
         .build_server(false)
         .out_dir("src/proto")
-        .compile(
+        .compile_with_config(
+            config,
             &[
                 "milvus-proto/proto/common.proto",
                 "milvus-proto/proto/milvus.proto",

docker-compose.yml

@@ -51,6 +51,8 @@ services:
     environment:
       ETCD_ENDPOINTS: etcd:2379
       MINIO_ADDRESS: minio:9000
+      COMMON_GRPC_RATE_LIMIT_ENABLE: "false"
+      COMMON_GRPC_RATE_LIMIT_QPSPERCALL: "1000"
     volumes:
       - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
     healthcheck:

examples/collection.rs

@@ -1,9 +1,10 @@
 use milvus::index::{IndexParams, IndexType};
 use milvus::options::LoadOptions;
+use milvus::proto::schema::DataType;
 use milvus::query::QueryOptions;
-use milvus::schema::{CollectionSchema, CollectionSchemaBuilder};
+use milvus::schema::{CollectionSchema, CollectionSchemaBuilder, FieldSchemaBuilder};
 use milvus::{
-    client::Client, collection::Collection, data::FieldColumn, error::Error, schema::FieldSchema,
+    client::Client, data::FieldColumn, error::Error,
 };
 use std::collections::HashMap;
 
@@ -20,16 +21,18 @@ async fn main() -> Result<(), Error> {
 
     let schema =
         CollectionSchemaBuilder::new("hello_milvus", "a guide example for milvus rust SDK")
-            .add_field(FieldSchema::new_primary_int64(
-                "id",
-                "primary key field",
-                true,
-            ))
-            .add_field(FieldSchema::new_float_vector(
-                DEFAULT_VEC_FIELD,
-                "feature field",
-                DIM,
-            ))
+            .add_field(FieldSchemaBuilder::new()
+                .with_name("id").with_dtype(DataType::Int64)
+                .with_description("primary key field")
+                .with_primary(true)
+                .build()
+            )
+            .add_field(FieldSchemaBuilder::new()
+                .with_name(DEFAULT_VEC_FIELD)
+                .with_dtype(DataType::FloatVector)
+                .with_dim(DIM)
+                .build()
+            )
             .build()?;
     client.create_collection(schema.clone(), None).await?;
 

src/collection.rs

@@ -18,6 +18,7 @@ use crate::config;
 use crate::data::FieldColumn;
 use crate::error::{Error as SuperError, Result};
 use crate::index::{IndexInfo, IndexParams};
+use crate::proto::common::KeyValuePair;
 use crate::proto::milvus::{
     CreateCollectionRequest, CreateIndexRequest, DescribeIndexRequest, DropCollectionRequest,
     DropIndexRequest, FlushRequest, GetCompactionStateRequest, GetCompactionStateResponse,
@@ -211,6 +212,11 @@ impl Client {
                 schema: buf.to_vec(),
                 shards_num: options.shard_num,
                 consistency_level: options.consistency_level as i32,
+                properties: options
+                    .properties
+                    .into_iter()
+                    .map(|(k, v)| KeyValuePair { key: k, value: v })
+                    .collect(),
                 ..Default::default()
             })
             .await?
@@ -406,6 +412,8 @@ impl Client {
                     replica_number: options.replica_number,
                     resource_groups: vec![],
                     refresh: false,
+                    load_fields: vec![],
+                    skip_load_dynamic_field: false,
                 })
                 .await?
                 .into_inner(),
@@ -637,6 +645,9 @@ impl Client {
             .manual_compaction(ManualCompactionRequest {
                 collection_id: collection.id,
                 timetravel: 0,
+                major_compaction: false,
+                collection_name: "".to_string(),
+                db_name: "".to_string(),
             })
             .await?
             .into_inner();
@@ -660,6 +671,7 @@ pub type ParamValue = serde_json::Value;
 pub use serde_json::json as ParamValue;
 
 // search result for a single vector
+#[derive(Debug)]
 pub struct SearchResult<'a> {
     pub size: i64,
     pub id: Vec<Value<'a>>,

src/data.rs

@@ -110,6 +110,8 @@ impl FieldColumn {
             ValueVec::String(v) => Value::String(Cow::Borrowed(v.get(idx)?.as_ref())),
             ValueVec::Json(v) => Value::Json(Cow::Borrowed(v.get(idx)?.as_ref())),
             ValueVec::Array(v) => Value::Array(Cow::Borrowed(v.get(idx)?)),
+            ValueVec::Bytes(v) => Value::Bytes(Cow::Borrowed(v.get(idx)?.as_ref())),
+            ValueVec::Geometry(v) => Value::Geometry(Cow::Borrowed(v.get(idx)?.as_ref())),
         })
     }
 
@@ -152,6 +154,8 @@ impl FieldColumn {
                 ValueVec::Json(_) => ValueVec::Json(Vec::new()),
                 ValueVec::Binary(_) => ValueVec::Binary(Vec::new()),
                 ValueVec::Array(_) => ValueVec::Array(Vec::new()),
+                ValueVec::Bytes(_) => ValueVec::Bytes(Vec::new()),
+                ValueVec::Geometry(_) => ValueVec::Geometry(Vec::new()),
             },
             is_dynamic: self.is_dynamic,
         }
@@ -204,8 +208,15 @@ impl From<FieldColumn> for schema::FieldData {
                     data: Some(VectorData::BinaryVector(v)),
                     dim: this.dim,
                 }),
+                ValueVec::Bytes(v) => Field::Scalars(ScalarField {
+                    data: Some(ScalarData::BytesData(schema::BytesArray { data: v })),
+                }),
+                ValueVec::Geometry(v) => Field::Scalars(ScalarField {
+                    data: Some(ScalarData::GeometryData(schema::GeometryArray { data: v })),
+                }),
             }),
             is_dynamic: false,
+            valid_data: vec![],
         }
     }
 }

src/index/mod.rs

@@ -2,7 +2,7 @@ use strum_macros::{Display, EnumString};
 
 use crate::proto::{
     common::{IndexState, KeyValuePair},
-    milvus::IndexDescription,
+    milvus::IndexDescription, schema,
 };
 use std::{collections::HashMap, str::FromStr};
 
@@ -40,6 +40,8 @@ pub enum IndexType {
     NGTPANNG,
     #[strum(serialize = "NGT_ONNG")]
     NGTONNG,
+    #[strum(serialize = "SPARSE_INVERTED_INDEX")]
+    SparseInvertedIndex,
 }
 
 #[derive(Debug, Clone, Copy, EnumString, Display)]
@@ -51,6 +53,9 @@ pub enum MetricType {
     TANIMOTO,
     SUBSTRUCTURE,
     SUPERSTRUCTURE,
+    COSINE,
+    // Only for sparse vector with BM25
+    BM25,
 }
 
 #[derive(Debug, Clone)]
@@ -164,3 +169,30 @@ impl From<IndexDescription> for IndexInfo {
         }
     }
 }
+
+#[derive(Debug, Clone, EnumString, Display)]
+pub enum FunctionType {
+    Unknown,
+    BM25,
+    TextEmbedding,
+}
+
+impl From<schema::FunctionType> for FunctionType {
+    fn from(value: schema::FunctionType) -> Self {
+        match value {
+            schema::FunctionType::Unknown => Self::Unknown,
+            schema::FunctionType::Bm25 => Self::BM25,
+            schema::FunctionType::TextEmbedding => Self::TextEmbedding,
+        }
+    }
+}
+
+impl Into<schema::FunctionType> for FunctionType  {
+    fn into(self) -> schema::FunctionType {
+        match self {
+            FunctionType::Unknown => schema::FunctionType::Unknown,
+            FunctionType::BM25 => schema::FunctionType::Bm25,
+            FunctionType::TextEmbedding => schema::FunctionType::TextEmbedding,
+        }
+    }
+}

src/mutate.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use prost::bytes::{BufMut, BytesMut};
 
 use crate::error::Result;
@@ -104,6 +106,7 @@ impl Client {
                 num_rows: row_num as u32,
                 fields_data: fields_data.into_iter().map(|f| f.into()).collect(),
                 hash_keys: Vec::new(),
+                schema_timestamp: 0,
             })
             .await?
             .into_inner();
@@ -133,6 +136,8 @@ impl Client {
                 expr: expr,
                 partition_name: options.partition_name.clone(),
                 hash_keys: Vec::new(),
+                consistency_level: crate::client::ConsistencyLevel::default() as i32,
+                expr_template_values: HashMap::new(),
             })
             .await?
             .into_inner();
@@ -214,6 +219,7 @@ impl Client {
                 num_rows: row_num as u32,
                 fields_data: fields_data.into_iter().map(|f| f.into()).collect(),
                 hash_keys: Vec::new(),
+                schema_timestamp: 0,
             })
             .await?
             .into_inner();

src/options.rs

@@ -1,16 +1,20 @@
+use std::collections::HashMap;
+
 use crate::proto::common::ConsistencyLevel;
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone)]
 pub struct CreateCollectionOptions {
     pub(crate) shard_num: i32,
     pub(crate) consistency_level: ConsistencyLevel,
+    pub(crate) properties: HashMap<String, String>,
 }
 
 impl Default for CreateCollectionOptions {
     fn default() -> Self {
         Self {
             shard_num: 0,
             consistency_level: ConsistencyLevel::Bounded,
+            properties: HashMap::new(),
         }
     }
 }
@@ -37,6 +41,16 @@ impl CreateCollectionOptions {
         self.consistency_level = consistency_level;
         self
     }
+
+    pub fn properties(mut self, properties: HashMap<String, String>) -> Self {
+        self.properties = properties;
+        self
+    }
+
+    pub fn add_property(mut self, key: &str, value: &str) -> Self {
+        self.properties.insert(key.to_owned(), value.to_owned());
+        self
+    }
 }
 
 #[derive(Debug, Clone, Copy)]