arc_vector_rust/

lib.rs

//! The ArcVector Vector Database client
//!
//! This library uses GRPC to connect to the ArcVector server and allows you to
//! access most if not all features. If you find a missing feature, please open
//! an [issue](https://github.com/arc_vector/rust-client/issues/new).
//!
//! If you use this library, you'll likely want to import the usual types and
//! functions:
//! ```
//!#[allow(unused_import)]
//! use arc_vector_rust::prelude::*;
//! ```
//!
//! To work with a ArcVector database, you'll first need to connect by creating a
//! [`ArcVectorClient`](crate::client::ArcVectorClient):
//! ```
//!# use arc_vector_rust::prelude::*;
//!# fn establish_connection(url: &str) -> anyhow::Result<ArcVectorClient> {
//! let mut config = ArcVectorClientConfig::from_url(url);
//! config.api_key = std::env::var("ARC_VECTOR_API_KEY").ok();
//! ArcVectorClient::new(Some(config))
//!# }
//! ```
//!
//! ArcVector works with *Collections* of *Points*. To add vector data, you first
//! create a collection:
//!
//! ```
//!# use arc_vector_rust::prelude::*;
//! use arc_vector_rust::arc_vector::{VectorParams, VectorsConfig};
//! use arc_vector_rust::arc_vector::vectors_config::Config;
//!# async fn create_collection(arc_vector_client: &ArcVectorClient)
//!# -> Result<(), Box<dyn std::error::Error>> {
//! let response = arc_vector_client
//!     .create_collection(&CreateCollection {
//!         collection_name: "my_collection".into(),
//!         vectors_config: Some(VectorsConfig {
//!             config: Some(Config::Params(VectorParams {
//!                 size: 512,
//!                 distance: Distance::Cosine as i32,
//!                 ..Default::default()
//!             })),
//!         }),
//!         ..Default::default()
//!     })
//!     .await?;
//!# Ok(())
//!# }
//! ```
//! The most interesting parts are the `collection_name` and the
//! `vectors_config.size` (the length of vectors to store) and `distance`
//! (which is the [`Distance`](crate::arc_vector::Distance) measure to gauge
//! similarity for the nearest neighbors search).
//!
//! Now we have a collection, we can insert (or rather upsert) points.
//! Points have an id, one or more vectors and a payload.
//! We can usually do that in bulk, but for this example, we'll add a
//! single point:
//! ```
//!# use arc_vector_rust::{prelude::*, arc_vector::PointId};
//!# async fn do_upsert(arc_vector_client: &ArcVectorClient)
//!# -> Result<(), Box<dyn std::error::Error>> {
//! let point = PointStruct {
//!     id: Some(PointId::from(42)), // unique u64 or String
//!     vectors: Some(vec![0.0_f32; 512].into()),
//!     payload: std::collections::HashMap::from([
//!         ("great".into(), Value::from(true)),
//!         ("level".into(), Value::from(9000)),
//!         ("text".into(), Value::from("Hi ArcVector!")),
//!         ("list".into(), Value::from(vec![1.234, 0.815])),
//!     ]),
//! };
//!
//! let response = arc_vector_client
//!     .upsert_points("my_collection", vec![point], None)
//!     .await?;
//!# Ok(())
//!# }
//! ```
//!
//! Finally, we can retrieve points in various ways, the canonical one being
//! a plain similarity search:
//! ```
//!# use arc_vector_rust::prelude::*;
//!# async fn search(arc_vector_client: &ArcVectorClient)
//!# -> Result<(), Box<dyn std::error::Error>> {
//! let response = arc_vector_client
//!     .search_points(&SearchPoints {
//!         collection_name: "my_collection".to_string(),
//!         vector: vec![0.0_f32; 512],
//!         limit: 4,
//!         with_payload: Some(true.into()),
//!         ..Default::default()
//!     })
//!     .await?;
//!# Ok(())
//!# }
//! ```
//!
//! You can also add a `filters: Some(filters)` field to the
//! [`SearchPoints`](crate::arc_vector::SearchPoints) argument to filter the
//! result. See the [`Filter`](crate::arc_vector::Filter) documentation for
//! details.

