Add apis for query and partition

What's changed:
1. client.rs: Add comment for list_aliases
2. collectoin.rs: Change the return type for describe_collection() from `collection` to `DescribeCollection`,containing more info than before,which now acts like pymilvus.
3. partitions.rs: Add load_partitions and release_partitions,following their options' struct.
4. query.rs: Add get() and it's following functions and struct,like extract_primary_field and pack_pks_expr,like pymilvus.Add options for query() to use and some build functions for QueryOptions.
5. test/: Change some settings and details to adapt to the changes above

Why Change Some Struct:
To catch up with pymivus,we need to approve optional params,so we need to change the option struct

Signed-off-by: Yinwei Li <yinwei.li@zilliz.com>

Author: Yinwei Li

src/client.rs

@@ -281,6 +281,15 @@ impl Client {
         Ok((resp.alias, resp.collection, resp.db_name))
     }
 
+    /// List a collection's aliases
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` containing the database name, collection name, and aliases.
     pub async fn list_aliases(
         &self,
         collection_name: &str,

src/collection.rs

@@ -60,6 +60,18 @@ pub struct Collection {
     // pub enable_dynamic_field: bool,
 }
 
+/// Return type for describe collection,containing enough messages
+pub struct DescribeCollection {
+    pub collection_name: String,
+    pub collection_id: i64,
+    pub shards_num: i32,
+    pub aliases: Vec<String>,
+    pub consistency_level: i32,
+    pub properties: Vec<proto::common::KeyValuePair>,
+    pub num_partitions: i64,
+    pub schema: crate::proto::schema::CollectionSchema,
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct CollectionCache {
     collections: dashmap::DashMap<String, Collection>,
@@ -275,7 +287,7 @@ impl Client {
     /// # Returns
     ///
     /// Returns a `Result` containing the `Collection` information if successful, or an error if the collection does not exist or cannot be accessed.
-    pub async fn describe_collection<S>(&self, name: S) -> Result<Collection>
+    pub async fn describe_collection<S>(&self, name: S) -> Result<DescribeCollection>
     where
         S: Into<String>,
     {
@@ -294,7 +306,16 @@ impl Client {
 
         status_to_result(&resp.status)?;
 
-        Ok(resp.into())
+        Ok(DescribeCollection {
+            collection_name: resp.collection_name,
+            collection_id: resp.collection_id,
+            shards_num: resp.shards_num,
+            aliases: resp.aliases,
+            consistency_level: resp.consistency_level,
+            properties: resp.properties,
+            num_partitions: resp.num_partitions,
+            schema: resp.schema.unwrap_or_default(),
+        })
     }
 
     /// Checks if a collection with the given name exists.
@@ -328,26 +349,30 @@ impl Client {
         Ok(res.value)
     }
 
-    
-    pub async fn rename_collection<S>(&self, name: S, new_name: S,options:Option<(String,String)>) -> Result<()>
+    pub async fn rename_collection<S>(
+        &self,
+        name: S,
+        new_name: S,
+        options: Option<(String, String)>,
+    ) -> Result<()>
     where
         S: Into<String>,
     {
         let name = name.into();
         let new_name = new_name.into();
-        let (db_name,new_db_name)=match options {
-            Some((db_name,new_db_name))=>(db_name,new_db_name),
-            None=>("".to_string(),"".to_string()),
+        let (db_name, new_db_name) = match options {
+            Some((db_name, new_db_name)) => (db_name, new_db_name),
+            None => ("".to_string(), "".to_string()),
         };
         let res = self
             .client
             .clone()
             .rename_collection(crate::proto::milvus::RenameCollectionRequest {
                 base: Some(MsgBase::new(MsgType::RenameCollection)),
-                db_name:db_name,
-                old_name:name,
-                new_name:new_name,
-                new_db_name:new_db_name,
+                db_name: db_name,
+                old_name: name,
+                new_name: new_name,
+                new_db_name: new_db_name,
             })
             .await?
             .into_inner();

src/partition.rs

@@ -10,7 +10,41 @@ use crate::{
     utils::status_to_result,
 };
 
+/// load_partitions' waitting time
+const WAIT_LOAD_DURATION_MS: u64 = 100;
+
+#[derive(Debug, Clone)]
+pub struct LoadPartitionsOption {
+    resource_groups: Vec<String>,
+    refresh: bool,
+    load_fields: Vec<String>,
+    skip_load_dynamic_field: bool,
+    load_params: HashMap<String, String>,
+}
+
+impl Default for LoadPartitionsOption {
+    fn default() -> Self {
+        LoadPartitionsOption {
+            resource_groups: Vec::new(),
+            refresh: false,
+            load_fields: Vec::new(),
+            skip_load_dynamic_field: false,
+            load_params: HashMap::new(),
+        }
+    }
+}
+
 impl Client {
+    /// Creates a new partition in the specified collection.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection where the partition will be created.
+    /// * `partition_name` - The name of the partition to be created.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` indicating success or failure.
     pub async fn create_partition(
         &self,
         collection_name: String,
@@ -30,6 +64,16 @@ impl Client {
         ))
     }
 
