From a8bf329dfcbe38546cc4f74d3d8510b3878056f3 Mon Sep 17 00:00:00 2001 From: timvisee Date: Wed, 15 Jan 2025 15:29:04 +0100 Subject: [PATCH] Add dense, sparse and multi dense vector builders and conversions --- src/builders/dense_vector_builder.rs | 48 +++++++++++++++ src/builders/mod.rs | 9 +++ src/builders/multi_dense_vector_builder.rs | 72 ++++++++++++++++++++++ src/builders/sparse_vector_builder.rs | 60 ++++++++++++++++++ src/qdrant_client/builders/vectors.rs | 59 ++++++++---------- 5 files changed, 213 insertions(+), 35 deletions(-) create mode 100644 src/builders/dense_vector_builder.rs create mode 100644 src/builders/multi_dense_vector_builder.rs create mode 100644 src/builders/sparse_vector_builder.rs diff --git a/src/builders/dense_vector_builder.rs b/src/builders/dense_vector_builder.rs new file mode 100644 index 0000000..fdf705c --- /dev/null +++ b/src/builders/dense_vector_builder.rs @@ -0,0 +1,48 @@ +use crate::qdrant::*; + +pub struct DenseVectorBuilder { + pub(crate) values: Vec, +} + +impl DenseVectorBuilder { + pub fn new(values: impl Into>) -> Self { + Self { + values: values.into(), + } + } + + #[allow(unused_mut)] + pub fn values(mut self, values: impl Into>) -> Self { + self.values = values.into(); + self + } + + /// Builds the desired type. Can often be omitted. + pub fn build(self) -> DenseVector { + DenseVector { data: self.values } + } +} + +impl From> for DenseVector { + fn from(values: Vec) -> Self { + DenseVectorBuilder::new(values).build() + } +} + +impl From for Vector { + fn from(dense_vector: DenseVector) -> Self { + crate::qdrant::vector::Vector::from(dense_vector).into() + } +} + +impl From for Vector { + fn from(dense_vector: DenseVectorBuilder) -> Self { + crate::qdrant::vector::Vector::from(dense_vector.build()).into() + } +} + +impl From for crate::qdrant::vector::Vector { + fn from(dense_vector: DenseVector) -> Self { + Self::Dense(dense_vector) + } +} diff --git a/src/builders/mod.rs b/src/builders/mod.rs index 61e80d1..1d94883 100644 --- a/src/builders/mod.rs +++ b/src/builders/mod.rs @@ -219,3 +219,12 @@ pub use count_points_builder::*; pub mod has_vector_condition_builder; pub use has_vector_condition_builder::HasVectorConditionBuilder; + +pub mod dense_vector_builder; +pub use dense_vector_builder::DenseVectorBuilder; + +pub mod sparse_vector_builder; +pub use sparse_vector_builder::SparseVectorBuilder; + +pub mod multi_dense_vector_builder; +pub use multi_dense_vector_builder::MultiDenseVectorBuilder; diff --git a/src/builders/multi_dense_vector_builder.rs b/src/builders/multi_dense_vector_builder.rs new file mode 100644 index 0000000..59ac8d9 --- /dev/null +++ b/src/builders/multi_dense_vector_builder.rs @@ -0,0 +1,72 @@ +use crate::qdrant::*; + +#[derive(Default)] +pub struct MultiDenseVectorBuilder { + pub(crate) vectors: Vec, +} + +impl MultiDenseVectorBuilder { + pub fn new(vectors: impl Into>) -> Self { + Self { + vectors: vectors.into(), + } + } + + pub fn single(vector: impl Into) -> Self { + Self::new(vec![vector.into()]) + } + + #[allow(unused_mut)] + pub fn vectors(mut self, vectors: impl Into>) -> Self { + self.vectors = vectors.into(); + self + } + + #[allow(unused_mut)] + pub fn add_vector(mut self, vector: impl Into) -> Self { + self.vectors.push(vector.into()); + self + } + + /// Builds the desired type. Can often be omitted. + pub fn build(self) -> MultiDenseVector { + MultiDenseVector { + vectors: self.vectors, + } + } +} + +impl From>> for MultiDenseVector { + fn from(vectors: Vec>) -> Self { + Self::from( + vectors + .into_iter() + .map(DenseVector::from) + .collect::>(), + ) + } +} + +impl From> for MultiDenseVector { + fn from(vectors: Vec) -> Self { + MultiDenseVectorBuilder::new(vectors).build() + } +} + +impl From for Vector { + fn from(dense_vector: MultiDenseVector) -> Self { + crate::qdrant::vector::Vector::from(dense_vector).into() + } +} + +impl From for Vector { + fn from(dense_vector: MultiDenseVectorBuilder) -> Self { + crate::qdrant::vector::Vector::from(dense_vector.build()).into() + } +} + +impl From for crate::qdrant::vector::Vector { + fn from(dense_vector: MultiDenseVector) -> Self { + Self::MultiDense(dense_vector) + } +} diff --git a/src/builders/sparse_vector_builder.rs b/src/builders/sparse_vector_builder.rs new file mode 100644 index 0000000..c225608 --- /dev/null +++ b/src/builders/sparse_vector_builder.rs @@ -0,0 +1,60 @@ +use crate::qdrant::*; + +#[derive(Default)] +pub struct SparseVectorBuilder { + pub(crate) indices: Vec, + pub(crate) values: Vec, +} + +impl SparseVectorBuilder { + pub fn new(indices: impl Into>, values: impl Into>) -> Self { + Self { + indices: indices.into(), + values: values.into(), + } + } + + #[allow(unused_mut)] + pub fn indices(mut self, indices: impl Into>) -> Self { + self.indices = indices.into(); + self + } + + #[allow(unused_mut)] + pub fn values(mut self, values: impl Into>) -> Self { + self.values = values.into(); + self + } + + /// Builds the desired type. Can often be omitted. + pub fn build(self) -> SparseVector { + SparseVector { + indices: self.indices, + values: self.values, + } + } +} + +impl From<(Vec, Vec)> for SparseVector { + fn from((indices, values): (Vec, Vec)) -> Self { + SparseVectorBuilder::new(indices, values).build() + } +} + +impl From for Vector { + fn from(dense_vector: SparseVector) -> Self { + crate::qdrant::vector::Vector::from(dense_vector).into() + } +} + +impl From for Vector { + fn from(dense_vector: SparseVectorBuilder) -> Self { + crate::qdrant::vector::Vector::from(dense_vector.build()).into() + } +} + +impl From for crate::qdrant::vector::Vector { + fn from(dense_vector: SparseVector) -> Self { + Self::Sparse(dense_vector) + } +} diff --git a/src/qdrant_client/builders/vectors.rs b/src/qdrant_client/builders/vectors.rs index 9f82378..457c61e 100644 --- a/src/qdrant_client/builders/vectors.rs +++ b/src/qdrant_client/builders/vectors.rs @@ -1,50 +1,27 @@ -use crate::qdrant::{DenseVector, MultiDenseVector, NamedVectors, SparseVector, Vector}; +use crate::qdrant::{ + DenseVectorBuilder, MultiDenseVector, NamedVectors, SparseVectorBuilder, Vector, +}; use crate::QdrantError; impl Vector { + #[inline] pub fn new(values: Vec) -> Self { Self::new_dense(values) } - pub fn new_dense(values: Vec) -> Self { - Vector { - vector: Some(crate::qdrant::vector::Vector::Dense(DenseVector { - data: values, - })), - // Deprecated - data: vec![], - indices: None, - vectors_count: None, - } + #[inline] + pub fn new_dense(values: impl Into>) -> Self { + DenseVectorBuilder::new(values.into()).build().into() } + #[inline] pub fn new_sparse(indices: impl Into>, values: impl Into>) -> Self { - Vector { - vector: Some(crate::qdrant::vector::Vector::Sparse(SparseVector { - values: values.into(), - indices: indices.into(), - })), - // Deprecated - data: vec![], - indices: None, - vectors_count: None, - } + SparseVectorBuilder::new(indices, values).build().into() } - pub fn new_multi(values: Vec>) -> Self { - let vectors = values - .into_iter() - .map(|data| DenseVector { data }) - .collect(); - Vector { - vector: Some(crate::qdrant::vector::Vector::MultiDense( - MultiDenseVector { vectors }, - )), - // Deprecated - data: vec![], - indices: None, - vectors_count: None, - } + #[inline] + pub fn new_multi(vectors: impl Into>>) -> Self { + MultiDenseVector::from(vectors.into()).into() } pub fn try_into_dense(self) -> Result, QdrantError> { @@ -120,3 +97,15 @@ impl NamedVectors { self } } + +impl From for Vector { + fn from(vector: crate::qdrant::vector::Vector) -> Self { + Vector { + vector: Some(vector), + // Deprecated + data: vec![], + indices: None, + vectors_count: None, + } + } +}