Implement describe_index_stats

pinecone_sdk/src/pinecone/data.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use crate::pinecone::PineconeClient;
 use crate::utils::errors::PineconeError;
 use once_cell::sync::Lazy;
@@ -8,7 +10,8 @@ use tonic::service::Interceptor;
 use tonic::transport::Channel;
 use tonic::{Request, Status};
 
-pub use pb::{UpsertResponse, Vector};
+pub use pb::{DescribeIndexStatsResponse, UpsertResponse, Vector};
+pub use prost_types::{value::Kind, Struct, Value};
 
 /// Generated protobuf module for data plane.
 pub mod pb {
@@ -53,7 +56,7 @@ impl Index {
  /// ### Example
  /// ```no_run
  /// use pinecone_sdk::pinecone::PineconeClient;
- /// use pinecone_sdk::pinecone::data::pb::Vector;
+ /// use pinecone_sdk::pinecone::data::Vector;
  /// # use pinecone_sdk::utils::errors::PineconeError;
  ///
  /// # #[tokio::main]
@@ -91,6 +94,54 @@ impl Index {
 
  Ok(response)
  }
+
+ /// The describe_index_stats operation returns statistics about the index.
+ ///
+ /// ### Arguments
+ /// * `filter: Option<BTreeMap<String, Value>>` - An optional filter to specify which vectors to return statistics for. Note that the filter is only supported by pod indexes.
+ ///
+ /// ### Return
+ /// * Returns a `Result<DescribeIndexStatsResponse, PineconeError>` object.
+ ///
+ /// ### Example
+ /// ```no_run
+ /// use pinecone_sdk::pinecone::PineconeClient;
+ /// use pinecone_sdk::pinecone::data::{Value, Kind};
+ /// use std::collections::BTreeMap;
+ /// # use pinecone_sdk::utils::errors::PineconeError;
+ ///
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<(), PineconeError>{
+ /// let pinecone = PineconeClient::new(None, None, None, None).unwrap();
+ ///
+ /// let mut index = pinecone.index("index-host").await.unwrap();
+ ///
+ /// let mut filter = BTreeMap::new();
+ /// filter.insert("field".to_string(), Value { kind: Some(Kind::StringValue("value".to_string())) });
+ ///
+ /// let response = index.describe_index_stats(None).await.unwrap();
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub async fn describe_index_stats(
+ &mut self,
+ filter: Option<BTreeMap<String, Value>>,
+ ) -> Result<DescribeIndexStatsResponse, PineconeError> {
+ let filter = match filter {
+ Some(filter) => Some(Struct { fields: filter }),
+ None => None,
+ };
+ let request = pb::DescribeIndexStatsRequest { filter };
+
+ let response = self
+ .connection
+ .describe_index_stats(request)
+ .await
+ .map_err(|e| PineconeError::DescribeIndexStatsError { source: e })?
+ .into_inner();
+
+ Ok(response)
+ }
 }
 
 impl PineconeClient {
@@ -172,14 +223,20 @@ impl PineconeClient {
 
  // connect to server
  let endpoint = Channel::from_shared(host)
- .map_err(|e| PineconeError::ConnectionError { source: Box::new(e) })?
+ .map_err(|e| PineconeError::ConnectionError {
+ source: Box::new(e),
+ })?
  .tls_config(tls_config)
- .map_err(|e| PineconeError::ConnectionError { source: Box::new(e) })?;
+ .map_err(|e| PineconeError::ConnectionError {
+ source: Box::new(e),
+ })?;
 
  let channel = endpoint
  .connect()
  .await
- .map_err(|e| PineconeError::ConnectionError { source: Box::new(e) })?;
+ .map_err(|e| PineconeError::ConnectionError {
+ source: Box::new(e),
+ })?;
 
  // add api key in metadata through interceptor
  let token: TonicMetadataVal<_> = self.api_key.parse().unwrap();

pinecone_sdk/src/utils/errors.rs

@@ -131,6 +131,12 @@ pub enum PineconeError {
  /// Source error
  source: tonic::Status,
  },
+
+ /// DescribeIndexStatsError: Failed to describe index stats.
+ DescribeIndexStatsError {
+ /// Source error
+ source: tonic::Status,
+ },
 }
 
 // Implement the conversion from OpenApiError to PineconeError for CreateIndexError.
@@ -198,67 +204,70 @@ impl std::fmt::Display for PineconeError {
  "Unknown response error: status: {}, message: {}",
  status, message
  )
- }
+ },
  PineconeError::ResourceAlreadyExistsError { source } => {
  write!(f, "Resource already exists error: {}", source)
- }
+ },
  PineconeError::UnprocessableEntityError { source } => {
  write!(f, "Unprocessable entity error: {}", source)
- }
+ },
  PineconeError::PendingCollectionError { source } => {
  write!(f, "Pending collection error: {}", source)
- }
+ },
  PineconeError::InternalServerError { source } => {
  write!(f, "Internal server error: {}", source)
- }
+ },
  PineconeError::ReqwestError { source } => {
  write!(f, "Reqwest error: {}", source.to_string())
- }
+ },
  PineconeError::SerdeError { source } => {
  write!(f, "Serde error: {}", source.to_string())
- }
+ },
  PineconeError::IoError { message } => {
  write!(f, "IO error: {}", message)
- }
+ },
  PineconeError::BadRequestError { source } => {
  write!(f, "Bad request error: {}", source)
- }
+ },
  PineconeError::UnauthorizedError { source } => {
  write!(f, "Unauthorized error: status: {}", source)
- }
+ },
  PineconeError::PodQuotaExceededError { source } => {
  write!(f, "Pod quota exceeded error: {}", source)
- }
+ },
  PineconeError::CollectionsQuotaExceededError { source } => {
  write!(f, "Collections quota exceeded error: {}", source)
- }
+ },
  PineconeError::InvalidCloudError { source } => {
  write!(f, "Invalid cloud error: status: {}", source)
- }
+ },
  PineconeError::InvalidRegionError { source } => {
  write!(f, "Invalid region error: {}", source)
- }
+ },
  PineconeError::CollectionNotFoundError { source } => {
  write!(f, "Collection not found error: {}", source)
- }
+ },
  PineconeError::IndexNotFoundError { source } => {
  write!(f, "Index not found error: status: {}", source)
- }
+ },
  PineconeError::APIKeyMissingError { message } => {
  write!(f, "API key missing error: {}", message)
- }
+ },
  PineconeError::InvalidHeadersError { message } => {
  write!(f, "Invalid headers error: {}", message)
- }
+ },
  PineconeError::TimeoutError { message } => {
  write!(f, "Timeout error: {}", message)
- }
+ },
  PineconeError::ConnectionError { source } => {
  write!(f, "Connection error: {}", source)
- }
+ },
  PineconeError::UpsertError { source } => {
  write!(f, "Upsert error: {}", source)
- }
+ },
+ PineconeError::DescribeIndexStatsError { source } => {
+ write!(f, "Describe index stats error: {}", source)
+ },
  }
  }
 }