+    /// Drops a partition from the specified collection.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection containing the partition.
+    /// * `partition_name` - The name of the partition to be dropped.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` indicating success or failure.
     pub async fn drop_partition(
         &self,
         collection_name: String,
@@ -49,6 +93,15 @@ impl Client {
         ))
     }
 
+    /// Retrieves a list of all partitions in the specified collection.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection to list partitions for.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` containing a vector of partition names if successful, or an error if the operation fails.
     pub async fn list_partitions(&self, collection_name: String) -> Result<Vec<String>> {
         let res = self
             .client
@@ -67,6 +120,16 @@ impl Client {
         Ok(res.partition_names)
     }
 
+    /// Checks if a partition exists in the specified collection.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection to check.
+    /// * `partition_name` - The name of the partition to check for existence.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` containing a boolean indicating whether the partition exists.
     pub async fn has_partition(
         &self,
         collection_name: String,
@@ -87,6 +150,16 @@ impl Client {
         Ok(res.value)
     }
 
+    /// Retrieves statistics for a specific partition.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection containing the partition.
+    /// * `partition_name` - The name of the partition to get statistics for.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` containing a HashMap of statistics key-value pairs.
     pub async fn get_partition_stats(
         &self,
         collection_name: String,
@@ -108,51 +181,127 @@ impl Client {
         Ok(res.stats.into_iter().map(|s| (s.key, s.value)).collect())
     }
 
-    // pub async fn load_partitions<S: Into<String>, I: IntoIterator<Item = S>>(
-    //     &self,
-    //     collection_name: S,
-    //     partition_names: I,
-    //     replica_number: i32,
-    // ) -> Result<()> {
-    //     let names: Vec<String> = partition_names.into_iter().map(|x| x.to_string()).collect();
-    //     status_to_result(&Some(
-    //         self.client
-    //             .clone()
-    //             .load_partitions(proto::milvus::LoadPartitionsRequest {
-    //                 base: Some(MsgBase::new(MsgType::LoadPartitions)),
-    //                 db_name: "".to_string(),
-    //                 collection_name: collection_name.into(),
-    //                 replica_number,
-    //                 partition_names: names.clone(),
-    //             })
-    //             .await?
-    //             .into_inner(),
-    //     ))?;
-
-    //     loop {
-    //         if self.get_loading_progress(&names).await? >= 100 {
-    //             return Ok(());
-    //         }
-
-    //         tokio::time::sleep(Duration::from_millis(config::WAIT_LOAD_DURATION_MS)).await;
-    //     }
-    // }
-
-    // pub async fn release_partitions<S: ToString, I: IntoIterator<Item = S>>(
-    //     &self,
-    //     partition_names: I,
-    // ) -> Result<()> {
-    //     status_to_result(&Some(
-    //         self.client
-    //             .clone()
-    //             .release_partitions(ReleasePartitionsRequest {
-    //                 base: Some(MsgBase::new(MsgType::ReleasePartitions)),
-    //                 db_name: "".to_string(),
-    //                 collection_name: self.schema().name.to_string(),
-    //                 partition_names: partition_names.into_iter().map(|x| x.to_string()).collect(),
-    //             })
-    //             .await?
-    //             .into_inner(),
-    //     ))
-    // }
+    /// Gets the loading progress for specified partitions.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection.
+    /// * `partition_names` - An iterator of partition names to check progress for.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` containing the loading progress percentage (0-100).
+    async fn get_loading_progress<'a, S, I>(
+        &self,
+        collection_name: S,
+        partition_names: I,
+    ) -> Result<i64>
+    where
+        S: Into<String>,
+        I: IntoIterator<Item = &'a String>,
+    {
+        let partition_names: Vec<String> = partition_names.into_iter().map(|x| x.into()).collect();
+        let resp = self
+            .client
+            .clone()
+            .get_loading_progress(crate::proto::milvus::GetLoadingProgressRequest {
+                base: Some(MsgBase::new(MsgType::LoadPartitions)),
+                db_name: "".to_string(),
+                collection_name: collection_name.into(),
+                partition_names: partition_names,
+            })
+            .await?
+            .into_inner();
+
+        status_to_result(&resp.status)?;
+        Ok(resp.progress)
+    }
+
+    /// Loads partitions into memory with configurable options.
+    ///
+    /// This method loads the specified partitions into memory and waits for the loading
+    /// process to complete. The method polls the loading progress until it reaches 100%.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection containing the partitions.
+    /// * `partition_names` - An iterator of partition names to load.
+    /// * `replica_number` - The number of replicas to load.
+    /// * `options` - Optional configuration for the loading process.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` indicating success or failure. The method will wait until
+    /// all partitions are fully loaded before returning.
+    pub async fn load_partitions<S: Into<String> + Copy, I: IntoIterator<Item = S>>(
+        &self,
+        collection_name: S,
+        partition_names: I,
+        replica_number: i32,
+        options: Option<LoadPartitionsOption>,
+    ) -> Result<()> {
+        let names: Vec<String> = partition_names.into_iter().map(|x| x.into()).collect();
+        let options = options.unwrap_or_default();
+
+        status_to_result(&Some(
+            self.client
+                .clone()
+                .load_partitions(proto::milvus::LoadPartitionsRequest {
+                    base: Some(MsgBase::new(MsgType::LoadPartitions)),
+                    db_name: "".to_string(),
+                    collection_name: collection_name.into(),
+                    replica_number,
+                    partition_names: names.clone(),
+                    resource_groups: options.resource_groups,
+                    refresh: options.refresh,
+                    load_fields: options.load_fields,
+                    skip_load_dynamic_field: options.skip_load_dynamic_field,
+                    load_params: options.load_params,
+                })
+                .await?
+                .into_inner(),
+        ))?;
+
+        loop {
+            if self.get_loading_progress(collection_name, &names).await? >= 100 {
+                return Ok(());
+            }
+
+            tokio::time::sleep(tokio::time::Duration::from_millis(WAIT_LOAD_DURATION_MS)).await;
+        }
+    }
+
+    /// Releases partitions from memory.
+    ///
+    /// This method releases the specified partitions from memory, freeing up
+    /// system resources. After releasing, the partitions will no longer be
+    /// available for queries until they are loaded again.
+    ///
+    /// # Arguments
+    ///
+    /// * `collection_name` - The name of the collection containing the partitions.
+    /// * `partition_names` - An iterator of partition names to release.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` indicating success or failure.
+    pub async fn release_partitions<S: Into<String>, I: IntoIterator<Item = S>>(
+        &self,
+        collection_name: S,
+        partition_names: I,
+    ) -> Result<()> {
+        let names: Vec<String> = partition_names.into_iter().map(|x| x.into()).collect();
+        status_to_result(&Some(
+            self.client
+                .clone()
+                .release_partitions(crate::proto::milvus::ReleasePartitionsRequest {
+                    base: Some(MsgBase::new(MsgType::ReleasePartitions)),
+                    db_name: "".to_string(),
+                    collection_name: collection_name.into(),
+                    partition_names: names,
+                })
+                .await?
+                .into_inner(),
+        ))
+    }
 }