Skip to content

Commit

Permalink
Add dense, sparse and multi dense vector builders and conversions
Browse files Browse the repository at this point in the history
  • Loading branch information
timvisee committed Jan 15, 2025
1 parent e5dec4c commit a8bf329
Show file tree
Hide file tree
Showing 5 changed files with 213 additions and 35 deletions.
48 changes: 48 additions & 0 deletions src/builders/dense_vector_builder.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
use crate::qdrant::*;

pub struct DenseVectorBuilder {
pub(crate) values: Vec<f32>,
}

impl DenseVectorBuilder {
pub fn new(values: impl Into<Vec<f32>>) -> Self {
Self {
values: values.into(),
}
}

#[allow(unused_mut)]
pub fn values(mut self, values: impl Into<Vec<f32>>) -> 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<Vec<f32>> for DenseVector {
fn from(values: Vec<f32>) -> Self {
DenseVectorBuilder::new(values).build()
}
}

impl From<DenseVector> for Vector {
fn from(dense_vector: DenseVector) -> Self {
crate::qdrant::vector::Vector::from(dense_vector).into()
}
}

impl From<DenseVectorBuilder> for Vector {
fn from(dense_vector: DenseVectorBuilder) -> Self {
crate::qdrant::vector::Vector::from(dense_vector.build()).into()
}
}

impl From<DenseVector> for crate::qdrant::vector::Vector {
fn from(dense_vector: DenseVector) -> Self {
Self::Dense(dense_vector)
}
}
9 changes: 9 additions & 0 deletions src/builders/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
72 changes: 72 additions & 0 deletions src/builders/multi_dense_vector_builder.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
use crate::qdrant::*;

#[derive(Default)]
pub struct MultiDenseVectorBuilder {
pub(crate) vectors: Vec<DenseVector>,
}

impl MultiDenseVectorBuilder {
pub fn new(vectors: impl Into<Vec<DenseVector>>) -> Self {
Self {
vectors: vectors.into(),
}
}

pub fn single(vector: impl Into<DenseVector>) -> Self {
Self::new(vec![vector.into()])
}

#[allow(unused_mut)]
pub fn vectors(mut self, vectors: impl Into<Vec<DenseVector>>) -> Self {
self.vectors = vectors.into();
self
}

#[allow(unused_mut)]
pub fn add_vector(mut self, vector: impl Into<DenseVector>) -> 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<Vec<Vec<f32>>> for MultiDenseVector {
fn from(vectors: Vec<Vec<f32>>) -> Self {
Self::from(
vectors
.into_iter()
.map(DenseVector::from)
.collect::<Vec<_>>(),
)
}
}

impl From<Vec<DenseVector>> for MultiDenseVector {
fn from(vectors: Vec<DenseVector>) -> Self {
MultiDenseVectorBuilder::new(vectors).build()
}
}

impl From<MultiDenseVector> for Vector {
fn from(dense_vector: MultiDenseVector) -> Self {
crate::qdrant::vector::Vector::from(dense_vector).into()
}
}

impl From<MultiDenseVectorBuilder> for Vector {
fn from(dense_vector: MultiDenseVectorBuilder) -> Self {
crate::qdrant::vector::Vector::from(dense_vector.build()).into()
}
}

impl From<MultiDenseVector> for crate::qdrant::vector::Vector {
fn from(dense_vector: MultiDenseVector) -> Self {
Self::MultiDense(dense_vector)
}
}
60 changes: 60 additions & 0 deletions src/builders/sparse_vector_builder.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
use crate::qdrant::*;

#[derive(Default)]
pub struct SparseVectorBuilder {
pub(crate) indices: Vec<u32>,
pub(crate) values: Vec<f32>,
}

impl SparseVectorBuilder {
pub fn new(indices: impl Into<Vec<u32>>, values: impl Into<Vec<f32>>) -> Self {
Self {
indices: indices.into(),
values: values.into(),
}
}

#[allow(unused_mut)]
pub fn indices(mut self, indices: impl Into<Vec<u32>>) -> Self {
self.indices = indices.into();
self
}

#[allow(unused_mut)]
pub fn values(mut self, values: impl Into<Vec<f32>>) -> 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<u32>, Vec<f32>)> for SparseVector {
fn from((indices, values): (Vec<u32>, Vec<f32>)) -> Self {
SparseVectorBuilder::new(indices, values).build()
}
}

impl From<SparseVector> for Vector {
fn from(dense_vector: SparseVector) -> Self {
crate::qdrant::vector::Vector::from(dense_vector).into()
}
}

impl From<SparseVectorBuilder> for Vector {
fn from(dense_vector: SparseVectorBuilder) -> Self {
crate::qdrant::vector::Vector::from(dense_vector.build()).into()
}
}

impl From<SparseVector> for crate::qdrant::vector::Vector {
fn from(dense_vector: SparseVector) -> Self {
Self::Sparse(dense_vector)
}
}
59 changes: 24 additions & 35 deletions src/qdrant_client/builders/vectors.rs
Original file line number Diff line number Diff line change
@@ -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<f32>) -> Self {
Self::new_dense(values)
}

pub fn new_dense(values: Vec<f32>) -> 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<Vec<f32>>) -> Self {
DenseVectorBuilder::new(values.into()).build().into()
}

#[inline]
pub fn new_sparse(indices: impl Into<Vec<u32>>, values: impl Into<Vec<f32>>) -> 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<Vec<f32>>) -> 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<Vec<Vec<f32>>>) -> Self {
MultiDenseVector::from(vectors.into()).into()
}

pub fn try_into_dense(self) -> Result<Vec<f32>, QdrantError> {
Expand Down Expand Up @@ -120,3 +97,15 @@ impl NamedVectors {
self
}
}

impl From<crate::qdrant::vector::Vector> for Vector {
fn from(vector: crate::qdrant::vector::Vector) -> Self {
Vector {
vector: Some(vector),
// Deprecated
data: vec![],
indices: None,
vectors_count: None,
}
}
}

0 comments on commit a8bf329

Please sign in to comment.