mod channel_pool;
pub mod client;
pub mod prelude;
// Do not lint/fmt code that is generated by tonic
#[allow(clippy::all)]
#[rustfmt::skip]
pub mod arc_vector;
pub mod filters;
#[cfg(feature = "serde")]
pub mod serde;

use arc_vector::{value::Kind::*, ListValue, RetrievedPoint, ScoredPoint, Struct, Value};

use std::error::Error;
use std::fmt::{Debug, Display, Formatter};

static NULL_VALUE: Value = Value {
    kind: Some(NullValue(0)),
};

macro_rules! get_payload {
    ($ty:ty) => {
        impl $ty {
            /// get a payload value for the specified key. If the key is not present,
            /// this will return a null value.
            ///
            /// # Examples:
            /// ```
            #[doc = concat!("use arc_vector_rust::arc_vector::", stringify!($ty), ";")]
            #[doc = concat!("let point = ", stringify!($ty), "::default();")]
            /// assert!(point.get("not_present").is_null());
            /// ````
            pub fn get(&self, key: &str) -> &Value {
                self.payload.get(key).unwrap_or(&NULL_VALUE)
            }
        }
    };
}

get_payload!(RetrievedPoint);
get_payload!(ScoredPoint);

macro_rules! extract {
    ($kind:ident, $check:ident) => {
        /// check if this value is a
        #[doc = stringify!($kind)]
        pub fn $check(&self) -> bool {
            matches!(self.kind, Some($kind(_)))
        }
    };
    ($kind:ident, $check:ident, $extract:ident, $ty:ty) => {
        extract!($kind, $check);

        /// extract the contents if this value is a
        #[doc = stringify!($kind)]
        pub fn $extract(&self) -> Option<$ty> {
            if let Some($kind(v)) = self.kind {
                Some(v)
            } else {
                None
            }
        }
    };
    ($kind:ident, $check:ident, $extract:ident, ref $ty:ty) => {
        extract!($kind, $check);

        /// extract the contents if this value is a
        #[doc = stringify!($kind)]
        pub fn $extract(&self) -> Option<&$ty> {
            if let Some($kind(v)) = &self.kind {
                Some(v)
            } else {
                None
            }
        }
    };
}

impl Value {
    extract!(NullValue, is_null);
    extract!(BoolValue, is_bool, as_bool, bool);
    extract!(IntegerValue, is_integer, as_integer, i64);
    extract!(DoubleValue, is_double, as_double, f64);
    extract!(StringValue, is_str, as_str, ref String);
    extract!(ListValue, is_list, as_list, ref [Value]);
    extract!(StructValue, is_struct, as_struct, ref Struct);

    #[cfg(feature = "serde")]
    /// convert this into a `serde_json::Value`
    ///
    /// # Examples:
    ///
    /// ```
    /// use serde_json::json;
    /// use arc_vector_rust::prelude::*;
    /// use arc_vector_rust::arc_vector::{value::Kind::*, Struct};
    /// let value = Value { kind: Some(StructValue(Struct {
    ///     fields: [
    ///         ("text".into(), Value { kind: Some(StringValue("Hi ArcVector!".into())) }),
    ///         ("int".into(), Value { kind: Some(IntegerValue(42))}),
    ///     ].into()
    /// }))};
    /// assert_eq!(value.into_json(), json!({
    ///    "text": "Hi ArcVector!",
    ///    "int": 42
    /// }));
    /// ```
    pub fn into_json(self) -> serde_json::Value {
        use serde_json::Value as JsonValue;
        match self.kind {
            Some(BoolValue(b)) => JsonValue::Bool(b),
            Some(IntegerValue(i)) => JsonValue::from(i),
            Some(DoubleValue(d)) => JsonValue::from(d),
            Some(StringValue(s)) => JsonValue::String(s),
            Some(ListValue(vs)) => vs.into_iter().map(Value::into_json).collect(),
            Some(StructValue(s)) => s
                .fields
                .into_iter()
                .map(|(k, v)| (k, v.into_json()))
                .collect(),
            Some(NullValue(_)) | None => JsonValue::Null,
        }
    }
}

#[cfg(feature = "serde")]
impl From<Value> for serde_json::Value {
    fn from(value: Value) -> Self {
        value.into_json()
    }
}

