From 1d723fe2481c5f374981b3a581857561b87f39b1 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 23 Sep 2024 15:50:33 +0800 Subject: [PATCH] refactor: make SQ build & preprocess faster (#596) * chore: rename inverted_index to sparse_inverted_index Signed-off-by: usamoi * refactor: make scalar quantization faster Signed-off-by: usamoi --------- Signed-off-by: usamoi --- Cargo.lock | 22 +- crates/base/src/index.rs | 64 +++--- crates/base/src/operator/mod.rs | 1 - crates/flat/src/lib.rs | 4 +- crates/hnsw/src/lib.rs | 6 +- crates/index/src/lib.rs | 4 +- crates/index/src/optimizing/index_source.rs | 16 +- crates/index/src/optimizing/indexing.rs | 5 +- crates/index/src/segment/sealed.rs | 2 +- crates/indexing/Cargo.toml | 2 +- crates/indexing/src/lib.rs | 6 +- crates/indexing/src/sealed.rs | 28 +-- crates/ivf/src/lib.rs | 4 +- crates/ivf/src/operator.rs | 16 +- crates/pyvectors/pyproject.toml | 1 + crates/quantization/src/lib.rs | 8 +- crates/quantization/src/product.rs | 8 +- crates/quantization/src/quantizer.rs | 10 +- crates/quantization/src/rabitq.rs | 16 +- crates/quantization/src/scalar.rs | 208 ++++++++++++------ crates/quantization/src/trivial.rs | 10 +- .../Cargo.toml | 2 +- .../src/lib.rs | 22 +- .../src/operator.rs | 44 ++-- crates/storage/src/lib.rs | 2 +- 25 files changed, 293 insertions(+), 218 deletions(-) rename crates/{inverted => sparse_inverted_index}/Cargo.toml (88%) rename crates/{inverted => sparse_inverted_index}/src/lib.rs (88%) rename crates/{inverted => sparse_inverted_index}/src/operator.rs (51%) diff --git a/Cargo.lock b/Cargo.lock index 10a31d200..30b8ba9f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1518,9 +1518,9 @@ dependencies = [ "base", "flat", "hnsw", - "inverted", "ivf", "quantization", + "sparse_inverted_index", "thiserror", ] @@ -1557,16 +1557,6 @@ dependencies = [ "ulock-sys", ] -[[package]] -name = "inverted" -version = "0.0.0" -dependencies = [ - "base", - "common", - "quantization", - "storage", -] - [[package]] name = "io-lifetimes" version = "1.0.11" @@ -3028,6 +3018,16 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "sparse_inverted_index" +version = "0.0.0" +dependencies = [ + "base", + "common", + "quantization", + "storage", +] + [[package]] name = "spin" version = "0.9.8" diff --git a/crates/base/src/index.rs b/crates/base/src/index.rs index 344d384c2..e79b8fa1e 100644 --- a/crates/base/src/index.rs +++ b/crates/base/src/index.rs @@ -104,51 +104,49 @@ pub struct IndexOptions { } impl IndexOptions { - fn validate_self_quantization( - &self, - quantization: &Option, - ) -> Result<(), ValidationError> { - match quantization { - None => Ok(()), - Some( - QuantizationOptions::Scalar(_) - | QuantizationOptions::Product(_) - | QuantizationOptions::Rabitq(_), - ) => { - if !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16) { - return Err(ValidationError::new( - "quantization is not support for vectors that are not dense vectors", - )); - } - Ok(()) - } - } - } fn validate_self(&self) -> Result<(), ValidationError> { match &self.indexing { IndexingOptions::Flat(FlatIndexingOptions { quantization }) => { - self.validate_self_quantization(quantization)?; + if quantization.is_some() + && !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16) + { + return Err(ValidationError::new( + "quantization is only supported for dense vectors", + )); + } } IndexingOptions::Ivf(IvfIndexingOptions { quantization, .. }) => { if !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16) { return Err(ValidationError::new( - "ivf is not support for vectors that are not dense vectors", + "ivf is only supported for dense vectors", + )); + } + if quantization.is_some() + && !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16) + { + return Err(ValidationError::new( + "quantization is only supported for dense vectors", )); } - self.validate_self_quantization(quantization)?; } IndexingOptions::Hnsw(HnswIndexingOptions { quantization, .. }) => { - self.validate_self_quantization(quantization)?; - } - IndexingOptions::InvertedIndex(_) => { - if !matches!(self.vector.d, DistanceKind::Dot) { + if quantization.is_some() + && !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16) + { return Err(ValidationError::new( - "inverted_index is not support for distance that is not negative dot product", + "quantization is only supported for dense vectors", )); } + } + IndexingOptions::SparseInvertedIndex(_) => { if !matches!(self.vector.v, VectorKind::SVecf32) { return Err(ValidationError::new( - "inverted_index is not support for vectors that are not sparse vectors", + "sparse_inverted_index is only supported for sparse vectors", + )); + } + if !matches!(self.vector.d, DistanceKind::Dot) { + return Err(ValidationError::new( + "sparse_inverted_index is only supported for dot distance", )); } } @@ -284,7 +282,7 @@ pub enum IndexingOptions { Flat(FlatIndexingOptions), Ivf(IvfIndexingOptions), Hnsw(HnswIndexingOptions), - InvertedIndex(InvertedIndexingOptions), + SparseInvertedIndex(SparseInvertedIndexIndexingOptions), } impl IndexingOptions { @@ -320,16 +318,16 @@ impl Validate for IndexingOptions { Self::Flat(x) => x.validate(), Self::Ivf(x) => x.validate(), Self::Hnsw(x) => x.validate(), - Self::InvertedIndex(x) => x.validate(), + Self::SparseInvertedIndex(x) => x.validate(), } } } #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct InvertedIndexingOptions {} +pub struct SparseInvertedIndexIndexingOptions {} -impl Default for InvertedIndexingOptions { +impl Default for SparseInvertedIndexIndexingOptions { fn default() -> Self { Self {} } diff --git a/crates/base/src/operator/mod.rs b/crates/base/src/operator/mod.rs index 4e21cf655..b3b140e7c 100644 --- a/crates/base/src/operator/mod.rs +++ b/crates/base/src/operator/mod.rs @@ -23,5 +23,4 @@ pub trait Operator: Copy + 'static + Send + Sync { fn distance(lhs: Borrowed<'_, Self>, rhs: Borrowed<'_, Self>) -> Distance; } -pub type Owned = ::Vector; pub type Borrowed<'a, T> = <::Vector as VectorOwned>::Borrowed<'a>; diff --git a/crates/flat/src/lib.rs b/crates/flat/src/lib.rs index 763bcbdec..a70995632 100644 --- a/crates/flat/src/lib.rs +++ b/crates/flat/src/lib.rs @@ -29,7 +29,7 @@ impl> Flat { pub fn create( path: impl AsRef, options: IndexOptions, - source: &(impl Vectors> + Collection + Source + Sync), + source: &(impl Vectors + Collection + Source + Sync), ) -> Self { let remapped = RemappedCollection::from_source(source); from_nothing(path, options, &remapped) @@ -83,7 +83,7 @@ impl> Flat { fn from_nothing>( path: impl AsRef, options: IndexOptions, - collection: &(impl Vectors> + Collection + Sync), + collection: &(impl Vectors + Collection + Sync), ) -> Flat { create_dir(path.as_ref()).unwrap(); let flat_indexing_options = options.indexing.clone().unwrap_flat(); diff --git a/crates/hnsw/src/lib.rs b/crates/hnsw/src/lib.rs index 4ca38d280..3de8afd00 100644 --- a/crates/hnsw/src/lib.rs +++ b/crates/hnsw/src/lib.rs @@ -46,7 +46,7 @@ impl> Hnsw { pub fn create( path: impl AsRef, options: IndexOptions, - source: &(impl Vectors> + Collection + Source + Sync), + source: &(impl Vectors + Collection + Source + Sync), ) -> Self { let remapped = RemappedCollection::from_source(source); if let Some(main) = source.get_main::() { @@ -116,7 +116,7 @@ impl> Hnsw { fn from_nothing>( path: impl AsRef, options: IndexOptions, - collection: &(impl Vectors> + Collection + Sync), + collection: &(impl Vectors + Collection + Sync), ) -> Hnsw { create_dir(path.as_ref()).unwrap(); let HnswIndexingOptions { @@ -198,7 +198,7 @@ fn from_nothing>( fn from_main>( path: impl AsRef, options: IndexOptions, - remapped: &RemappedCollection, impl Vectors> + Collection + Sync>, + remapped: &RemappedCollection + Collection + Sync>, main: &Hnsw, ) -> Hnsw { create_dir(path.as_ref()).unwrap(); diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs index 9a92b5b74..6d791ee0d 100644 --- a/crates/index/src/lib.rs +++ b/crates/index/src/lib.rs @@ -316,7 +316,7 @@ impl Index { } pub fn create_sealed_segment( &self, - source: &(impl Vectors> + Collection + Source + Sync), + source: &(impl Vectors + Collection + Source + Sync), sealed_segment_ids: &[NonZeroU128], growing_segment_ids: &[NonZeroU128], ) -> Option>> { @@ -444,7 +444,7 @@ impl IndexView { } pub fn insert( &self, - vector: Owned, + vector: O::Vector, pointer: Pointer, ) -> Result, InsertError> { if self.options.vector.dims != vector.as_borrowed().dims() { diff --git a/crates/index/src/optimizing/index_source.rs b/crates/index/src/optimizing/index_source.rs index ce98f2e80..55554bf1b 100644 --- a/crates/index/src/optimizing/index_source.rs +++ b/crates/index/src/optimizing/index_source.rs @@ -2,7 +2,7 @@ use crate::delete::Delete; use crate::Op; use crate::{GrowingSegment, SealedSegment}; use base::index::IndexOptions; -use base::operator::{Borrowed, Owned}; +use base::operator::Borrowed; use base::search::*; use std::any::Any; use std::fmt::Debug; @@ -17,7 +17,7 @@ pub struct IndexSource { _phantom: PhantomData V>, } -impl IndexSource, O> { +impl IndexSource { pub fn new( options: IndexOptions, sealed: Option>>, @@ -34,7 +34,7 @@ impl IndexSource, O> { } } -impl Vectors> for IndexSource, O> { +impl Vectors for IndexSource { fn dims(&self) -> u32 { self.dims } @@ -61,7 +61,7 @@ impl Vectors> for IndexSource, O> { } } -impl Collection for IndexSource, O> { +impl Collection for IndexSource { fn payload(&self, mut index: u32) -> Payload { for x in self.sealed.iter() { if index < x.len() { @@ -79,7 +79,7 @@ impl Collection for IndexSource, O> { } } -impl Source for IndexSource, O> { +impl Source for IndexSource { fn get_main(&self) -> Option<&T> { let x = self.sealed.as_ref()?; Some( @@ -104,7 +104,7 @@ pub struct RoGrowingCollection { _phantom: PhantomData V>, } -impl Debug for RoGrowingCollection, O> { +impl Debug for RoGrowingCollection { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("RoGrowingCollection") .field("growing", &self.growing) @@ -113,7 +113,7 @@ impl Debug for RoGrowingCollection, O> { } } -impl Vectors> for RoGrowingCollection, O> { +impl Vectors for RoGrowingCollection { fn dims(&self) -> u32 { self.dims } @@ -133,7 +133,7 @@ impl Vectors> for RoGrowingCollection, O> { } } -impl Collection for RoGrowingCollection, O> { +impl Collection for RoGrowingCollection { fn payload(&self, mut index: u32) -> Payload { for x in self.growing.iter() { if index < x.len() { diff --git a/crates/index/src/optimizing/indexing.rs b/crates/index/src/optimizing/indexing.rs index 1cd23cfeb..01d1a4335 100644 --- a/crates/index/src/optimizing/indexing.rs +++ b/crates/index/src/optimizing/indexing.rs @@ -1,14 +1,13 @@ use crate::optimizing::index_source::IndexSource; use crate::Index; use crate::Op; -use base::operator::Owned; use std::sync::Arc; pub fn scan( index: Arc>, capacity: u32, delete_threshold: f64, -) -> Option, O>> { +) -> Option> { let (sealed, growing) = 'a: { let protect = index.protect.lock(); // approach 1: merge small segments to a big segment @@ -87,7 +86,7 @@ pub fn scan( )) } -pub fn make(index: Arc>, source: IndexSource, O>) { +pub fn make(index: Arc>, source: IndexSource) { let _ = index.create_sealed_segment( &source, &source.sealed.iter().map(|x| x.id()).collect::>(), diff --git a/crates/index/src/segment/sealed.rs b/crates/index/src/segment/sealed.rs index de8d91631..4d7a0fc71 100644 --- a/crates/index/src/segment/sealed.rs +++ b/crates/index/src/segment/sealed.rs @@ -37,7 +37,7 @@ impl SealedSegment { path: PathBuf, id: NonZeroU128, options: IndexOptions, - source: &(impl Vectors> + Collection + Source + Sync), + source: &(impl Vectors + Collection + Source + Sync), ) -> Arc { let indexing = SealedIndexing::create(&path, options, source); Arc::new(Self { diff --git a/crates/indexing/Cargo.toml b/crates/indexing/Cargo.toml index 0d5ed0976..d2a0f28fd 100644 --- a/crates/indexing/Cargo.toml +++ b/crates/indexing/Cargo.toml @@ -11,9 +11,9 @@ base = { path = "../base" } # algorithms flat = { path = "../flat" } hnsw = { path = "../hnsw" } -inverted = { path = "../inverted" } ivf = { path = "../ivf" } quantization = { path = "../quantization" } +sparse_inverted_index = { path = "../sparse_inverted_index" } [lints] workspace = true diff --git a/crates/indexing/src/lib.rs b/crates/indexing/src/lib.rs index 19af8cd36..6b1632838 100644 --- a/crates/indexing/src/lib.rs +++ b/crates/indexing/src/lib.rs @@ -4,16 +4,16 @@ use quantization::rabitq::OperatorRabitqQuantization; pub use sealed::SealedIndexing; use base::operator::Operator; -use inverted::operator::OperatorInvertedIndex; use ivf::operator::OperatorIvf; use quantization::product::OperatorProductQuantization; use quantization::scalar::OperatorScalarQuantization; +use sparse_inverted_index::operator::OperatorSparseInvertedIndex; pub trait OperatorIndexing where Self: Operator, Self: OperatorIvf, - Self: OperatorInvertedIndex, + Self: OperatorSparseInvertedIndex, Self: OperatorScalarQuantization, Self: OperatorProductQuantization, Self: OperatorRabitqQuantization, @@ -24,7 +24,7 @@ impl OperatorIndexing for T where Self: Operator, Self: OperatorIvf, - Self: OperatorInvertedIndex, + Self: OperatorSparseInvertedIndex, Self: OperatorScalarQuantization, Self: OperatorProductQuantization, Self: OperatorRabitqQuantization, diff --git a/crates/indexing/src/sealed.rs b/crates/indexing/src/sealed.rs index a01729eb1..da5accb53 100644 --- a/crates/indexing/src/sealed.rs +++ b/crates/indexing/src/sealed.rs @@ -4,12 +4,12 @@ use base::operator::*; use base::search::*; use flat::Flat; use hnsw::Hnsw; -use inverted::InvertedIndex; use ivf::Ivf; use quantization::product::ProductQuantizer; use quantization::rabitq::RabitqQuantizer; use quantization::scalar::ScalarQuantizer; use quantization::trivial::TrivialQuantizer; +use sparse_inverted_index::SparseInvertedIndex; use std::any::Any; use std::path::Path; @@ -26,14 +26,14 @@ pub enum SealedIndexing { HnswSq(Hnsw>), HnswPq(Hnsw>), HnswRq(Hnsw>), - InvertedIndex(InvertedIndex), + SparseInvertedIndex(SparseInvertedIndex), } impl SealedIndexing { pub fn create( path: impl AsRef, options: IndexOptions, - source: &(impl Vectors> + Collection + Source + Sync), + source: &(impl Vectors + Collection + Source + Sync), ) -> Self { match options.indexing { IndexingOptions::Flat(FlatIndexingOptions { @@ -78,8 +78,8 @@ impl SealedIndexing { Self::HnswRq(Hnsw::create(path, options, source)) } }, - IndexingOptions::InvertedIndex(_) => { - Self::InvertedIndex(InvertedIndex::create(path, options, source)) + IndexingOptions::SparseInvertedIndex(_) => { + Self::SparseInvertedIndex(SparseInvertedIndex::create(path, options, source)) } } } @@ -110,7 +110,9 @@ impl SealedIndexing { Some(QuantizationOptions::Product(_)) => Self::HnswPq(Hnsw::open(path)), Some(QuantizationOptions::Rabitq(_)) => Self::HnswRq(Hnsw::open(path)), }, - IndexingOptions::InvertedIndex(_) => Self::InvertedIndex(InvertedIndex::open(path)), + IndexingOptions::SparseInvertedIndex(_) => { + Self::SparseInvertedIndex(SparseInvertedIndex::open(path)) + } } } @@ -132,7 +134,7 @@ impl SealedIndexing { SealedIndexing::HnswPq(x) => x.vbase(vector, opts), SealedIndexing::HnswSq(x) => x.vbase(vector, opts), SealedIndexing::HnswRq(x) => x.vbase(vector, opts), - SealedIndexing::InvertedIndex(x) => x.vbase(vector, opts), + SealedIndexing::SparseInvertedIndex(x) => x.vbase(vector, opts), } } @@ -150,12 +152,12 @@ impl SealedIndexing { SealedIndexing::HnswPq(x) => x, SealedIndexing::HnswSq(x) => x, SealedIndexing::HnswRq(x) => x, - SealedIndexing::InvertedIndex(x) => x, + SealedIndexing::SparseInvertedIndex(x) => x, } } } -impl Vectors> for SealedIndexing { +impl Vectors for SealedIndexing { fn dims(&self) -> u32 { match self { SealedIndexing::Flat(x) => x.dims(), @@ -170,7 +172,7 @@ impl Vectors> for SealedIndexing { SealedIndexing::HnswPq(x) => x.dims(), SealedIndexing::HnswSq(x) => x.dims(), SealedIndexing::HnswRq(x) => x.dims(), - SealedIndexing::InvertedIndex(x) => x.dims(), + SealedIndexing::SparseInvertedIndex(x) => x.dims(), } } @@ -188,7 +190,7 @@ impl Vectors> for SealedIndexing { SealedIndexing::HnswPq(x) => x.len(), SealedIndexing::HnswSq(x) => x.len(), SealedIndexing::HnswRq(x) => x.len(), - SealedIndexing::InvertedIndex(x) => x.len(), + SealedIndexing::SparseInvertedIndex(x) => x.len(), } } @@ -206,7 +208,7 @@ impl Vectors> for SealedIndexing { SealedIndexing::HnswSq(x) => x.vector(i), SealedIndexing::HnswPq(x) => x.vector(i), SealedIndexing::HnswRq(x) => x.vector(i), - SealedIndexing::InvertedIndex(x) => x.vector(i), + SealedIndexing::SparseInvertedIndex(x) => x.vector(i), } } } @@ -226,7 +228,7 @@ impl Collection for SealedIndexing { SealedIndexing::HnswPq(x) => x.payload(i), SealedIndexing::HnswSq(x) => x.payload(i), SealedIndexing::HnswRq(x) => x.payload(i), - SealedIndexing::InvertedIndex(x) => x.payload(i), + SealedIndexing::SparseInvertedIndex(x) => x.payload(i), } } } diff --git a/crates/ivf/src/lib.rs b/crates/ivf/src/lib.rs index 7a82d2306..682e190b2 100644 --- a/crates/ivf/src/lib.rs +++ b/crates/ivf/src/lib.rs @@ -39,7 +39,7 @@ impl> Ivf { pub fn create( path: impl AsRef, options: IndexOptions, - source: &(impl Vectors> + Collection + Source + Sync), + source: &(impl Vectors + Collection + Source + Sync), ) -> Self { let remapped = RemappedCollection::from_source(source); from_nothing(path, options, &remapped) @@ -122,7 +122,7 @@ impl> Ivf { fn from_nothing>( path: impl AsRef, options: IndexOptions, - collection: &(impl Vectors> + Collection + Sync), + collection: &(impl Vectors + Collection + Sync), ) -> Ivf { create_dir(path.as_ref()).unwrap(); let IvfIndexingOptions { diff --git a/crates/ivf/src/operator.rs b/crates/ivf/src/operator.rs index 06d6da3e5..ac72d4854 100644 --- a/crates/ivf/src/operator.rs +++ b/crates/ivf/src/operator.rs @@ -13,7 +13,7 @@ pub trait OperatorIvf: OperatorStorage { fn interpret(vector: Borrowed<'_, Self>) -> &[Self::Scalar]; fn project>(quantizer: &Q, slice: &[Self::Scalar]) -> Vec; const SUPPORT_RESIDUAL: bool; - fn residual(lhs: Borrowed<'_, Self>, rhs: &[Self::Scalar]) -> Owned; + fn residual(lhs: Borrowed<'_, Self>, rhs: &[Self::Scalar]) -> Self::Vector; } impl OperatorIvf for BVectorDot { @@ -28,7 +28,7 @@ impl OperatorIvf for BVectorDot { unimplemented!() } const SUPPORT_RESIDUAL: bool = false; - fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Owned { + fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Self::Vector { unimplemented!() } } @@ -45,7 +45,7 @@ impl OperatorIvf for BVectorJaccard { unimplemented!() } const SUPPORT_RESIDUAL: bool = false; - fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Owned { + fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Self::Vector { unimplemented!() } } @@ -62,7 +62,7 @@ impl OperatorIvf for BVectorHamming { unimplemented!() } const SUPPORT_RESIDUAL: bool = false; - fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Owned { + fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Self::Vector { unimplemented!() } } @@ -79,7 +79,7 @@ impl OperatorIvf for SVectDot { unimplemented!() } const SUPPORT_RESIDUAL: bool = false; - fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Owned { + fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Self::Vector { unimplemented!() } } @@ -96,7 +96,7 @@ impl OperatorIvf for SVectL2 { unimplemented!() } const SUPPORT_RESIDUAL: bool = false; - fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Owned { + fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[Self::Scalar]) -> Self::Vector { unimplemented!() } } @@ -118,7 +118,7 @@ impl OperatorIvf for VectDot { quantizer.project(VectBorrowed::new(centroid)).into_vec() } const SUPPORT_RESIDUAL: bool = false; - fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[S]) -> Owned { + fn residual(_lhs: Borrowed<'_, Self>, _rhs: &[S]) -> Self::Vector { unimplemented!() } } @@ -140,7 +140,7 @@ impl OperatorIvf for VectL2 { quantizer.project(VectBorrowed::new(vector)).into_vec() } const SUPPORT_RESIDUAL: bool = true; - fn residual(lhs: Borrowed<'_, Self>, rhs: &[S]) -> Owned { + fn residual(lhs: Borrowed<'_, Self>, rhs: &[S]) -> Self::Vector { lhs.operator_sub(VectBorrowed::new(rhs)) } } diff --git a/crates/pyvectors/pyproject.toml b/crates/pyvectors/pyproject.toml index 4de29d904..f4828dba4 100644 --- a/crates/pyvectors/pyproject.toml +++ b/crates/pyvectors/pyproject.toml @@ -7,3 +7,4 @@ module-name = "vectors" [project] name = "vectors" +version = "0.0.0" diff --git a/crates/quantization/src/lib.rs b/crates/quantization/src/lib.rs index 412fd4753..a7f756d47 100644 --- a/crates/quantization/src/lib.rs +++ b/crates/quantization/src/lib.rs @@ -42,8 +42,8 @@ impl> Quantization { path: impl AsRef, vector_options: VectorOptions, quantization_options: Option, - vectors: &(impl Vectors> + Sync), - transform: impl Fn(Borrowed<'_, O>) -> Owned + Copy + Send + Sync, + vectors: &(impl Vectors + Sync), + transform: impl Fn(Borrowed<'_, O>) -> O::Vector + Copy + Send + Sync, ) -> Self { std::fs::create_dir(path.as_ref()).unwrap(); let quantizer = Json::create( @@ -107,7 +107,7 @@ impl> Quantization { &self.quantizer } - pub fn project(&self, vector: Borrowed<'_, O>) -> Owned { + pub fn project(&self, vector: Borrowed<'_, O>) -> O::Vector { Q::project(&self.quantizer, vector) } @@ -123,7 +123,7 @@ impl> Quantization { Q::flat_rerank_preprocess(&self.quantizer, vector, opts) } - pub fn process(&self, vectors: &impl Vectors>, lut: &Q::Lut, u: u32) -> Distance { + pub fn process(&self, vectors: &impl Vectors, lut: &Q::Lut, u: u32) -> Distance { let locate = |i| { let code_size = self.quantizer.code_size() as usize; let start = i as usize * code_size; diff --git a/crates/quantization/src/product.rs b/crates/quantization/src/product.rs index 00f6b37cb..8808cbbf3 100644 --- a/crates/quantization/src/product.rs +++ b/crates/quantization/src/product.rs @@ -41,8 +41,8 @@ impl Quantizer for ProductQuantizer { fn train( vector_options: VectorOptions, options: Option, - vectors: &(impl Vectors> + Sync), - transform: impl Fn(Borrowed<'_, O>) -> Owned + Copy + Sync, + vectors: &(impl Vectors + Sync), + transform: impl Fn(Borrowed<'_, O>) -> O::Vector + Copy + Sync, ) -> Self { let dims = vector_options.dims; let options = if let Some(QuantizationOptions::Product(x)) = options { @@ -116,7 +116,7 @@ impl Quantizer for ProductQuantizer { } } - fn fscan_encode(&self, vectors: [Owned; 32]) -> Vec { + fn fscan_encode(&self, vectors: [O::Vector; 32]) -> Vec { let dims = self.dims; let ratio = self.ratio; let bits = self.bits; @@ -141,7 +141,7 @@ impl Quantizer for ProductQuantizer { } } - fn project(&self, vector: Borrowed<'_, O>) -> Owned { + fn project(&self, vector: Borrowed<'_, O>) -> O::Vector { vector.own() } diff --git a/crates/quantization/src/quantizer.rs b/crates/quantization/src/quantizer.rs index 1c2800ab6..8434968cc 100644 --- a/crates/quantization/src/quantizer.rs +++ b/crates/quantization/src/quantizer.rs @@ -1,7 +1,7 @@ use base::distance::Distance; use base::index::{QuantizationOptions, SearchOptions, VectorOptions}; +use base::operator::Borrowed; use base::operator::Operator; -use base::operator::{Borrowed, Owned}; use base::search::{RerankerPop, RerankerPush, Vectors}; use serde::{Deserialize, Serialize}; use std::ops::Range; @@ -12,16 +12,16 @@ pub trait Quantizer: fn train( vector_options: VectorOptions, options: Option, - vectors: &(impl Vectors> + Sync), - transform: impl Fn(Borrowed<'_, O>) -> Owned + Copy + Sync, + vectors: &(impl Vectors + Sync), + transform: impl Fn(Borrowed<'_, O>) -> O::Vector + Copy + Sync, ) -> Self; fn encode(&self, vector: Borrowed<'_, O>) -> Vec; - fn fscan_encode(&self, vectors: [Owned; 32]) -> Vec; + fn fscan_encode(&self, vectors: [O::Vector; 32]) -> Vec; fn code_size(&self) -> u32; fn fcode_size(&self) -> u32; - fn project(&self, vector: Borrowed<'_, O>) -> Owned; + fn project(&self, vector: Borrowed<'_, O>) -> O::Vector; type Lut; fn preprocess(&self, vector: Borrowed<'_, O>) -> Self::Lut; diff --git a/crates/quantization/src/rabitq.rs b/crates/quantization/src/rabitq.rs index f570301f8..3e669681b 100644 --- a/crates/quantization/src/rabitq.rs +++ b/crates/quantization/src/rabitq.rs @@ -31,8 +31,8 @@ impl Quantizer for RabitqQuantizer { fn train( vector_options: VectorOptions, _: Option, - _: &(impl Vectors> + Sync), - _: impl Fn(Borrowed<'_, O>) -> Owned + Copy + Sync, + _: &(impl Vectors + Sync), + _: impl Fn(Borrowed<'_, O>) -> O::Vector + Copy + Sync, ) -> Self { let dims = vector_options.dims; let projection = { @@ -78,7 +78,7 @@ impl Quantizer for RabitqQuantizer { result } - fn fscan_encode(&self, vectors: [Owned; 32]) -> Vec { + fn fscan_encode(&self, vectors: [O::Vector; 32]) -> Vec { let dims = self.dims; let coded = vectors.map(|vector| O::code(vector.as_borrowed())); let codes = coded.clone().map(|(_, _, _, _, e)| { @@ -123,7 +123,7 @@ impl Quantizer for RabitqQuantizer { O::process(lut, c) } - fn project(&self, vector: Borrowed<'_, O>) -> Owned { + fn project(&self, vector: Borrowed<'_, O>) -> O::Vector { O::project(&self.projection, vector) } @@ -263,7 +263,7 @@ pub trait OperatorRabitqQuantization: Operator { fn code(vector: Borrowed<'_, Self>) -> (f32, f32, f32, f32, Vec); - fn project(projection: &[Vec], vector: Borrowed<'_, Self>) -> Owned; + fn project(projection: &[Vec], vector: Borrowed<'_, Self>) -> Self::Vector; type Lut; fn preprocess(vector: Borrowed<'_, Self>) -> Self::Lut; @@ -320,7 +320,7 @@ impl OperatorRabitqQuantization for VectL2 { (sum_of_x2, factor_ppc, factor_ip, factor_err, code) } - fn project(projection: &[Vec], vector: Borrowed<'_, Self>) -> Owned { + fn project(projection: &[Vec], vector: Borrowed<'_, Self>) -> Self::Vector { let slice = (0..projection.len()) .map(|i| S::from_f32(S::reduce_sum_of_xy(&projection[i], vector.slice()))) .collect(); @@ -465,7 +465,7 @@ impl OperatorRabitqQuantization for VectDot { (sum_of_x2, factor_ppc, factor_ip, factor_err, code) } - fn project(projection: &[Vec], vector: Borrowed<'_, Self>) -> Owned { + fn project(projection: &[Vec], vector: Borrowed<'_, Self>) -> Self::Vector { let slice = (0..projection.len()) .map(|i| S::from_f32(S::reduce_sum_of_xy(&projection[i], vector.slice()))) .collect(); @@ -574,7 +574,7 @@ macro_rules! unimpl_operator_rabitq_quantization { unimplemented!() } - fn project(_: &[Vec], _: Borrowed<'_, Self>) -> Owned { + fn project(_: &[Vec], _: Borrowed<'_, Self>) -> Self::Vector { unimplemented!() } diff --git a/crates/quantization/src/scalar.rs b/crates/quantization/src/scalar.rs index fab3b28fb..9f6918081 100644 --- a/crates/quantization/src/scalar.rs +++ b/crates/quantization/src/scalar.rs @@ -18,21 +18,22 @@ use base::search::RerankerPop; use base::search::RerankerPush; use base::search::Vectors; use base::vector::*; -use common::vec2::Vec2; +use rayon::iter::IntoParallelIterator; +use rayon::iter::ParallelIterator; use serde::Deserialize; use serde::Serialize; use std::cmp::Reverse; use std::marker::PhantomData; use std::ops::Range; +use stoppable_rayon as rayon; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(bound = "")] pub struct ScalarQuantizer { dims: u32, bits: u32, - max: Vec, min: Vec, - centroids: Vec2, + max: Vec, _phantom: PhantomData O>, } @@ -40,8 +41,8 @@ impl Quantizer for ScalarQuantizer { fn train( vector_options: VectorOptions, options: Option, - vectors: &impl Vectors>, - transform: impl Fn(Borrowed<'_, O>) -> Owned + Copy, + vectors: &(impl Vectors + Sync), + transform: impl Fn(Borrowed<'_, O>) -> O::Vector + Copy + Sync, ) -> Self { let options = if let Some(QuantizationOptions::Scalar(x)) = options { x @@ -50,32 +51,46 @@ impl Quantizer for ScalarQuantizer { }; let dims = vector_options.dims; let bits = options.bits; - let mut max = vec![f32::NEG_INFINITY; dims as usize]; - let mut min = vec![f32::INFINITY; dims as usize]; let n = vectors.len(); - for i in 0..n { - let vector = transform(vectors.vector(i)); - let vector = vector.as_borrowed(); - for j in 0..dims { - min[j as usize] = min[j as usize].min(O::get(vector, j).to_f32()); - max[j as usize] = max[j as usize].max(O::get(vector, j).to_f32()); - } - } - let mut centroids = Vec2::zeros((1 << bits, dims as usize)); - for i in 0..dims { - let bas = min[i as usize]; - let del = max[i as usize] - min[i as usize]; - for j in 0_usize..(1 << bits) { - let val = j as f32 / ((1 << bits) - 1) as f32; - centroids[(j, i as usize)] = bas + val * del; - } - } + let (min, max) = (0..n) + .into_par_iter() + .fold( + || { + ( + vec![f32::INFINITY; dims as usize], + vec![f32::NEG_INFINITY; dims as usize], + ) + }, + |(mut min, mut max), i| { + let vector = transform(vectors.vector(i)); + let vector = vector.as_borrowed(); + for j in 0..dims { + min[j as usize] = min[j as usize].min(O::get(vector, j).to_f32()); + max[j as usize] = max[j as usize].max(O::get(vector, j).to_f32()); + } + (min, max) + }, + ) + .reduce( + || { + ( + vec![f32::INFINITY; dims as usize], + vec![f32::NEG_INFINITY; dims as usize], + ) + }, + |(mut min, mut max), (rmin, rmax)| { + for j in 0..dims { + min[j as usize] = min[j as usize].min(rmin[j as usize]); + max[j as usize] = max[j as usize].max(rmax[j as usize]); + } + (min, max) + }, + ); Self { dims, bits, - max, min, - centroids, + max, _phantom: PhantomData, } } @@ -83,13 +98,16 @@ impl Quantizer for ScalarQuantizer { fn encode(&self, vector: Borrowed<'_, O>) -> Vec { let dims = self.dims; let bits = self.bits; + let min = self.min.as_slice(); + let max = self.max.as_slice(); let code_size = (dims * bits).div_ceil(8); let mut code = Vec::with_capacity(dims as usize); for i in 0..dims { - let del = self.max[i as usize] - self.min[i as usize]; - let w = (((O::get(vector, i).to_f32() - self.min[i as usize]) / del).to_f32() - * (((1 << bits) - 1) as f32)) as u32; - code.push(w.clamp(0, 255) as u8); + let val = O::get(vector, i).to_f32(); + let bas = min[i as usize]; + let del = (max[i as usize] - min[i as usize]) / ((1 << bits) - 1) as f32; + let j = ((val - bas) / del).round_ties_even() as u32; + code.push(j.clamp(0, (1 << bits) - 1) as u8); } match bits { 1 => InfiniteByteChunks::new(code.into_iter()) @@ -109,19 +127,20 @@ impl Quantizer for ScalarQuantizer { } } - fn fscan_encode(&self, vectors: [Owned; 32]) -> Vec { + fn fscan_encode(&self, vectors: [O::Vector; 32]) -> Vec { let dims = self.dims; let bits = self.bits; + let min = self.min.as_slice(); + let max = self.max.as_slice(); if bits == 4 { let codes = vectors.map(|vector| { let mut code = Vec::with_capacity(dims as usize); for i in 0..dims { - let del = self.max[i as usize] - self.min[i as usize]; - let w = (((O::get(vector.as_borrowed(), i).to_f32() - self.min[i as usize]) - / del) - .to_f32() - * (((1 << bits) - 1) as f32)) as u32; - code.push(w.clamp(0, 255) as u8); + let val = O::get(vector.as_borrowed(), i).to_f32(); + let bas = min[i as usize]; + let del = (max[i as usize] - min[i as usize]) / ((1 << bits) - 1) as f32; + let j = ((val - bas) / del).round_ties_even() as u32; + code.push(j.clamp(0, (1 << bits) - 1) as u8); } code }); @@ -143,14 +162,14 @@ impl Quantizer for ScalarQuantizer { } } - fn project(&self, vector: Borrowed<'_, O>) -> Owned { + fn project(&self, vector: Borrowed<'_, O>) -> O::Vector { vector.own() } type Lut = Vec; fn preprocess(&self, vector: Borrowed<'_, O>) -> Self::Lut { - O::preprocess(self.dims, self.bits, &self.max, &self.min, vector) + O::preprocess(self.dims, self.bits, &self.min, &self.max, vector) } fn process(&self, lut: &Self::Lut, code: &[u8], _: Borrowed<'_, O>) -> Distance { @@ -165,7 +184,7 @@ impl Quantizer for ScalarQuantizer { ); fn fscan_preprocess(&self, vector: Borrowed<'_, O>) -> Self::FLut { - O::fscan_preprocess(self.dims, self.bits, &self.max, &self.min, vector) + O::fscan_preprocess(self.dims, self.bits, &self.min, &self.max, vector) } fn fscan_process(&self, flut: &Self::FLut, code: &[u8]) -> [Distance; 32] { @@ -291,8 +310,8 @@ pub trait OperatorScalarQuantization: Operator { fn preprocess( dims: u32, bits: u32, - max: &[f32], min: &[f32], + max: &[f32], vector: Borrowed<'_, Self>, ) -> Vec; fn process(dims: u32, bits: u32, lut: &[f32], code: &[u8]) -> Distance; @@ -300,8 +319,8 @@ pub trait OperatorScalarQuantization: Operator { fn fscan_preprocess( dims: u32, bits: u32, - max: &[f32], min: &[f32], + max: &[f32], vector: Borrowed<'_, Self>, ) -> (u32, f32, f32, Vec); fn fscan_process(flut: &(u32, f32, f32, Vec), code: &[u8]) -> [Distance; 32]; @@ -316,22 +335,46 @@ impl OperatorScalarQuantization for VectDot { fn preprocess( dims: u32, bits: u32, - max: &[f32], min: &[f32], + max: &[f32], vector: Borrowed<'_, Self>, ) -> Vec { - let mut xy = Vec::with_capacity(dims as _); - for i in 0..dims { - let bas = min[i as usize]; - let del = max[i as usize] - min[i as usize]; - xy.extend((0..1 << bits).map(|k| { - let x = vector.slice()[i as usize].to_f32(); - let val = k as f32 / ((1 << bits) - 1) as f32; - let y = bas + val * del; - x * y - })); + #[inline(never)] + fn internal( + dims: usize, + min: &[f32], + max: &[f32], + vector: &[S], + ) -> Vec { + assert!(dims <= 65535); + assert!(dims == min.len()); + assert!(dims == max.len()); + assert!(dims == vector.len()); + let mut table = Vec::::with_capacity(dims * (1 << BITS)); + for i in 0..dims { + let bas = min[i]; + let del = (max[i] - min[i]) / ((1 << BITS) - 1) as f32; + for j in 0..1 << BITS { + let x = vector[i].to_f32(); + let y = bas + (j as f32) * del; + let value = x * y; + unsafe { + table.as_mut_ptr().add(i * (1 << BITS) + j).write(value); + } + } + } + unsafe { + table.set_len(dims * (1 << BITS)); + } + table + } + match bits { + 1 => internal::<1, _>(dims as _, min, max, vector.slice()), + 2 => internal::<2, _>(dims as _, min, max, vector.slice()), + 4 => internal::<4, _>(dims as _, min, max, vector.slice()), + 8 => internal::<8, _>(dims as _, min, max, vector.slice()), + _ => unreachable!(), } - xy } fn process(dims: u32, bits: u32, lut: &[f32], rhs: &[u8]) -> Distance { fn internal(dims: u32, t: &[f32], f: impl Fn(usize) -> usize) -> Distance { @@ -359,11 +402,11 @@ impl OperatorScalarQuantization for VectDot { fn fscan_preprocess( dims: u32, bits: u32, - max: &[f32], min: &[f32], + max: &[f32], vector: Borrowed<'_, Self>, ) -> (u32, f32, f32, Vec) { - let (k, b, t) = quantize::<255>(&Self::preprocess(dims, bits, max, min, vector)); + let (k, b, t) = quantize::<255>(&Self::preprocess(dims, bits, min, max, vector)); (dims, k, b, t) } fn fscan_process(flut: &(u32, f32, f32, Vec), codes: &[u8]) -> [Distance; 32] { @@ -382,23 +425,46 @@ impl OperatorScalarQuantization for VectL2 { fn preprocess( dims: u32, bits: u32, - max: &[f32], min: &[f32], + max: &[f32], vector: Borrowed<'_, Self>, ) -> Vec { - let mut d2 = Vec::with_capacity(dims as _); - for i in 0..dims { - let bas = min[i as usize]; - let del = max[i as usize] - min[i as usize]; - d2.extend((0..1 << bits).map(|k| { - let x = vector.slice()[i as usize].to_f32(); - let val = k as f32 / ((1 << bits) - 1) as f32; - let y = bas + val * del; - let d = x - y; - d * d - })); + #[inline(never)] + fn internal( + dims: usize, + min: &[f32], + max: &[f32], + vector: &[S], + ) -> Vec { + assert!(dims <= 65535); + assert!(dims == min.len()); + assert!(dims == max.len()); + assert!(dims == vector.len()); + let mut table = Vec::::with_capacity(dims * (1 << BITS)); + for i in 0..dims { + let bas = min[i]; + let del = (max[i] - min[i]) / ((1 << BITS) - 1) as f32; + for j in 0..1 << BITS { + let x = vector[i].to_f32(); + let y = bas + (j as f32) * del; + let value = (x - y) * (x - y); + unsafe { + table.as_mut_ptr().add(i * (1 << BITS) + j).write(value); + } + } + } + unsafe { + table.set_len(dims * (1 << BITS)); + } + table + } + match bits { + 1 => internal::<1, _>(dims as _, min, max, vector.slice()), + 2 => internal::<2, _>(dims as _, min, max, vector.slice()), + 4 => internal::<4, _>(dims as _, min, max, vector.slice()), + 8 => internal::<8, _>(dims as _, min, max, vector.slice()), + _ => unreachable!(), } - d2 } fn process(dims: u32, bits: u32, lut: &[f32], rhs: &[u8]) -> Distance { fn internal(dims: u32, t: &[f32], f: impl Fn(usize) -> usize) -> Distance { @@ -426,11 +492,11 @@ impl OperatorScalarQuantization for VectL2 { fn fscan_preprocess( dims: u32, bits: u32, - max: &[f32], min: &[f32], + max: &[f32], vector: Borrowed<'_, Self>, ) -> (u32, f32, f32, Vec) { - let (k, b, t) = quantize::<255>(&Self::preprocess(dims, bits, max, min, vector)); + let (k, b, t) = quantize::<255>(&Self::preprocess(dims, bits, min, max, vector)); (dims, k, b, t) } fn fscan_process(flut: &(u32, f32, f32, Vec), codes: &[u8]) -> [Distance; 32] { diff --git a/crates/quantization/src/trivial.rs b/crates/quantization/src/trivial.rs index 076fb93c9..270c4a0f4 100644 --- a/crates/quantization/src/trivial.rs +++ b/crates/quantization/src/trivial.rs @@ -24,8 +24,8 @@ impl Quantizer for TrivialQuantizer { fn train( _: VectorOptions, _: Option, - _: &impl Vectors>, - _: impl Fn(Borrowed<'_, O>) -> Owned + Copy, + _: &impl Vectors, + _: impl Fn(Borrowed<'_, O>) -> O::Vector + Copy, ) -> Self { Self { _maker: PhantomData, @@ -36,7 +36,7 @@ impl Quantizer for TrivialQuantizer { Vec::new() } - fn fscan_encode(&self, _: [Owned; 32]) -> Vec { + fn fscan_encode(&self, _: [O::Vector; 32]) -> Vec { Vec::new() } @@ -48,11 +48,11 @@ impl Quantizer for TrivialQuantizer { 0 } - fn project(&self, vector: Borrowed<'_, O>) -> Owned { + fn project(&self, vector: Borrowed<'_, O>) -> O::Vector { vector.own() } - type Lut = Owned; + type Lut = O::Vector; fn preprocess(&self, vector: Borrowed<'_, O>) -> Self::Lut { vector.own() diff --git a/crates/inverted/Cargo.toml b/crates/sparse_inverted_index/Cargo.toml similarity index 88% rename from crates/inverted/Cargo.toml rename to crates/sparse_inverted_index/Cargo.toml index 080323f8a..925e0fe7c 100644 --- a/crates/inverted/Cargo.toml +++ b/crates/sparse_inverted_index/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "inverted" +name = "sparse_inverted_index" version.workspace = true edition.workspace = true diff --git a/crates/inverted/src/lib.rs b/crates/sparse_inverted_index/src/lib.rs similarity index 88% rename from crates/inverted/src/lib.rs rename to crates/sparse_inverted_index/src/lib.rs index 7b3553e81..ba2f424ac 100644 --- a/crates/inverted/src/lib.rs +++ b/crates/sparse_inverted_index/src/lib.rs @@ -2,11 +2,11 @@ pub mod operator; -use self::operator::OperatorInvertedIndex; +use self::operator::OperatorSparseInvertedIndex; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::index::{IndexOptions, SearchOptions}; -use base::operator::{Borrowed, Owned}; +use base::operator::Borrowed; use base::scalar::ScalarLike; use base::search::{Collection, Element, Payload, Source, Vectors}; use common::json::Json; @@ -21,7 +21,7 @@ use std::path::Path; const ZERO: f32 = 0.0f32; #[allow(dead_code)] -pub struct InvertedIndex { +pub struct SparseInvertedIndex { storage: O::Storage, payloads: MmapArray, indexes: Json>, @@ -29,11 +29,11 @@ pub struct InvertedIndex { scores: Json>, } -impl InvertedIndex { +impl SparseInvertedIndex { pub fn create( path: impl AsRef, options: IndexOptions, - source: &(impl Vectors> + Collection + Source), + source: &(impl Vectors + Collection + Source), ) -> Self { let remapped = RemappedCollection::from_source(source); from_nothing(path, options, &remapped) @@ -87,11 +87,11 @@ impl InvertedIndex { } } -fn from_nothing( +fn from_nothing( path: impl AsRef, opts: IndexOptions, - collection: &(impl Vectors> + Collection), -) -> InvertedIndex { + collection: &(impl Vectors + Collection), +) -> SparseInvertedIndex { create_dir(path.as_ref()).expect("failed to create path for inverted index"); let mut token_collection = vec![Vec::new(); opts.vector.dims as usize]; @@ -110,7 +110,7 @@ fn from_nothing( let json_index = Json::create(path.as_ref().join("indexes"), indexes); let json_offset = Json::create(path.as_ref().join("offsets"), offsets); let json_score = Json::create(path.as_ref().join("scores"), scores); - InvertedIndex { + SparseInvertedIndex { storage, payloads, indexes: json_index, @@ -119,13 +119,13 @@ fn from_nothing( } } -fn open(path: impl AsRef) -> InvertedIndex { +fn open(path: impl AsRef) -> SparseInvertedIndex { let storage = O::Storage::open(path.as_ref().join("storage")); let payloads = MmapArray::open(path.as_ref().join("payloads")); let offsets = Json::open(path.as_ref().join("offsets")); let indexes = Json::open(path.as_ref().join("indexes")); let scores = Json::open(path.as_ref().join("scores")); - InvertedIndex { + SparseInvertedIndex { storage, payloads, indexes, diff --git a/crates/inverted/src/operator.rs b/crates/sparse_inverted_index/src/operator.rs similarity index 51% rename from crates/inverted/src/operator.rs rename to crates/sparse_inverted_index/src/operator.rs index a7111c857..b25659287 100644 --- a/crates/inverted/src/operator.rs +++ b/crates/sparse_inverted_index/src/operator.rs @@ -1,43 +1,53 @@ -use base::{operator::*, scalar::ScalarLike}; +use base::operator::*; +use base::scalar::ScalarLike; use std::iter::{zip, Empty}; use storage::OperatorStorage; -pub trait OperatorInvertedIndex: OperatorStorage { +pub trait OperatorSparseInvertedIndex: OperatorStorage { fn to_index_vec(vec: Borrowed<'_, Self>) -> impl Iterator + '_; } -impl OperatorInvertedIndex for SVectDot { - fn to_index_vec(vec: Borrowed<'_, Self>) -> impl Iterator + '_ { - zip(vec.indexes().iter().copied(), vec.values().iter().copied()) +impl OperatorSparseInvertedIndex for SVectDot { + fn to_index_vec(vector: Borrowed<'_, Self>) -> impl Iterator + '_ { + zip( + vector.indexes().iter().copied(), + vector.values().iter().copied().map(S::to_f32), + ) } } -macro_rules! unimpl_operator_inverted_index { - ($t:ty) => { - impl OperatorInvertedIndex for $t { - fn to_index_vec(_: Borrowed<'_, Self>) -> impl Iterator + '_ { - #![allow(unreachable_code)] - unimplemented!() as Empty<(u32, f32)> - } - } - }; +impl OperatorSparseInvertedIndex for SVectL2 { + fn to_index_vec(_: Borrowed<'_, Self>) -> impl Iterator + '_ { + #![allow(unreachable_code)] + unimplemented!() as Empty<(u32, f32)> + } } -impl OperatorInvertedIndex for VectDot { +impl OperatorSparseInvertedIndex for VectDot { fn to_index_vec(_: Borrowed<'_, Self>) -> impl Iterator + '_ { #![allow(unreachable_code)] unimplemented!() as Empty<(u32, f32)> } } -impl OperatorInvertedIndex for VectL2 { +impl OperatorSparseInvertedIndex for VectL2 { fn to_index_vec(_: Borrowed<'_, Self>) -> impl Iterator + '_ { #![allow(unreachable_code)] unimplemented!() as Empty<(u32, f32)> } } -unimpl_operator_inverted_index!(SVectL2); +macro_rules! unimpl_operator_inverted_index { + ($t:ty) => { + impl OperatorSparseInvertedIndex for $t { + #![allow(unreachable_code)] + fn to_index_vec(_: Borrowed<'_, Self>) -> impl Iterator + '_ { + unimplemented!() as Empty<(u32, f32)> + } + } + }; +} + unimpl_operator_inverted_index!(BVectorDot); unimpl_operator_inverted_index!(BVectorJaccard); unimpl_operator_inverted_index!(BVectorHamming); diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs index d397c636e..93f13832f 100644 --- a/crates/storage/src/lib.rs +++ b/crates/storage/src/lib.rs @@ -14,7 +14,7 @@ pub trait Storage: Vectors { } pub trait OperatorStorage: Operator { - type Storage: Storage> + Send + Sync; + type Storage: Storage + Send + Sync; } impl OperatorStorage for SVectDot {