diff --git a/.gitignore b/.gitignore index 8cba6e3..37e65bd 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,7 @@ Cargo.lock /target #Cargo.lock /.idea/ + +/volumes + +/.vscode/ \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 838d2c2..2cc4887 100644 --- a/Cargo.toml +++ b/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" diff --git a/README.md b/README.md index 92832b7..1c7e420 100644 --- a/README.md +++ b/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 ``` diff --git a/build.rs b/build.rs index 003123a..13526dc 100644 --- a/build.rs +++ b/build.rs @@ -1,8 +1,12 @@ fn main() -> Result<(), Box> { + 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", diff --git a/docker-compose.yml b/docker-compose.yml index 70c75d8..7189e25 100644 --- a/docker-compose.yml +++ b/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: diff --git a/examples/collection.rs b/examples/collection.rs index 97cbe74..a993259 100644 --- a/examples/collection.rs +++ b/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?; diff --git a/milvus-proto b/milvus-proto index 02cbad3..0a19881 160000 --- a/milvus-proto +++ b/milvus-proto @@ -1 +1 @@ -Subproject commit 02cbad30332f52be598373b3f0c5968270f753f8 +Subproject commit 0a1988183e535eadf47fb78cea2a3f8322c2d45b diff --git a/src/collection.rs b/src/collection.rs index d70748f..9c1ad88 100644 --- a/src/collection.rs +++ b/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>, diff --git a/src/data.rs b/src/data.rs index 7258bf6..048a067 100644 --- a/src/data.rs +++ b/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 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![], } } } diff --git a/src/index/mod.rs b/src/index/mod.rs index 3eac2af..0412b08 100644 --- a/src/index/mod.rs +++ b/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 for IndexInfo { } } } + +#[derive(Debug, Clone, EnumString, Display)] +pub enum FunctionType { + Unknown, + BM25, + TextEmbedding, +} + +impl From 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 for FunctionType { + fn into(self) -> schema::FunctionType { + match self { + FunctionType::Unknown => schema::FunctionType::Unknown, + FunctionType::BM25 => schema::FunctionType::Bm25, + FunctionType::TextEmbedding => schema::FunctionType::TextEmbedding, + } + } +} \ No newline at end of file diff --git a/src/mutate.rs b/src/mutate.rs index 17680f6..1c59d69 100644 --- a/src/mutate.rs +++ b/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(); diff --git a/src/options.rs b/src/options.rs index 097740a..d069b8c 100644 --- a/src/options.rs +++ b/src/options.rs @@ -1,9 +1,12 @@ +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, } impl Default for CreateCollectionOptions { @@ -11,6 +14,7 @@ impl Default for CreateCollectionOptions { 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) -> 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)] diff --git a/src/proto/milvus.proto.common.rs b/src/proto/milvus.proto.common.rs index 161a5bb..d6866a6 100644 --- a/src/proto/milvus.proto.common.rs +++ b/src/proto/milvus.proto.common.rs @@ -12,6 +12,11 @@ pub struct Status { pub retriable: bool, #[prost(string, tag = "5")] pub detail: ::prost::alloc::string::String, + #[prost(map = "string, string", tag = "6")] + pub extra_info: ::std::collections::HashMap< + ::prost::alloc::string::String, + ::prost::alloc::string::String, + >, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -42,7 +47,9 @@ pub struct PlaceholderValue { pub tag: ::prost::alloc::string::String, #[prost(enumeration = "PlaceholderType", tag = "2")] pub r#type: i32, - /// values is a 2d-array, every array contains a vector + /// values is a 2d-array of nq rows, every row contains a query vector. + /// for dense vector, all rows are of the same length; for sparse vector, + /// the length of each row may vary depending on their number of non-zeros. #[prost(bytes = "vec", repeated, tag = "3")] pub values: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, } @@ -88,6 +95,8 @@ pub struct ReplicateInfo { pub is_replicate: bool, #[prost(uint64, tag = "2")] pub msg_timestamp: u64, + #[prost(string, tag = "3")] + pub replicate_id: ::prost::alloc::string::String, } /// Don't Modify This. @czs #[allow(clippy::derive_partial_eq_without_eq)] @@ -167,6 +176,17 @@ pub struct ServerInfo { ::prost::alloc::string::String, >, } +/// NodeInfo is used to describe the node information. +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct NodeInfo { + #[prost(int64, tag = "1")] + pub node_id: i64, + #[prost(string, tag = "2")] + pub address: ::prost::alloc::string::String, + #[prost(string, tag = "3")] + pub hostname: ::prost::alloc::string::String, +} /// Deprecated #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] @@ -229,6 +249,11 @@ pub enum ErrorCode { NotReadyServe = 56, /// Coord is switching from standby mode to active mode NotReadyCoordActivating = 57, + CreatePrivilegeGroupFailure = 58, + DropPrivilegeGroupFailure = 59, + ListPrivilegeGroupsFailure = 60, + OperatePrivilegeGroupFailure = 61, + SchemaMismatch = 62, /// Service availability. /// NA: Not Available. DataCoordNa = 100, @@ -299,6 +324,11 @@ impl ErrorCode { ErrorCode::TimeTickLongDelay => "TimeTickLongDelay", ErrorCode::NotReadyServe => "NotReadyServe", ErrorCode::NotReadyCoordActivating => "NotReadyCoordActivating", + ErrorCode::CreatePrivilegeGroupFailure => "CreatePrivilegeGroupFailure", + ErrorCode::DropPrivilegeGroupFailure => "DropPrivilegeGroupFailure", + ErrorCode::ListPrivilegeGroupsFailure => "ListPrivilegeGroupsFailure", + ErrorCode::OperatePrivilegeGroupFailure => "OperatePrivilegeGroupFailure", + ErrorCode::SchemaMismatch => "SchemaMismatch", ErrorCode::DataCoordNa => "DataCoordNA", ErrorCode::DdRequestRace => "DDRequestRace", } @@ -363,6 +393,11 @@ impl ErrorCode { "TimeTickLongDelay" => Some(Self::TimeTickLongDelay), "NotReadyServe" => Some(Self::NotReadyServe), "NotReadyCoordActivating" => Some(Self::NotReadyCoordActivating), + "CreatePrivilegeGroupFailure" => Some(Self::CreatePrivilegeGroupFailure), + "DropPrivilegeGroupFailure" => Some(Self::DropPrivilegeGroupFailure), + "ListPrivilegeGroupsFailure" => Some(Self::ListPrivilegeGroupsFailure), + "OperatePrivilegeGroupFailure" => Some(Self::OperatePrivilegeGroupFailure), + "SchemaMismatch" => Some(Self::SchemaMismatch), "DataCoordNA" => Some(Self::DataCoordNa), "DDRequestRace" => Some(Self::DdRequestRace), _ => None, @@ -453,12 +488,50 @@ impl SegmentState { } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] +pub enum SegmentLevel { + /// zero value for legacy logic + Legacy = 0, + /// L0 segment, contains delta data for current channel + L0 = 1, + /// L1 segment, normal segment, with no extra compaction attribute + L1 = 2, + /// L2 segment, segment with extra data distribution info + L2 = 3, +} +impl SegmentLevel { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + SegmentLevel::Legacy => "Legacy", + SegmentLevel::L0 => "L0", + SegmentLevel::L1 => "L1", + SegmentLevel::L2 => "L2", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "Legacy" => Some(Self::Legacy), + "L0" => Some(Self::L0), + "L1" => Some(Self::L1), + "L2" => Some(Self::L2), + _ => None, + } + } +} +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] pub enum PlaceholderType { None = 0, BinaryVector = 100, FloatVector = 101, Float16Vector = 102, BFloat16Vector = 103, + SparseFloatVector = 104, + Int8Vector = 105, Int64 = 5, VarChar = 21, } @@ -474,6 +547,8 @@ impl PlaceholderType { PlaceholderType::FloatVector => "FloatVector", PlaceholderType::Float16Vector => "Float16Vector", PlaceholderType::BFloat16Vector => "BFloat16Vector", + PlaceholderType::SparseFloatVector => "SparseFloatVector", + PlaceholderType::Int8Vector => "Int8Vector", PlaceholderType::Int64 => "Int64", PlaceholderType::VarChar => "VarChar", } @@ -486,6 +561,8 @@ impl PlaceholderType { "FloatVector" => Some(Self::FloatVector), "Float16Vector" => Some(Self::Float16Vector), "BFloat16Vector" => Some(Self::BFloat16Vector), + "SparseFloatVector" => Some(Self::SparseFloatVector), + "Int8Vector" => Some(Self::Int8Vector), "Int64" => Some(Self::Int64), "VarChar" => Some(Self::VarChar), _ => None, @@ -512,6 +589,7 @@ pub enum MsgType { RenameCollection = 112, DescribeAlias = 113, ListAliases = 114, + AlterCollectionField = 115, /// DEFINITION REQUESTS: PARTITION CreatePartition = 200, DropPartition = 201, @@ -535,12 +613,20 @@ pub enum MsgType { DescribeIndex = 301, DropIndex = 302, GetIndexStatistics = 303, + AlterIndex = 304, /// MANIPULATION REQUESTS Insert = 400, Delete = 401, Flush = 402, ResendSegmentStats = 403, Upsert = 404, + /// streaming service new msg type for internal usage compatible + ManualFlush = 405, + /// streaming service new msg type for internal usage compatible + FlushSegment = 406, + /// streaming service new msg type for internal usage compatible + CreateSegment = 407, + Import = 408, /// QUERY Search = 500, SearchResult = 501, @@ -580,6 +666,7 @@ pub enum MsgType { Connect = 1209, ListClientInfos = 1210, AllocTimestamp = 1211, + Replicate = 1212, /// Credential CreateCredential = 1500, GetCredential = 1501, @@ -597,6 +684,11 @@ pub enum MsgType { SelectGrant = 1607, RefreshPolicyInfoCache = 1608, ListPolicy = 1609, + CreatePrivilegeGroup = 1610, + DropPrivilegeGroup = 1611, + ListPrivilegeGroups = 1612, + OperatePrivilegeGroup = 1613, + OperatePrivilegeV2 = 1614, /// Resource group CreateResourceGroup = 1700, DropResourceGroup = 1701, @@ -604,10 +696,14 @@ pub enum MsgType { DescribeResourceGroup = 1703, TransferNode = 1704, TransferReplica = 1705, + UpdateResourceGroups = 1706, /// Database group CreateDatabase = 1801, DropDatabase = 1802, ListDatabases = 1803, + AlterDatabase = 1804, + DescribeDatabase = 1805, + AddCollectionField = 1900, } impl MsgType { /// String value of the enum field names used in the ProtoBuf definition. @@ -632,6 +728,7 @@ impl MsgType { MsgType::RenameCollection => "RenameCollection", MsgType::DescribeAlias => "DescribeAlias", MsgType::ListAliases => "ListAliases", + MsgType::AlterCollectionField => "AlterCollectionField", MsgType::CreatePartition => "CreatePartition", MsgType::DropPartition => "DropPartition", MsgType::HasPartition => "HasPartition", @@ -652,11 +749,16 @@ impl MsgType { MsgType::DescribeIndex => "DescribeIndex", MsgType::DropIndex => "DropIndex", MsgType::GetIndexStatistics => "GetIndexStatistics", + MsgType::AlterIndex => "AlterIndex", MsgType::Insert => "Insert", MsgType::Delete => "Delete", MsgType::Flush => "Flush", MsgType::ResendSegmentStats => "ResendSegmentStats", MsgType::Upsert => "Upsert", + MsgType::ManualFlush => "ManualFlush", + MsgType::FlushSegment => "FlushSegment", + MsgType::CreateSegment => "CreateSegment", + MsgType::Import => "Import", MsgType::Search => "Search", MsgType::SearchResult => "SearchResult", MsgType::GetIndexState => "GetIndexState", @@ -692,6 +794,7 @@ impl MsgType { MsgType::Connect => "Connect", MsgType::ListClientInfos => "ListClientInfos", MsgType::AllocTimestamp => "AllocTimestamp", + MsgType::Replicate => "Replicate", MsgType::CreateCredential => "CreateCredential", MsgType::GetCredential => "GetCredential", MsgType::DeleteCredential => "DeleteCredential", @@ -707,15 +810,24 @@ impl MsgType { MsgType::SelectGrant => "SelectGrant", MsgType::RefreshPolicyInfoCache => "RefreshPolicyInfoCache", MsgType::ListPolicy => "ListPolicy", + MsgType::CreatePrivilegeGroup => "CreatePrivilegeGroup", + MsgType::DropPrivilegeGroup => "DropPrivilegeGroup", + MsgType::ListPrivilegeGroups => "ListPrivilegeGroups", + MsgType::OperatePrivilegeGroup => "OperatePrivilegeGroup", + MsgType::OperatePrivilegeV2 => "OperatePrivilegeV2", MsgType::CreateResourceGroup => "CreateResourceGroup", MsgType::DropResourceGroup => "DropResourceGroup", MsgType::ListResourceGroups => "ListResourceGroups", MsgType::DescribeResourceGroup => "DescribeResourceGroup", MsgType::TransferNode => "TransferNode", MsgType::TransferReplica => "TransferReplica", + MsgType::UpdateResourceGroups => "UpdateResourceGroups", MsgType::CreateDatabase => "CreateDatabase", MsgType::DropDatabase => "DropDatabase", MsgType::ListDatabases => "ListDatabases", + MsgType::AlterDatabase => "AlterDatabase", + MsgType::DescribeDatabase => "DescribeDatabase", + MsgType::AddCollectionField => "AddCollectionField", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -737,6 +849,7 @@ impl MsgType { "RenameCollection" => Some(Self::RenameCollection), "DescribeAlias" => Some(Self::DescribeAlias), "ListAliases" => Some(Self::ListAliases), + "AlterCollectionField" => Some(Self::AlterCollectionField), "CreatePartition" => Some(Self::CreatePartition), "DropPartition" => Some(Self::DropPartition), "HasPartition" => Some(Self::HasPartition), @@ -757,11 +870,16 @@ impl MsgType { "DescribeIndex" => Some(Self::DescribeIndex), "DropIndex" => Some(Self::DropIndex), "GetIndexStatistics" => Some(Self::GetIndexStatistics), + "AlterIndex" => Some(Self::AlterIndex), "Insert" => Some(Self::Insert), "Delete" => Some(Self::Delete), "Flush" => Some(Self::Flush), "ResendSegmentStats" => Some(Self::ResendSegmentStats), "Upsert" => Some(Self::Upsert), + "ManualFlush" => Some(Self::ManualFlush), + "FlushSegment" => Some(Self::FlushSegment), + "CreateSegment" => Some(Self::CreateSegment), + "Import" => Some(Self::Import), "Search" => Some(Self::Search), "SearchResult" => Some(Self::SearchResult), "GetIndexState" => Some(Self::GetIndexState), @@ -797,6 +915,7 @@ impl MsgType { "Connect" => Some(Self::Connect), "ListClientInfos" => Some(Self::ListClientInfos), "AllocTimestamp" => Some(Self::AllocTimestamp), + "Replicate" => Some(Self::Replicate), "CreateCredential" => Some(Self::CreateCredential), "GetCredential" => Some(Self::GetCredential), "DeleteCredential" => Some(Self::DeleteCredential), @@ -812,15 +931,24 @@ impl MsgType { "SelectGrant" => Some(Self::SelectGrant), "RefreshPolicyInfoCache" => Some(Self::RefreshPolicyInfoCache), "ListPolicy" => Some(Self::ListPolicy), + "CreatePrivilegeGroup" => Some(Self::CreatePrivilegeGroup), + "DropPrivilegeGroup" => Some(Self::DropPrivilegeGroup), + "ListPrivilegeGroups" => Some(Self::ListPrivilegeGroups), + "OperatePrivilegeGroup" => Some(Self::OperatePrivilegeGroup), + "OperatePrivilegeV2" => Some(Self::OperatePrivilegeV2), "CreateResourceGroup" => Some(Self::CreateResourceGroup), "DropResourceGroup" => Some(Self::DropResourceGroup), "ListResourceGroups" => Some(Self::ListResourceGroups), "DescribeResourceGroup" => Some(Self::DescribeResourceGroup), "TransferNode" => Some(Self::TransferNode), "TransferReplica" => Some(Self::TransferReplica), + "UpdateResourceGroups" => Some(Self::UpdateResourceGroups), "CreateDatabase" => Some(Self::CreateDatabase), "DropDatabase" => Some(Self::DropDatabase), "ListDatabases" => Some(Self::ListDatabases), + "AlterDatabase" => Some(Self::AlterDatabase), + "DescribeDatabase" => Some(Self::DescribeDatabase), + "AddCollectionField" => Some(Self::AddCollectionField), _ => None, } } @@ -1041,6 +1169,34 @@ pub enum ObjectPrivilege { PrivilegeShowPartitions = 41, PrivilegeHasPartition = 42, PrivilegeGetFlushState = 43, + PrivilegeCreateAlias = 44, + PrivilegeDropAlias = 45, + PrivilegeDescribeAlias = 46, + PrivilegeListAliases = 47, + PrivilegeUpdateResourceGroups = 48, + PrivilegeAlterDatabase = 49, + PrivilegeDescribeDatabase = 50, + PrivilegeBackupRbac = 51, + PrivilegeRestoreRbac = 52, + PrivilegeGroupReadOnly = 53, + PrivilegeGroupReadWrite = 54, + PrivilegeGroupAdmin = 55, + PrivilegeCreatePrivilegeGroup = 56, + PrivilegeDropPrivilegeGroup = 57, + PrivilegeListPrivilegeGroups = 58, + PrivilegeOperatePrivilegeGroup = 59, + PrivilegeGroupClusterReadOnly = 60, + PrivilegeGroupClusterReadWrite = 61, + PrivilegeGroupClusterAdmin = 62, + PrivilegeGroupDatabaseReadOnly = 63, + PrivilegeGroupDatabaseReadWrite = 64, + PrivilegeGroupDatabaseAdmin = 65, + PrivilegeGroupCollectionReadOnly = 66, + PrivilegeGroupCollectionReadWrite = 67, + PrivilegeGroupCollectionAdmin = 68, + PrivilegeGetImportProgress = 69, + PrivilegeListImport = 70, + PrivilegeAddCollectionField = 71, } impl ObjectPrivilege { /// String value of the enum field names used in the ProtoBuf definition. @@ -1097,6 +1253,56 @@ impl ObjectPrivilege { ObjectPrivilege::PrivilegeShowPartitions => "PrivilegeShowPartitions", ObjectPrivilege::PrivilegeHasPartition => "PrivilegeHasPartition", ObjectPrivilege::PrivilegeGetFlushState => "PrivilegeGetFlushState", + ObjectPrivilege::PrivilegeCreateAlias => "PrivilegeCreateAlias", + ObjectPrivilege::PrivilegeDropAlias => "PrivilegeDropAlias", + ObjectPrivilege::PrivilegeDescribeAlias => "PrivilegeDescribeAlias", + ObjectPrivilege::PrivilegeListAliases => "PrivilegeListAliases", + ObjectPrivilege::PrivilegeUpdateResourceGroups => { + "PrivilegeUpdateResourceGroups" + } + ObjectPrivilege::PrivilegeAlterDatabase => "PrivilegeAlterDatabase", + ObjectPrivilege::PrivilegeDescribeDatabase => "PrivilegeDescribeDatabase", + ObjectPrivilege::PrivilegeBackupRbac => "PrivilegeBackupRBAC", + ObjectPrivilege::PrivilegeRestoreRbac => "PrivilegeRestoreRBAC", + ObjectPrivilege::PrivilegeGroupReadOnly => "PrivilegeGroupReadOnly", + ObjectPrivilege::PrivilegeGroupReadWrite => "PrivilegeGroupReadWrite", + ObjectPrivilege::PrivilegeGroupAdmin => "PrivilegeGroupAdmin", + ObjectPrivilege::PrivilegeCreatePrivilegeGroup => { + "PrivilegeCreatePrivilegeGroup" + } + ObjectPrivilege::PrivilegeDropPrivilegeGroup => "PrivilegeDropPrivilegeGroup", + ObjectPrivilege::PrivilegeListPrivilegeGroups => { + "PrivilegeListPrivilegeGroups" + } + ObjectPrivilege::PrivilegeOperatePrivilegeGroup => { + "PrivilegeOperatePrivilegeGroup" + } + ObjectPrivilege::PrivilegeGroupClusterReadOnly => { + "PrivilegeGroupClusterReadOnly" + } + ObjectPrivilege::PrivilegeGroupClusterReadWrite => { + "PrivilegeGroupClusterReadWrite" + } + ObjectPrivilege::PrivilegeGroupClusterAdmin => "PrivilegeGroupClusterAdmin", + ObjectPrivilege::PrivilegeGroupDatabaseReadOnly => { + "PrivilegeGroupDatabaseReadOnly" + } + ObjectPrivilege::PrivilegeGroupDatabaseReadWrite => { + "PrivilegeGroupDatabaseReadWrite" + } + ObjectPrivilege::PrivilegeGroupDatabaseAdmin => "PrivilegeGroupDatabaseAdmin", + ObjectPrivilege::PrivilegeGroupCollectionReadOnly => { + "PrivilegeGroupCollectionReadOnly" + } + ObjectPrivilege::PrivilegeGroupCollectionReadWrite => { + "PrivilegeGroupCollectionReadWrite" + } + ObjectPrivilege::PrivilegeGroupCollectionAdmin => { + "PrivilegeGroupCollectionAdmin" + } + ObjectPrivilege::PrivilegeGetImportProgress => "PrivilegeGetImportProgress", + ObjectPrivilege::PrivilegeListImport => "PrivilegeListImport", + ObjectPrivilege::PrivilegeAddCollectionField => "PrivilegeAddCollectionField", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -1148,6 +1354,46 @@ impl ObjectPrivilege { "PrivilegeShowPartitions" => Some(Self::PrivilegeShowPartitions), "PrivilegeHasPartition" => Some(Self::PrivilegeHasPartition), "PrivilegeGetFlushState" => Some(Self::PrivilegeGetFlushState), + "PrivilegeCreateAlias" => Some(Self::PrivilegeCreateAlias), + "PrivilegeDropAlias" => Some(Self::PrivilegeDropAlias), + "PrivilegeDescribeAlias" => Some(Self::PrivilegeDescribeAlias), + "PrivilegeListAliases" => Some(Self::PrivilegeListAliases), + "PrivilegeUpdateResourceGroups" => Some(Self::PrivilegeUpdateResourceGroups), + "PrivilegeAlterDatabase" => Some(Self::PrivilegeAlterDatabase), + "PrivilegeDescribeDatabase" => Some(Self::PrivilegeDescribeDatabase), + "PrivilegeBackupRBAC" => Some(Self::PrivilegeBackupRbac), + "PrivilegeRestoreRBAC" => Some(Self::PrivilegeRestoreRbac), + "PrivilegeGroupReadOnly" => Some(Self::PrivilegeGroupReadOnly), + "PrivilegeGroupReadWrite" => Some(Self::PrivilegeGroupReadWrite), + "PrivilegeGroupAdmin" => Some(Self::PrivilegeGroupAdmin), + "PrivilegeCreatePrivilegeGroup" => Some(Self::PrivilegeCreatePrivilegeGroup), + "PrivilegeDropPrivilegeGroup" => Some(Self::PrivilegeDropPrivilegeGroup), + "PrivilegeListPrivilegeGroups" => Some(Self::PrivilegeListPrivilegeGroups), + "PrivilegeOperatePrivilegeGroup" => { + Some(Self::PrivilegeOperatePrivilegeGroup) + } + "PrivilegeGroupClusterReadOnly" => Some(Self::PrivilegeGroupClusterReadOnly), + "PrivilegeGroupClusterReadWrite" => { + Some(Self::PrivilegeGroupClusterReadWrite) + } + "PrivilegeGroupClusterAdmin" => Some(Self::PrivilegeGroupClusterAdmin), + "PrivilegeGroupDatabaseReadOnly" => { + Some(Self::PrivilegeGroupDatabaseReadOnly) + } + "PrivilegeGroupDatabaseReadWrite" => { + Some(Self::PrivilegeGroupDatabaseReadWrite) + } + "PrivilegeGroupDatabaseAdmin" => Some(Self::PrivilegeGroupDatabaseAdmin), + "PrivilegeGroupCollectionReadOnly" => { + Some(Self::PrivilegeGroupCollectionReadOnly) + } + "PrivilegeGroupCollectionReadWrite" => { + Some(Self::PrivilegeGroupCollectionReadWrite) + } + "PrivilegeGroupCollectionAdmin" => Some(Self::PrivilegeGroupCollectionAdmin), + "PrivilegeGetImportProgress" => Some(Self::PrivilegeGetImportProgress), + "PrivilegeListImport" => Some(Self::PrivilegeListImport), + "PrivilegeAddCollectionField" => Some(Self::PrivilegeAddCollectionField), _ => None, } } diff --git a/src/proto/milvus.proto.milvus.rs b/src/proto/milvus.proto.milvus.rs index 9a077ab..d2cd089 100644 --- a/src/proto/milvus.proto.milvus.rs +++ b/src/proto/milvus.proto.milvus.rs @@ -142,6 +142,24 @@ pub struct AlterCollectionRequest { pub collection_id: i64, #[prost(message, repeated, tag = "5")] pub properties: ::prost::alloc::vec::Vec, + #[prost(string, repeated, tag = "6")] + pub delete_keys: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AlterCollectionFieldRequest { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(string, tag = "2")] + pub db_name: ::prost::alloc::string::String, + /// The unique collection name in milvus.(Required) + #[prost(string, tag = "3")] + pub collection_name: ::prost::alloc::string::String, + #[prost(string, tag = "4")] + pub field_name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "5")] + pub properties: ::prost::alloc::vec::Vec, } /// * /// Check collection exist in milvus or not. @@ -243,6 +261,12 @@ pub struct DescribeCollectionResponse { pub db_name: ::prost::alloc::string::String, #[prost(int64, tag = "15")] pub num_partitions: i64, + #[prost(int64, tag = "16")] + pub db_id: i64, + #[prost(uint64, tag = "17")] + pub request_time: u64, + #[prost(uint64, tag = "18")] + pub update_timestamp: u64, } /// * /// Load collection data into query nodes, then you can do vector search on this collection. @@ -266,6 +290,12 @@ pub struct LoadCollectionRequest { /// Whether to enable refresh mode. #[prost(bool, tag = "6")] pub refresh: bool, + /// Field Partial Load fields list + #[prost(string, repeated, tag = "7")] + pub load_fields: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + /// Field Partial load skip load dynamic fields + #[prost(bool, tag = "8")] + pub skip_load_dynamic_field: bool, } /// * /// Release collection data from query nodes, then you can't do vector search on this collection. @@ -469,6 +499,12 @@ pub struct LoadPartitionsRequest { /// Whether to enable refresh mode. #[prost(bool, tag = "7")] pub refresh: bool, + /// Field Partial Load fields list + #[prost(string, repeated, tag = "8")] + pub load_fields: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + /// Field Partial load skip load dynamic fields + #[prost(bool, tag = "9")] + pub skip_load_dynamic_field: bool, } /// /// Release specific partitions data of one collection from query nodes. @@ -645,6 +681,8 @@ pub struct AlterIndexRequest { pub index_name: ::prost::alloc::string::String, #[prost(message, repeated, tag = "5")] pub extra_params: ::prost::alloc::vec::Vec, + #[prost(string, repeated, tag = "6")] + pub delete_keys: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } /// /// Get created index information. @@ -802,6 +840,23 @@ pub struct InsertRequest { pub hash_keys: ::prost::alloc::vec::Vec, #[prost(uint32, tag = "7")] pub num_rows: u32, + #[prost(uint64, tag = "8")] + pub schema_timestamp: u64, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AddCollectionFieldRequest { + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(string, tag = "2")] + pub db_name: ::prost::alloc::string::String, + #[prost(string, tag = "3")] + pub collection_name: ::prost::alloc::string::String, + #[prost(int64, tag = "4")] + pub collection_id: i64, + /// The serialized `schema.FieldSchema` + #[prost(bytes = "vec", tag = "5")] + pub schema: ::prost::alloc::vec::Vec, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -820,6 +875,8 @@ pub struct UpsertRequest { pub hash_keys: ::prost::alloc::vec::Vec, #[prost(uint32, tag = "7")] pub num_rows: u32, + #[prost(uint64, tag = "8")] + pub schema_timestamp: u64, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -861,6 +918,38 @@ pub struct DeleteRequest { pub expr: ::prost::alloc::string::String, #[prost(uint32, repeated, tag = "6")] pub hash_keys: ::prost::alloc::vec::Vec, + #[prost(enumeration = "super::common::ConsistencyLevel", tag = "7")] + pub consistency_level: i32, + #[prost(map = "string, message", tag = "8")] + pub expr_template_values: ::std::collections::HashMap< + ::prost::alloc::string::String, + super::schema::TemplateValue, + >, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SubSearchRequest { + /// must + #[prost(string, tag = "1")] + pub dsl: ::prost::alloc::string::String, + /// serialized `PlaceholderGroup` + /// + /// must + #[prost(bytes = "vec", tag = "2")] + pub placeholder_group: ::prost::alloc::vec::Vec, + /// must + #[prost(enumeration = "super::common::DslType", tag = "3")] + pub dsl_type: i32, + /// must + #[prost(message, repeated, tag = "4")] + pub search_params: ::prost::alloc::vec::Vec, + #[prost(int64, tag = "5")] + pub nq: i64, + #[prost(map = "string, message", tag = "6")] + pub expr_template_values: ::std::collections::HashMap< + ::prost::alloc::string::String, + super::schema::TemplateValue, + >, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -907,6 +996,13 @@ pub struct SearchRequest { pub use_default_consistency: bool, #[prost(bool, tag = "16")] pub search_by_primary_keys: bool, + #[prost(message, repeated, tag = "17")] + pub sub_reqs: ::prost::alloc::vec::Vec, + #[prost(map = "string, message", tag = "18")] + pub expr_template_values: ::std::collections::HashMap< + ::prost::alloc::string::String, + super::schema::TemplateValue, + >, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -927,10 +1023,13 @@ pub struct SearchResults { pub results: ::core::option::Option, #[prost(string, tag = "3")] pub collection_name: ::prost::alloc::string::String, + /// for session-like operation like iterator + #[prost(uint64, tag = "4")] + pub session_ts: u64, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct SearchRequestV2 { +pub struct HybridSearchRequest { /// must #[prost(message, optional, tag = "1")] pub base: ::core::option::Option, @@ -997,6 +1096,11 @@ pub struct FlushResponse { /// hybrid ts for geting flush tate #[prost(map = "string, uint64", tag = "6")] pub coll_flush_ts: ::std::collections::HashMap<::prost::alloc::string::String, u64>, + #[prost(map = "string, message", tag = "7")] + pub channel_cps: ::std::collections::HashMap< + ::prost::alloc::string::String, + super::msg::MsgPosition, + >, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1027,6 +1131,11 @@ pub struct QueryRequest { pub consistency_level: i32, #[prost(bool, tag = "12")] pub use_default_consistency: bool, + #[prost(map = "string, message", tag = "13")] + pub expr_template_values: ::std::collections::HashMap< + ::prost::alloc::string::String, + super::schema::TemplateValue, + >, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1039,6 +1148,30 @@ pub struct QueryResults { pub collection_name: ::prost::alloc::string::String, #[prost(string, repeated, tag = "4")] pub output_fields: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + /// for session-like operation like iterator + #[prost(uint64, tag = "5")] + pub session_ts: u64, + #[prost(string, tag = "6")] + pub primary_field_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct QueryCursor { + #[prost(uint64, tag = "1")] + pub session_ts: u64, + #[prost(oneof = "query_cursor::CursorPk", tags = "2, 3")] + pub cursor_pk: ::core::option::Option, +} +/// Nested message and enum types in `QueryCursor`. +pub mod query_cursor { + #[allow(clippy::derive_partial_eq_without_eq)] + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum CursorPk { + #[prost(string, tag = "2")] + StrPk(::prost::alloc::string::String), + #[prost(int64, tag = "3")] + IntPk(i64), + } } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1136,6 +1269,10 @@ pub struct PersistentSegmentInfo { pub num_rows: i64, #[prost(enumeration = "super::common::SegmentState", tag = "5")] pub state: i32, + #[prost(enumeration = "super::common::SegmentLevel", tag = "6")] + pub level: i32, + #[prost(bool, tag = "7")] + pub is_sorted: bool, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1182,6 +1319,10 @@ pub struct QuerySegmentInfo { pub state: i32, #[prost(int64, repeated, tag = "10")] pub node_ids: ::prost::alloc::vec::Vec, + #[prost(enumeration = "super::common::SegmentLevel", tag = "11")] + pub level: i32, + #[prost(bool, tag = "12")] + pub is_sorted: bool, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1297,6 +1438,12 @@ pub struct ManualCompactionRequest { pub collection_id: i64, #[prost(uint64, tag = "2")] pub timetravel: u64, + #[prost(bool, tag = "3")] + pub major_compaction: bool, + #[prost(string, tag = "4")] + pub collection_name: ::prost::alloc::string::String, + #[prost(string, tag = "5")] + pub db_name: ::prost::alloc::string::String, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1413,8 +1560,12 @@ pub struct ImportRequest { /// import options, bucket, etc. #[prost(message, repeated, tag = "6")] pub options: ::prost::alloc::vec::Vec, + /// target database #[prost(string, tag = "7")] pub db_name: ::prost::alloc::string::String, + /// serialized `schema.ClusteringInfo` + #[prost(bytes = "vec", tag = "8")] + pub clustering_info: ::prost::alloc::vec::Vec, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1645,6 +1796,60 @@ pub struct DropRoleRequest { /// role name #[prost(string, tag = "2")] pub role_name: ::prost::alloc::string::String, + /// force to drop the role even if there is permission binding + #[prost(bool, tag = "3")] + pub force_drop: bool, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct CreatePrivilegeGroupRequest { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + /// group name + #[prost(string, tag = "2")] + pub group_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DropPrivilegeGroupRequest { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + /// group name + #[prost(string, tag = "2")] + pub group_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ListPrivilegeGroupsRequest { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ListPrivilegeGroupsResponse { + #[prost(message, optional, tag = "1")] + pub status: ::core::option::Option, + #[prost(message, repeated, tag = "2")] + pub privilege_groups: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct OperatePrivilegeGroupRequest { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + /// group name + #[prost(string, tag = "2")] + pub group_name: ::prost::alloc::string::String, + /// privileges + #[prost(message, repeated, tag = "3")] + pub privileges: ::prost::alloc::vec::Vec, + /// operation type + #[prost(enumeration = "OperatePrivilegeGroupType", tag = "4")] + pub r#type: i32, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1664,6 +1869,14 @@ pub struct OperateUserRoleRequest { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PrivilegeGroupInfo { + #[prost(string, tag = "1")] + pub group_name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "2")] + pub privileges: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct SelectRoleRequest { /// Not useful for now #[prost(message, optional, tag = "1")] @@ -1801,6 +2014,81 @@ pub struct OperatePrivilegeRequest { /// operation type #[prost(enumeration = "OperatePrivilegeType", tag = "3")] pub r#type: i32, + /// version + #[prost(string, tag = "4")] + pub version: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct OperatePrivilegeV2Request { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + /// role + #[prost(message, optional, tag = "2")] + pub role: ::core::option::Option, + /// privilege + #[prost(message, optional, tag = "3")] + pub grantor: ::core::option::Option, + /// operation type + #[prost(enumeration = "OperatePrivilegeType", tag = "4")] + pub r#type: i32, + /// db name + #[prost(string, tag = "5")] + pub db_name: ::prost::alloc::string::String, + /// collection name + #[prost(string, tag = "6")] + pub collection_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct UserInfo { + #[prost(string, tag = "1")] + pub user: ::prost::alloc::string::String, + #[prost(string, tag = "2")] + pub password: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "3")] + pub roles: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RbacMeta { + /// user + #[prost(message, repeated, tag = "1")] + pub users: ::prost::alloc::vec::Vec, + /// role + #[prost(message, repeated, tag = "2")] + pub roles: ::prost::alloc::vec::Vec, + /// (role, object, previledge) + #[prost(message, repeated, tag = "3")] + pub grants: ::prost::alloc::vec::Vec, + /// privilege group info + #[prost(message, repeated, tag = "4")] + pub privilege_groups: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct BackupRbacMetaRequest { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct BackupRbacMetaResponse { + #[prost(message, optional, tag = "1")] + pub status: ::core::option::Option, + #[prost(message, optional, tag = "2")] + pub rbac_meta: ::core::option::Option, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RestoreRbacMetaRequest { + /// Not useful for now + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(message, optional, tag = "2")] + pub rbac_meta: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1885,6 +2173,19 @@ pub struct CreateResourceGroupRequest { pub base: ::core::option::Option, #[prost(string, tag = "2")] pub resource_group: ::prost::alloc::string::String, + #[prost(message, optional, tag = "3")] + pub config: ::core::option::Option, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct UpdateResourceGroupsRequest { + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(map = "string, message", tag = "2")] + pub resource_groups: ::std::collections::HashMap< + ::prost::alloc::string::String, + super::rg::ResourceGroupConfig, + >, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1981,6 +2282,12 @@ pub struct ResourceGroup { ::prost::alloc::string::String, i32, >, + /// resource group configuration. + #[prost(message, optional, tag = "7")] + pub config: ::core::option::Option, + /// query node belong to this resource group now. + #[prost(message, repeated, tag = "8")] + pub nodes: ::prost::alloc::vec::Vec, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -2063,6 +2370,8 @@ pub struct CreateDatabaseRequest { pub base: ::core::option::Option, #[prost(string, tag = "2")] pub db_name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "3")] + pub properties: ::prost::alloc::vec::Vec, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -2087,6 +2396,44 @@ pub struct ListDatabasesResponse { pub db_names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, #[prost(uint64, repeated, tag = "3")] pub created_timestamp: ::prost::alloc::vec::Vec, + #[prost(int64, repeated, tag = "4")] + pub db_ids: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AlterDatabaseRequest { + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(string, tag = "2")] + pub db_name: ::prost::alloc::string::String, + #[prost(string, tag = "3")] + pub db_id: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "4")] + pub properties: ::prost::alloc::vec::Vec, + #[prost(string, repeated, tag = "5")] + pub delete_keys: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DescribeDatabaseRequest { + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(string, tag = "2")] + pub db_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DescribeDatabaseResponse { + #[prost(message, optional, tag = "1")] + pub status: ::core::option::Option, + #[prost(string, tag = "2")] + pub db_name: ::prost::alloc::string::String, + #[prost(int64, tag = "3")] + pub db_id: i64, + #[prost(uint64, tag = "4")] + pub created_timestamp: u64, + #[prost(message, repeated, tag = "5")] + pub properties: ::prost::alloc::vec::Vec, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -2114,6 +2461,54 @@ pub struct ReplicateMessageResponse { #[prost(string, tag = "2")] pub position: ::prost::alloc::string::String, } +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ImportAuthPlaceholder { + #[prost(string, tag = "1")] + pub db_name: ::prost::alloc::string::String, + #[prost(string, tag = "2")] + pub collection_name: ::prost::alloc::string::String, + #[prost(string, tag = "3")] + pub partition_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct GetImportProgressAuthPlaceholder { + #[prost(string, tag = "1")] + pub db_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ListImportsAuthPlaceholder { + #[prost(string, tag = "3")] + pub db_name: ::prost::alloc::string::String, + #[prost(string, tag = "1")] + pub collection_name: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RunAnalyzerRequset { + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(string, tag = "2")] + pub analyzer_params: ::prost::alloc::string::String, + #[prost(bytes = "vec", repeated, tag = "3")] + pub placeholder: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AnalyzerResult { + #[prost(string, repeated, tag = "1")] + pub tokens: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RunAnalyzerResponse { + #[prost(message, optional, tag = "1")] + pub status: ::core::option::Option, + #[prost(message, repeated, tag = "2")] + pub results: ::prost::alloc::vec::Vec, +} /// Deprecated: use GetLoadingProgress rpc instead #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] @@ -2145,6 +2540,34 @@ impl ShowType { } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] +pub enum OperatePrivilegeGroupType { + AddPrivilegesToGroup = 0, + RemovePrivilegesFromGroup = 1, +} +impl OperatePrivilegeGroupType { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + OperatePrivilegeGroupType::AddPrivilegesToGroup => "AddPrivilegesToGroup", + OperatePrivilegeGroupType::RemovePrivilegesFromGroup => { + "RemovePrivilegesFromGroup" + } + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "AddPrivilegesToGroup" => Some(Self::AddPrivilegesToGroup), + "RemovePrivilegesFromGroup" => Some(Self::RemovePrivilegesFromGroup), + _ => None, + } + } +} +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] pub enum OperateUserRoleType { AddUserToRole = 0, RemoveUserFromRole = 1, @@ -2171,6 +2594,35 @@ impl OperateUserRoleType { } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] +pub enum PrivilegeLevel { + Cluster = 0, + Database = 1, + Collection = 2, +} +impl PrivilegeLevel { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + PrivilegeLevel::Cluster => "Cluster", + PrivilegeLevel::Database => "Database", + PrivilegeLevel::Collection => "Collection", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "Cluster" => Some(Self::Cluster), + "Database" => Some(Self::Database), + "Collection" => Some(Self::Collection), + _ => None, + } + } +} +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] pub enum OperatePrivilegeType { Grant = 0, Revoke = 1, @@ -2462,6 +2914,25 @@ pub mod milvus_service_client { ); self.inner.unary(request.into_request(), path, codec).await } + pub async fn alter_collection_field( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/AlterCollectionField", + ); + self.inner.unary(request.into_request(), path, codec).await + } pub async fn create_partition( &mut self, request: impl tonic::IntoRequest, @@ -2945,9 +3416,9 @@ pub mod milvus_service_client { ); self.inner.unary(request.into_request(), path, codec).await } - pub async fn search_v2( + pub async fn hybrid_search( &mut self, - request: impl tonic::IntoRequest, + request: impl tonic::IntoRequest, ) -> Result, tonic::Status> { self.inner .ready() @@ -2960,7 +3431,7 @@ pub mod milvus_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/milvus.proto.milvus.MilvusService/SearchV2", + "/milvus.proto.milvus.MilvusService/HybridSearch", ); self.inner.unary(request.into_request(), path, codec).await } @@ -3040,6 +3511,25 @@ pub mod milvus_service_client { ); self.inner.unary(request.into_request(), path, codec).await } + pub async fn add_collection_field( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/AddCollectionField", + ); + self.inner.unary(request.into_request(), path, codec).await + } pub async fn get_flush_state( &mut self, request: impl tonic::IntoRequest, @@ -3542,6 +4032,25 @@ pub mod milvus_service_client { ); self.inner.unary(request.into_request(), path, codec).await } + pub async fn operate_privilege_v2( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/OperatePrivilegeV2", + ); + self.inner.unary(request.into_request(), path, codec).await + } pub async fn select_grant( &mut self, request: impl tonic::IntoRequest, @@ -3637,6 +4146,25 @@ pub mod milvus_service_client { ); self.inner.unary(request.into_request(), path, codec).await } + pub async fn update_resource_groups( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/UpdateResourceGroups", + ); + self.inner.unary(request.into_request(), path, codec).await + } pub async fn transfer_node( &mut self, request: impl tonic::IntoRequest, @@ -3878,6 +4406,44 @@ pub mod milvus_service_client { ); self.inner.unary(request.into_request(), path, codec).await } + pub async fn alter_database( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/AlterDatabase", + ); + self.inner.unary(request.into_request(), path, codec).await + } + pub async fn describe_database( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/DescribeDatabase", + ); + self.inner.unary(request.into_request(), path, codec).await + } pub async fn replicate_message( &mut self, request: impl tonic::IntoRequest, @@ -3897,6 +4463,139 @@ pub mod milvus_service_client { ); self.inner.unary(request.into_request(), path, codec).await } + pub async fn backup_rbac( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/BackupRBAC", + ); + self.inner.unary(request.into_request(), path, codec).await + } + pub async fn restore_rbac( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/RestoreRBAC", + ); + self.inner.unary(request.into_request(), path, codec).await + } + pub async fn create_privilege_group( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/CreatePrivilegeGroup", + ); + self.inner.unary(request.into_request(), path, codec).await + } + pub async fn drop_privilege_group( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/DropPrivilegeGroup", + ); + self.inner.unary(request.into_request(), path, codec).await + } + pub async fn list_privilege_groups( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/ListPrivilegeGroups", + ); + self.inner.unary(request.into_request(), path, codec).await + } + pub async fn operate_privilege_group( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/OperatePrivilegeGroup", + ); + self.inner.unary(request.into_request(), path, codec).await + } + pub async fn run_analyzer( + &mut self, + request: impl tonic::IntoRequest, + ) -> Result, tonic::Status> { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/milvus.proto.milvus.MilvusService/RunAnalyzer", + ); + self.inner.unary(request.into_request(), path, codec).await + } } } /// Generated client implementations. diff --git a/src/proto/milvus.proto.msg.rs b/src/proto/milvus.proto.msg.rs index 1132b94..c9e8b6c 100644 --- a/src/proto/milvus.proto.msg.rs +++ b/src/proto/milvus.proto.msg.rs @@ -61,6 +61,8 @@ pub struct DeleteRequest { pub num_rows: i64, #[prost(message, optional, tag = "12")] pub primary_keys: ::core::option::Option, + #[prost(int64, tag = "13")] + pub segment_id: i64, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -170,6 +172,54 @@ pub struct DataNodeTtMsg { #[prost(message, repeated, tag = "4")] pub segments_stats: ::prost::alloc::vec::Vec, } +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ReplicateMsg { + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(bool, tag = "2")] + pub is_end: bool, + #[prost(bool, tag = "3")] + pub is_cluster: bool, + #[prost(string, tag = "4")] + pub database: ::prost::alloc::string::String, + #[prost(string, tag = "5")] + pub collection: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ImportFile { + #[prost(int64, tag = "1")] + pub id: i64, + /// A singular row-based file or multiple column-based files. + #[prost(string, repeated, tag = "2")] + pub paths: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ImportMsg { + #[prost(message, optional, tag = "1")] + pub base: ::core::option::Option, + #[prost(string, tag = "2")] + pub db_name: ::prost::alloc::string::String, + #[prost(string, tag = "3")] + pub collection_name: ::prost::alloc::string::String, + #[prost(int64, tag = "4")] + pub collection_id: i64, + #[prost(int64, repeated, tag = "5")] + pub partition_i_ds: ::prost::alloc::vec::Vec, + #[prost(map = "string, string", tag = "6")] + pub options: ::std::collections::HashMap< + ::prost::alloc::string::String, + ::prost::alloc::string::String, + >, + #[prost(message, repeated, tag = "7")] + pub files: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "8")] + pub schema: ::core::option::Option, + #[prost(int64, tag = "9")] + pub job_id: i64, +} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum InsertDataVersion { diff --git a/src/proto/milvus.proto.rg.rs b/src/proto/milvus.proto.rg.rs new file mode 100644 index 0000000..846183c --- /dev/null +++ b/src/proto/milvus.proto.rg.rs @@ -0,0 +1,40 @@ +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ResourceGroupLimit { + /// preserve for other limit. + #[prost(int32, tag = "1")] + pub node_num: i32, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ResourceGroupTransfer { + /// resource groups can be transfered with current resource group. + #[prost(string, tag = "1")] + pub resource_group: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ResourceGroupNodeFilter { + /// node in resource group must match node labels requirements + #[prost(message, repeated, tag = "1")] + pub node_labels: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ResourceGroupConfig { + /// requests node num in resource group, if node num is less than requests.nodeNum, it will be transfer from other resource group. + #[prost(message, optional, tag = "1")] + pub requests: ::core::option::Option, + /// limited node num in resource group, if node num is more than limits.nodeNum, it will be transfer to other resource group. + #[prost(message, optional, tag = "2")] + pub limits: ::core::option::Option, + /// missing node should be transfer from given resource group at high priority in repeated list. + #[prost(message, repeated, tag = "3")] + pub transfer_from: ::prost::alloc::vec::Vec, + /// redundant node should be transfer to given resource group at high priority in repeated list. + #[prost(message, repeated, tag = "4")] + pub transfer_to: ::prost::alloc::vec::Vec, + /// node in resource group must match node filters + #[prost(message, optional, tag = "5")] + pub node_filter: ::core::option::Option, +} diff --git a/src/proto/milvus.proto.schema.rs b/src/proto/milvus.proto.schema.rs index d28ba36..742044d 100644 --- a/src/proto/milvus.proto.schema.rs +++ b/src/proto/milvus.proto.schema.rs @@ -36,6 +36,35 @@ pub struct FieldSchema { /// enable logic partitions #[prost(bool, tag = "13")] pub is_partition_key: bool, + #[prost(bool, tag = "14")] + pub is_clustering_key: bool, + /// enable set null as field value + #[prost(bool, tag = "15")] + pub nullable: bool, + #[prost(bool, tag = "16")] + pub is_function_output: bool, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct FunctionSchema { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(int64, tag = "2")] + pub id: i64, + #[prost(string, tag = "3")] + pub description: ::prost::alloc::string::String, + #[prost(enumeration = "FunctionType", tag = "4")] + pub r#type: i32, + #[prost(string, repeated, tag = "5")] + pub input_field_names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(int64, repeated, tag = "6")] + pub input_field_ids: ::prost::alloc::vec::Vec, + #[prost(string, repeated, tag = "7")] + pub output_field_names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(int64, repeated, tag = "8")] + pub output_field_ids: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "9")] + pub params: ::prost::alloc::vec::Vec, } /// * /// @brief Collection schema @@ -55,6 +84,12 @@ pub struct CollectionSchema { /// mark whether this table has the dynamic field function enabled. #[prost(bool, tag = "5")] pub enable_dynamic_field: bool, + #[prost(message, repeated, tag = "6")] + pub properties: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "7")] + pub functions: ::prost::alloc::vec::Vec, + #[prost(string, tag = "8")] + pub db_name: ::prost::alloc::string::String, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -115,6 +150,12 @@ pub struct JsonArray { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct GeometryArray { + #[prost(bytes = "vec", repeated, tag = "1")] + pub data: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ValueField { #[prost(oneof = "value_field::Data", tags = "1, 2, 3, 4, 5, 6, 7")] pub data: ::core::option::Option, @@ -143,7 +184,7 @@ pub mod value_field { #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct ScalarField { - #[prost(oneof = "scalar_field::Data", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9")] + #[prost(oneof = "scalar_field::Data", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10")] pub data: ::core::option::Option, } /// Nested message and enum types in `ScalarField`. @@ -169,14 +210,27 @@ pub mod scalar_field { ArrayData(super::ArrayArray), #[prost(message, tag = "9")] JsonData(super::JsonArray), + #[prost(message, tag = "10")] + GeometryData(super::GeometryArray), } } +/// beta, api may change +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SparseFloatArray { + #[prost(bytes = "vec", repeated, tag = "1")] + pub contents: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, + /// dim is the max dimension of the current batch of vectors + #[prost(int64, tag = "2")] + pub dim: i64, +} #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct VectorField { + /// For sparse vector, dim is the max dimension of the current batch of vectors #[prost(int64, tag = "1")] pub dim: i64, - #[prost(oneof = "vector_field::Data", tags = "2, 3, 4, 5")] + #[prost(oneof = "vector_field::Data", tags = "2, 3, 4, 5, 6, 7")] pub data: ::core::option::Option, } /// Nested message and enum types in `VectorField`. @@ -192,6 +246,10 @@ pub mod vector_field { Float16Vector(::prost::alloc::vec::Vec), #[prost(bytes, tag = "5")] Bfloat16Vector(::prost::alloc::vec::Vec), + #[prost(message, tag = "6")] + SparseFloatVector(super::SparseFloatArray), + #[prost(bytes, tag = "7")] + Int8Vector(::prost::alloc::vec::Vec), } } #[allow(clippy::derive_partial_eq_without_eq)] @@ -205,6 +263,8 @@ pub struct FieldData { pub field_id: i64, #[prost(bool, tag = "6")] pub is_dynamic: bool, + #[prost(bool, repeated, tag = "7")] + pub valid_data: ::prost::alloc::vec::Vec, #[prost(oneof = "field_data::Field", tags = "3, 4")] pub field: ::core::option::Option, } @@ -238,6 +298,14 @@ pub mod i_ds { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct SearchIteratorV2Results { + #[prost(string, tag = "1")] + pub token: ::prost::alloc::string::String, + #[prost(float, tag = "2")] + pub last_bound: f32, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct SearchResultData { #[prost(int64, tag = "1")] pub num_queries: i64, @@ -255,6 +323,97 @@ pub struct SearchResultData { pub output_fields: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, #[prost(message, optional, tag = "8")] pub group_by_field_value: ::core::option::Option, + #[prost(int64, tag = "9")] + pub all_search_count: i64, + #[prost(float, repeated, tag = "10")] + pub distances: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "11")] + pub search_iterator_v2_results: ::core::option::Option, + #[prost(float, repeated, tag = "12")] + pub recalls: ::prost::alloc::vec::Vec, + #[prost(string, tag = "13")] + pub primary_field_name: ::prost::alloc::string::String, +} +/// vector field clustering info +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct VectorClusteringInfo { + /// for multi vectors + #[prost(string, tag = "1")] + pub field: ::prost::alloc::string::String, + #[prost(message, optional, tag = "2")] + pub centroid: ::core::option::Option, +} +/// Scalar field clustering info +/// todo more definitions: min/max, etc +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ScalarClusteringInfo { + #[prost(string, tag = "1")] + pub field: ::prost::alloc::string::String, +} +/// clustering distribution info of a certain data unit, it can be segment, partition, etc. +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ClusteringInfo { + #[prost(message, repeated, tag = "1")] + pub vector_clustering_infos: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "2")] + pub scalar_clustering_infos: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TemplateValue { + #[prost(oneof = "template_value::Val", tags = "1, 2, 3, 4, 5")] + pub val: ::core::option::Option, +} +/// Nested message and enum types in `TemplateValue`. +pub mod template_value { + #[allow(clippy::derive_partial_eq_without_eq)] + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Val { + #[prost(bool, tag = "1")] + BoolVal(bool), + #[prost(int64, tag = "2")] + Int64Val(i64), + #[prost(double, tag = "3")] + FloatVal(f64), + #[prost(string, tag = "4")] + StringVal(::prost::alloc::string::String), + #[prost(message, tag = "5")] + ArrayVal(super::TemplateArrayValue), + } +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TemplateArrayValue { + #[prost(oneof = "template_array_value::Data", tags = "1, 2, 3, 4, 5, 6")] + pub data: ::core::option::Option, +} +/// Nested message and enum types in `TemplateArrayValue`. +pub mod template_array_value { + #[allow(clippy::derive_partial_eq_without_eq)] + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Data { + #[prost(message, tag = "1")] + BoolData(super::BoolArray), + #[prost(message, tag = "2")] + LongData(super::LongArray), + #[prost(message, tag = "3")] + DoubleData(super::DoubleArray), + #[prost(message, tag = "4")] + StringData(super::StringArray), + #[prost(message, tag = "5")] + ArrayData(super::TemplateArrayValueArray), + #[prost(message, tag = "6")] + JsonData(super::JsonArray), + } +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TemplateArrayValueArray { + #[prost(message, repeated, tag = "1")] + pub data: ::prost::alloc::vec::Vec, } /// * /// @brief Field data type @@ -274,10 +433,14 @@ pub enum DataType { VarChar = 21, Array = 22, Json = 23, + Geometry = 24, + Text = 25, BinaryVector = 100, FloatVector = 101, Float16Vector = 102, BFloat16Vector = 103, + SparseFloatVector = 104, + Int8Vector = 105, } impl DataType { /// String value of the enum field names used in the ProtoBuf definition. @@ -298,10 +461,14 @@ impl DataType { DataType::VarChar => "VarChar", DataType::Array => "Array", DataType::Json => "JSON", + DataType::Geometry => "Geometry", + DataType::Text => "Text", DataType::BinaryVector => "BinaryVector", DataType::FloatVector => "FloatVector", DataType::Float16Vector => "Float16Vector", DataType::BFloat16Vector => "BFloat16Vector", + DataType::SparseFloatVector => "SparseFloatVector", + DataType::Int8Vector => "Int8Vector", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -319,10 +486,43 @@ impl DataType { "VarChar" => Some(Self::VarChar), "Array" => Some(Self::Array), "JSON" => Some(Self::Json), + "Geometry" => Some(Self::Geometry), + "Text" => Some(Self::Text), "BinaryVector" => Some(Self::BinaryVector), "FloatVector" => Some(Self::FloatVector), "Float16Vector" => Some(Self::Float16Vector), "BFloat16Vector" => Some(Self::BFloat16Vector), + "SparseFloatVector" => Some(Self::SparseFloatVector), + "Int8Vector" => Some(Self::Int8Vector), + _ => None, + } + } +} +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum FunctionType { + Unknown = 0, + Bm25 = 1, + TextEmbedding = 2, +} +impl FunctionType { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + FunctionType::Unknown => "Unknown", + FunctionType::Bm25 => "BM25", + FunctionType::TextEmbedding => "TextEmbedding", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "Unknown" => Some(Self::Unknown), + "BM25" => Some(Self::Bm25), + "TextEmbedding" => Some(Self::TextEmbedding), _ => None, } } diff --git a/src/proto/mod.rs b/src/proto/mod.rs index 224a71b..fecf15e 100644 --- a/src/proto/mod.rs +++ b/src/proto/mod.rs @@ -26,6 +26,8 @@ pub mod milvus; pub mod msg; #[path = "milvus.proto.schema.rs"] pub mod schema; +#[path = "milvus.proto.rg.rs"] +pub mod rg; impl MsgBase { pub fn new(msg_type: MsgType) -> Self { diff --git a/src/query.rs b/src/query.rs index 044e4c6..ac5e7b1 100644 --- a/src/query.rs +++ b/src/query.rs @@ -218,6 +218,7 @@ impl Client { not_return_all_meta: false, consistency_level: ConsistencyLevel::default() as _, use_default_consistency: false, + expr_template_values: HashMap::new(), }) .await? .into_inner(); @@ -295,6 +296,8 @@ impl Client { consistency_level: ConsistencyLevel::default() as _, use_default_consistency: false, search_by_primary_keys: false, + expr_template_values: HashMap::new(), + sub_reqs: Vec::new(), }) .await? .into_inner(); @@ -373,6 +376,7 @@ fn get_place_holder_value(vectors: Vec) -> Result { match vectors[0] { Value::FloatArray(_) => place_holder.r#type = PlaceholderType::FloatVector as _, Value::Binary(_) => place_holder.r#type = PlaceholderType::BinaryVector as _, + Value::String(_) => place_holder.r#type = PlaceholderType::VarChar as _, _ => { return Err(SuperError::from(crate::collection::Error::IllegalType( "place holder".to_string(), @@ -391,6 +395,7 @@ fn get_place_holder_value(vectors: Vec) -> Result { place_holder.values.push(bytes) } (Value::Binary(d), Value::Binary(_)) => place_holder.values.push(d.to_vec()), + (Value::String(d), Value::String(_)) => place_holder.values.push(d.as_bytes().to_vec()), _ => { return Err(SuperError::from(crate::collection::Error::IllegalType( "place holder".to_string(), diff --git a/src/schema.rs b/src/schema.rs index 73460b2..b12330c 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -14,9 +14,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::error; use crate::error::Result; use crate::proto::schema::FieldState; +use crate::{error, index::FunctionType}; use prost::alloc::vec::Vec; use prost::encoding::bool; use thiserror::Error as ThisError; @@ -153,8 +153,10 @@ pub struct FieldSchema { pub is_primary: bool, pub auto_id: bool, pub chunk_size: usize, - pub dim: i64, // only for BinaryVector and FloatVector - pub max_length: i32, // only for VarChar + pub dim: i64, // only for BinaryVector and FloatVector + pub max_length: i32, // only for VarChar + pub enable_analyzer: Option, // for BM25 tokenizer + pub enable_match: Option, // for BM25 match } impl FieldSchema { @@ -168,6 +170,8 @@ impl FieldSchema { chunk_size: 0, dim: 0, max_length: 0, + enable_analyzer: None, + enable_match: None, } } } @@ -202,6 +206,8 @@ impl From for FieldSchema { _ => dim, }) as _, dim, + enable_analyzer: None, + enable_match: None, } } } @@ -217,6 +223,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -230,6 +238,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -243,6 +253,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -256,6 +268,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -269,6 +283,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -282,6 +298,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -300,6 +318,8 @@ impl FieldSchema { max_length, chunk_size: 1, dim: 1, + enable_analyzer: None, + enable_match: None, } } @@ -313,6 +333,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -326,6 +348,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -339,6 +363,8 @@ impl FieldSchema { chunk_size: 1, dim: 1, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -356,6 +382,8 @@ impl FieldSchema { auto_id: false, chunk_size: 1, dim: 1, + enable_analyzer: None, + enable_match: None, } } @@ -373,6 +401,8 @@ impl FieldSchema { is_primary: false, auto_id: false, max_length: 0, + enable_analyzer: None, + enable_match: None, } } @@ -390,13 +420,30 @@ impl FieldSchema { is_primary: false, auto_id: false, max_length: 0, + enable_analyzer: None, + enable_match: None, + } + } + + pub fn new_sparse_float_vector(name: &str, description: &str) -> Self { + Self { + name: name.to_owned(), + description: description.to_owned(), + dtype: DataType::SparseFloatVector, + chunk_size: 0, + dim: 0, + is_primary: false, + auto_id: false, + max_length: 0, + enable_analyzer: None, + enable_match: None, } } } impl From for schema::FieldSchema { fn from(fld: FieldSchema) -> schema::FieldSchema { - let params = match fld.dtype { + let mut params = match fld.dtype { DataType::BinaryVector | DataType::FloatVector => vec![KeyValuePair { key: "dim".to_string(), value: fld.dim.to_string(), @@ -408,6 +455,20 @@ impl From for schema::FieldSchema { _ => Vec::new(), }; + if let Some(enable_analyzer) = fld.enable_analyzer { + params.push(KeyValuePair { + key: "enable_analyzer".to_string(), + value: enable_analyzer.to_string(), + }); + } + + if let Some(enable_match) = fld.enable_match { + params.push(KeyValuePair { + key: "enable_match".to_string(), + value: enable_match.to_string(), + }); + } + schema::FieldSchema { field_id: 0, name: fld.name.into(), @@ -422,6 +483,187 @@ impl From for schema::FieldSchema { default_value: None, is_dynamic: false, is_partition_key: false, + is_clustering_key: false, + nullable: false, + is_function_output: false, + } + } +} + +#[derive(Debug, Clone)] +pub struct FieldSchemaBuilder { + name: String, + description: String, + dtype: DataType, + chunk_size: usize, + dim: i64, + is_primary: bool, + auto_id: bool, + max_length: i32, + enable_analyzer: Option, + enable_match: Option, +} + +impl FieldSchemaBuilder { + pub fn new() -> Self { + Self { + name: String::new(), + description: String::new(), + dtype: DataType::None, + is_primary: false, + auto_id: false, + chunk_size: 1, + dim: 1, + max_length: 0, + enable_analyzer: None, + enable_match: None, + } + } + + pub fn with_name(mut self, name: &str) -> Self { + self.name = name.to_owned(); + self + } + + pub fn with_description(mut self, description: &str) -> Self { + self.description = description.to_owned(); + self + } + + pub fn with_dtype(mut self, dtype: DataType) -> Self { + self.dtype = dtype; + self + } + + pub fn with_chunk_size(mut self, chunk_size: usize) -> Self { + self.chunk_size = chunk_size; + self + } + + pub fn with_dim(mut self, dim: i64) -> Self { + self.dim = dim; + self + } + + pub fn with_primary(mut self, is_primary: bool) -> Self { + self.is_primary = is_primary; + self + } + + pub fn with_auto_id(mut self, auto_id: bool) -> Self { + self.auto_id = auto_id; + self + } + + pub fn with_max_length(mut self, max_length: i32) -> Self { + self.max_length = max_length; + self + } + + pub fn enable_analyzer(mut self, enable_analyzer: bool) -> Self { + self.enable_analyzer = Some(enable_analyzer); + self + } + + pub fn with_enable_match(mut self, enable_match: bool) -> Self { + self.enable_match = Some(enable_match); + self + } + + pub fn build(self) -> FieldSchema { + FieldSchema { + name: self.name, + description: self.description, + dtype: self.dtype, + is_primary: self.is_primary, + auto_id: self.auto_id, + chunk_size: self.chunk_size, + dim: self.dim, + max_length: self.max_length, + enable_analyzer: self.enable_analyzer, + enable_match: self.enable_match, + } + } +} + +#[derive(Debug, Clone)] +pub struct FunctionSchema { + pub name: String, + pub typ: FunctionType, + pub input_field_names: Vec, + pub output_field_names: Vec, +} + +impl From for schema::FunctionSchema { + fn from(value: FunctionSchema) -> Self { + Self { + name: value.name, + id: 0, + description: "".into(), + r#type: Into::::into(value.typ) as i32, + input_field_names: value.input_field_names, + input_field_ids: Vec::new(), + output_field_names: value.output_field_names, + output_field_ids: Vec::new(), + params: Vec::new(), + } + } +} + +impl From for FunctionSchema { + fn from(value: schema::FunctionSchema) -> Self { + Self { + name: value.name, + typ: schema::FunctionType::from_i32(value.r#type).unwrap().into(), + input_field_names: value.input_field_names, + output_field_names: value.output_field_names, + } + } +} + +pub struct FunctionSchemaBuilder { + pub name: String, + pub typ: FunctionType, + pub input_field_names: Vec, + pub output_field_names: Vec, +} + +impl FunctionSchemaBuilder { + pub fn new() -> Self { + Self { + name: String::new(), + typ: FunctionType::Unknown, + input_field_names: Vec::new(), + output_field_names: Vec::new(), + } + } + + pub fn with_name(mut self, name: &str) -> Self { + self.name = name.to_owned(); + self + } + + pub fn with_typ(mut self, typ: FunctionType) -> Self { + self.typ = typ; + self + } + + pub fn with_input_field_names(mut self, input_field_names: Vec) -> Self { + self.input_field_names = input_field_names; + self + } + + pub fn with_output_field_names(mut self, output_field_names: Vec) -> Self { + self.output_field_names = output_field_names; + self + } + + pub fn build(self) -> FunctionSchema { + FunctionSchema { + name: self.name, + typ: self.typ, + input_field_names: self.input_field_names, + output_field_names: self.output_field_names, } } } @@ -432,6 +674,7 @@ pub struct CollectionSchema { pub(crate) description: String, pub(crate) fields: Vec, pub(crate) enable_dynamic_field: bool, + pub(crate) functions: Vec, } impl CollectionSchema { @@ -486,6 +729,9 @@ impl From for schema::CollectionSchema { description: col.description, fields: col.fields.into_iter().map(Into::into).collect(), enable_dynamic_field: col.enable_dynamic_field, + properties: Vec::new(), + db_name: "".to_string(), + functions: col.functions.into_iter().map(Into::into).collect(), } } } @@ -497,6 +743,7 @@ impl From for CollectionSchema { name: v.name, description: v.description, enable_dynamic_field: v.enable_dynamic_field, + functions: v.functions.into_iter().map(Into::into).collect(), } } } @@ -507,6 +754,7 @@ pub struct CollectionSchemaBuilder { description: String, inner: Vec, enable_dynamic_field: bool, + functions: Vec, } impl CollectionSchemaBuilder { @@ -516,6 +764,7 @@ impl CollectionSchemaBuilder { description: description.to_owned(), inner: Vec::new(), enable_dynamic_field: false, + functions: Vec::new(), } } @@ -576,6 +825,11 @@ impl CollectionSchemaBuilder { self } + pub fn add_function(&mut self, schema: FunctionSchema) -> &mut Self { + self.functions.push(schema); + self + } + pub fn build(&mut self) -> Result { let mut has_primary = false; @@ -596,7 +850,8 @@ impl CollectionSchemaBuilder { fields: this.inner.into(), name: this.name, description: this.description, - enable_dynamic_field: self.enable_dynamic_field, + enable_dynamic_field: this.enable_dynamic_field, + functions: this.functions.into(), }) } } diff --git a/src/value.rs b/src/value.rs index f2c72e1..d837c01 100644 --- a/src/value.rs +++ b/src/value.rs @@ -8,6 +8,7 @@ use crate::proto::{ }, }; +#[derive(Debug)] pub enum Value<'a> { None, Bool(bool), @@ -22,6 +23,8 @@ pub enum Value<'a> { String(Cow<'a, str>), Json(Cow<'a, [u8]>), Array(Cow<'a, proto::schema::ScalarField>), + Bytes(Cow<'a, [u8]>), + Geometry(Cow<'a, [u8]>), } macro_rules! impl_from_for_field_data_column { @@ -60,6 +63,8 @@ impl Value<'_> { Value::FloatArray(_) => DataType::FloatVector, Value::Binary(_) => DataType::BinaryVector, Value::Array(_) => DataType::Array, + Value::Bytes(_) => DataType::Int8Vector, + Value::Geometry(_) => DataType::Geometry, } } } @@ -112,6 +117,8 @@ pub enum ValueVec { String(Vec), Json(Vec>), Array(Vec), + Bytes(Vec>), + Geometry(Vec>), } macro_rules! impl_from_for_value_vec { @@ -189,6 +196,10 @@ impl ValueVec { DataType::FloatVector => Self::Float(Vec::new()), DataType::Float16Vector => Self::Binary(Vec::new()), DataType::BFloat16Vector => Self::Binary(Vec::new()), + DataType::SparseFloatVector => Self::Bytes(Vec::new()), + DataType::Int8Vector => Self::Binary(Vec::new()), + DataType::Geometry => Self::Geometry(Vec::new()), + DataType::Text => Self::String(Vec::new()), } } @@ -227,6 +238,8 @@ impl ValueVec { ValueVec::String(v) => v.len(), ValueVec::Json(v) => v.len(), ValueVec::Array(v) => v.len(), + ValueVec::Bytes(v) => v.len(), + ValueVec::Geometry(v) => v.len(), } } @@ -242,6 +255,8 @@ impl ValueVec { ValueVec::String(v) => v.clear(), ValueVec::Json(v) => v.clear(), ValueVec::Array(v) => v.clear(), + ValueVec::Bytes(v) => v.clear(), + ValueVec::Geometry(v) => v.clear(), } } } @@ -259,7 +274,8 @@ impl From for ValueVec { ScalarData::StringData(v) => Self::String(v.data), ScalarData::JsonData(v) => Self::Json(v.data), ScalarData::ArrayData(v) => Self::Array(v.data), - ScalarData::BytesData(_) => unimplemented!(), // Self::Bytes(v.data), + ScalarData::BytesData(v) => Self::Bytes(v.data), + ScalarData::GeometryData(v) => Self::Geometry(v.data), }, None => Self::None, }, @@ -270,6 +286,8 @@ impl From for ValueVec { VectorData::BinaryVector(v) => Self::Binary(v), VectorData::Bfloat16Vector(v) => Self::Binary(v), VectorData::Float16Vector(v) => Self::Binary(v), + VectorData::Int8Vector(v) => Self::Binary(v), + VectorData::SparseFloatVector(v) => Self::Bytes(v.contents), }, None => Self::None, }, diff --git a/tests/alias.rs b/tests/alias.rs index 51b562e..8f614c0 100644 --- a/tests/alias.rs +++ b/tests/alias.rs @@ -10,7 +10,7 @@ async fn clean_test_collection(client: Client, collection_name: &str) -> Result< #[tokio::test] async fn test_create_alias() -> Result<()> { let alias = "test_create_alias"; - let (client, schema) = create_test_collection(true).await?; + let (client, schema) = create_test_collection(true, None).await?; client.create_alias(schema.name(), alias).await?; client.drop_alias(alias).await?; clean_test_collection(client, schema.name()).await?; @@ -20,9 +20,9 @@ async fn test_create_alias() -> Result<()> { #[tokio::test] async fn test_alter_alias() -> Result<()> { let alias = "test_alter_alias"; - let (client1, schema1) = create_test_collection(true).await?; + let (client1, schema1) = create_test_collection(true, None).await?; client1.create_alias(schema1.name(), alias).await?; - let (client2, schema2) = create_test_collection(true).await?; + let (client2, schema2) = create_test_collection(true, None).await?; client2.alter_alias(schema2.name(), alias).await?; client2.drop_alias(alias).await?; clean_test_collection(client1, schema1.name()).await?; diff --git a/tests/client.rs b/tests/client.rs index 8720e4a..7c0f47e 100644 --- a/tests/client.rs +++ b/tests/client.rs @@ -17,6 +17,7 @@ use milvus::client::*; use milvus::error::Result; use milvus::options::CreateCollectionOptions; +use milvus::proto::schema::DataType; use milvus::schema::*; mod common; @@ -73,8 +74,25 @@ async fn create_has_drop_collection() -> Result<()> { let mut schema = CollectionSchemaBuilder::new(NAME, "hello world"); let schema = schema - .add_field(FieldSchema::new_int64("i64_field", "")) - .add_field(FieldSchema::new_bool("bool_field", "")) + .add_field( + FieldSchemaBuilder::new() + .with_name("i64_field") + .with_dtype(DataType::Int64) + .build(), + ) + .add_field( + FieldSchemaBuilder::new() + .with_name("bool_field") + .with_dtype(DataType::Bool) + .build(), + ) + .add_field( + FieldSchemaBuilder::new() + .with_name("vec_field") + .with_dtype(DataType::FloatVector) + .with_dim(128) + .build(), + ) .set_primary_key("i64_field")? .enable_auto_id()? .build()?; @@ -107,8 +125,8 @@ async fn create_alter_drop_alias() -> Result<()> { let client = Client::new(URL).await?; - let (_, schema1) = create_test_collection(true).await?; - let (_, schema2) = create_test_collection(true).await?; + let (_, schema1) = create_test_collection(true, None).await?; + let (_, schema2) = create_test_collection(true, None).await?; client.create_alias(schema1.name(), &alias0).await?; assert!(client.has_collection(alias0).await?); diff --git a/tests/collection.rs b/tests/collection.rs index afb58f7..92d5600 100644 --- a/tests/collection.rs +++ b/tests/collection.rs @@ -14,230 +14,367 @@ // See the License for the specific language governing permissions and // limitations under the License. -use milvus::client::ConsistencyLevel; -use milvus::collection::{Collection, ParamValue}; +use milvus::collection::ParamValue; use milvus::data::FieldColumn; use milvus::error::Result; -use milvus::index::{IndexParams, IndexType, MetricType}; -use milvus::mutate::InsertOptions; +use milvus::index::{FunctionType, IndexParams, IndexType, MetricType}; use milvus::options::LoadOptions; +use milvus::proto::schema::DataType; use milvus::query::{QueryOptions, SearchOptions}; +use milvus::schema::{ + CollectionSchema, CollectionSchemaBuilder, FieldSchemaBuilder, FunctionSchemaBuilder, +}; use std::collections::HashMap; mod common; use common::*; -use milvus::value::ValueVec; +use milvus::value::{Value, ValueVec}; + +const METRIC_TYPE_LIST: &[MetricType; 2] = &[MetricType::L2, MetricType::COSINE]; #[tokio::test] async fn manual_compaction_empty_collection() -> Result<()> { - let (client, schema) = create_test_collection(true).await?; + let (client, schema) = create_test_collection(true, None).await?; let resp = client.manual_compaction(schema.name()).await?; assert_eq!(0, resp.plan_count); + + client.drop_collection(schema.name()).await?; Ok(()) } #[tokio::test] async fn collection_upsert() -> Result<()> { - let (client, schema) = create_test_collection(false).await?; - let pk_data = gen_random_int64_vector(2000); - let vec_data = gen_random_f32_vector(DEFAULT_DIM * 2000); - let pk_col = FieldColumn::new(schema.get_field("id").unwrap(), pk_data); - let vec_col = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), vec_data); - client - .upsert(schema.name(), vec![pk_col, vec_col], None) - .await?; - let index_params = IndexParams::new( - DEFAULT_INDEX_NAME.to_owned(), - IndexType::IvfFlat, - milvus::index::MetricType::L2, - HashMap::from([("nlist".to_owned(), "32".to_owned())]), - ); - client - .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) - .await?; - client - .load_collection(schema.name(), Some(LoadOptions::default())) - .await?; + for metric_type in METRIC_TYPE_LIST { + let (client, schema) = create_test_collection(false, None).await?; + let pk_data = gen_random_int64_vector(2000); + let vec_data = gen_random_f32_vector(DEFAULT_DIM * 2000); + let pk_col = FieldColumn::new(schema.get_field("id").unwrap(), pk_data); + let vec_col = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), vec_data); + client + .upsert(schema.name(), vec![pk_col, vec_col], None) + .await?; + let index_params = IndexParams::new( + DEFAULT_INDEX_NAME.to_owned(), + IndexType::IvfFlat, + *metric_type, + HashMap::from([("nlist".to_owned(), "32".to_owned())]), + ); + client + .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) + .await?; + client + .load_collection(schema.name(), Some(LoadOptions::default())) + .await?; + + let options = QueryOptions::default(); + let options = options.output_fields(vec![String::from("count(*)")]); + let result = client.query(schema.name(), "", &options).await?; + if let ValueVec::Long(vec) = &result[0].value { + assert_eq!(2000, vec[0]); + } else { + panic!("invalid result"); + } - let options = QueryOptions::default(); - let options = options.output_fields(vec![String::from("count(*)")]); - let result = client.query(schema.name(), "", &options).await?; - if let ValueVec::Long(vec) = &result[0].value { - assert_eq!(2000, vec[0]); - } else { - panic!("invalid result"); + client.drop_collection(schema.name()).await?; } Ok(()) } #[tokio::test] async fn collection_basic() -> Result<()> { - let (client, schema) = create_test_collection(true).await?; - - let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); - - let embed_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); - - client - .insert(schema.name(), vec![embed_column], None) - .await?; - client.flush(schema.name()).await?; - let index_params = IndexParams::new( - DEFAULT_INDEX_NAME.to_owned(), - IndexType::IvfFlat, - milvus::index::MetricType::L2, - HashMap::from([("nlist".to_owned(), "32".to_owned())]), - ); - client - .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) - .await?; - client - .load_collection(schema.name(), Some(LoadOptions::default())) - .await?; - - let options = QueryOptions::default(); - let result = client.query(schema.name(), "id > 0", &options).await?; - - println!( - "result num: {}", - result.first().map(|c| c.len()).unwrap_or(0), - ); - - client.drop_collection(schema.name()).await?; + for metric_type in METRIC_TYPE_LIST { + let (client, schema) = create_test_collection(true, None).await?; + + let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); + + let embed_column = + FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); + + client + .insert(schema.name(), vec![embed_column], None) + .await?; + client.flush(schema.name()).await?; + let index_params = IndexParams::new( + DEFAULT_INDEX_NAME.to_owned(), + IndexType::IvfFlat, + *metric_type, + HashMap::from([("nlist".to_owned(), "32".to_owned())]), + ); + client + .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) + .await?; + client + .load_collection(schema.name(), Some(LoadOptions::default())) + .await?; + + let options = QueryOptions::default(); + let result = client.query(schema.name(), "id > 0", &options).await?; + + println!( + "result num: {}", + result.first().map(|c| c.len()).unwrap_or(0), + ); + + client.drop_collection(schema.name()).await?; + } Ok(()) } #[tokio::test] async fn collection_index() -> Result<()> { - let (client, schema) = create_test_collection(true).await?; - - let feature = gen_random_f32_vector(DEFAULT_DIM * 2000); - - let feature_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), feature); - - client - .insert(schema.name(), vec![feature_column], None) - .await?; - client.flush(schema.name()).await?; - - let index_params = IndexParams::new( - DEFAULT_INDEX_NAME.to_owned(), - IndexType::IvfFlat, - milvus::index::MetricType::L2, - HashMap::from([("nlist".to_owned(), "32".to_owned())]), - ); - client - .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params.clone()) - .await?; - let index_list = client - .describe_index(schema.name(), DEFAULT_VEC_FIELD) - .await?; - assert!(index_list.len() == 1, "{}", index_list.len()); - let index = &index_list[0]; - - assert_eq!(index.params().name(), index_params.name()); - assert_eq!(index.params().extra_params(), index_params.extra_params()); - - client.drop_index(schema.name(), DEFAULT_VEC_FIELD).await?; - client.drop_collection(schema.name()).await?; + for metric_type in METRIC_TYPE_LIST { + let (client, schema) = create_test_collection(true, None).await?; + + let feature = gen_random_f32_vector(DEFAULT_DIM * 2000); + + let feature_column = + FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), feature); + + client + .insert(schema.name(), vec![feature_column], None) + .await?; + client.flush(schema.name()).await?; + + let index_params = IndexParams::new( + DEFAULT_INDEX_NAME.to_owned(), + IndexType::IvfFlat, + *metric_type, + HashMap::from([("nlist".to_owned(), "32".to_owned())]), + ); + client + .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params.clone()) + .await?; + let index_list = client + .describe_index(schema.name(), DEFAULT_VEC_FIELD) + .await?; + assert!(index_list.len() == 1, "{}", index_list.len()); + let index = &index_list[0]; + + assert_eq!(index.params().name(), index_params.name()); + assert_eq!(index.params().extra_params(), index_params.extra_params()); + + client.drop_index(schema.name(), DEFAULT_VEC_FIELD).await?; + client.drop_collection(schema.name()).await?; + } Ok(()) } #[tokio::test] async fn collection_search() -> Result<()> { - let (client, schema) = create_test_collection(true).await?; - - let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); - let embed_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); - - client - .insert(schema.name(), vec![embed_column], None) - .await?; - client.flush(schema.name()).await?; - let index_params = IndexParams::new( - "ivf_flat".to_owned(), - IndexType::IvfFlat, - MetricType::L2, - HashMap::from_iter([("nlist".to_owned(), 32.to_string())]), - ); - client - .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) - .await?; - client.flush(schema.name()).await?; - client - .load_collection(schema.name(), Some(LoadOptions::default())) - .await?; - - let mut option = SearchOptions::with_limit(10) - .metric_type(MetricType::L2) - .output_fields(vec!["id".to_owned()]); - option = option.add_param("nprobe", ParamValue!(16)); - let query_vec = gen_random_f32_vector(DEFAULT_DIM); - - let result = client - .search( - schema.name(), - vec![query_vec.into()], - DEFAULT_VEC_FIELD, - &option, - ) - .await?; - - assert_eq!(result[0].size, 10); - - client.drop_collection(schema.name()).await?; + for metric_type in METRIC_TYPE_LIST { + let (client, schema) = create_test_collection(true, None).await?; + + let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); + let embed_column = + FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); + + client + .insert(schema.name(), vec![embed_column], None) + .await?; + + let index_params = IndexParams::new( + "ivf_flat".to_owned(), + IndexType::IvfFlat, + *metric_type, + HashMap::from_iter([("nlist".to_owned(), 32.to_string())]), + ); + client + .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) + .await?; + client.flush(schema.name()).await?; + client + .load_collection(schema.name(), Some(LoadOptions::default())) + .await?; + + let mut option = SearchOptions::with_limit(10) + .metric_type(*metric_type) + .output_fields(vec!["id".to_owned()]); + option = option.add_param("nprobe", ParamValue!(16)); + let query_vec = gen_random_f32_vector(DEFAULT_DIM); + + let result = client + .search( + schema.name(), + vec![query_vec.into()], + DEFAULT_VEC_FIELD, + &option, + ) + .await?; + + assert_eq!(result[0].size, 10); + + client.drop_collection(schema.name()).await?; + } Ok(()) } #[tokio::test] async fn collection_range_search() -> Result<()> { - let (client, schema) = create_test_collection(true).await?; + for metric_type in METRIC_TYPE_LIST { + let (client, schema) = create_test_collection(true, None).await?; + + let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); + let embed_column = + FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); + + client + .insert(schema.name(), vec![embed_column], None) + .await?; + client.flush(schema.name()).await?; + let index_params = IndexParams::new( + "ivf_flat".to_owned(), + IndexType::IvfFlat, + *metric_type, + HashMap::from_iter([("nlist".to_owned(), 32.to_string())]), + ); + client + .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) + .await?; + + client + .load_collection(schema.name(), Some(LoadOptions::default())) + .await?; + + let radius_limit: f32 = match metric_type { + MetricType::L2 => 20.0, + MetricType::COSINE => 0.2, + _ => unimplemented!(), + }; + let range_filter_limit: f32 = match metric_type { + MetricType::L2 => 10.0, + MetricType::COSINE => 0.9, + _ => unimplemented!(), + }; + + let mut option = SearchOptions::with_limit(5) + .metric_type(*metric_type) + .output_fields(vec!["id".to_owned()]); + option = option.add_param("nprobe", ParamValue!(16)); + option = option.radius(radius_limit).range_filter(range_filter_limit); + let query_vec = gen_random_f32_vector(DEFAULT_DIM); + + let result = client + .search( + schema.name(), + vec![query_vec.into()], + DEFAULT_VEC_FIELD, + &option, + ) + .await?; + for record in &result { + for value in &record.score { + match metric_type { + MetricType::L2 => { + assert!(*value >= range_filter_limit && *value <= radius_limit) + } + MetricType::COSINE => { + assert!(*value >= radius_limit && *value <= range_filter_limit) + } + _ => unimplemented!(), + } + } + } - let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); - let embed_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); + client.drop_collection(schema.name()).await?; + } + Ok(()) +} +#[tokio::test] +async fn test_text_search() -> Result<()> { + let schema_fn: CollectionSchemaFn = |name: &str| -> Result { + Ok(CollectionSchemaBuilder::new(name, "") + .add_field( + FieldSchemaBuilder::new() + .with_name("id") + .with_dtype(DataType::Int64) + .with_primary(true) + .with_auto_id(true) + .build(), + ) + .add_field( + FieldSchemaBuilder::new() + .with_name("text") + .with_dtype(DataType::VarChar) + .with_max_length(1024) + .enable_analyzer(true) + .build(), + ) + .add_field( + FieldSchemaBuilder::new() + .with_name("sparse") + .with_dtype(DataType::SparseFloatVector) + .build(), + ) + .add_function( + FunctionSchemaBuilder::new() + .with_name("text2vec") + .with_typ(FunctionType::BM25) + .with_input_field_names(vec!["text".to_owned()]) + .with_output_field_names(vec!["sparse".to_owned()]) + .build(), + ) + .build()?) + }; + let (client, schema) = create_test_collection(true, Some(schema_fn)).await?; + + let index_params = IndexParams::new( + "sparse".to_owned(), + IndexType::SparseInvertedIndex, + MetricType::BM25, + HashMap::new(), + ); client - .insert(schema.name(), vec![embed_column], None) + .create_index(schema.name(), "sparse", index_params) .await?; - client.flush(schema.name()).await?; - let index_params = IndexParams::new( - "ivf_flat".to_owned(), - IndexType::IvfFlat, - MetricType::L2, - HashMap::from_iter([("nlist".to_owned(), 32.to_string())]), + + let text_column = FieldColumn::new( + schema.get_field("text").unwrap(), + vec![ + "information retrieval is a field of study.".to_owned(), + "information retrieval focuses on finding relevant information in large datasets." + .to_owned(), + "data mining and information retrieval overlap in research.".to_owned(), + ], ); client - .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) + .insert(schema.name(), vec![text_column], None) .await?; + client.flush(schema.name()).await?; client .load_collection(schema.name(), Some(LoadOptions::default())) .await?; - let radius_limit: f32 = 20.0; - let range_filter_limit: f32 = 10.0; - - let mut option = SearchOptions::with_limit(5) - .metric_type(MetricType::L2) - .output_fields(vec!["id".to_owned()]); - option = option.add_param("nprobe", ParamValue!(16)); - option = option.radius(radius_limit).range_filter(range_filter_limit); - let query_vec = gen_random_f32_vector(DEFAULT_DIM); - + let mut options = SearchOptions::default(); + options = options.limit(3); + options = options.output_fields(vec!["text".into()]); + options = options.add_param("drop_ratio_search", ParamValue!(0.2)); + options = options.metric_type(MetricType::BM25); let result = client .search( schema.name(), - vec![query_vec.into()], - DEFAULT_VEC_FIELD, - &option, + vec!["whats the focus of information retrieval?".into()], + "sparse", + &options, ) .await?; + assert!(result.len() == 1); for record in &result { for value in &record.score { - assert!(*value >= range_filter_limit && *value <= radius_limit); + assert!(*value >= 0.2); + } + } + + for col in &result[0].field { + assert!(col.name == "text"); + + if let ValueVec::String(vec) = &col.value { + assert_eq!(&vec[0], "information retrieval is a field of study."); } + break; } client.drop_collection(schema.name()).await?; diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 2e17886..0b44cd2 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -1,35 +1,53 @@ use milvus::client::*; use milvus::error::Result; use milvus::options::CreateCollectionOptions; -use milvus::schema::{CollectionSchema, CollectionSchemaBuilder, FieldSchema}; +use milvus::proto::schema::DataType; +use milvus::schema::{CollectionSchema, CollectionSchemaBuilder, FieldSchemaBuilder}; use rand::Rng; pub const DEFAULT_DIM: i64 = 128; pub const DEFAULT_VEC_FIELD: &str = "feature"; pub const DEFAULT_INDEX_NAME: &str = "feature_index"; -pub const URL: &str = "http://localhost:19530"; +pub const URL: &str = "http://127.0.0.1:19530"; -pub async fn create_test_collection(autoid: bool) -> Result<(Client, CollectionSchema)> { - let collection_name = gen_random_name(); - let collection_name = format!("{}_{}", "test_collection", collection_name); +pub type CollectionSchemaFn = fn(&str) -> Result; + +pub async fn create_test_collection( + autoid: bool, + schema: Option, +) -> Result<(Client, CollectionSchema)> { + let collection_name = format!("{}_{}", "test_collection", gen_random_name()); let client = Client::new(URL).await?; - let schema = CollectionSchemaBuilder::new(&collection_name, "") - .add_field(FieldSchema::new_primary_int64("id", "", autoid)) - .add_field(FieldSchema::new_float_vector( - DEFAULT_VEC_FIELD, - "", - DEFAULT_DIM, - )) - .build()?; + let schema = match schema { + Some(schema_fn) => schema_fn(&collection_name)?, + None => CollectionSchemaBuilder::new(&collection_name, "") + .add_field( + FieldSchemaBuilder::new() + .with_name("id") + .with_dtype(DataType::Int64) + .with_primary(true) + .with_auto_id(autoid) + .build(), + ) + .add_field( + FieldSchemaBuilder::new() + .with_name(DEFAULT_VEC_FIELD) + .with_dtype(DataType::FloatVector) + .with_dim(DEFAULT_DIM) + .build(), + ) + .build()?, + }; if client.has_collection(&collection_name).await? { client.drop_collection(&collection_name).await?; } client .create_collection( schema.clone(), - Some(CreateCollectionOptions::with_consistency_level( - ConsistencyLevel::Eventually, - )), + Some( + CreateCollectionOptions::with_consistency_level(ConsistencyLevel::Eventually) + .add_property("collection.insertRate.max.mb", "2000000"), + ), ) .await?; Ok((client, schema))