impl Display for Value {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match &self.kind {
            Some(BoolValue(b)) => write!(f, "{}", b),
            Some(IntegerValue(i)) => write!(f, "{}", i),
            Some(DoubleValue(v)) => write!(f, "{}", v),
            Some(StringValue(s)) => write!(f, "{:?}", s),
            Some(ListValue(vs)) => {
                let mut i = vs.values.iter();
                write!(f, "[")?;
                if let Some(first) = i.next() {
                    write!(f, "{}", first)?;
                    for v in i {
                        write!(f, ",{}", v)?;
                    }
                }
                write!(f, "]")
            }
            Some(StructValue(s)) => {
                let mut i = s.fields.iter();
                write!(f, "{{")?;
                if let Some((key, value)) = i.next() {
                    write!(f, "{:?}:{}", key, value)?;
                    for (key, value) in i {
                        write!(f, ",{:?}:{}", key, value)?;
                    }
                }
                write!(f, "}}")
            }
            _ => write!(f, "null"),
        }
    }
}

pub mod error {
    use std::marker::PhantomData;

    /// An error for failed conversions (e.g. calling `String::try_from(v)`
    /// on an integer [`Value`](crate::Value))
    pub struct NotA<T> {
        marker: PhantomData<T>,
    }

    impl<T> Default for NotA<T> {
        fn default() -> Self {
            NotA {
                marker: PhantomData,
            }
        }
    }
}

use error::NotA;

macro_rules! not_a {
    ($ty:ty) => {
        impl Error for NotA<$ty> {}

        impl Debug for NotA<$ty> {
            fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
                write!(f, "{}", self)
            }
        }

        impl Display for NotA<$ty> {
            fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
                f.write_str(concat!("not a ", stringify!($ty)))
            }
        }
    };
}

macro_rules! impl_try_from {
    ($ty:ty, $key:ident) => {
        not_a!($ty);

        impl std::convert::TryFrom<Value> for $ty {
            type Error = NotA<$ty>;

            fn try_from(v: Value) -> Result<Self, NotA<$ty>> {
                if let Some($key(t)) = v.kind {
                    Ok(t)
                } else {
                    Err(NotA::default())
                }
            }
        }
    };
}

impl_try_from!(bool, BoolValue);
impl_try_from!(i64, IntegerValue);
impl_try_from!(f64, DoubleValue);
impl_try_from!(String, StringValue);

not_a!(ListValue);
not_a!(Struct);

impl Value {
    /// try to get an iterator over the items of the contained list value, if any
    pub fn iter_list(&self) -> Result<impl Iterator<Item = &Value>, NotA<ListValue>> {
        if let Some(ListValue(values)) = &self.kind {
            Ok(values.iter())
        } else {
            Err(NotA::default())
        }
    }

    /// try to get a field from the struct if this value contains one
    pub fn get_struct(&self, key: &str) -> Result<&Value, NotA<Struct>> {
        if let Some(StructValue(Struct { fields })) = &self.kind {
            Ok(fields.get(key).unwrap_or(&NULL_VALUE))
        } else {
            Err(NotA::default())
        }
    }
}

impl std::ops::Deref for ListValue {
    type Target = [Value];

    fn deref(&self) -> &[Value] {
        &self.values
    }
}

impl IntoIterator for ListValue {
    type Item = Value;

    type IntoIter = std::vec::IntoIter<Value>;

    fn into_iter(self) -> Self::IntoIter {
        self.values.into_iter()
    }
}

impl ListValue {
    pub fn iter(&self) -> std::slice::Iter<'_, Value> {
        self.values.iter()
    }
}

#[cfg(test)]
mod tests {
    use crate::prelude::*;
    use crate::arc_vector::value::Kind::*;
    use crate::arc_vector::vectors_config::Config;
    use crate::arc_vector::{
        CreateFieldIndexCollection, FieldType, ListValue, Struct, Value, VectorParams,
        VectorsConfig,
    };
    use std::collections::HashMap;