@@ -290,6 +299,7 @@ impl std::error::Error for PineconeError {
  PineconeError::TimeoutError { message: _ } => None,
  PineconeError::ConnectionError { source } => Some(source.as_ref()),
  PineconeError::UpsertError { source } => Some(source),
+ PineconeError::DescribeIndexStatsError { source } => Some(source),
  }
  }
 }

pinecone_sdk/tests/integration_test.rs

@@ -1,9 +1,10 @@
 use openapi::models::index_model::Metric as OpenApiMetric;
 use openapi::models::serverless_spec::Cloud as OpenApiCloud;
 use pinecone_sdk::pinecone::control::{Cloud, Metric, WaitPolicy};
-use pinecone_sdk::pinecone::data::Vector;
+use pinecone_sdk::pinecone::data::{Kind, Value, Vector};
 use pinecone_sdk::pinecone::PineconeClient;
 use pinecone_sdk::utils::errors::PineconeError;
+use std::collections::BTreeMap;
 use std::time::Duration;
 use std::vec;
 
@@ -463,3 +464,61 @@ async fn test_upsert() -> Result<(), PineconeError> {
 
  Ok(())
 }
+
+#[tokio::test]
+async fn test_describe_index_stats_with_filter() -> Result<(), PineconeError> {
+ let pinecone = PineconeClient::new(None, None, None, None).unwrap();
+
+ let host = pinecone
+ .describe_index(&get_pod_index())
+ .await
+ .unwrap()
+ .host;
+
+ let mut index = pinecone
+ .index(host.as_str())
+ .await
+ .expect("Failed to target index");
+
+ let mut filter = BTreeMap::new();
+ filter.insert(
+ "id".to_string(),
+ Value {
+ kind: Some(Kind::BoolValue(false)),
+ },
+ );
+
+ let describe_index_stats_response = index
+ .describe_index_stats(Some(filter))
+ .await
+ .expect("Failed to describe index stats");
+
+ assert_eq!(describe_index_stats_response.dimension, 12);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_describe_index_stats_no_filter() -> Result<(), PineconeError> {
+ let pinecone = PineconeClient::new(None, None, None, None).unwrap();
+
+ let host = pinecone
+ .describe_index(&get_serverless_index())
+ .await
+ .unwrap()
+ .host;
+
+ let mut index = pinecone
+ .index(host.as_str())
+ .await
+ .expect("Failed to target index");
+
+ let describe_index_stats_response = index
+ .describe_index_stats(None)
+ .await
+ .expect("Failed to describe index stats");
+
+ assert_eq!(describe_index_stats_response.dimension, 4);
+
+ Ok(())
+}