    #[test]
    fn display() {
        let value = Value {
            kind: Some(StructValue(Struct {
                fields: [
                    ("text", StringValue("Hi ArcVector!".into())),
                    ("int", IntegerValue(42)),
                    ("float", DoubleValue(1.23)),
                    (
                        "list",
                        ListValue(ListValue {
                            values: vec![Value {
                                kind: Some(NullValue(0)),
                            }],
                        }),
                    ),
                    (
                        "struct",
                        StructValue(Struct {
                            fields: [(
                                "bool".into(),
                                Value {
                                    kind: Some(BoolValue(true)),
                                },
                            )]
                            .into(),
                        }),
                    ),
                ]
                .into_iter()
                .map(|(k, v)| (k.into(), Value { kind: Some(v) }))
                .collect(),
            })),
        };
        let text = format!("{}", value);
        assert!([
            "\"float\":1.23",
            "\"list\":[null]",
            "\"struct\":{\"bool\":true}",
            "\"int\":42",
            "\"text\":\"Hi ArcVector!\""
        ]
        .into_iter()
        .all(|item| text.contains(item)));
    }

    #[tokio::test]
    async fn test_arc_vector_queries() -> anyhow::Result<()> {
        let config = ArcVectorClientConfig::from_url("http://localhost:6334");
        let client = ArcVectorClient::new(Some(config))?;

        let health = client.health_check().await?;
        println!("{:?}", health);

        let collections_list = client.list_collections().await?;
        println!("{:?}", collections_list);

        let collection_name = "test";
        client.delete_collection(collection_name).await?;

        client
            .create_collection(&CreateCollection {
                collection_name: collection_name.into(),
                vectors_config: Some(VectorsConfig {
                    config: Some(Config::Params(VectorParams {
                        size: 10,
                        distance: Distance::Cosine.into(),
                        hnsw_config: None,
                        quantization_config: None,
                        on_disk: None,
                    })),
                }),
                ..Default::default()
            })
            .await?;

        let collection_info = client.collection_info(collection_name).await?;
        println!("{:#?}", collection_info);

        let mut sub_payload = Payload::new();
        sub_payload.insert("foo", "Not bar");

        let payload: Payload = vec![
            ("foo", "Bar".into()),
            ("bar", 12.into()),
            ("sub_payload", sub_payload.into()),
        ]
        .into_iter()
        .collect::<HashMap<_, Value>>()
        .into();

        let points = vec![PointStruct::new(0, vec![12.; 10], payload)];
        client
            .upsert_points_blocking(collection_name, points, None)
            .await?;

        let search_result = client
            .search_points(&SearchPoints {
                collection_name: collection_name.into(),
                vector: vec![11.; 10],
                filter: None,
                limit: 10,
                with_payload: Some(true.into()),
                params: None,
                score_threshold: None,
                offset: None,
                vector_name: None,
                with_vectors: None,
                read_consistency: None,
            })
            .await?;

        eprintln!("search_result = {:#?}", search_result);

        // Override payload of the existing point
        let new_payload: Payload = vec![("foo", "BAZ".into())]
            .into_iter()
            .collect::<HashMap<_, Value>>()
            .into();
        client
            .set_payload(collection_name, &vec![0.into()].into(), new_payload, None)
            .await?;

        // Delete some payload fields
        client
            .delete_payload_blocking(
                collection_name,
                &vec![0.into()].into(),
                vec!["sub_payload".to_string()],
                None,
            )
            .await?;

        // retrieve points
        let points = client
            .get_points(collection_name, &[0.into()], Some(true), Some(true), None)
            .await?;

        assert_eq!(points.result.len(), 1);
        let point = points.result[0].clone();
        assert!(point.payload.contains_key("foo"));
        assert!(!point.payload.contains_key("sub_payload"));

        client
            .delete_points(collection_name, &vec![0.into()].into(), None)
            .await?;

        // Access raw point api with client
        client
            .with_points_client(|mut client| async move {
                client
                    .create_field_index(CreateFieldIndexCollection {
                        collection_name: collection_name.to_string(),
                        wait: None,
                        field_name: "foo".to_string(),
                        field_type: Some(FieldType::Keyword as i32),
                        field_index_params: None,
                        ordering: None,
                    })
                    .await
            })
            .await?;

        client.create_snapshot(collection_name).await?;
        #[cfg(feature = "download_snapshots")]
        client
            .download_snapshot("test.tar", collection_name, None, None)
            .await?;

        Ok(())
    }
}