From 6959df460fb0f3d312eabf37a8903b436c5b496d Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 3 Jan 2023 10:54:38 +0100 Subject: [PATCH 01/55] integrate quantized data to storages --- .gitignore | 2 +- Cargo.lock | 9 +++ lib/segment/Cargo.toml | 2 +- lib/segment/src/index/hnsw_index/hnsw.rs | 22 +++++-- .../vector_storage/memmap_vector_storage.rs | 22 +++++++ .../src/vector_storage/mmap_vectors.rs | 58 ++++++++++++++++++- lib/segment/src/vector_storage/mod.rs | 1 + .../quantized_vector_storage.rs | 47 +++++++++++++++ .../src/vector_storage/vector_storage_base.rs | 22 +------ 9 files changed, 155 insertions(+), 30 deletions(-) create mode 100644 lib/segment/src/vector_storage/quantized_vector_storage.rs diff --git a/.gitignore b/.gitignore index 9e2354f452a..6f77fd375bb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ /target /.idea -/storage +/storage* /snapshots *.log \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 7f70fef6031..55ae6fc883c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2723,6 +2723,14 @@ dependencies = [ "uuid", ] +[[package]] +name = "quantization" +version = "0.1.0" +source = "git+https://github.com/qdrant/quantization.git#13b734be5ed139ac4aff70e1f93b32d0ec46cf67" +dependencies = [ + "cc", +] + [[package]] name = "quick-error" version = "1.2.3" @@ -3239,6 +3247,7 @@ dependencies = [ "ordered-float", "parking_lot", "pprof", + "quantization", "rand", "rand_distr", "rayon", diff --git a/lib/segment/Cargo.toml b/lib/segment/Cargo.toml index 4d1bc4106f1..ef47999cb57 100644 --- a/lib/segment/Cargo.toml +++ b/lib/segment/Cargo.toml @@ -45,7 +45,7 @@ tar = "0.4.38" fs_extra = "1.2.0" semver = "1.0.16" tinyvec = { version = "1.6.0", features = ["alloc"] } - +quantization = { git = "https://github.com/qdrant/quantization.git" } [[bench]] name = "vector_search" diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index ffb80772b3e..0ebe5631640 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -4,7 +4,7 @@ use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; use std::sync::Arc; -use atomic_refcell::AtomicRefCell; +use atomic_refcell::{AtomicRef, AtomicRefCell}; use log::debug; use parking_lot::Mutex; use rand::thread_rng; @@ -29,7 +29,7 @@ use crate::index::{PayloadIndex, VectorIndex}; use crate::telemetry::VectorIndexSearchesTelemetry; use crate::types::Condition::Field; use crate::types::{FieldCondition, Filter, HnswConfig, SearchParams, VECTOR_ELEMENT_SIZE}; -use crate::vector_storage::{ScoredPointOffset, VectorStorageSS}; +use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorage, VectorStorageSS}; const HNSW_USE_HEURISTIC: bool = true; const BYTES_IN_KB: usize = 1024; @@ -156,7 +156,7 @@ impl HNSWIndex { check_process_stopped(stopped)?; let vector = vector_storage.get_vector(block_point_id).unwrap(); - let raw_scorer = vector_storage.raw_scorer(vector); + let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector); let block_condition_checker = BuildConditionChecker { filter_list: block_filter_list, current_point: block_point_id, @@ -185,7 +185,7 @@ impl HNSWIndex { let ef = max(req_ef, top); let vector_storage = self.vector_storage.borrow(); - let raw_scorer = vector_storage.raw_scorer(vector.to_owned()); + let raw_scorer = Self::get_raw_scorer(&vector_storage, vector); let payload_index = self.payload_index.borrow(); let filter_context = filter.map(|f| payload_index.filter_context(f)); @@ -211,6 +211,18 @@ impl HNSWIndex { .map(|vector| self.search_with_graph(vector, filter, top, params)) .collect() } + + fn get_raw_scorer<'a>( + vector_storage: &'a AtomicRef, + vector: &[VectorElementType], + ) -> Box { + // TODO: provide option for quantization disabling + if let Some(raw_scorer) = vector_storage.quantized_raw_scorer(vector) { + raw_scorer + } else { + vector_storage.raw_scorer(vector.to_owned()) + } + } } impl VectorIndex for HNSWIndex { @@ -336,7 +348,7 @@ impl VectorIndex for HNSWIndex { ids.into_par_iter().try_for_each(|vector_id| { check_process_stopped(stopped)?; let vector = vector_storage.get_vector(vector_id).unwrap(); - let raw_scorer = vector_storage.raw_scorer(vector); + let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector); let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), None); graph_layers_builder.link_new_point(vector_id, points_scorer); diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index d72b99ca0d5..861b15304ec 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use atomic_refcell::AtomicRefCell; +use super::quantized_vector_storage::QuantizedRawScorer; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::{check_process_stopped, OperationResult}; @@ -253,6 +254,27 @@ where }) } + fn quantized_raw_scorer( + &self, + vector: &[VectorElementType], + ) -> Option> { + let mmap_store = self.mmap_store.as_ref().unwrap(); + if let Some(quantized_data) = &mmap_store.quantized_vectors { + if let Some(deleted_ram) = &mmap_store.deleted_ram { + let query = TMetric::preprocess(&vector).unwrap_or(vector.to_owned()); + Some(Box::new(QuantizedRawScorer { + query: quantized_data.encode_query(&query), + quantized_data, + deleted: deleted_ram, + })) + } else { + None + } + } else { + None + } + } + fn score_points( &self, vector: &[VectorElementType], diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 315e864d9ba..112ead7c8e4 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -4,14 +4,16 @@ use std::mem::{size_of, transmute}; use std::path::Path; use std::sync::Arc; +use bitvec::vec::BitVec; use memmap2::{Mmap, MmapMut, MmapOptions}; use parking_lot::{RwLock, RwLockReadGuard}; +use super::quantized_vector_storage::EncodedVectors; use crate::common::error_logging::LogError; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; -use crate::entry::entry_point::OperationResult; -use crate::types::PointOffsetType; +use crate::entry::entry_point::{OperationError, OperationResult}; +use crate::types::{Distance, PointOffsetType}; const HEADER_SIZE: usize = 4; const DELETED_HEADER: &[u8; 4] = b"drop"; @@ -24,6 +26,8 @@ pub struct MmapVectors { mmap: Mmap, deleted_mmap: Arc>, pub deleted_count: usize, + pub deleted_ram: Option, + pub quantized_vectors: Option, } fn open_read(path: &Path) -> OperationResult { @@ -77,9 +81,49 @@ impl MmapVectors { mmap, deleted_mmap: Arc::new(RwLock::new(deleted_mmap)), deleted_count, + deleted_ram: None, + quantized_vectors: None, }) } + fn enable_deleted_ram(&mut self) { + if self.deleted_ram.is_some() { + return; + } + + let mut deleted = BitVec::new(); + deleted.resize(self.num_vectors, false); + for i in 0..self.num_vectors { + deleted.set(i, self.deleted(i as PointOffsetType).unwrap_or_default()); + } + self.deleted_ram = Some(deleted); + } + + #[allow(dead_code)] + fn quantize(&mut self, distance: Distance) -> OperationResult<()> { + if self.quantized_vectors.is_some() { + return Ok(()); + } + + self.enable_deleted_ram(); + self.quantized_vectors = Some( + EncodedVectors::encode( + (0..self.num_vectors as u32).map(|i| { + let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); + self.raw_vector_offset(offset) + }), + Vec::new(), + match distance { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + ) + .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, + ); + Ok(()) + } + pub fn data_offset(&self, key: PointOffsetType) -> Option { let vector_data_length = self.dim * size_of::(); let offset = (key as usize) * vector_data_length + HEADER_SIZE; @@ -113,7 +157,11 @@ impl MmapVectors { } pub fn deleted(&self, key: PointOffsetType) -> Option { - Self::check_deleted(&self.deleted_mmap.read(), key) + if let Some(deleted_ram) = &self.deleted_ram { + deleted_ram.get(key as usize).map(|x| *x) + } else { + Self::check_deleted(&self.deleted_mmap.read(), key) + } } /// Creates returns owned vector (copy of internal vector) @@ -131,6 +179,10 @@ impl MmapVectors { let mut deleted_mmap = self.deleted_mmap.write(); let flag = deleted_mmap.get_mut((key as usize) + HEADER_SIZE).unwrap(); + if let Some(deleted_ram) = &mut self.deleted_ram { + deleted_ram.set(key as usize, true); + } + if *flag == 0 { *flag = 1; self.deleted_count += 1; diff --git a/lib/segment/src/vector_storage/mod.rs b/lib/segment/src/vector_storage/mod.rs index 1de30c9b5e0..e6179a69b61 100644 --- a/lib/segment/src/vector_storage/mod.rs +++ b/lib/segment/src/vector_storage/mod.rs @@ -1,6 +1,7 @@ pub mod chunked_vectors; pub mod memmap_vector_storage; mod mmap_vectors; +mod quantized_vector_storage; pub mod simple_vector_storage; mod vector_storage_base; diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs new file mode 100644 index 00000000000..7a6f25aefaf --- /dev/null +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -0,0 +1,47 @@ +use bitvec::vec::BitVec; + +use super::{RawScorer, ScoredPointOffset}; +use crate::types::{PointOffsetType, ScoreType}; + +pub type EncodedVectors = quantization::encoder::EncodedVectors>; + +pub struct QuantizedRawScorer<'a> { + pub query: quantization::encoder::EncodedQuery, + pub deleted: &'a BitVec, + pub quantized_data: &'a EncodedVectors, +} + +impl RawScorer for QuantizedRawScorer<'_> { + fn score_points(&self, points: &[PointOffsetType], scores: &mut [ScoredPointOffset]) -> usize { + let mut size: usize = 0; + for point_id in points.iter().copied() { + if self.deleted[point_id as usize] { + continue; + } + scores[size] = ScoredPointOffset { + idx: point_id, + score: self + .quantized_data + .score_point(&self.query, point_id as usize), + }; + size += 1; + if size == scores.len() { + return size; + } + } + size + } + + fn check_point(&self, point: PointOffsetType) -> bool { + (point as usize) < self.deleted.len() && !self.deleted[point as usize] + } + + fn score_point(&self, point: PointOffsetType) -> ScoreType { + self.quantized_data.score_point(&self.query, point as usize) + } + + fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType { + self.quantized_data + .score_internal(point_a as usize, point_b as usize) + } +} diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index d0d6f0d6e5b..c4ec3954dc9 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -75,6 +75,8 @@ pub trait VectorStorage { /// Generate a `RawScorer` object which contains all required context for searching similar vector fn raw_scorer(&self, vector: Vec) -> Box; + fn quantized_raw_scorer(&self, vector: &[VectorElementType]) + -> Option>; fn score_points( &self, @@ -102,24 +104,4 @@ pub trait VectorStorage { } } -trait SuperVectorStorage {} - pub type VectorStorageSS = dyn VectorStorage + Sync + Send; - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_ordering() { - assert!( - ScoredPointOffset { - idx: 10, - score: 0.9 - } > ScoredPointOffset { - idx: 20, - score: 0.6 - } - ) - } -} From 63691b4ef56503ab81f2984c99877a336762c112 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 3 Jan 2023 10:56:35 +0100 Subject: [PATCH 02/55] revert gitignore --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 6f77fd375bb..9e2354f452a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ /target /.idea -/storage* +/storage /snapshots *.log \ No newline at end of file From 6abff6d56b2c3dedfc6a89218009c77c32ce0048 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 3 Jan 2023 11:00:45 +0100 Subject: [PATCH 03/55] are you happy clippy --- .../vector_storage/memmap_vector_storage.rs | 2 +- .../vector_storage/simple_vector_storage.rs | 42 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 861b15304ec..0f9597eeea1 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -261,7 +261,7 @@ where let mmap_store = self.mmap_store.as_ref().unwrap(); if let Some(quantized_data) = &mmap_store.quantized_vectors { if let Some(deleted_ram) = &mmap_store.deleted_ram { - let query = TMetric::preprocess(&vector).unwrap_or(vector.to_owned()); + let query = TMetric::preprocess(vector).unwrap_or_else(|| vector.to_owned()); Some(Box::new(QuantizedRawScorer { query: quantized_data.encode_query(&query), quantized_data, diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index be7e65c2f9f..0ada55e8c59 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -12,6 +12,7 @@ use rocksdb::DB; use serde::{Deserialize, Serialize}; use super::chunked_vectors::ChunkedVectors; +use super::quantized_vector_storage::{EncodedVectors, QuantizedRawScorer}; use super::vector_storage_base::VectorStorage; use crate::common::rocksdb_wrapper::DatabaseColumnWrapper; use crate::common::Flusher; @@ -30,6 +31,7 @@ pub struct SimpleVectorStorage { vectors: ChunkedVectors, deleted: BitVec, deleted_count: usize, + quantized_vectors: Option, db_wrapper: DatabaseColumnWrapper, } @@ -129,6 +131,7 @@ pub fn open_simple_vector_storage( vectors, deleted, deleted_count, + quantized_vectors: None, db_wrapper, }))), Distance::Euclid => Ok(Arc::new(AtomicRefCell::new(SimpleVectorStorage::< @@ -139,6 +142,7 @@ pub fn open_simple_vector_storage( vectors, deleted, deleted_count, + quantized_vectors: None, db_wrapper, }))), Distance::Dot => Ok(Arc::new(AtomicRefCell::new(SimpleVectorStorage::< @@ -149,6 +153,7 @@ pub fn open_simple_vector_storage( vectors, deleted, deleted_count, + quantized_vectors: None, db_wrapper, }))), } @@ -173,6 +178,27 @@ where Ok(()) } + + #[allow(dead_code)] + fn quantize(&mut self) -> OperationResult<()> { + if self.quantized_vectors.is_some() { + return Ok(()); + } + + self.quantized_vectors = Some( + EncodedVectors::encode( + (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), + Vec::new(), + match TMetric::distance() { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + ) + .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, + ); + Ok(()) + } } impl VectorStorage for SimpleVectorStorage @@ -277,6 +303,22 @@ where }) } + fn quantized_raw_scorer( + &self, + vector: &[VectorElementType], + ) -> Option> { + if let Some(quantized_data) = &self.quantized_vectors { + let query = TMetric::preprocess(vector).unwrap_or_else(|| vector.to_owned()); + Some(Box::new(QuantizedRawScorer { + query: quantized_data.encode_query(&query), + quantized_data, + deleted: &self.deleted, + })) + } else { + None + } + } + fn score_points( &self, vector: &[VectorElementType], From d70caab9d5adf02b0564b7320045e03c11ff1069 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 4 Jan 2023 10:17:07 +0100 Subject: [PATCH 04/55] quantize in optimizer --- lib/api/src/grpc/qdrant.rs | 7 ++-- .../segment_constructor/segment_builder.rs | 1 + .../vector_storage/memmap_vector_storage.rs | 5 +++ .../src/vector_storage/mmap_vectors.rs | 3 +- .../vector_storage/simple_vector_storage.rs | 41 +++++++++---------- .../src/vector_storage/vector_storage_base.rs | 1 + 6 files changed, 32 insertions(+), 26 deletions(-) diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 2e328d8cd86..9037243ec3c 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -1805,12 +1805,13 @@ pub struct ScrollPoints { #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct LookupLocation { - #[prost(string, tag="1")] + #[prost(string, tag = "1")] pub collection_name: ::prost::alloc::string::String, /// Which vector to use for search, if not specified - use default vector - #[prost(string, optional, tag="2")] + #[prost(string, optional, tag = "2")] pub vector_name: ::core::option::Option<::prost::alloc::string::String>, } +#[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct RecommendPoints { /// name of the collection @@ -1847,7 +1848,7 @@ pub struct RecommendPoints { #[prost(message, optional, tag = "12")] pub with_vectors: ::core::option::Option, /// Name of the collection to use for points lookup, if not specified - use current collection - #[prost(message, optional, tag="13")] + #[prost(message, optional, tag = "13")] pub lookup_from: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 97d98377ea7..c84e16a44b2 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -103,6 +103,7 @@ impl SegmentBuilder { vector_storage.update_from(&**other_vector_storage, stopped)?; internal_id_iter = Some(new_internal_range.zip(other_vector_storage.iter_ids())); + vector_storage.quantize()?; } if internal_id_iter.is_none() { return Err(OperationError::service_error( diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 0f9597eeea1..70d7642ae9c 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -275,6 +275,11 @@ where } } + fn quantize(&mut self) -> OperationResult<()> { + let mmap_store = self.mmap_store.as_mut().unwrap(); + mmap_store.quantize(TMetric::distance()) + } + fn score_points( &self, vector: &[VectorElementType], diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 112ead7c8e4..4e16eda625b 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -99,8 +99,7 @@ impl MmapVectors { self.deleted_ram = Some(deleted); } - #[allow(dead_code)] - fn quantize(&mut self, distance: Distance) -> OperationResult<()> { + pub fn quantize(&mut self, distance: Distance) -> OperationResult<()> { if self.quantized_vectors.is_some() { return Ok(()); } diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 0ada55e8c59..89499bb7ddc 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -178,27 +178,6 @@ where Ok(()) } - - #[allow(dead_code)] - fn quantize(&mut self) -> OperationResult<()> { - if self.quantized_vectors.is_some() { - return Ok(()); - } - - self.quantized_vectors = Some( - EncodedVectors::encode( - (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), - Vec::new(), - match TMetric::distance() { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - ) - .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, - ); - Ok(()) - } } impl VectorStorage for SimpleVectorStorage @@ -319,6 +298,26 @@ where } } + fn quantize(&mut self) -> OperationResult<()> { + if self.quantized_vectors.is_some() { + return Ok(()); + } + + self.quantized_vectors = Some( + EncodedVectors::encode( + (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), + Vec::new(), + match TMetric::distance() { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + ) + .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, + ); + Ok(()) + } + fn score_points( &self, vector: &[VectorElementType], diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index c4ec3954dc9..ef34ef56e5d 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -77,6 +77,7 @@ pub trait VectorStorage { fn raw_scorer(&self, vector: Vec) -> Box; fn quantized_raw_scorer(&self, vector: &[VectorElementType]) -> Option>; + fn quantize(&mut self) -> OperationResult<()>; fn score_points( &self, From 6749bebc1d1aff12d0a150673e46c5ef316108f7 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 4 Jan 2023 10:18:33 +0100 Subject: [PATCH 05/55] provide flag --- .../collection_manager/optimizers/segment_optimizer.rs | 2 +- lib/segment/src/segment_constructor/segment_builder.rs | 6 ++++-- lib/segment/tests/segment_builder_test.rs | 8 ++++---- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs index e84de4b91b0..0c4a6807016 100644 --- a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs @@ -248,7 +248,7 @@ pub trait SegmentOptimizer { match segment { LockedSegment::Original(segment_arc) => { let segment_guard = segment_arc.read(); - segment_builder.update_from(&segment_guard, stopped)?; + segment_builder.update_from(&segment_guard, stopped, false)?; } LockedSegment::Proxy(_) => panic!("Attempt to optimize segment which is already currently under optimization. Should never happen"), } diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index c84e16a44b2..5cf0efd8017 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -51,7 +51,7 @@ impl SegmentBuilder { /// /// * `bool` - if `true` - data successfully added, if `false` - process was interrupted /// - pub fn update_from(&mut self, other: &Segment, stopped: &AtomicBool) -> OperationResult { + pub fn update_from(&mut self, other: &Segment, stopped: &AtomicBool, quantize: bool) -> OperationResult { match &mut self.segment { None => Err(OperationError::service_error( "Segment building error: created segment not found", @@ -103,7 +103,9 @@ impl SegmentBuilder { vector_storage.update_from(&**other_vector_storage, stopped)?; internal_id_iter = Some(new_internal_range.zip(other_vector_storage.iter_ids())); - vector_storage.quantize()?; + if quantize { + vector_storage.quantize()?; + } } if internal_id_iter.is_none() { return Err(OperationError::service_error( diff --git a/lib/segment/tests/segment_builder_test.rs b/lib/segment/tests/segment_builder_test.rs index 990c02f4073..340716b31e3 100644 --- a/lib/segment/tests/segment_builder_test.rs +++ b/lib/segment/tests/segment_builder_test.rs @@ -38,9 +38,9 @@ mod tests { .upsert_vector(100, 3.into(), &only_default_vector(&[0., 0., 0., 0.])) .unwrap(); - builder.update_from(&segment1, &stopped).unwrap(); - builder.update_from(&segment2, &stopped).unwrap(); - builder.update_from(&segment2, &stopped).unwrap(); + builder.update_from(&segment1, &stopped, false).unwrap(); + builder.update_from(&segment2, &stopped, false).unwrap(); + builder.update_from(&segment2, &stopped, false).unwrap(); // Check what happens if segment building fails here @@ -94,7 +94,7 @@ mod tests { let mut builder = SegmentBuilder::new(dir.path(), temp_dir.path(), &segment_config).unwrap(); - builder.update_from(segment, &stopped).unwrap(); + builder.update_from(segment, &stopped, false).unwrap(); let now = Instant::now(); From d851e1dab852c08a773b10ce3bb0767159392356 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 5 Jan 2023 11:50:36 +0100 Subject: [PATCH 06/55] fix segfault --- lib/segment/src/vector_storage/mmap_vectors.rs | 8 -------- lib/segment/src/vector_storage/simple_vector_storage.rs | 4 ---- 2 files changed, 12 deletions(-) diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 4e16eda625b..55e342b6553 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -87,10 +87,6 @@ impl MmapVectors { } fn enable_deleted_ram(&mut self) { - if self.deleted_ram.is_some() { - return; - } - let mut deleted = BitVec::new(); deleted.resize(self.num_vectors, false); for i in 0..self.num_vectors { @@ -100,10 +96,6 @@ impl MmapVectors { } pub fn quantize(&mut self, distance: Distance) -> OperationResult<()> { - if self.quantized_vectors.is_some() { - return Ok(()); - } - self.enable_deleted_ram(); self.quantized_vectors = Some( EncodedVectors::encode( diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 89499bb7ddc..c4f79d369d1 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -299,10 +299,6 @@ where } fn quantize(&mut self) -> OperationResult<()> { - if self.quantized_vectors.is_some() { - return Ok(()); - } - self.quantized_vectors = Some( EncodedVectors::encode( (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), From f736ce4ddcd7eed2cb4b664763039bebb2260d2e Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Fri, 6 Jan 2023 13:31:52 +0100 Subject: [PATCH 07/55] skip quantization flag, update scores --- lib/api/src/grpc/conversions.rs | 2 ++ lib/api/src/grpc/proto/points.proto | 5 ++++ lib/api/src/grpc/qdrant.rs | 4 +++ lib/segment/src/index/hnsw_index/hnsw.rs | 30 +++++++++++++------ .../segment_constructor/segment_builder.rs | 7 ++++- lib/segment/src/types.rs | 6 +++- lib/segment/tests/exact_search_test.rs | 2 ++ lib/segment/tests/filtrable_hnsw_test.rs | 2 +- 8 files changed, 46 insertions(+), 12 deletions(-) diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index e7ddf1b77ba..cfac705bd16 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -333,6 +333,7 @@ impl From for segment::types::SearchParams { Self { hnsw_ef: params.hnsw_ef.map(|x| x as usize), exact: params.exact.unwrap_or(false), + ignore_quantization: params.ignore_quantization.unwrap_or(false), } } } @@ -342,6 +343,7 @@ impl From for SearchParams { Self { hnsw_ef: params.hnsw_ef.map(|x| x as u64), exact: Some(params.exact), + ignore_quantization: Some(params.ignore_quantization), } } } diff --git a/lib/api/src/grpc/proto/points.proto b/lib/api/src/grpc/proto/points.proto index 85b531bb996..7a1cac6ecca 100644 --- a/lib/api/src/grpc/proto/points.proto +++ b/lib/api/src/grpc/proto/points.proto @@ -137,6 +137,11 @@ message SearchParams { Search without approximation. If set to true, search may run long but with exact results. */ optional bool exact = 2; + + /* + If set to true, search will ignore quantized vector data + */ + optional bool ignore_quantization = 3; } message SearchPoints { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 9037243ec3c..1b4245fa0ed 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -1737,6 +1737,10 @@ pub struct SearchParams { /// Search without approximation. If set to true, search may run long but with exact results. #[prost(bool, optional, tag = "2")] pub exact: ::core::option::Option, + /// + /// If set to true, search will ignore quantized vector data + #[prost(bool, optional, tag = "3")] + pub ignore_quantization: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index 0ebe5631640..6d7a3022f38 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -156,7 +156,7 @@ impl HNSWIndex { check_process_stopped(stopped)?; let vector = vector_storage.get_vector(block_point_id).unwrap(); - let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector); + let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector, false).0; let block_condition_checker = BuildConditionChecker { filter_list: block_filter_list, current_point: block_point_id, @@ -185,7 +185,9 @@ impl HNSWIndex { let ef = max(req_ef, top); let vector_storage = self.vector_storage.borrow(); - let raw_scorer = Self::get_raw_scorer(&vector_storage, vector); + let ignore_quantization = params.map(|p| p.ignore_quantization).unwrap_or(false); + let (raw_scorer, quantized) = + Self::get_raw_scorer(&vector_storage, vector, ignore_quantization); let payload_index = self.payload_index.borrow(); let filter_context = filter.map(|f| payload_index.filter_context(f)); @@ -193,7 +195,16 @@ impl HNSWIndex { let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), filter_context.as_deref()); if let Some(graph) = &self.graph { - graph.search(top, ef, points_scorer) + let mut search_result = graph.search(top, ef, points_scorer); + if quantized { + let raw_scorer = vector_storage.raw_scorer(vector.to_owned()); + search_result.iter_mut().for_each(|scored_point| { + scored_point.score = raw_scorer.score_point(scored_point.idx); + }); + search_result + } else { + search_result + } } else { Vec::new() } @@ -215,13 +226,14 @@ impl HNSWIndex { fn get_raw_scorer<'a>( vector_storage: &'a AtomicRef, vector: &[VectorElementType], - ) -> Box { - // TODO: provide option for quantization disabling + ignore_quantization: bool, + ) -> (Box, bool) { if let Some(raw_scorer) = vector_storage.quantized_raw_scorer(vector) { - raw_scorer - } else { - vector_storage.raw_scorer(vector.to_owned()) + if !ignore_quantization { + return (raw_scorer, true); + } } + (vector_storage.raw_scorer(vector.to_owned()), false) } } @@ -348,7 +360,7 @@ impl VectorIndex for HNSWIndex { ids.into_par_iter().try_for_each(|vector_id| { check_process_stopped(stopped)?; let vector = vector_storage.get_vector(vector_id).unwrap(); - let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector); + let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector, false).0; let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), None); graph_layers_builder.link_new_point(vector_id, points_scorer); diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 5cf0efd8017..92bb83b230c 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -51,7 +51,12 @@ impl SegmentBuilder { /// /// * `bool` - if `true` - data successfully added, if `false` - process was interrupted /// - pub fn update_from(&mut self, other: &Segment, stopped: &AtomicBool, quantize: bool) -> OperationResult { + pub fn update_from( + &mut self, + other: &Segment, + stopped: &AtomicBool, + quantize: bool, + ) -> OperationResult { match &mut self.segment { None => Err(OperationError::service_error( "Segment building error: created segment not found", diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index eada6afbcdd..11f61ff388e 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -229,7 +229,7 @@ pub struct SegmentInfo { } /// Additional parameters of the search -#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Default)] #[serde(rename_all = "snake_case")] pub struct SearchParams { /// Params relevant to HNSW index @@ -239,6 +239,10 @@ pub struct SearchParams { /// Search without approximation. If set to true, search may run long but with exact results. #[serde(default)] pub exact: bool, + + /// If set to true, search will ignore quantized vector data + #[serde(default)] + pub ignore_quantization: bool, } /// Vector index configuration of the segment diff --git a/lib/segment/tests/exact_search_test.rs b/lib/segment/tests/exact_search_test.rs index 788c612c84f..ca865b7f39e 100644 --- a/lib/segment/tests/exact_search_test.rs +++ b/lib/segment/tests/exact_search_test.rs @@ -146,6 +146,7 @@ mod tests { Some(&SearchParams { hnsw_ef: Some(ef), exact: true, + ..Default::default() }), ); let plain_result = segment.vector_data[DEFAULT_VECTOR_NAME] @@ -180,6 +181,7 @@ mod tests { Some(&SearchParams { hnsw_ef: Some(ef), exact: true, + ..Default::default() }), ); let plain_result = segment.vector_data[DEFAULT_VECTOR_NAME] diff --git a/lib/segment/tests/filtrable_hnsw_test.rs b/lib/segment/tests/filtrable_hnsw_test.rs index 86a2560f5b1..383cc618fe8 100644 --- a/lib/segment/tests/filtrable_hnsw_test.rs +++ b/lib/segment/tests/filtrable_hnsw_test.rs @@ -163,7 +163,7 @@ mod tests { top, Some(&SearchParams { hnsw_ef: Some(ef), - exact: false, + ..Default::default() }), ); From 817efd1d0b00000589d924b562e448130f014883 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 17 Jan 2023 06:11:51 +0000 Subject: [PATCH 08/55] use quantization flag --- lib/api/src/grpc/proto/collections.proto | 1 + lib/api/src/grpc/qdrant.rs | 2 ++ lib/collection/benches/batch_search_bench.rs | 1 + lib/collection/src/collection_manager/fixtures.rs | 2 ++ .../src/collection_manager/optimizers/indexing_optimizer.rs | 2 ++ .../src/collection_manager/optimizers/vacuum_optimizer.rs | 1 + lib/collection/src/config.rs | 2 ++ lib/collection/src/operations/config_diff.rs | 1 + lib/collection/src/operations/conversions.rs | 2 ++ lib/collection/src/operations/types.rs | 3 +++ lib/collection/src/tests/snapshot_test.rs | 1 + lib/collection/tests/common/mod.rs | 1 + lib/collection/tests/multi_vec_test.rs | 2 ++ lib/segment/src/segment.rs | 5 +++++ .../src/segment_constructor/segment_constructor_base.rs | 1 + .../src/segment_constructor/simple_segment_constructor.rs | 3 +++ lib/segment/src/telemetry.rs | 1 + lib/segment/src/types.rs | 2 ++ lib/segment/tests/exact_search_test.rs | 1 + lib/segment/tests/filtrable_hnsw_test.rs | 1 + lib/segment/tests/fixtures/segment.rs | 3 +++ lib/segment/tests/payload_index_test.rs | 1 + lib/segment/tests/segment_builder_test.rs | 1 + lib/storage/tests/alias_tests.rs | 1 + src/consensus.rs | 1 + 25 files changed, 42 insertions(+) diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index fe68a358ce4..5c04d7647fd 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -4,6 +4,7 @@ package qdrant; message VectorParams { uint64 size = 1; // Size of the vectors Distance distance = 2; // Distance function used for comparing vectors + optional bool use_quantization = 3; } message VectorParamsMap { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 1b4245fa0ed..833cfdacbd6 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -7,6 +7,8 @@ pub struct VectorParams { /// Distance function used for comparing vectors #[prost(enumeration = "Distance", tag = "2")] pub distance: i32, + #[prost(bool, optional, tag = "3")] + pub use_quantization: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/lib/collection/benches/batch_search_bench.rs b/lib/collection/benches/batch_search_bench.rs index ec2bef0fb83..e7f8fd7038f 100644 --- a/lib/collection/benches/batch_search_bench.rs +++ b/lib/collection/benches/batch_search_bench.rs @@ -61,6 +61,7 @@ fn batch_search_bench(c: &mut Criterion) { vectors: VectorParams { size: NonZeroU64::new(100).unwrap(), distance: Distance::Dot, + use_quantization: false, } .into(), shard_number: NonZeroU32::new(1).expect("Shard number can not be zero"), diff --git a/lib/collection/src/collection_manager/fixtures.rs b/lib/collection/src/collection_manager/fixtures.rs index 9de173c9736..ec3e746dd6a 100644 --- a/lib/collection/src/collection_manager/fixtures.rs +++ b/lib/collection/src/collection_manager/fixtures.rs @@ -176,6 +176,7 @@ pub(crate) fn get_merge_optimizer( vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(dim as u64).unwrap(), distance: Distance::Dot, + use_quantization: false, }), shard_number: NonZeroU32::new(1).unwrap(), on_disk_payload: false, @@ -203,6 +204,7 @@ pub(crate) fn get_indexing_optimizer( vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(dim as u64).unwrap(), distance: Distance::Dot, + use_quantization: false, }), shard_number: NonZeroU32::new(1).unwrap(), on_disk_payload: false, diff --git a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs index 520afdfc493..05a27da676d 100644 --- a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs @@ -308,6 +308,7 @@ mod tests { VectorParams { size: NonZeroU64::new(params.size as u64).unwrap(), distance: params.distance, + use_quantization: params.use_quantization, }, ) }) @@ -422,6 +423,7 @@ mod tests { ) .unwrap(), distance: segment_config.vector_data[DEFAULT_VECTOR_NAME].distance, + use_quantization: segment_config.vector_data[DEFAULT_VECTOR_NAME].use_quantization, }), shard_number: NonZeroU32::new(1).unwrap(), replication_factor: NonZeroU32::new(1).unwrap(), diff --git a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs index 4f9ce1dd53b..a6ac6177d2a 100644 --- a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs @@ -219,6 +219,7 @@ mod tests { vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, + use_quantization: false, }), shard_number: NonZeroU32::new(1).unwrap(), on_disk_payload: false, diff --git a/lib/collection/src/config.rs b/lib/collection/src/config.rs index 88974943349..e97e22f57e8 100644 --- a/lib/collection/src/config.rs +++ b/lib/collection/src/config.rs @@ -161,6 +161,7 @@ impl CollectionParams { VectorDataConfig { size: params.size.get() as usize, distance: params.distance, + use_quantization: params.use_quantization, }, ); map @@ -173,6 +174,7 @@ impl CollectionParams { VectorDataConfig { size: params.size.get() as usize, distance: params.distance, + use_quantization: params.use_quantization, }, ) }) diff --git a/lib/collection/src/operations/config_diff.rs b/lib/collection/src/operations/config_diff.rs index afb8bad20e0..83123ea49ad 100644 --- a/lib/collection/src/operations/config_diff.rs +++ b/lib/collection/src/operations/config_diff.rs @@ -240,6 +240,7 @@ mod tests { vectors: VectorParams { size: NonZeroU64::new(128).unwrap(), distance: Distance::Cosine, + use_quantization: false, } .into(), shard_number: NonZeroU32::new(1).unwrap(), diff --git a/lib/collection/src/operations/conversions.rs b/lib/collection/src/operations/conversions.rs index 357d7b9fe56..79cfc90d8dd 100644 --- a/lib/collection/src/operations/conversions.rs +++ b/lib/collection/src/operations/conversions.rs @@ -265,6 +265,7 @@ impl TryFrom for VectorParams { Status::invalid_argument("VectorParams size must be greater than zero") })?, distance: from_grpc_dist(vector_params.distance)?, + use_quantization: vector_params.use_quantization.unwrap_or(false), }) } } @@ -635,6 +636,7 @@ impl From for api::grpc::qdrant::VectorParams { Distance::Dot => api::grpc::qdrant::Distance::Dot, } .into(), + use_quantization: Some(value.use_quantization), } } } diff --git a/lib/collection/src/operations/types.rs b/lib/collection/src/operations/types.rs index d4cc3378d8d..caa7aa5848a 100644 --- a/lib/collection/src/operations/types.rs +++ b/lib/collection/src/operations/types.rs @@ -641,6 +641,9 @@ pub struct VectorParams { pub size: NonZeroU64, /// Type of distance function used for measuring distance between vectors pub distance: Distance, + /// + #[serde(default)] + pub use_quantization: bool, } /// Vector params separator for single and multiple vector modes diff --git a/lib/collection/src/tests/snapshot_test.rs b/lib/collection/src/tests/snapshot_test.rs index 119e7b3b31f..018f9d5ef47 100644 --- a/lib/collection/src/tests/snapshot_test.rs +++ b/lib/collection/src/tests/snapshot_test.rs @@ -43,6 +43,7 @@ async fn test_snapshot_collection() { vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, + use_quantization: false, }), shard_number: NonZeroU32::new(4).unwrap(), replication_factor: NonZeroU32::new(3).unwrap(), diff --git a/lib/collection/tests/common/mod.rs b/lib/collection/tests/common/mod.rs index 57a3d198ea6..d6047f0369e 100644 --- a/lib/collection/tests/common/mod.rs +++ b/lib/collection/tests/common/mod.rs @@ -43,6 +43,7 @@ pub async fn simple_collection_fixture(collection_path: &Path, shard_number: u32 vectors: VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, + use_quantization: false, } .into(), shard_number: NonZeroU32::new(shard_number).expect("Shard number can not be zero"), diff --git a/lib/collection/tests/multi_vec_test.rs b/lib/collection/tests/multi_vec_test.rs index 262f165bd86..1038e65ffbb 100644 --- a/lib/collection/tests/multi_vec_test.rs +++ b/lib/collection/tests/multi_vec_test.rs @@ -41,10 +41,12 @@ pub async fn multi_vec_collection_fixture(collection_path: &Path, shard_number: let vector_params1 = VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, + use_quantization: false, }; let vector_params2 = VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, + use_quantization: false, }; let mut vectors_config = BTreeMap::new(); diff --git a/lib/segment/src/segment.rs b/lib/segment/src/segment.rs index 9201c635faf..28f2e99f199 100644 --- a/lib/segment/src/segment.rs +++ b/lib/segment/src/segment.rs @@ -1237,6 +1237,7 @@ mod tests { VectorDataConfig { size: dim, distance: Distance::Dot, + use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1306,6 +1307,7 @@ mod tests { VectorDataConfig { size: dim, distance: Distance::Dot, + use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1394,6 +1396,7 @@ mod tests { VectorDataConfig { size: 2, distance: Distance::Dot, + use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1477,6 +1480,7 @@ mod tests { VectorDataConfig { size: 2, distance: Distance::Dot, + use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1529,6 +1533,7 @@ mod tests { VectorDataConfig { size: 2, distance: Distance::Dot, + use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 47f3f843623..7e1f7e2c696 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -264,6 +264,7 @@ fn load_segment_state_v3(segment_path: &Path) -> OperationResult { let vector_data = VectorDataConfig { size: state.config.vector_size, distance: state.config.distance, + use_quantization: false, }; SegmentState { version: state.version, diff --git a/lib/segment/src/segment_constructor/simple_segment_constructor.rs b/lib/segment/src/segment_constructor/simple_segment_constructor.rs index 291c8d75f88..b03bf06560f 100644 --- a/lib/segment/src/segment_constructor/simple_segment_constructor.rs +++ b/lib/segment/src/segment_constructor/simple_segment_constructor.rs @@ -26,6 +26,7 @@ pub fn build_simple_segment( VectorDataConfig { size: dim, distance, + use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -47,6 +48,7 @@ pub fn build_multivec_segment( VectorDataConfig { size: dim1, distance, + use_quantization: false, }, ); vectors_config.insert( @@ -54,6 +56,7 @@ pub fn build_multivec_segment( VectorDataConfig { size: dim2, distance, + use_quantization: false, }, ); diff --git a/lib/segment/src/telemetry.rs b/lib/segment/src/telemetry.rs index b8f580a155a..7ebb168459f 100644 --- a/lib/segment/src/telemetry.rs +++ b/lib/segment/src/telemetry.rs @@ -105,6 +105,7 @@ impl Anonymize for VectorDataConfig { VectorDataConfig { size: self.size.anonymize(), distance: self.distance, + use_quantization: self.use_quantization, } } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 11f61ff388e..4e1ef62b67e 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -388,6 +388,8 @@ pub struct VectorDataConfig { pub size: usize, /// Type of distance function used for measuring distance between vectors pub distance: Distance, + #[serde(default)] + pub use_quantization: bool, } /// Default value based on diff --git a/lib/segment/tests/exact_search_test.rs b/lib/segment/tests/exact_search_test.rs index ca865b7f39e..b7e77359af6 100644 --- a/lib/segment/tests/exact_search_test.rs +++ b/lib/segment/tests/exact_search_test.rs @@ -45,6 +45,7 @@ mod tests { VectorDataConfig { size: dim, distance, + use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/tests/filtrable_hnsw_test.rs b/lib/segment/tests/filtrable_hnsw_test.rs index 383cc618fe8..2f77657dc29 100644 --- a/lib/segment/tests/filtrable_hnsw_test.rs +++ b/lib/segment/tests/filtrable_hnsw_test.rs @@ -45,6 +45,7 @@ mod tests { VectorDataConfig { size: dim, distance, + use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/tests/fixtures/segment.rs b/lib/segment/tests/fixtures/segment.rs index 39e1681e3ff..b025fa47b0d 100644 --- a/lib/segment/tests/fixtures/segment.rs +++ b/lib/segment/tests/fixtures/segment.rs @@ -117,6 +117,7 @@ pub fn build_segment_3(path: &Path) -> Segment { VectorDataConfig { size: 4, distance: Distance::Dot, + use_quantization: false, }, ), ( @@ -124,6 +125,7 @@ pub fn build_segment_3(path: &Path) -> Segment { VectorDataConfig { size: 1, distance: Distance::Dot, + use_quantization: false, }, ), ( @@ -131,6 +133,7 @@ pub fn build_segment_3(path: &Path) -> Segment { VectorDataConfig { size: 4, distance: Distance::Euclid, + use_quantization: false, }, ), ]), diff --git a/lib/segment/tests/payload_index_test.rs b/lib/segment/tests/payload_index_test.rs index 76a7dd516bb..8c6ce0328e8 100644 --- a/lib/segment/tests/payload_index_test.rs +++ b/lib/segment/tests/payload_index_test.rs @@ -32,6 +32,7 @@ mod tests { VectorDataConfig { size: dim, distance: Distance::Dot, + use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/tests/segment_builder_test.rs b/lib/segment/tests/segment_builder_test.rs index 340716b31e3..bc1386ec568 100644 --- a/lib/segment/tests/segment_builder_test.rs +++ b/lib/segment/tests/segment_builder_test.rs @@ -84,6 +84,7 @@ mod tests { VectorDataConfig { size: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].size, distance: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].distance, + use_quantization: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].use_quantization, }, )]), index: Indexes::Hnsw(Default::default()), diff --git a/lib/storage/tests/alias_tests.rs b/lib/storage/tests/alias_tests.rs index 6aa08a00635..a13be7546c1 100644 --- a/lib/storage/tests/alias_tests.rs +++ b/lib/storage/tests/alias_tests.rs @@ -73,6 +73,7 @@ mod tests { vectors: VectorParams { size: NonZeroU64::new(10).unwrap(), distance: Distance::Cosine, + use_quantization: false, } .into(), hnsw_config: None, diff --git a/src/consensus.rs b/src/consensus.rs index f4021619750..a654a2a69df 100644 --- a/src/consensus.rs +++ b/src/consensus.rs @@ -912,6 +912,7 @@ mod tests { vectors: VectorParams { size: NonZeroU64::new(10).unwrap(), distance: Distance::Cosine, + use_quantization: false, } .into(), hnsw_config: None, From 0e4490d44c1152f4c48a7610444db9ec82dc6001 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 17 Jan 2023 06:13:18 +0000 Subject: [PATCH 09/55] are you happy fmt --- .../src/collection_manager/optimizers/indexing_optimizer.rs | 3 ++- lib/segment/tests/segment_builder_test.rs | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs index 05a27da676d..1034da86e02 100644 --- a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs @@ -423,7 +423,8 @@ mod tests { ) .unwrap(), distance: segment_config.vector_data[DEFAULT_VECTOR_NAME].distance, - use_quantization: segment_config.vector_data[DEFAULT_VECTOR_NAME].use_quantization, + use_quantization: segment_config.vector_data[DEFAULT_VECTOR_NAME] + .use_quantization, }), shard_number: NonZeroU32::new(1).unwrap(), replication_factor: NonZeroU32::new(1).unwrap(), diff --git a/lib/segment/tests/segment_builder_test.rs b/lib/segment/tests/segment_builder_test.rs index bc1386ec568..1c3785457ef 100644 --- a/lib/segment/tests/segment_builder_test.rs +++ b/lib/segment/tests/segment_builder_test.rs @@ -84,7 +84,8 @@ mod tests { VectorDataConfig { size: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].size, distance: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].distance, - use_quantization: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].use_quantization, + use_quantization: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME] + .use_quantization, }, )]), index: Indexes::Hnsw(Default::default()), From dfd69deb2f9dccf652b23e692f428628e9c86eb0 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 17 Jan 2023 07:51:48 +0000 Subject: [PATCH 10/55] use quantization flag --- .../optimizers/segment_optimizer.rs | 4 +++- lib/segment/src/segment.rs | 15 +++++++++++++++ .../src/segment_constructor/segment_builder.rs | 10 +--------- lib/segment/tests/segment_builder_test.rs | 8 ++++---- 4 files changed, 23 insertions(+), 14 deletions(-) diff --git a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs index 0c4a6807016..64770a5cbf7 100644 --- a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs @@ -248,7 +248,7 @@ pub trait SegmentOptimizer { match segment { LockedSegment::Original(segment_arc) => { let segment_guard = segment_arc.read(); - segment_builder.update_from(&segment_guard, stopped, false)?; + segment_builder.update_from(&segment_guard, stopped)?; } LockedSegment::Proxy(_) => panic!("Attempt to optimize segment which is already currently under optimization. Should never happen"), } @@ -296,6 +296,8 @@ pub trait SegmentOptimizer { self.check_cancellation(stopped)?; } + optimized_segment.update_quantization()?; + Ok(optimized_segment) } diff --git a/lib/segment/src/segment.rs b/lib/segment/src/segment.rs index 28f2e99f199..39473d05fdc 100644 --- a/lib/segment/src/segment.rs +++ b/lib/segment/src/segment.rs @@ -470,6 +470,21 @@ impl Segment { } Ok(()) } + + pub fn update_quantization(&self) -> OperationResult<()> { + for (vector_name, vector_data) in self.vector_data.iter() { + let use_quantization = self + .segment_config + .vector_data + .get(vector_name) + .map(|config| config.use_quantization) + .unwrap_or(false); + if use_quantization { + vector_data.vector_storage.borrow_mut().quantize()?; + } + } + Ok(()) + } } /// This is a basic implementation of `SegmentEntry`, diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 92bb83b230c..97d98377ea7 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -51,12 +51,7 @@ impl SegmentBuilder { /// /// * `bool` - if `true` - data successfully added, if `false` - process was interrupted /// - pub fn update_from( - &mut self, - other: &Segment, - stopped: &AtomicBool, - quantize: bool, - ) -> OperationResult { + pub fn update_from(&mut self, other: &Segment, stopped: &AtomicBool) -> OperationResult { match &mut self.segment { None => Err(OperationError::service_error( "Segment building error: created segment not found", @@ -108,9 +103,6 @@ impl SegmentBuilder { vector_storage.update_from(&**other_vector_storage, stopped)?; internal_id_iter = Some(new_internal_range.zip(other_vector_storage.iter_ids())); - if quantize { - vector_storage.quantize()?; - } } if internal_id_iter.is_none() { return Err(OperationError::service_error( diff --git a/lib/segment/tests/segment_builder_test.rs b/lib/segment/tests/segment_builder_test.rs index 1c3785457ef..84d51ef6d3c 100644 --- a/lib/segment/tests/segment_builder_test.rs +++ b/lib/segment/tests/segment_builder_test.rs @@ -38,9 +38,9 @@ mod tests { .upsert_vector(100, 3.into(), &only_default_vector(&[0., 0., 0., 0.])) .unwrap(); - builder.update_from(&segment1, &stopped, false).unwrap(); - builder.update_from(&segment2, &stopped, false).unwrap(); - builder.update_from(&segment2, &stopped, false).unwrap(); + builder.update_from(&segment1, &stopped).unwrap(); + builder.update_from(&segment2, &stopped).unwrap(); + builder.update_from(&segment2, &stopped).unwrap(); // Check what happens if segment building fails here @@ -96,7 +96,7 @@ mod tests { let mut builder = SegmentBuilder::new(dir.path(), temp_dir.path(), &segment_config).unwrap(); - builder.update_from(segment, &stopped, false).unwrap(); + builder.update_from(segment, &stopped).unwrap(); let now = Instant::now(); From 82dc7d212607f30bdadb5be6858165a163949259 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 18 Jan 2023 10:28:49 +0000 Subject: [PATCH 11/55] quantized search test --- .../tests/hnsw_quantized_search_test.rs | 205 ++++++++++++++++++ 1 file changed, 205 insertions(+) create mode 100644 lib/segment/tests/hnsw_quantized_search_test.rs diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs new file mode 100644 index 00000000000..2fd285a8610 --- /dev/null +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -0,0 +1,205 @@ +#[cfg(test)] +mod tests { + use std::collections::{HashMap, BTreeSet}; + use std::sync::atomic::AtomicBool; + + use itertools::Itertools; + use rand::{thread_rng, Rng}; + use segment::data_types::vectors::{only_default_vector, DEFAULT_VECTOR_NAME}; + use segment::entry::entry_point::SegmentEntry; + use segment::fixtures::payload_fixtures::{random_int_payload, random_vector}; + use segment::index::hnsw_index::graph_links::GraphLinksRam; + use segment::index::hnsw_index::hnsw::HNSWIndex; + use segment::index::{PayloadIndex, VectorIndex}; + use segment::segment_constructor::build_segment; + use segment::types::{ + Condition, Distance, FieldCondition, Filter, HnswConfig, Indexes, Payload, + PayloadSchemaType, PointOffsetType, Range, SearchParams, SegmentConfig, SeqNumberType, + StorageType, VectorDataConfig, + }; + use segment::vector_storage::ScoredPointOffset; + use serde_json::json; + use tempfile::Builder; + + fn sames_count( + a: &[Vec], + b: &[Vec], + ) -> usize { + a.iter().collect::>() + .intersection(&b.iter().collect()) + .count() + } + + #[test] + fn hnsw_quantized_search_test() { + let stopped = AtomicBool::new(false); + + let dim = 32; + let m = 8; + let num_vectors: u64 = 5_000; + let ef = 32; + let ef_construct = 16; + let distance = Distance::Cosine; + let full_scan_threshold = 16; // KB + let indexing_threshold = 500; // num vectors + let num_payload_values = 2; + + let mut rnd = thread_rng(); + + let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); + let hnsw_dir = Builder::new().prefix("hnsw_dir").tempdir().unwrap(); + + let config = SegmentConfig { + vector_data: HashMap::from([( + DEFAULT_VECTOR_NAME.to_owned(), + VectorDataConfig { + size: dim, + distance, + use_quantization: true, + }, + )]), + index: Indexes::Plain {}, + storage_type: StorageType::InMemory, + payload_storage_type: Default::default(), + }; + + let int_key = "int"; + + let mut segment = build_segment(dir.path(), &config).unwrap(); + for n in 0..num_vectors { + let idx = n.into(); + let vector = random_vector(&mut rnd, dim); + + let int_payload = random_int_payload(&mut rnd, num_payload_values..=num_payload_values); + let payload: Payload = json!({int_key:int_payload,}).into(); + + segment + .upsert_vector(n as SeqNumberType, idx, &only_default_vector(&vector)) + .unwrap(); + segment + .set_full_payload(n as SeqNumberType, idx, &payload) + .unwrap(); + } + segment.update_quantization().unwrap(); + + let payload_index_ptr = segment.payload_index.clone(); + + let hnsw_config = HnswConfig { + m, + ef_construct, + full_scan_threshold, + max_indexing_threads: 2, + on_disk: Some(false), + payload_m: None, + }; + + let mut hnsw_index = HNSWIndex::::open( + hnsw_dir.path(), + segment.vector_data[DEFAULT_VECTOR_NAME] + .vector_storage + .clone(), + payload_index_ptr.clone(), + hnsw_config, + ) + .unwrap(); + + hnsw_index.build_index(&stopped).unwrap(); + + payload_index_ptr + .borrow_mut() + .set_indexed(int_key, PayloadSchemaType::Integer.into()) + .unwrap(); + let borrowed_payload_index = payload_index_ptr.borrow(); + let blocks = borrowed_payload_index + .payload_blocks(int_key, indexing_threshold) + .collect_vec(); + for block in blocks.iter() { + assert!( + block.condition.range.is_some(), + "only range conditions should be generated for this type of payload" + ); + } + + let mut coverage: HashMap = Default::default(); + for block in &blocks { + let px = payload_index_ptr.borrow(); + let filter = Filter::new_must(Condition::Field(block.condition.clone())); + let points = px.query_points(&filter); + for point in points { + coverage.insert(point, coverage.get(&point).unwrap_or(&0) + 1); + } + } + let expected_blocks = num_vectors as usize / indexing_threshold * 2; + + println!("blocks.len() = {:#?}", blocks.len()); + assert!( + (blocks.len() as i64 - expected_blocks as i64).abs() <= 3, + "real number of payload blocks is too far from expected" + ); + + assert_eq!( + coverage.len(), + num_vectors as usize, + "not all points are covered by payload blocks" + ); + + hnsw_index.build_index(&stopped).unwrap(); + + let top = 3; + let attempts = 50; + let mut sames: usize = 0; + for _i in 0..attempts { + let query = random_vector(&mut rnd, dim); + + let index_result = hnsw_index.search( + &[&query], + None, + top, + Some(&SearchParams { + hnsw_ef: Some(ef), + ..Default::default() + }), + ); + let plain_result = segment.vector_data[DEFAULT_VECTOR_NAME] + .vector_index + .borrow() + .search(&[&query], None, top, None); + sames += sames_count(&index_result, &plain_result); + + let range_size = 40; + let left_range = rnd.gen_range(0..400); + let right_range = left_range + range_size; + + let filter = Filter::new_must(Condition::Field(FieldCondition::new_range( + int_key.to_owned(), + Range { + lt: None, + gt: None, + gte: Some(left_range as f64), + lte: Some(right_range as f64), + }, + ))); + + let filter_query = Some(&filter); + let index_result = hnsw_index.search( + &[&query], + filter_query, + top, + Some(&SearchParams { + hnsw_ef: Some(ef), + ..Default::default() + }), + ); + let plain_result = segment.vector_data[DEFAULT_VECTOR_NAME] + .vector_index + .borrow() + .search(&[&query], filter_query, top, None); + sames += sames_count(&index_result, &plain_result); + } + println!("sames = {}, attempts = {}, top = {}", sames, attempts, top); + assert!( + 4 * sames > attempts * 2 * top, + "not all results are the same for index and plain search" + ); + } +} From 708c4cc7398bc278568b1c3fdd40f02d83de0126 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 18 Jan 2023 10:30:29 +0000 Subject: [PATCH 12/55] are you happy fmt --- lib/segment/tests/hnsw_quantized_search_test.rs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 2fd285a8610..91ac6a8462b 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use std::collections::{HashMap, BTreeSet}; + use std::collections::{BTreeSet, HashMap}; use std::sync::atomic::AtomicBool; use itertools::Itertools; @@ -21,11 +21,9 @@ mod tests { use serde_json::json; use tempfile::Builder; - fn sames_count( - a: &[Vec], - b: &[Vec], - ) -> usize { - a.iter().collect::>() + fn sames_count(a: &[Vec], b: &[Vec]) -> usize { + a.iter() + .collect::>() .intersection(&b.iter().collect()) .count() } From ff4d359e80ab3dc6d76fc9748ff357a764cb20bc Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 18 Jan 2023 11:40:25 +0000 Subject: [PATCH 13/55] refactor test, refactor scorer choosing --- lib/segment/src/index/hnsw_index/hnsw.rs | 11 +- .../tests/hnsw_quantized_search_test.rs | 128 +++--------------- 2 files changed, 28 insertions(+), 111 deletions(-) diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index 6d7a3022f38..1b685c143ae 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -228,12 +228,15 @@ impl HNSWIndex { vector: &[VectorElementType], ignore_quantization: bool, ) -> (Box, bool) { - if let Some(raw_scorer) = vector_storage.quantized_raw_scorer(vector) { - if !ignore_quantization { - return (raw_scorer, true); + if ignore_quantization { + (vector_storage.raw_scorer(vector.to_owned()), false) + } else { + if let Some(quantized_raw_scorer) = vector_storage.quantized_raw_scorer(vector) { + (quantized_raw_scorer, true) + } else { + (vector_storage.raw_scorer(vector.to_owned()), false) } } - (vector_storage.raw_scorer(vector.to_owned()), false) } } diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 91ac6a8462b..b68ea671cf1 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -3,28 +3,27 @@ mod tests { use std::collections::{BTreeSet, HashMap}; use std::sync::atomic::AtomicBool; - use itertools::Itertools; - use rand::{thread_rng, Rng}; + use rand::thread_rng; use segment::data_types::vectors::{only_default_vector, DEFAULT_VECTOR_NAME}; use segment::entry::entry_point::SegmentEntry; - use segment::fixtures::payload_fixtures::{random_int_payload, random_vector}; + use segment::fixtures::payload_fixtures::random_vector; use segment::index::hnsw_index::graph_links::GraphLinksRam; use segment::index::hnsw_index::hnsw::HNSWIndex; - use segment::index::{PayloadIndex, VectorIndex}; + use segment::index::VectorIndex; use segment::segment_constructor::build_segment; use segment::types::{ - Condition, Distance, FieldCondition, Filter, HnswConfig, Indexes, Payload, - PayloadSchemaType, PointOffsetType, Range, SearchParams, SegmentConfig, SeqNumberType, + Distance, HnswConfig, Indexes, + SearchParams, SegmentConfig, SeqNumberType, StorageType, VectorDataConfig, }; use segment::vector_storage::ScoredPointOffset; - use serde_json::json; use tempfile::Builder; fn sames_count(a: &[Vec], b: &[Vec]) -> usize { - a.iter() + a[0].iter() + .map(|x| x.idx) .collect::>() - .intersection(&b.iter().collect()) + .intersection(&b[0].iter().map(|x| x.idx).collect()) .count() } @@ -32,15 +31,12 @@ mod tests { fn hnsw_quantized_search_test() { let stopped = AtomicBool::new(false); - let dim = 32; - let m = 8; - let num_vectors: u64 = 5_000; - let ef = 32; - let ef_construct = 16; + let dim = 128; + let m = 16; + let num_vectors: u64 = 100_000; + let ef = 100; + let ef_construct = 100; let distance = Distance::Cosine; - let full_scan_threshold = 16; // KB - let indexing_threshold = 500; // num vectors - let num_payload_values = 2; let mut rnd = thread_rng(); @@ -61,31 +57,20 @@ mod tests { payload_storage_type: Default::default(), }; - let int_key = "int"; - let mut segment = build_segment(dir.path(), &config).unwrap(); for n in 0..num_vectors { let idx = n.into(); let vector = random_vector(&mut rnd, dim); - - let int_payload = random_int_payload(&mut rnd, num_payload_values..=num_payload_values); - let payload: Payload = json!({int_key:int_payload,}).into(); - segment .upsert_vector(n as SeqNumberType, idx, &only_default_vector(&vector)) .unwrap(); - segment - .set_full_payload(n as SeqNumberType, idx, &payload) - .unwrap(); } segment.update_quantization().unwrap(); - let payload_index_ptr = segment.payload_index.clone(); - let hnsw_config = HnswConfig { m, ef_construct, - full_scan_threshold, + full_scan_threshold: usize::MAX, max_indexing_threads: 2, on_disk: Some(false), payload_m: None, @@ -96,55 +81,16 @@ mod tests { segment.vector_data[DEFAULT_VECTOR_NAME] .vector_storage .clone(), - payload_index_ptr.clone(), + segment.payload_index.clone(), hnsw_config, ) .unwrap(); hnsw_index.build_index(&stopped).unwrap(); + println!("\nbuild_index finished"); - payload_index_ptr - .borrow_mut() - .set_indexed(int_key, PayloadSchemaType::Integer.into()) - .unwrap(); - let borrowed_payload_index = payload_index_ptr.borrow(); - let blocks = borrowed_payload_index - .payload_blocks(int_key, indexing_threshold) - .collect_vec(); - for block in blocks.iter() { - assert!( - block.condition.range.is_some(), - "only range conditions should be generated for this type of payload" - ); - } - - let mut coverage: HashMap = Default::default(); - for block in &blocks { - let px = payload_index_ptr.borrow(); - let filter = Filter::new_must(Condition::Field(block.condition.clone())); - let points = px.query_points(&filter); - for point in points { - coverage.insert(point, coverage.get(&point).unwrap_or(&0) + 1); - } - } - let expected_blocks = num_vectors as usize / indexing_threshold * 2; - - println!("blocks.len() = {:#?}", blocks.len()); - assert!( - (blocks.len() as i64 - expected_blocks as i64).abs() <= 3, - "real number of payload blocks is too far from expected" - ); - - assert_eq!( - coverage.len(), - num_vectors as usize, - "not all points are covered by payload blocks" - ); - - hnsw_index.build_index(&stopped).unwrap(); - - let top = 3; - let attempts = 50; + let top = 10; + let attempts = 10; let mut sames: usize = 0; for _i in 0..attempts { let query = random_vector(&mut rnd, dim); @@ -163,41 +109,9 @@ mod tests { .borrow() .search(&[&query], None, top, None); sames += sames_count(&index_result, &plain_result); - - let range_size = 40; - let left_range = rnd.gen_range(0..400); - let right_range = left_range + range_size; - - let filter = Filter::new_must(Condition::Field(FieldCondition::new_range( - int_key.to_owned(), - Range { - lt: None, - gt: None, - gte: Some(left_range as f64), - lte: Some(right_range as f64), - }, - ))); - - let filter_query = Some(&filter); - let index_result = hnsw_index.search( - &[&query], - filter_query, - top, - Some(&SearchParams { - hnsw_ef: Some(ef), - ..Default::default() - }), - ); - let plain_result = segment.vector_data[DEFAULT_VECTOR_NAME] - .vector_index - .borrow() - .search(&[&query], filter_query, top, None); - sames += sames_count(&index_result, &plain_result); } - println!("sames = {}, attempts = {}, top = {}", sames, attempts, top); - assert!( - 4 * sames > attempts * 2 * top, - "not all results are the same for index and plain search" - ); + let acc = 100.0 * sames as f64 / (attempts * top) as f64; + println!("sames = {:}, attempts = {:}, top = {:}, acc = {:}", sames, attempts, top, acc); + assert!(acc > 40.0); } } From d1892f71499519ea31f5e69b68f8c9941c44560a Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 18 Jan 2023 11:40:53 +0000 Subject: [PATCH 14/55] are you happy fmt --- lib/segment/tests/hnsw_quantized_search_test.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index b68ea671cf1..c3f98ed29f1 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -12,9 +12,8 @@ mod tests { use segment::index::VectorIndex; use segment::segment_constructor::build_segment; use segment::types::{ - Distance, HnswConfig, Indexes, - SearchParams, SegmentConfig, SeqNumberType, - StorageType, VectorDataConfig, + Distance, HnswConfig, Indexes, SearchParams, SegmentConfig, SeqNumberType, StorageType, + VectorDataConfig, }; use segment::vector_storage::ScoredPointOffset; use tempfile::Builder; @@ -111,7 +110,10 @@ mod tests { sames += sames_count(&index_result, &plain_result); } let acc = 100.0 * sames as f64 / (attempts * top) as f64; - println!("sames = {:}, attempts = {:}, top = {:}, acc = {:}", sames, attempts, top, acc); + println!( + "sames = {:}, attempts = {:}, top = {:}, acc = {:}", + sames, attempts, top, acc + ); assert!(acc > 40.0); } } From 59113ab9b1860c34d55dc522570eef791580b003 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 18 Jan 2023 17:16:28 +0000 Subject: [PATCH 15/55] run quantization on segment builder --- .../collection_manager/optimizers/segment_optimizer.rs | 2 -- lib/segment/src/segment_constructor/segment_builder.rs | 8 ++++++-- lib/segment/src/vector_storage/simple_vector_storage.rs | 1 + 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs index 64770a5cbf7..e84de4b91b0 100644 --- a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs @@ -296,8 +296,6 @@ pub trait SegmentOptimizer { self.check_cancellation(stopped)?; } - optimized_segment.update_quantization()?; - Ok(optimized_segment) } diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 97d98377ea7..fc303abdfee 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -187,6 +187,8 @@ impl SegmentBuilder { check_process_stopped(stopped)?; } + segment.update_quantization()?; + for vector_data in segment.vector_data.values_mut() { vector_data.vector_index.borrow_mut().build_index(stopped)?; } @@ -199,11 +201,13 @@ impl SegmentBuilder { fs::rename(&self.temp_path, &self.destination_path) .describe("Moving segment data after optimization")?; - load_segment(&self.destination_path)?.ok_or_else(|| { + let loaded_segment = load_segment(&self.destination_path)?.ok_or_else(|| { OperationError::service_error(&format!( "Segment loading error: {}", self.destination_path.display() )) - }) + })?; + loaded_segment.update_quantization()?; + Ok(loaded_segment) } } diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index c4f79d369d1..b9349bf4962 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -299,6 +299,7 @@ where } fn quantize(&mut self) -> OperationResult<()> { + log::info!("Quantizing vectors..."); self.quantized_vectors = Some( EncodedVectors::encode( (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), From 9f47e8d3145242081ea51f79ff2d097fde2c8b8b Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 18 Jan 2023 17:45:00 +0000 Subject: [PATCH 16/55] decrease testing parameters --- lib/segment/tests/hnsw_quantized_search_test.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index c3f98ed29f1..8f1b9a3fd2b 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -32,9 +32,9 @@ mod tests { let dim = 128; let m = 16; - let num_vectors: u64 = 100_000; - let ef = 100; - let ef_construct = 100; + let num_vectors: u64 = 5_000; + let ef = 64; + let ef_construct = 64; let distance = Distance::Cosine; let mut rnd = thread_rng(); @@ -86,7 +86,6 @@ mod tests { .unwrap(); hnsw_index.build_index(&stopped).unwrap(); - println!("\nbuild_index finished"); let top = 10; let attempts = 10; From fa3c3a83ec13495ab3b19bf9cde75ff6752c76db Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 19 Jan 2023 15:26:45 +0000 Subject: [PATCH 17/55] simplify segment --- lib/segment/src/segment.rs | 15 --------------- .../segment_constructor/segment_builder.rs | 19 +++++++++++++++++-- .../tests/hnsw_quantized_search_test.rs | 8 +++++++- 3 files changed, 24 insertions(+), 18 deletions(-) diff --git a/lib/segment/src/segment.rs b/lib/segment/src/segment.rs index 39473d05fdc..28f2e99f199 100644 --- a/lib/segment/src/segment.rs +++ b/lib/segment/src/segment.rs @@ -470,21 +470,6 @@ impl Segment { } Ok(()) } - - pub fn update_quantization(&self) -> OperationResult<()> { - for (vector_name, vector_data) in self.vector_data.iter() { - let use_quantization = self - .segment_config - .vector_data - .get(vector_name) - .map(|config| config.use_quantization) - .unwrap_or(false); - if use_quantization { - vector_data.vector_storage.borrow_mut().quantize()?; - } - } - Ok(()) - } } /// This is a basic implementation of `SegmentEntry`, diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index fc303abdfee..0197c30a9b4 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -187,7 +187,7 @@ impl SegmentBuilder { check_process_stopped(stopped)?; } - segment.update_quantization()?; + Self::update_quantization(&segment)?; for vector_data in segment.vector_data.values_mut() { vector_data.vector_index.borrow_mut().build_index(stopped)?; @@ -207,7 +207,22 @@ impl SegmentBuilder { self.destination_path.display() )) })?; - loaded_segment.update_quantization()?; + Self::update_quantization(&loaded_segment)?; Ok(loaded_segment) } + + fn update_quantization(segment: &Segment) -> OperationResult<()> { + for (vector_name, vector_data) in segment.vector_data.iter() { + let use_quantization = segment + .segment_config + .vector_data + .get(vector_name) + .map(|config| config.use_quantization) + .unwrap_or(false); + if use_quantization { + vector_data.vector_storage.borrow_mut().quantize()?; + } + } + Ok(()) + } } diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 8f1b9a3fd2b..1799bd0cba8 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -64,7 +64,13 @@ mod tests { .upsert_vector(n as SeqNumberType, idx, &only_default_vector(&vector)) .unwrap(); } - segment.update_quantization().unwrap(); + segment.vector_data.values_mut().for_each(|vector_storage| { + vector_storage + .vector_storage + .borrow_mut() + .quantize() + .unwrap() + }); let hnsw_config = HnswConfig { m, From 2926b0c3ac93e58affc3bfcb1e0ab46a549ed088 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 19 Jan 2023 15:26:55 +0000 Subject: [PATCH 18/55] update version --- Cargo.lock | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 55ae6fc883c..0eb559552fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2726,9 +2726,12 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#13b734be5ed139ac4aff70e1f93b32d0ec46cf67" +source = "git+https://github.com/qdrant/quantization.git#328f37c6d951d2ae483f56bf1f609ab13b6ab737" dependencies = [ + "atomicwrites", "cc", + "serde", + "serde_json", ] [[package]] From 1b1d5b14a7435ef362082b46b6c38ef05bc91b2b Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 19 Jan 2023 18:44:36 +0000 Subject: [PATCH 19/55] remove use_quantization flag --- lib/api/src/grpc/proto/collections.proto | 1 - lib/api/src/grpc/qdrant.rs | 2 -- lib/collection/benches/batch_search_bench.rs | 1 - lib/collection/src/collection_manager/fixtures.rs | 2 -- .../collection_manager/optimizers/indexing_optimizer.rs | 3 --- .../collection_manager/optimizers/vacuum_optimizer.rs | 1 - lib/collection/src/config.rs | 2 -- lib/collection/src/operations/config_diff.rs | 1 - lib/collection/src/operations/conversions.rs | 2 -- lib/collection/src/operations/types.rs | 3 --- lib/collection/src/tests/snapshot_test.rs | 1 - lib/collection/tests/common/mod.rs | 1 - lib/collection/tests/multi_vec_test.rs | 2 -- lib/segment/src/segment.rs | 5 ----- lib/segment/src/segment_constructor/segment_builder.rs | 9 ++------- .../src/segment_constructor/segment_constructor_base.rs | 1 - .../segment_constructor/simple_segment_constructor.rs | 3 --- lib/segment/src/telemetry.rs | 1 - lib/segment/src/types.rs | 2 -- lib/segment/tests/exact_search_test.rs | 1 - lib/segment/tests/filtrable_hnsw_test.rs | 1 - lib/segment/tests/fixtures/segment.rs | 3 --- lib/segment/tests/hnsw_quantized_search_test.rs | 1 - lib/segment/tests/payload_index_test.rs | 1 - lib/segment/tests/segment_builder_test.rs | 2 -- lib/storage/tests/alias_tests.rs | 1 - src/consensus.rs | 1 - 27 files changed, 2 insertions(+), 52 deletions(-) diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index 5c04d7647fd..fe68a358ce4 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -4,7 +4,6 @@ package qdrant; message VectorParams { uint64 size = 1; // Size of the vectors Distance distance = 2; // Distance function used for comparing vectors - optional bool use_quantization = 3; } message VectorParamsMap { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 833cfdacbd6..1b4245fa0ed 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -7,8 +7,6 @@ pub struct VectorParams { /// Distance function used for comparing vectors #[prost(enumeration = "Distance", tag = "2")] pub distance: i32, - #[prost(bool, optional, tag = "3")] - pub use_quantization: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/lib/collection/benches/batch_search_bench.rs b/lib/collection/benches/batch_search_bench.rs index e7f8fd7038f..ec2bef0fb83 100644 --- a/lib/collection/benches/batch_search_bench.rs +++ b/lib/collection/benches/batch_search_bench.rs @@ -61,7 +61,6 @@ fn batch_search_bench(c: &mut Criterion) { vectors: VectorParams { size: NonZeroU64::new(100).unwrap(), distance: Distance::Dot, - use_quantization: false, } .into(), shard_number: NonZeroU32::new(1).expect("Shard number can not be zero"), diff --git a/lib/collection/src/collection_manager/fixtures.rs b/lib/collection/src/collection_manager/fixtures.rs index ec3e746dd6a..9de173c9736 100644 --- a/lib/collection/src/collection_manager/fixtures.rs +++ b/lib/collection/src/collection_manager/fixtures.rs @@ -176,7 +176,6 @@ pub(crate) fn get_merge_optimizer( vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(dim as u64).unwrap(), distance: Distance::Dot, - use_quantization: false, }), shard_number: NonZeroU32::new(1).unwrap(), on_disk_payload: false, @@ -204,7 +203,6 @@ pub(crate) fn get_indexing_optimizer( vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(dim as u64).unwrap(), distance: Distance::Dot, - use_quantization: false, }), shard_number: NonZeroU32::new(1).unwrap(), on_disk_payload: false, diff --git a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs index 1034da86e02..520afdfc493 100644 --- a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs @@ -308,7 +308,6 @@ mod tests { VectorParams { size: NonZeroU64::new(params.size as u64).unwrap(), distance: params.distance, - use_quantization: params.use_quantization, }, ) }) @@ -423,8 +422,6 @@ mod tests { ) .unwrap(), distance: segment_config.vector_data[DEFAULT_VECTOR_NAME].distance, - use_quantization: segment_config.vector_data[DEFAULT_VECTOR_NAME] - .use_quantization, }), shard_number: NonZeroU32::new(1).unwrap(), replication_factor: NonZeroU32::new(1).unwrap(), diff --git a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs index a6ac6177d2a..4f9ce1dd53b 100644 --- a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs @@ -219,7 +219,6 @@ mod tests { vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, - use_quantization: false, }), shard_number: NonZeroU32::new(1).unwrap(), on_disk_payload: false, diff --git a/lib/collection/src/config.rs b/lib/collection/src/config.rs index e97e22f57e8..88974943349 100644 --- a/lib/collection/src/config.rs +++ b/lib/collection/src/config.rs @@ -161,7 +161,6 @@ impl CollectionParams { VectorDataConfig { size: params.size.get() as usize, distance: params.distance, - use_quantization: params.use_quantization, }, ); map @@ -174,7 +173,6 @@ impl CollectionParams { VectorDataConfig { size: params.size.get() as usize, distance: params.distance, - use_quantization: params.use_quantization, }, ) }) diff --git a/lib/collection/src/operations/config_diff.rs b/lib/collection/src/operations/config_diff.rs index 83123ea49ad..afb8bad20e0 100644 --- a/lib/collection/src/operations/config_diff.rs +++ b/lib/collection/src/operations/config_diff.rs @@ -240,7 +240,6 @@ mod tests { vectors: VectorParams { size: NonZeroU64::new(128).unwrap(), distance: Distance::Cosine, - use_quantization: false, } .into(), shard_number: NonZeroU32::new(1).unwrap(), diff --git a/lib/collection/src/operations/conversions.rs b/lib/collection/src/operations/conversions.rs index 79cfc90d8dd..357d7b9fe56 100644 --- a/lib/collection/src/operations/conversions.rs +++ b/lib/collection/src/operations/conversions.rs @@ -265,7 +265,6 @@ impl TryFrom for VectorParams { Status::invalid_argument("VectorParams size must be greater than zero") })?, distance: from_grpc_dist(vector_params.distance)?, - use_quantization: vector_params.use_quantization.unwrap_or(false), }) } } @@ -636,7 +635,6 @@ impl From for api::grpc::qdrant::VectorParams { Distance::Dot => api::grpc::qdrant::Distance::Dot, } .into(), - use_quantization: Some(value.use_quantization), } } } diff --git a/lib/collection/src/operations/types.rs b/lib/collection/src/operations/types.rs index caa7aa5848a..d4cc3378d8d 100644 --- a/lib/collection/src/operations/types.rs +++ b/lib/collection/src/operations/types.rs @@ -641,9 +641,6 @@ pub struct VectorParams { pub size: NonZeroU64, /// Type of distance function used for measuring distance between vectors pub distance: Distance, - /// - #[serde(default)] - pub use_quantization: bool, } /// Vector params separator for single and multiple vector modes diff --git a/lib/collection/src/tests/snapshot_test.rs b/lib/collection/src/tests/snapshot_test.rs index 018f9d5ef47..119e7b3b31f 100644 --- a/lib/collection/src/tests/snapshot_test.rs +++ b/lib/collection/src/tests/snapshot_test.rs @@ -43,7 +43,6 @@ async fn test_snapshot_collection() { vectors: VectorsConfig::Single(VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, - use_quantization: false, }), shard_number: NonZeroU32::new(4).unwrap(), replication_factor: NonZeroU32::new(3).unwrap(), diff --git a/lib/collection/tests/common/mod.rs b/lib/collection/tests/common/mod.rs index d6047f0369e..57a3d198ea6 100644 --- a/lib/collection/tests/common/mod.rs +++ b/lib/collection/tests/common/mod.rs @@ -43,7 +43,6 @@ pub async fn simple_collection_fixture(collection_path: &Path, shard_number: u32 vectors: VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, - use_quantization: false, } .into(), shard_number: NonZeroU32::new(shard_number).expect("Shard number can not be zero"), diff --git a/lib/collection/tests/multi_vec_test.rs b/lib/collection/tests/multi_vec_test.rs index 1038e65ffbb..262f165bd86 100644 --- a/lib/collection/tests/multi_vec_test.rs +++ b/lib/collection/tests/multi_vec_test.rs @@ -41,12 +41,10 @@ pub async fn multi_vec_collection_fixture(collection_path: &Path, shard_number: let vector_params1 = VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, - use_quantization: false, }; let vector_params2 = VectorParams { size: NonZeroU64::new(4).unwrap(), distance: Distance::Dot, - use_quantization: false, }; let mut vectors_config = BTreeMap::new(); diff --git a/lib/segment/src/segment.rs b/lib/segment/src/segment.rs index 28f2e99f199..9201c635faf 100644 --- a/lib/segment/src/segment.rs +++ b/lib/segment/src/segment.rs @@ -1237,7 +1237,6 @@ mod tests { VectorDataConfig { size: dim, distance: Distance::Dot, - use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1307,7 +1306,6 @@ mod tests { VectorDataConfig { size: dim, distance: Distance::Dot, - use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1396,7 +1394,6 @@ mod tests { VectorDataConfig { size: 2, distance: Distance::Dot, - use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1480,7 +1477,6 @@ mod tests { VectorDataConfig { size: 2, distance: Distance::Dot, - use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -1533,7 +1529,6 @@ mod tests { VectorDataConfig { size: 2, distance: Distance::Dot, - use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 0197c30a9b4..600a516c680 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -212,13 +212,8 @@ impl SegmentBuilder { } fn update_quantization(segment: &Segment) -> OperationResult<()> { - for (vector_name, vector_data) in segment.vector_data.iter() { - let use_quantization = segment - .segment_config - .vector_data - .get(vector_name) - .map(|config| config.use_quantization) - .unwrap_or(false); + for (_vector_name, vector_data) in segment.vector_data.iter() { + let use_quantization = false; if use_quantization { vector_data.vector_storage.borrow_mut().quantize()?; } diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 7e1f7e2c696..47f3f843623 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -264,7 +264,6 @@ fn load_segment_state_v3(segment_path: &Path) -> OperationResult { let vector_data = VectorDataConfig { size: state.config.vector_size, distance: state.config.distance, - use_quantization: false, }; SegmentState { version: state.version, diff --git a/lib/segment/src/segment_constructor/simple_segment_constructor.rs b/lib/segment/src/segment_constructor/simple_segment_constructor.rs index b03bf06560f..291c8d75f88 100644 --- a/lib/segment/src/segment_constructor/simple_segment_constructor.rs +++ b/lib/segment/src/segment_constructor/simple_segment_constructor.rs @@ -26,7 +26,6 @@ pub fn build_simple_segment( VectorDataConfig { size: dim, distance, - use_quantization: false, }, )]), index: Indexes::Plain {}, @@ -48,7 +47,6 @@ pub fn build_multivec_segment( VectorDataConfig { size: dim1, distance, - use_quantization: false, }, ); vectors_config.insert( @@ -56,7 +54,6 @@ pub fn build_multivec_segment( VectorDataConfig { size: dim2, distance, - use_quantization: false, }, ); diff --git a/lib/segment/src/telemetry.rs b/lib/segment/src/telemetry.rs index 7ebb168459f..b8f580a155a 100644 --- a/lib/segment/src/telemetry.rs +++ b/lib/segment/src/telemetry.rs @@ -105,7 +105,6 @@ impl Anonymize for VectorDataConfig { VectorDataConfig { size: self.size.anonymize(), distance: self.distance, - use_quantization: self.use_quantization, } } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 4e1ef62b67e..11f61ff388e 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -388,8 +388,6 @@ pub struct VectorDataConfig { pub size: usize, /// Type of distance function used for measuring distance between vectors pub distance: Distance, - #[serde(default)] - pub use_quantization: bool, } /// Default value based on diff --git a/lib/segment/tests/exact_search_test.rs b/lib/segment/tests/exact_search_test.rs index b7e77359af6..ca865b7f39e 100644 --- a/lib/segment/tests/exact_search_test.rs +++ b/lib/segment/tests/exact_search_test.rs @@ -45,7 +45,6 @@ mod tests { VectorDataConfig { size: dim, distance, - use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/tests/filtrable_hnsw_test.rs b/lib/segment/tests/filtrable_hnsw_test.rs index 2f77657dc29..383cc618fe8 100644 --- a/lib/segment/tests/filtrable_hnsw_test.rs +++ b/lib/segment/tests/filtrable_hnsw_test.rs @@ -45,7 +45,6 @@ mod tests { VectorDataConfig { size: dim, distance, - use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/tests/fixtures/segment.rs b/lib/segment/tests/fixtures/segment.rs index b025fa47b0d..39e1681e3ff 100644 --- a/lib/segment/tests/fixtures/segment.rs +++ b/lib/segment/tests/fixtures/segment.rs @@ -117,7 +117,6 @@ pub fn build_segment_3(path: &Path) -> Segment { VectorDataConfig { size: 4, distance: Distance::Dot, - use_quantization: false, }, ), ( @@ -125,7 +124,6 @@ pub fn build_segment_3(path: &Path) -> Segment { VectorDataConfig { size: 1, distance: Distance::Dot, - use_quantization: false, }, ), ( @@ -133,7 +131,6 @@ pub fn build_segment_3(path: &Path) -> Segment { VectorDataConfig { size: 4, distance: Distance::Euclid, - use_quantization: false, }, ), ]), diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 1799bd0cba8..205ed938716 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -48,7 +48,6 @@ mod tests { VectorDataConfig { size: dim, distance, - use_quantization: true, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/tests/payload_index_test.rs b/lib/segment/tests/payload_index_test.rs index 8c6ce0328e8..76a7dd516bb 100644 --- a/lib/segment/tests/payload_index_test.rs +++ b/lib/segment/tests/payload_index_test.rs @@ -32,7 +32,6 @@ mod tests { VectorDataConfig { size: dim, distance: Distance::Dot, - use_quantization: false, }, )]), index: Indexes::Plain {}, diff --git a/lib/segment/tests/segment_builder_test.rs b/lib/segment/tests/segment_builder_test.rs index 84d51ef6d3c..990c02f4073 100644 --- a/lib/segment/tests/segment_builder_test.rs +++ b/lib/segment/tests/segment_builder_test.rs @@ -84,8 +84,6 @@ mod tests { VectorDataConfig { size: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].size, distance: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME].distance, - use_quantization: segment.segment_config.vector_data[DEFAULT_VECTOR_NAME] - .use_quantization, }, )]), index: Indexes::Hnsw(Default::default()), diff --git a/lib/storage/tests/alias_tests.rs b/lib/storage/tests/alias_tests.rs index a13be7546c1..6aa08a00635 100644 --- a/lib/storage/tests/alias_tests.rs +++ b/lib/storage/tests/alias_tests.rs @@ -73,7 +73,6 @@ mod tests { vectors: VectorParams { size: NonZeroU64::new(10).unwrap(), distance: Distance::Cosine, - use_quantization: false, } .into(), hnsw_config: None, diff --git a/src/consensus.rs b/src/consensus.rs index a654a2a69df..f4021619750 100644 --- a/src/consensus.rs +++ b/src/consensus.rs @@ -912,7 +912,6 @@ mod tests { vectors: VectorParams { size: NonZeroU64::new(10).unwrap(), distance: Distance::Cosine, - use_quantization: false, } .into(), hnsw_config: None, From 5ca3e1254cedeaaf2cab3440d45fb77da9b00eae Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 19 Jan 2023 21:10:36 +0000 Subject: [PATCH 20/55] provide quantization config --- lib/api/src/grpc/proto/collections.proto | 6 ++++++ lib/api/src/grpc/qdrant.rs | 10 ++++++++++ lib/collection/benches/batch_search_bench.rs | 1 + lib/collection/src/collection_manager/fixtures.rs | 2 ++ .../optimizers/indexing_optimizer.rs | 13 ++++++++++++- .../optimizers/merge_optimizer.rs | 9 ++++++++- .../optimizers/segment_optimizer.rs | 7 ++++++- .../optimizers/vacuum_optimizer.rs | 10 +++++++++- lib/collection/src/config.rs | 4 +++- lib/collection/src/operations/conversions.rs | 8 +++++++- lib/collection/src/optimizers_builder.rs | 6 +++++- lib/collection/src/shards/local_shard.rs | 4 ++++ lib/collection/src/telemetry.rs | 1 + lib/collection/src/tests/snapshot_test.rs | 1 + lib/collection/tests/common/mod.rs | 1 + lib/collection/tests/multi_vec_test.rs | 1 + lib/segment/src/segment.rs | 10 +++++----- .../src/segment_constructor/segment_builder.rs | 9 +++++---- .../segment_constructor/segment_constructor_base.rs | 1 + .../simple_segment_constructor.rs | 6 ++---- lib/segment/src/telemetry.rs | 1 + lib/segment/src/types.rs | 11 ++++++++++- lib/segment/tests/exact_search_test.rs | 2 +- lib/segment/tests/filtrable_hnsw_test.rs | 2 +- lib/segment/tests/fixtures/segment.rs | 3 +-- lib/segment/tests/hnsw_quantized_search_test.rs | 2 +- lib/segment/tests/payload_index_test.rs | 2 +- lib/segment/tests/segment_builder_test.rs | 3 +-- .../src/content_manager/collection_meta_ops.rs | 3 +++ lib/storage/src/content_manager/conversions.rs | 4 ++++ lib/storage/src/content_manager/toc.rs | 2 ++ lib/storage/tests/alias_tests.rs | 1 + src/consensus.rs | 1 + src/migrations/single_to_cluster.rs | 1 + 34 files changed, 119 insertions(+), 29 deletions(-) diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index fe68a358ce4..1955d4525b2 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -154,6 +154,10 @@ message OptimizersConfigDiff { optional uint64 max_optimization_threads = 8; } +message QuantizationConfig { + bool enable = 1; +} + message CreateCollection { string collection_name = 1; // Name of the collection reserved 2; // Deprecated @@ -167,6 +171,7 @@ message CreateCollection { optional VectorsConfig vectors_config = 10; // Configuration for vectors optional uint32 replication_factor = 11; // Number of replicas of each shard that network tries to maintain, default = 1 optional uint32 write_consistency_factor = 12; // How many replicas should apply the operation for us to consider it successful, default = 1 + optional QuantizationConfig quantization_config = 13; } message UpdateCollection { @@ -206,6 +211,7 @@ message CollectionConfig { HnswConfigDiff hnsw_config = 2; // Configuration of vector index OptimizersConfigDiff optimizer_config = 3; // Configuration of the optimizers WalConfigDiff wal_config = 4; // Configuration of the Write-Ahead-Log + optional QuantizationConfig quantization_config = 5; } enum TokenizerType { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 1b4245fa0ed..e17983d657e 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -171,6 +171,12 @@ pub struct OptimizersConfigDiff { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct QuantizationConfig { + #[prost(bool, tag = "1")] + pub enable: bool, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct CreateCollection { /// Name of the collection #[prost(string, tag = "1")] @@ -202,6 +208,8 @@ pub struct CreateCollection { /// How many replicas should apply the operation for us to consider it successful, default = 1 #[prost(uint32, optional, tag = "12")] pub write_consistency_factor: ::core::option::Option, + #[prost(message, optional, tag = "13")] + pub quantization_config: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -283,6 +291,8 @@ pub struct CollectionConfig { /// Configuration of the Write-Ahead-Log #[prost(message, optional, tag = "4")] pub wal_config: ::core::option::Option, + #[prost(message, optional, tag = "5")] + pub quantization_config: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/lib/collection/benches/batch_search_bench.rs b/lib/collection/benches/batch_search_bench.rs index ec2bef0fb83..e28c4bee039 100644 --- a/lib/collection/benches/batch_search_bench.rs +++ b/lib/collection/benches/batch_search_bench.rs @@ -83,6 +83,7 @@ fn batch_search_bench(c: &mut Criterion) { }, wal_config, hnsw_config: Default::default(), + quantization_config: Default::default(), }; let shared_config = Arc::new(RwLock::new(collection_config)); diff --git a/lib/collection/src/collection_manager/fixtures.rs b/lib/collection/src/collection_manager/fixtures.rs index 9de173c9736..504faee0f7a 100644 --- a/lib/collection/src/collection_manager/fixtures.rs +++ b/lib/collection/src/collection_manager/fixtures.rs @@ -183,6 +183,7 @@ pub(crate) fn get_merge_optimizer( write_consistency_factor: NonZeroU32::new(1).unwrap(), }, Default::default(), + Default::default(), ) } @@ -210,5 +211,6 @@ pub(crate) fn get_indexing_optimizer( write_consistency_factor: NonZeroU32::new(1).unwrap(), }, Default::default(), + Default::default(), ) } diff --git a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs index 520afdfc493..d817d4ccb4d 100644 --- a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs @@ -6,7 +6,9 @@ use parking_lot::Mutex; use segment::common::operation_time_statistics::{ OperationDurationStatistics, OperationDurationsAggregator, }; -use segment::types::{HnswConfig, Indexes, SegmentType, StorageType, VECTOR_ELEMENT_SIZE}; +use segment::types::{ + HnswConfig, Indexes, QuantizationConfig, SegmentType, StorageType, VECTOR_ELEMENT_SIZE, +}; use crate::collection_manager::holders::segment_holder::{ LockedSegmentHolder, SegmentHolder, SegmentId, @@ -28,6 +30,7 @@ pub struct IndexingOptimizer { collection_temp_dir: PathBuf, collection_params: CollectionParams, hnsw_config: HnswConfig, + quantization_config: Option, telemetry_durations_aggregator: Arc>, } @@ -38,6 +41,7 @@ impl IndexingOptimizer { collection_temp_dir: PathBuf, collection_params: CollectionParams, hnsw_config: HnswConfig, + quantization_config: Option, ) -> Self { IndexingOptimizer { thresholds_config, @@ -45,6 +49,7 @@ impl IndexingOptimizer { collection_temp_dir, collection_params, hnsw_config, + quantization_config, telemetry_durations_aggregator: OperationDurationsAggregator::new(), } } @@ -224,6 +229,10 @@ impl SegmentOptimizer for IndexingOptimizer { self.hnsw_config } + fn quantization_config(&self) -> Option { + self.quantization_config + } + fn threshold_config(&self) -> &OptimizerThresholds { &self.thresholds_config } @@ -329,6 +338,7 @@ mod tests { on_disk_payload: false, }, Default::default(), + Default::default(), ); let locked_holder: Arc> = Arc::new(RwLock::new(holder)); @@ -429,6 +439,7 @@ mod tests { on_disk_payload: false, }, Default::default(), + Default::default(), ); let locked_holder: Arc> = Arc::new(RwLock::new(holder)); diff --git a/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs b/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs index 45b5ec80224..156ae11064f 100644 --- a/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs @@ -7,7 +7,7 @@ use parking_lot::Mutex; use segment::common::operation_time_statistics::{ OperationDurationStatistics, OperationDurationsAggregator, }; -use segment::types::{HnswConfig, SegmentType, VECTOR_ELEMENT_SIZE}; +use segment::types::{HnswConfig, QuantizationConfig, SegmentType, VECTOR_ELEMENT_SIZE}; use crate::collection_manager::holders::segment_holder::{ LockedSegment, LockedSegmentHolder, SegmentId, @@ -30,6 +30,7 @@ pub struct MergeOptimizer { collection_temp_dir: PathBuf, collection_params: CollectionParams, hnsw_config: HnswConfig, + quantization_config: Option, telemetry_durations_aggregator: Arc>, } @@ -42,6 +43,7 @@ impl MergeOptimizer { collection_temp_dir: PathBuf, collection_params: CollectionParams, hnsw_config: HnswConfig, + quantization_config: Option, ) -> Self { MergeOptimizer { max_segments, @@ -50,6 +52,7 @@ impl MergeOptimizer { collection_temp_dir, collection_params, hnsw_config, + quantization_config, telemetry_durations_aggregator: OperationDurationsAggregator::new(), } } @@ -72,6 +75,10 @@ impl SegmentOptimizer for MergeOptimizer { self.hnsw_config } + fn quantization_config(&self) -> Option { + self.quantization_config + } + fn threshold_config(&self) -> &OptimizerThresholds { &self.thresholds_config } diff --git a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs index e84de4b91b0..15ef6e6366b 100644 --- a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs @@ -14,7 +14,7 @@ use segment::segment_constructor::build_segment; use segment::segment_constructor::segment_builder::SegmentBuilder; use segment::types::{ HnswConfig, Indexes, PayloadFieldSchema, PayloadKeyType, PayloadStorageType, PointIdType, - SegmentConfig, StorageType, VECTOR_ELEMENT_SIZE, + QuantizationConfig, SegmentConfig, StorageType, VECTOR_ELEMENT_SIZE, }; use crate::collection_manager::holders::proxy_segment::ProxySegment; @@ -54,6 +54,9 @@ pub trait SegmentOptimizer { /// Get HNSW config fn hnsw_config(&self) -> HnswConfig; + /// Get quantization config + fn quantization_config(&self) -> Option; + /// Get thresholds configuration for the current optimizer fn threshold_config(&self) -> &OptimizerThresholds; @@ -79,6 +82,7 @@ pub trait SegmentOptimizer { true => PayloadStorageType::OnDisk, false => PayloadStorageType::InMemory, }, + quantization_config: None, }; Ok(LockedSegment::new(build_segment( self.collection_path(), @@ -132,6 +136,7 @@ pub trait SegmentOptimizer { true => PayloadStorageType::OnDisk, false => PayloadStorageType::InMemory, }, + quantization_config: self.quantization_config(), }; Ok(SegmentBuilder::new( diff --git a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs index 4f9ce1dd53b..6113a374864 100644 --- a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs @@ -7,7 +7,7 @@ use parking_lot::Mutex; use segment::common::operation_time_statistics::{ OperationDurationStatistics, OperationDurationsAggregator, }; -use segment::types::{HnswConfig, SegmentType}; +use segment::types::{HnswConfig, QuantizationConfig, SegmentType}; use crate::collection_manager::holders::segment_holder::{ LockedSegment, LockedSegmentHolder, SegmentId, @@ -27,6 +27,7 @@ pub struct VacuumOptimizer { collection_temp_dir: PathBuf, collection_params: CollectionParams, hnsw_config: HnswConfig, + quantization_config: Option, telemetry_durations_aggregator: Arc>, } @@ -40,6 +41,7 @@ impl VacuumOptimizer { collection_temp_dir: PathBuf, collection_params: CollectionParams, hnsw_config: HnswConfig, + quantization_config: Option, ) -> Self { VacuumOptimizer { deleted_threshold, @@ -49,6 +51,7 @@ impl VacuumOptimizer { collection_temp_dir, collection_params, hnsw_config, + quantization_config, telemetry_durations_aggregator: OperationDurationsAggregator::new(), } } @@ -104,6 +107,10 @@ impl SegmentOptimizer for VacuumOptimizer { self.hnsw_config } + fn quantization_config(&self) -> Option { + self.quantization_config + } + fn threshold_config(&self) -> &OptimizerThresholds { &self.thresholds_config } @@ -226,6 +233,7 @@ mod tests { write_consistency_factor: NonZeroU32::new(1).unwrap(), }, Default::default(), + Default::default(), ); let suggested_to_optimize = diff --git a/lib/collection/src/config.rs b/lib/collection/src/config.rs index 88974943349..fd4c716eda2 100644 --- a/lib/collection/src/config.rs +++ b/lib/collection/src/config.rs @@ -9,7 +9,7 @@ use atomicwrites::OverwriteBehavior::AllowOverwrite; use schemars::JsonSchema; use segment::common::anonymize::Anonymize; use segment::data_types::vectors::DEFAULT_VECTOR_NAME; -use segment::types::{HnswConfig, VectorDataConfig}; +use segment::types::{HnswConfig, QuantizationConfig, VectorDataConfig}; use serde::{Deserialize, Serialize}; use wal::WalOptions; @@ -103,6 +103,8 @@ pub struct CollectionConfig { pub hnsw_config: HnswConfig, pub optimizer_config: OptimizersConfig, pub wal_config: WalConfig, + #[serde(default)] + pub quantization_config: Option, } impl CollectionConfig { diff --git a/lib/collection/src/operations/conversions.rs b/lib/collection/src/operations/conversions.rs index 357d7b9fe56..0e236f7025f 100644 --- a/lib/collection/src/operations/conversions.rs +++ b/lib/collection/src/operations/conversions.rs @@ -4,7 +4,7 @@ use std::num::{NonZeroU32, NonZeroU64}; use api::grpc::conversions::{from_grpc_dist, payload_to_proto, proto_to_payloads}; use itertools::Itertools; use segment::data_types::vectors::{NamedVector, VectorStruct, DEFAULT_VECTOR_NAME}; -use segment::types::Distance; +use segment::types::{Distance, QuantizationConfig}; use tonic::Status; use super::config_diff::CollectionParamsDiff; @@ -176,6 +176,9 @@ impl From for api::grpc::qdrant::CollectionInfo { wal_capacity_mb: Some(config.wal_config.wal_capacity_mb as u64), wal_segments_ahead: Some(config.wal_config.wal_segments_ahead as u64), }), + quantization_config: config + .quantization_config + .map(|x| api::grpc::qdrant::QuantizationConfig { enable: x.enable }), }), payload_schema: payload_schema .into_iter() @@ -336,6 +339,9 @@ impl TryFrom for CollectionConfig { None => return Err(Status::invalid_argument("Malformed WalConfig type")), Some(wal_config) => wal_config.into(), }, + quantization_config: config + .quantization_config + .map(|x| QuantizationConfig { enable: x.enable }), }) } } diff --git a/lib/collection/src/optimizers_builder.rs b/lib/collection/src/optimizers_builder.rs index d26952c1ff8..36250f1415a 100644 --- a/lib/collection/src/optimizers_builder.rs +++ b/lib/collection/src/optimizers_builder.rs @@ -2,7 +2,7 @@ use std::path::Path; use std::sync::Arc; use schemars::JsonSchema; -use segment::types::HnswConfig; +use segment::types::{HnswConfig, QuantizationConfig}; use serde::{Deserialize, Serialize}; use crate::collection_manager::optimizers::indexing_optimizer::IndexingOptimizer; @@ -100,6 +100,7 @@ pub fn build_optimizers( collection_params: &CollectionParams, optimizers_config: &OptimizersConfig, hnsw_config: &HnswConfig, + quantization_config: &Option, ) -> Arc>> { let segments_path = shard_path.join("segments"); let temp_segments_path = shard_path.join("temp_segments"); @@ -118,6 +119,7 @@ pub fn build_optimizers( temp_segments_path.clone(), collection_params.clone(), *hnsw_config, + quantization_config.clone(), )), Arc::new(IndexingOptimizer::new( threshold_config.clone(), @@ -125,6 +127,7 @@ pub fn build_optimizers( temp_segments_path.clone(), collection_params.clone(), *hnsw_config, + quantization_config.clone(), )), Arc::new(VacuumOptimizer::new( optimizers_config.deleted_threshold, @@ -134,6 +137,7 @@ pub fn build_optimizers( temp_segments_path, collection_params.clone(), *hnsw_config, + quantization_config.clone(), )), ]) } diff --git a/lib/collection/src/shards/local_shard.rs b/lib/collection/src/shards/local_shard.rs index a377372f820..039473c7f0c 100644 --- a/lib/collection/src/shards/local_shard.rs +++ b/lib/collection/src/shards/local_shard.rs @@ -221,6 +221,7 @@ impl LocalShard { &collection_config.params, &collection_config.optimizer_config, &collection_config.hnsw_config, + &collection_config.quantization_config, ); drop(collection_config); // release `shared_config` from borrow checker @@ -309,6 +310,7 @@ impl LocalShard { true => PayloadStorageType::OnDisk, false => PayloadStorageType::InMemory, }, + quantization_config: Default::default(), }; let segment = thread::Builder::new() .name(format!("shard-build-{}-{}", collection_id, id)) @@ -344,6 +346,7 @@ impl LocalShard { &config.params, &config.optimizer_config, &config.hnsw_config, + &config.quantization_config, ); drop(config); // release `shared_config` from borrow checker @@ -415,6 +418,7 @@ impl LocalShard { &config.params, &config.optimizer_config, &config.hnsw_config, + &config.quantization_config, ); update_handler.optimizers = new_optimizers; update_handler.flush_interval_sec = config.optimizer_config.flush_interval_sec; diff --git a/lib/collection/src/telemetry.rs b/lib/collection/src/telemetry.rs index 81ec0a3eca1..9ac4d680e88 100644 --- a/lib/collection/src/telemetry.rs +++ b/lib/collection/src/telemetry.rs @@ -34,6 +34,7 @@ impl Anonymize for CollectionConfig { hnsw_config: self.hnsw_config, optimizer_config: self.optimizer_config.clone(), wal_config: self.wal_config.clone(), + quantization_config: self.quantization_config.clone(), } } } diff --git a/lib/collection/src/tests/snapshot_test.rs b/lib/collection/src/tests/snapshot_test.rs index 119e7b3b31f..30e3c35b5e2 100644 --- a/lib/collection/src/tests/snapshot_test.rs +++ b/lib/collection/src/tests/snapshot_test.rs @@ -55,6 +55,7 @@ async fn test_snapshot_collection() { optimizer_config: TEST_OPTIMIZERS_CONFIG.clone(), wal_config, hnsw_config: Default::default(), + quantization_config: Default::default(), }; let snapshots_path = Builder::new().prefix("test_snapshots").tempdir().unwrap(); diff --git a/lib/collection/tests/common/mod.rs b/lib/collection/tests/common/mod.rs index 57a3d198ea6..bb518bdc745 100644 --- a/lib/collection/tests/common/mod.rs +++ b/lib/collection/tests/common/mod.rs @@ -56,6 +56,7 @@ pub async fn simple_collection_fixture(collection_path: &Path, shard_number: u32 optimizer_config: TEST_OPTIMIZERS_CONFIG.clone(), wal_config, hnsw_config: Default::default(), + quantization_config: Default::default(), }; let snapshot_path = collection_path.join("snapshots"); diff --git a/lib/collection/tests/multi_vec_test.rs b/lib/collection/tests/multi_vec_test.rs index 262f165bd86..03d970cb7ad 100644 --- a/lib/collection/tests/multi_vec_test.rs +++ b/lib/collection/tests/multi_vec_test.rs @@ -65,6 +65,7 @@ pub async fn multi_vec_collection_fixture(collection_path: &Path, shard_number: optimizer_config: TEST_OPTIMIZERS_CONFIG.clone(), wal_config, hnsw_config: Default::default(), + quantization_config: Default::default(), }; let snapshot_path = collection_path.join("snapshots"); diff --git a/lib/segment/src/segment.rs b/lib/segment/src/segment.rs index 9201c635faf..e7b5e9e4df8 100644 --- a/lib/segment/src/segment.rs +++ b/lib/segment/src/segment.rs @@ -1241,7 +1241,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut segment = build_segment(dir.path(), &config).unwrap(); @@ -1310,7 +1310,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut segment = build_segment(dir.path(), &config).unwrap(); @@ -1398,7 +1398,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut segment = build_segment(segment_base_dir.path(), &config).unwrap(); @@ -1481,7 +1481,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut segment = build_segment(segment_base_dir.path(), &config).unwrap(); @@ -1533,7 +1533,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut segment = build_segment(segment_base_dir.path(), &config).unwrap(); diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 600a516c680..a927a4dba1d 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -212,10 +212,11 @@ impl SegmentBuilder { } fn update_quantization(segment: &Segment) -> OperationResult<()> { - for (_vector_name, vector_data) in segment.vector_data.iter() { - let use_quantization = false; - if use_quantization { - vector_data.vector_storage.borrow_mut().quantize()?; + if let Some(quantization) = &segment.config().quantization_config { + if quantization.enable { + for vector_data in segment.vector_data.values() { + vector_data.vector_storage.borrow_mut().quantize()?; + } } } Ok(()) diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 47f3f843623..9d770e15255 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -272,6 +272,7 @@ fn load_segment_state_v3(segment_path: &Path) -> OperationResult { index: state.config.index, storage_type: state.config.storage_type, payload_storage_type: state.config.payload_storage_type, + quantization_config: None, }, } }) diff --git a/lib/segment/src/segment_constructor/simple_segment_constructor.rs b/lib/segment/src/segment_constructor/simple_segment_constructor.rs index 291c8d75f88..0e7860c277b 100644 --- a/lib/segment/src/segment_constructor/simple_segment_constructor.rs +++ b/lib/segment/src/segment_constructor/simple_segment_constructor.rs @@ -29,8 +29,7 @@ pub fn build_simple_segment( }, )]), index: Indexes::Plain {}, - storage_type: Default::default(), - payload_storage_type: Default::default(), + ..Default::default() }, ) } @@ -62,8 +61,7 @@ pub fn build_multivec_segment( &SegmentConfig { vector_data: vectors_config, index: Indexes::Plain {}, - storage_type: Default::default(), - payload_storage_type: Default::default(), + ..Default::default() }, ) } diff --git a/lib/segment/src/telemetry.rs b/lib/segment/src/telemetry.rs index b8f580a155a..47eb457746d 100644 --- a/lib/segment/src/telemetry.rs +++ b/lib/segment/src/telemetry.rs @@ -96,6 +96,7 @@ impl Anonymize for SegmentConfig { index: self.index, storage_type: self.storage_type, payload_storage_type: self.payload_storage_type, + quantization_config: self.quantization_config, } } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 11f61ff388e..b450b51c544 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -289,6 +289,12 @@ fn default_max_indexing_threads() -> usize { 0 } +#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Copy, Clone, PartialEq, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub struct QuantizationConfig { + pub enable: bool, +} + pub const DEFAULT_HNSW_EF_CONSTRUCT: usize = 100; impl Default for HnswConfig { @@ -367,7 +373,7 @@ impl Default for PayloadStorageType { } } -#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)] +#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone)] #[serde(rename_all = "snake_case")] pub struct SegmentConfig { pub vector_data: HashMap, @@ -378,6 +384,9 @@ pub struct SegmentConfig { /// Defines payload storage type #[serde(default)] pub payload_storage_type: PayloadStorageType, + + #[serde(default)] + pub quantization_config: Option, } /// Config of single vector data storage diff --git a/lib/segment/tests/exact_search_test.rs b/lib/segment/tests/exact_search_test.rs index ca865b7f39e..78c0fc3b8c5 100644 --- a/lib/segment/tests/exact_search_test.rs +++ b/lib/segment/tests/exact_search_test.rs @@ -49,7 +49,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let int_key = "int"; diff --git a/lib/segment/tests/filtrable_hnsw_test.rs b/lib/segment/tests/filtrable_hnsw_test.rs index 383cc618fe8..2eb0ccb1171 100644 --- a/lib/segment/tests/filtrable_hnsw_test.rs +++ b/lib/segment/tests/filtrable_hnsw_test.rs @@ -49,7 +49,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let int_key = "int"; diff --git a/lib/segment/tests/fixtures/segment.rs b/lib/segment/tests/fixtures/segment.rs index 39e1681e3ff..9b271a4e563 100644 --- a/lib/segment/tests/fixtures/segment.rs +++ b/lib/segment/tests/fixtures/segment.rs @@ -135,8 +135,7 @@ pub fn build_segment_3(path: &Path) -> Segment { ), ]), index: Indexes::Plain {}, - storage_type: Default::default(), - payload_storage_type: Default::default(), + ..Default::default() }, ) .unwrap(); diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 205ed938716..04c8135ab68 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -52,7 +52,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut segment = build_segment(dir.path(), &config).unwrap(); diff --git a/lib/segment/tests/payload_index_test.rs b/lib/segment/tests/payload_index_test.rs index 76a7dd516bb..bb0bad68327 100644 --- a/lib/segment/tests/payload_index_test.rs +++ b/lib/segment/tests/payload_index_test.rs @@ -36,7 +36,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut plain_segment = build_segment(path_plain, &config).unwrap(); diff --git a/lib/segment/tests/segment_builder_test.rs b/lib/segment/tests/segment_builder_test.rs index 990c02f4073..4cf857e99f6 100644 --- a/lib/segment/tests/segment_builder_test.rs +++ b/lib/segment/tests/segment_builder_test.rs @@ -87,8 +87,7 @@ mod tests { }, )]), index: Indexes::Hnsw(Default::default()), - storage_type: Default::default(), - payload_storage_type: Default::default(), + ..Default::default() }; let mut builder = diff --git a/lib/storage/src/content_manager/collection_meta_ops.rs b/lib/storage/src/content_manager/collection_meta_ops.rs index 948a43d250a..e7cabaa72cc 100644 --- a/lib/storage/src/content_manager/collection_meta_ops.rs +++ b/lib/storage/src/content_manager/collection_meta_ops.rs @@ -7,6 +7,7 @@ use collection::shards::shard::{PeerId, ShardId}; use collection::shards::transfer::shard_transfer::{ShardTransfer, ShardTransferKey}; use collection::shards::{replica_set, CollectionId}; use schemars::JsonSchema; +use segment::types::QuantizationConfig; use serde::{Deserialize, Serialize}; use crate::content_manager::shard_distribution::ShardDistributionProposal; @@ -120,6 +121,8 @@ pub struct CreateCollection { pub wal_config: Option, /// Custom params for Optimizers. If none - values from service configuration file are used. pub optimizers_config: Option, + #[serde(default)] + pub quantization_config: Option, } /// Operation for creating new collection and (optionally) specify index params diff --git a/lib/storage/src/content_manager/conversions.rs b/lib/storage/src/content_manager/conversions.rs index fc2f0a8c361..2a3b656880a 100644 --- a/lib/storage/src/content_manager/conversions.rs +++ b/lib/storage/src/content_manager/conversions.rs @@ -1,6 +1,7 @@ use std::collections::BTreeMap; use collection::operations::types::VectorsConfig; +use segment::types::QuantizationConfig; use tonic::Status; use crate::content_manager::collection_meta_ops::{ @@ -56,6 +57,9 @@ impl TryFrom for CollectionMetaOperations { on_disk_payload: value.on_disk_payload, replication_factor: value.replication_factor, write_consistency_factor: value.write_consistency_factor, + quantization_config: value + .quantization_config + .map(|v| QuantizationConfig { enable: v.enable }), }, ))) } diff --git a/lib/storage/src/content_manager/toc.rs b/lib/storage/src/content_manager/toc.rs index 588d94947df..e49ebda72cd 100644 --- a/lib/storage/src/content_manager/toc.rs +++ b/lib/storage/src/content_manager/toc.rs @@ -258,6 +258,7 @@ impl TableOfContent { optimizers_config: optimizers_config_diff, replication_factor, write_consistency_factor, + quantization_config, } = operation; self.collections @@ -320,6 +321,7 @@ impl TableOfContent { params: collection_params, optimizer_config: optimizers_config, hnsw_config, + quantization_config, }; let collection = Collection::new( collection_name.to_string(), diff --git a/lib/storage/tests/alias_tests.rs b/lib/storage/tests/alias_tests.rs index 6aa08a00635..964481cfebf 100644 --- a/lib/storage/tests/alias_tests.rs +++ b/lib/storage/tests/alias_tests.rs @@ -82,6 +82,7 @@ mod tests { on_disk_payload: None, replication_factor: None, write_consistency_factor: None, + quantization_config: None, }, )), None, diff --git a/src/consensus.rs b/src/consensus.rs index f4021619750..5399d3c7b52 100644 --- a/src/consensus.rs +++ b/src/consensus.rs @@ -921,6 +921,7 @@ mod tests { on_disk_payload: None, replication_factor: None, write_consistency_factor: None, + quantization_config: None, }, )), None, diff --git a/src/migrations/single_to_cluster.rs b/src/migrations/single_to_cluster.rs index c866da5bb4b..694d09b9322 100644 --- a/src/migrations/single_to_cluster.rs +++ b/src/migrations/single_to_cluster.rs @@ -47,6 +47,7 @@ pub async fn handle_existing_collections( hnsw_config: Some(collection_state.config.hnsw_config.into()), wal_config: Some(collection_state.config.wal_config.into()), optimizers_config: Some(collection_state.config.optimizer_config.into()), + quantization_config: collection_state.config.quantization_config.into(), }, ); From 40a896013099b753458dd98018f74611f56facf7 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 23 Jan 2023 18:41:11 +0000 Subject: [PATCH 21/55] quantization version up --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 0eb559552fa..e3defa0f2e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2726,7 +2726,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#328f37c6d951d2ae483f56bf1f609ab13b6ab737" +source = "git+https://github.com/qdrant/quantization.git#b12910ff528dccffb59c5636c02e97b1987419a2" dependencies = [ "atomicwrites", "cc", From 77263c13a80d07238bce9f9c1f5becd1f9c909ec Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 25 Jan 2023 23:08:13 +0000 Subject: [PATCH 22/55] euclid dist --- Cargo.lock | 2 +- lib/segment/src/vector_storage/mmap_vectors.rs | 1 + lib/segment/src/vector_storage/simple_vector_storage.rs | 1 + 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index e3defa0f2e9..b252ce65eff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2726,7 +2726,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#b12910ff528dccffb59c5636c02e97b1987419a2" +source = "git+https://github.com/qdrant/quantization.git#8f6da740db3a22dc0a31fdc3fd9adf9d25f58a51" dependencies = [ "atomicwrites", "cc", diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 55e342b6553..2f1e39a8d72 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -109,6 +109,7 @@ impl MmapVectors { Distance::Euclid => quantization::encoder::SimilarityType::L2, Distance::Dot => quantization::encoder::SimilarityType::Dot, }, + distance == Distance::Euclid, ) .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, ); diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index b9349bf4962..892ffc4a953 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -309,6 +309,7 @@ where Distance::Euclid => quantization::encoder::SimilarityType::L2, Distance::Dot => quantization::encoder::SimilarityType::Dot, }, + TMetric::distance() == Distance::Euclid, ) .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, ); From ae44533364b37f7e2d234f52ba6e0cfc3e3c7624 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 25 Jan 2023 23:11:33 +0000 Subject: [PATCH 23/55] add euclid test --- lib/segment/tests/hnsw_quantized_search_test.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 04c8135ab68..e4c7b5cb232 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -26,8 +26,7 @@ mod tests { .count() } - #[test] - fn hnsw_quantized_search_test() { + fn hnsw_quantized_search_test(distance: Distance) { let stopped = AtomicBool::new(false); let dim = 128; @@ -35,7 +34,6 @@ mod tests { let num_vectors: u64 = 5_000; let ef = 64; let ef_construct = 64; - let distance = Distance::Cosine; let mut rnd = thread_rng(); @@ -120,4 +118,14 @@ mod tests { ); assert!(acc > 40.0); } + + #[test] + fn hnsw_quantized_search_cosine_test() { + hnsw_quantized_search_test(Distance::Cosine); + } + + #[test] + fn hnsw_quantized_search_euclid_test() { + hnsw_quantized_search_test(Distance::Euclid); + } } From 8c36fa59cb91ba7f1b6725a71df4d87fbf1c9b5b Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 26 Jan 2023 05:51:31 +0000 Subject: [PATCH 24/55] saveload --- .../segment_constructor/segment_builder.rs | 19 +++++-- .../segment_constructor_base.rs | 52 ++++++++++++++----- .../vector_storage/memmap_vector_storage.rs | 9 +++- .../src/vector_storage/mmap_vectors.rs | 47 +++++++++++------ .../vector_storage/simple_vector_storage.rs | 36 ++++++++----- .../src/vector_storage/vector_storage_base.rs | 5 +- .../tests/hnsw_quantized_search_test.rs | 5 +- 7 files changed, 120 insertions(+), 53 deletions(-) diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index a927a4dba1d..793540b5335 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -4,6 +4,7 @@ use std::fs; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; +use super::{get_vector_storage_path, QUANTIZED_DATA_PATH, QUANTIZED_META_PATH}; use crate::common::error_logging::LogError; use crate::entry::entry_point::{ check_process_stopped, OperationError, OperationResult, SegmentEntry, @@ -187,7 +188,7 @@ impl SegmentBuilder { check_process_stopped(stopped)?; } - Self::update_quantization(&segment)?; + Self::update_quantization(&segment, stopped)?; for vector_data in segment.vector_data.values_mut() { vector_data.vector_index.borrow_mut().build_index(stopped)?; @@ -207,15 +208,23 @@ impl SegmentBuilder { self.destination_path.display() )) })?; - Self::update_quantization(&loaded_segment)?; Ok(loaded_segment) } - fn update_quantization(segment: &Segment) -> OperationResult<()> { + fn update_quantization(segment: &Segment, stopped: &AtomicBool) -> OperationResult<()> { if let Some(quantization) = &segment.config().quantization_config { if quantization.enable { - for vector_data in segment.vector_data.values() { - vector_data.vector_storage.borrow_mut().quantize()?; + let segment_path = segment.current_path.as_path().clone(); + for (vector_name, vector_data) in &segment.vector_data { + check_process_stopped(stopped)?; + + let vector_storage_path = get_vector_storage_path(segment_path, &vector_name); + let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); + let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); + vector_data + .vector_storage + .borrow_mut() + .quantize(&quantized_meta_path, &quantized_data_path)?; } } } diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 9d770e15255..87046dc2568 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use std::fs::{create_dir_all, File}; use std::io::Read; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; use atomic_refcell::AtomicRefCell; @@ -32,22 +32,40 @@ use crate::vector_storage::memmap_vector_storage::open_memmap_vector_storage; use crate::vector_storage::simple_vector_storage::open_simple_vector_storage; use crate::vector_storage::VectorStorageSS; +pub const PAYLOAD_INDEX_PATH: &str = "payload_index"; +pub const VECTOR_STORAGE_PATH: &str = "vector_storage"; +pub const VECTOR_INDEX_PATH: &str = "vector_index"; +pub const QUANTIZED_DATA_PATH: &str = "quantized.data"; +pub const QUANTIZED_META_PATH: &str = "quantized.meta"; + fn sp(t: T) -> Arc> { Arc::new(AtomicRefCell::new(t)) } +fn get_vector_name_with_prefix(prefix: &str, vector_name: &str) -> String { + if !vector_name.is_empty() { + format!("{}-{}", prefix, vector_name) + } else { + prefix.to_owned() + } +} + +pub fn get_vector_storage_path(segment_path: &Path, vector_name: &str) -> PathBuf { + segment_path.join(get_vector_name_with_prefix( + VECTOR_STORAGE_PATH, + vector_name, + )) +} + +pub fn get_vector_index_path(segment_path: &Path, vector_name: &str) -> PathBuf { + segment_path.join(get_vector_name_with_prefix(VECTOR_INDEX_PATH, vector_name)) +} + fn create_segment( version: SeqNumberType, segment_path: &Path, config: &SegmentConfig, ) -> OperationResult { - let get_vector_name_with_prefix = |prefix: &str, vector_name: &str| { - if !vector_name.is_empty() { - format!("{}-{}", prefix, vector_name) - } else { - prefix.to_owned() - } - }; let vector_db_names: Vec = config .vector_data .keys() @@ -63,7 +81,7 @@ fn create_segment( let id_tracker = sp(SimpleIdTracker::open(database.clone())?); - let payload_index_path = segment_path.join("payload_index"); + let payload_index_path = segment_path.join(PAYLOAD_INDEX_PATH); let payload_index: Arc> = sp(StructPayloadIndex::open( payload_storage, id_tracker.clone(), @@ -72,10 +90,8 @@ fn create_segment( let mut vector_data = HashMap::new(); for (vector_name, vector_config) in &config.vector_data { - let vector_storage_path = - segment_path.join(get_vector_name_with_prefix("vector_storage", vector_name)); - let vector_index_path = - segment_path.join(get_vector_name_with_prefix("vector_index", vector_name)); + let vector_storage_path = get_vector_storage_path(&segment_path, vector_name); + let vector_index_path = get_vector_index_path(&segment_path, vector_name); let vector_storage: Arc> = match config.storage_type { StorageType::InMemory => { @@ -94,6 +110,16 @@ fn create_segment( )?, }; + if let Some(quantization_config) = config.quantization_config { + if quantization_config.enable { + let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); + let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); + vector_storage + .borrow_mut() + .load_quantization(&quantized_meta_path, &quantized_data_path)?; + } + } + let vector_index: Arc> = match config.index { Indexes::Plain { .. } => sp(PlainIndex::new( vector_storage.clone(), diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 70d7642ae9c..1a6755f1583 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -275,9 +275,14 @@ where } } - fn quantize(&mut self) -> OperationResult<()> { + fn quantize(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { let mmap_store = self.mmap_store.as_mut().unwrap(); - mmap_store.quantize(TMetric::distance()) + mmap_store.quantize(TMetric::distance(), meta_path, data_path) + } + + fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { + let mmap_store = self.mmap_store.as_mut().unwrap(); + mmap_store.load_quantization(meta_path, data_path) } fn score_points( diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 2f1e39a8d72..18020e490aa 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -95,24 +95,37 @@ impl MmapVectors { self.deleted_ram = Some(deleted); } - pub fn quantize(&mut self, distance: Distance) -> OperationResult<()> { + pub fn quantize( + &mut self, + distance: Distance, + meta_path: &Path, + data_path: &Path, + ) -> OperationResult<()> { self.enable_deleted_ram(); - self.quantized_vectors = Some( - EncodedVectors::encode( - (0..self.num_vectors as u32).map(|i| { - let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); - self.raw_vector_offset(offset) - }), - Vec::new(), - match distance { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - distance == Distance::Euclid, - ) - .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, - ); + let quantized_vectors = EncodedVectors::encode( + (0..self.num_vectors as u32).map(|i| { + let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); + self.raw_vector_offset(offset) + }), + Vec::new(), + match distance { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + distance == Distance::Euclid, + ) + .map_err(|e| { + OperationError::service_error(&format!("Cannot quantize vector data: {}", e)) + })?; + quantized_vectors.save(data_path, meta_path)?; + self.quantized_vectors = Some(quantized_vectors); + Ok(()) + } + + pub fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { + self.enable_deleted_ram(); + self.quantized_vectors = Some(EncodedVectors::load(data_path, meta_path)?); Ok(()) } diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 892ffc4a953..5430d404615 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -1,6 +1,7 @@ use std::marker::PhantomData; use std::mem::size_of; use std::ops::Range; +use std::path::Path; use std::sync::atomic::AtomicBool; use std::sync::Arc; @@ -298,21 +299,28 @@ where } } - fn quantize(&mut self) -> OperationResult<()> { + fn quantize(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { log::info!("Quantizing vectors..."); - self.quantized_vectors = Some( - EncodedVectors::encode( - (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), - Vec::new(), - match TMetric::distance() { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - TMetric::distance() == Distance::Euclid, - ) - .map_err(|_| OperationError::service_error("cannot quantize vector data"))?, - ); + let quantized_vectors = EncodedVectors::encode( + (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), + Vec::new(), + match TMetric::distance() { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + TMetric::distance() == Distance::Euclid, + ) + .map_err(|e| { + OperationError::service_error(&format!("Cannot quantize vector data: {}", e)) + })?; + quantized_vectors.save(data_path, meta_path)?; + self.quantized_vectors = Some(quantized_vectors); + Ok(()) + } + + fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { + self.quantized_vectors = Some(EncodedVectors::load(data_path, meta_path)?); Ok(()) } diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index ef34ef56e5d..4d4a970ce21 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -1,5 +1,6 @@ use std::cmp::Ordering; use std::ops::Range; +use std::path::Path; use std::sync::atomic::AtomicBool; use ordered_float::OrderedFloat; @@ -75,9 +76,11 @@ pub trait VectorStorage { /// Generate a `RawScorer` object which contains all required context for searching similar vector fn raw_scorer(&self, vector: Vec) -> Box; + fn quantized_raw_scorer(&self, vector: &[VectorElementType]) -> Option>; - fn quantize(&mut self) -> OperationResult<()>; + fn quantize(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; + fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; fn score_points( &self, diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index e4c7b5cb232..fca866bccd7 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -28,6 +28,9 @@ mod tests { fn hnsw_quantized_search_test(distance: Distance) { let stopped = AtomicBool::new(false); + let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); + let quantized_meta_path = dir.path().join("quantized.meta"); + let quantized_data_path = dir.path().join("quantized.data"); let dim = 128; let m = 16; @@ -65,7 +68,7 @@ mod tests { vector_storage .vector_storage .borrow_mut() - .quantize() + .quantize(&quantized_meta_path, &quantized_data_path) .unwrap() }); From e08fafb6a530c2b4a7987da0275fc9fe19f1d141 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 26 Jan 2023 08:33:26 +0000 Subject: [PATCH 25/55] fix initialization bugs --- lib/collection/src/shards/local_shard.rs | 2 +- .../src/segment_constructor/segment_constructor_base.rs | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/lib/collection/src/shards/local_shard.rs b/lib/collection/src/shards/local_shard.rs index 039473c7f0c..5f44845c1d7 100644 --- a/lib/collection/src/shards/local_shard.rs +++ b/lib/collection/src/shards/local_shard.rs @@ -310,7 +310,7 @@ impl LocalShard { true => PayloadStorageType::OnDisk, false => PayloadStorageType::InMemory, }, - quantization_config: Default::default(), + quantization_config: config.quantization_config.clone(), }; let segment = thread::Builder::new() .name(format!("shard-build-{}-{}", collection_id, id)) diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 87046dc2568..c0ea3426eb2 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -114,9 +114,11 @@ fn create_segment( if quantization_config.enable { let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); - vector_storage - .borrow_mut() - .load_quantization(&quantized_meta_path, &quantized_data_path)?; + if quantized_meta_path.exists() && quantized_data_path.exists() { + vector_storage + .borrow_mut() + .load_quantization(&quantized_meta_path, &quantized_data_path)?; + } } } From 48ff551dcd7e42bb87b374b3d6a8f6003671395c Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 26 Jan 2023 10:02:37 +0000 Subject: [PATCH 26/55] quantization lib version up --- Cargo.lock | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b252ce65eff..91ee35bd2ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2726,9 +2726,8 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#8f6da740db3a22dc0a31fdc3fd9adf9d25f58a51" +source = "git+https://github.com/qdrant/quantization.git#781a2cf448100a5c86eb38a7c228b080597fc0fa" dependencies = [ - "atomicwrites", "cc", "serde", "serde_json", From e4be928bce174ca4c267140042efd6f3a6ef05e1 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 26 Jan 2023 14:23:12 +0400 Subject: [PATCH 27/55] fix arm build --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 91ee35bd2ed..bccaf707e95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2726,7 +2726,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#781a2cf448100a5c86eb38a7c228b080597fc0fa" +source = "git+https://github.com/qdrant/quantization.git#43db76514c51dc65eb57252cfeb29cd62568a10e" dependencies = [ "cc", "serde", From d628c522735eb0f26b03797b277f258ad0747d27 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 26 Jan 2023 10:57:35 +0000 Subject: [PATCH 28/55] refactor scorer selecting --- lib/segment/src/index/hnsw_index/hnsw.rs | 47 +++++++++++++----------- 1 file changed, 25 insertions(+), 22 deletions(-) diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index 1b685c143ae..9643a9a1eef 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -4,7 +4,7 @@ use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; use std::sync::Arc; -use atomic_refcell::{AtomicRef, AtomicRefCell}; +use atomic_refcell::AtomicRefCell; use log::debug; use parking_lot::Mutex; use rand::thread_rng; @@ -29,7 +29,7 @@ use crate::index::{PayloadIndex, VectorIndex}; use crate::telemetry::VectorIndexSearchesTelemetry; use crate::types::Condition::Field; use crate::types::{FieldCondition, Filter, HnswConfig, SearchParams, VECTOR_ELEMENT_SIZE}; -use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorage, VectorStorageSS}; +use crate::vector_storage::{ScoredPointOffset, VectorStorageSS}; const HNSW_USE_HEURISTIC: bool = true; const BYTES_IN_KB: usize = 1024; @@ -156,7 +156,13 @@ impl HNSWIndex { check_process_stopped(stopped)?; let vector = vector_storage.get_vector(block_point_id).unwrap(); - let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector, false).0; + let raw_scorer = if let Some(quantized_raw_scorer) = + vector_storage.quantized_raw_scorer(&vector) + { + quantized_raw_scorer + } else { + vector_storage.raw_scorer(vector.to_owned()) + }; let block_condition_checker = BuildConditionChecker { filter_list: block_filter_list, current_point: block_point_id, @@ -186,8 +192,15 @@ impl HNSWIndex { let vector_storage = self.vector_storage.borrow(); let ignore_quantization = params.map(|p| p.ignore_quantization).unwrap_or(false); - let (raw_scorer, quantized) = - Self::get_raw_scorer(&vector_storage, vector, ignore_quantization); + let (raw_scorer, quantized) = if ignore_quantization { + (vector_storage.raw_scorer(vector.to_owned()), false) + } else { + if let Some(quantized_raw_scorer) = vector_storage.quantized_raw_scorer(vector) { + (quantized_raw_scorer, true) + } else { + (vector_storage.raw_scorer(vector.to_owned()), false) + } + }; let payload_index = self.payload_index.borrow(); let filter_context = filter.map(|f| payload_index.filter_context(f)); @@ -222,22 +235,6 @@ impl HNSWIndex { .map(|vector| self.search_with_graph(vector, filter, top, params)) .collect() } - - fn get_raw_scorer<'a>( - vector_storage: &'a AtomicRef, - vector: &[VectorElementType], - ignore_quantization: bool, - ) -> (Box, bool) { - if ignore_quantization { - (vector_storage.raw_scorer(vector.to_owned()), false) - } else { - if let Some(quantized_raw_scorer) = vector_storage.quantized_raw_scorer(vector) { - (quantized_raw_scorer, true) - } else { - (vector_storage.raw_scorer(vector.to_owned()), false) - } - } - } } impl VectorIndex for HNSWIndex { @@ -363,7 +360,13 @@ impl VectorIndex for HNSWIndex { ids.into_par_iter().try_for_each(|vector_id| { check_process_stopped(stopped)?; let vector = vector_storage.get_vector(vector_id).unwrap(); - let raw_scorer = Self::get_raw_scorer(&vector_storage, &vector, false).0; + let raw_scorer = if let Some(quantized_raw_scorer) = + vector_storage.quantized_raw_scorer(&vector) + { + quantized_raw_scorer + } else { + vector_storage.raw_scorer(vector) + }; let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), None); graph_layers_builder.link_new_point(vector_id, points_scorer); From 6333fff291989b35086a5a2d14b9b4da5dbeee80 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 26 Jan 2023 20:24:40 +0000 Subject: [PATCH 29/55] quant lib version up --- Cargo.lock | 2 +- lib/segment/src/vector_storage/quantized_vector_storage.rs | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bccaf707e95..d04c3a54032 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2726,7 +2726,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#43db76514c51dc65eb57252cfeb29cd62568a10e" +source = "git+https://github.com/qdrant/quantization.git#c34a173afd90d9daf6bd774298c8c19b1ae8583e" dependencies = [ "cc", "serde", diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs index 7a6f25aefaf..b1a4bd1a095 100644 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -22,7 +22,7 @@ impl RawScorer for QuantizedRawScorer<'_> { idx: point_id, score: self .quantized_data - .score_point(&self.query, point_id as usize), + .score_point(&self.query, point_id), }; size += 1; if size == scores.len() { @@ -37,11 +37,10 @@ impl RawScorer for QuantizedRawScorer<'_> { } fn score_point(&self, point: PointOffsetType) -> ScoreType { - self.quantized_data.score_point(&self.query, point as usize) + self.quantized_data.score_point(&self.query, point) } fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType { - self.quantized_data - .score_internal(point_a as usize, point_b as usize) + self.quantized_data.score_internal(point_a, point_b) } } From f7e54eb3b49fdb437e0d51153aadafb4e90965b5 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 26 Jan 2023 20:25:30 +0000 Subject: [PATCH 30/55] are you happy fmt --- lib/segment/src/vector_storage/quantized_vector_storage.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs index b1a4bd1a095..8efb135c422 100644 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -20,9 +20,7 @@ impl RawScorer for QuantizedRawScorer<'_> { } scores[size] = ScoredPointOffset { idx: point_id, - score: self - .quantized_data - .score_point(&self.query, point_id), + score: self.quantized_data.score_point(&self.query, point_id), }; size += 1; if size == scores.len() { From f36c4436cfb65bf89c83657cac6b7dfb7610faf9 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 30 Jan 2023 08:28:22 +0000 Subject: [PATCH 31/55] are you happy fmt --- lib/segment/src/vector_storage/mmap_vectors.rs | 2 +- lib/segment/src/vector_storage/vector_storage_base.rs | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 58409717e47..e406c5dc2c2 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -13,8 +13,8 @@ use crate::common::error_logging::LogError; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::{OperationError, OperationResult}; -use crate::types::{Distance, PointOffsetType}; use crate::madvise; +use crate::types::{Distance, PointOffsetType}; const HEADER_SIZE: usize = 4; const DELETED_HEADER: &[u8; 4] = b"drop"; diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index 82ca3ecb009..5b07f78083e 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -1,7 +1,6 @@ use std::cmp::Ordering; use std::ops::Range; -use std::path::Path; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; use ordered_float::OrderedFloat; From 6773f03f9eaaa42303e8ba3a4679ca9dd79a6ed2 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 30 Jan 2023 08:47:12 +0000 Subject: [PATCH 32/55] are you happy clippy --- .../collection_manager/optimizers/indexing_optimizer.rs | 2 +- .../src/collection_manager/optimizers/merge_optimizer.rs | 2 +- .../src/collection_manager/optimizers/vacuum_optimizer.rs | 2 +- lib/segment/src/index/hnsw_index/hnsw.rs | 8 +++----- lib/segment/src/segment.rs | 1 + lib/segment/src/segment_constructor/segment_builder.rs | 4 ++-- .../src/segment_constructor/segment_constructor_base.rs | 6 +++--- lib/segment/src/telemetry.rs | 2 +- lib/segment/src/types.rs | 2 +- lib/segment/src/vector_storage/mmap_vectors.rs | 4 +--- lib/segment/src/vector_storage/simple_vector_storage.rs | 4 +--- src/migrations/single_to_cluster.rs | 2 +- 12 files changed, 17 insertions(+), 22 deletions(-) diff --git a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs index f9a8bbc644f..63070202ea4 100644 --- a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs @@ -230,7 +230,7 @@ impl SegmentOptimizer for IndexingOptimizer { } fn quantization_config(&self) -> Option { - self.quantization_config + self.quantization_config.clone() } fn threshold_config(&self) -> &OptimizerThresholds { diff --git a/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs b/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs index 156ae11064f..d70cc8d79e3 100644 --- a/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/merge_optimizer.rs @@ -76,7 +76,7 @@ impl SegmentOptimizer for MergeOptimizer { } fn quantization_config(&self) -> Option { - self.quantization_config + self.quantization_config.clone() } fn threshold_config(&self) -> &OptimizerThresholds { diff --git a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs index 6113a374864..a0a2dd84ec9 100644 --- a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs @@ -108,7 +108,7 @@ impl SegmentOptimizer for VacuumOptimizer { } fn quantization_config(&self) -> Option { - self.quantization_config + self.quantization_config.clone() } fn threshold_config(&self) -> &OptimizerThresholds { diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index 7369d226761..35e1e43d5be 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -194,12 +194,10 @@ impl HNSWIndex { let ignore_quantization = params.map(|p| p.ignore_quantization).unwrap_or(false); let (raw_scorer, quantized) = if ignore_quantization { (vector_storage.raw_scorer(vector.to_owned()), false) + } else if let Some(quantized_raw_scorer) = vector_storage.quantized_raw_scorer(vector) { + (quantized_raw_scorer, true) } else { - if let Some(quantized_raw_scorer) = vector_storage.quantized_raw_scorer(vector) { - (quantized_raw_scorer, true) - } else { - (vector_storage.raw_scorer(vector.to_owned()), false) - } + (vector_storage.raw_scorer(vector.to_owned()), false) }; let payload_index = self.payload_index.borrow(); diff --git a/lib/segment/src/segment.rs b/lib/segment/src/segment.rs index ffa9bf3aae7..4702f311dd1 100644 --- a/lib/segment/src/segment.rs +++ b/lib/segment/src/segment.rs @@ -1616,6 +1616,7 @@ mod tests { index: Indexes::Plain {}, storage_type: StorageType::InMemory, payload_storage_type: Default::default(), + quantization_config: None, }; let mut segment = build_segment(dir.path(), &config).unwrap(); diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 581f6b0c562..d5e4ff4da8a 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -212,11 +212,11 @@ impl SegmentBuilder { fn update_quantization(segment: &Segment, stopped: &AtomicBool) -> OperationResult<()> { if let Some(quantization) = &segment.config().quantization_config { if quantization.enable { - let segment_path = segment.current_path.as_path().clone(); + let segment_path = segment.current_path.as_path(); for (vector_name, vector_data) in &segment.vector_data { check_process_stopped(stopped)?; - let vector_storage_path = get_vector_storage_path(segment_path, &vector_name); + let vector_storage_path = get_vector_storage_path(segment_path, vector_name); let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); vector_data diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index e1da23978ca..093ffcfcae3 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -90,8 +90,8 @@ fn create_segment( let mut vector_data = HashMap::new(); for (vector_name, vector_config) in &config.vector_data { - let vector_storage_path = get_vector_storage_path(&segment_path, vector_name); - let vector_index_path = get_vector_index_path(&segment_path, vector_name); + let vector_storage_path = get_vector_storage_path(segment_path, vector_name); + let vector_index_path = get_vector_index_path(segment_path, vector_name); let vector_storage: Arc> = match config.storage_type { StorageType::InMemory => { @@ -110,7 +110,7 @@ fn create_segment( )?, }; - if let Some(quantization_config) = config.quantization_config { + if let Some(quantization_config) = &config.quantization_config { if quantization_config.enable { let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); diff --git a/lib/segment/src/telemetry.rs b/lib/segment/src/telemetry.rs index b9a86f0b411..f81a978c3e4 100644 --- a/lib/segment/src/telemetry.rs +++ b/lib/segment/src/telemetry.rs @@ -110,7 +110,7 @@ impl Anonymize for SegmentConfig { index: self.index, storage_type: self.storage_type, payload_storage_type: self.payload_storage_type, - quantization_config: self.quantization_config, + quantization_config: self.quantization_config.clone(), } } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index a794d1ca898..526158f3e9a 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -289,7 +289,7 @@ fn default_max_indexing_threads() -> usize { 0 } -#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Copy, Clone, PartialEq, Eq, Hash)] +#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] pub struct QuantizationConfig { pub enable: bool, diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index e406c5dc2c2..74716cb57a7 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -120,9 +120,7 @@ impl MmapVectors { }, distance == Distance::Euclid, ) - .map_err(|e| { - OperationError::service_error(&format!("Cannot quantize vector data: {}", e)) - })?; + .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; quantized_vectors.save(data_path, meta_path)?; self.quantized_vectors = Some(quantized_vectors); Ok(()) diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index c0a5bac2810..f32099f1a7b 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -311,9 +311,7 @@ where }, TMetric::distance() == Distance::Euclid, ) - .map_err(|e| { - OperationError::service_error(&format!("Cannot quantize vector data: {}", e)) - })?; + .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; quantized_vectors.save(data_path, meta_path)?; self.quantized_vectors = Some(quantized_vectors); Ok(()) diff --git a/src/migrations/single_to_cluster.rs b/src/migrations/single_to_cluster.rs index 9947dcb5a2b..6ec8f041002 100644 --- a/src/migrations/single_to_cluster.rs +++ b/src/migrations/single_to_cluster.rs @@ -48,7 +48,7 @@ pub async fn handle_existing_collections( wal_config: Some(collection_state.config.wal_config.into()), optimizers_config: Some(collection_state.config.optimizer_config.into()), init_from: None, - quantization_config: collection_state.config.quantization_config.into(), + quantization_config: collection_state.config.quantization_config, }, ); From 6958d047226264117492e56e2c69cb9c96814c28 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 30 Jan 2023 09:39:02 +0000 Subject: [PATCH 33/55] add save/load test for simple storage --- .../vector_storage/simple_vector_storage.rs | 54 +++++++++++++++++++ .../src/vector_storage/vector_storage_base.rs | 3 ++ .../tests/hnsw_quantized_search_test.rs | 5 +- 3 files changed, 58 insertions(+), 4 deletions(-) diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index f32099f1a7b..8b2122eff08 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -453,4 +453,58 @@ mod tests { assert!(!all_ids1.contains(&top_idx)) } + + #[test] + fn test_score_quantized_points() { + let dir = Builder::new().prefix("storage_dir").tempdir().unwrap(); + let db = open_db(dir.path(), &[DB_VECTOR_CF]).unwrap(); + let distance = Distance::Dot; + let dim = 4; + let storage = open_simple_vector_storage(db, DB_VECTOR_CF, dim, distance).unwrap(); + let mut borrowed_storage = storage.borrow_mut(); + + let vec0 = vec![1.0, 0.0, 1.0, 1.0]; + let vec1 = vec![1.0, 0.0, 1.0, 0.0]; + let vec2 = vec![1.0, 1.0, 1.0, 1.0]; + let vec3 = vec![1.0, 1.0, 0.0, 1.0]; + let vec4 = vec![1.0, 0.0, 0.0, 0.0]; + + borrowed_storage.put_vector(vec0).unwrap(); + borrowed_storage.put_vector(vec1).unwrap(); + borrowed_storage.put_vector(vec2).unwrap(); + borrowed_storage.put_vector(vec3).unwrap(); + borrowed_storage.put_vector(vec4).unwrap(); + + let quantized_meta_path = dir.path().join("quantized.meta"); + let quantized_data_path = dir.path().join("quantized.data"); + + borrowed_storage + .quantize(&quantized_meta_path, &quantized_data_path) + .unwrap(); + + let query = vec![0.0, 1.0, 1.1, 1.0]; + + { + let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_quant = borrowed_storage.raw_scorer(query.clone()); + for i in 0..5 { + let orig = scorer_orig.score_point(i); + let quant = scorer_quant.score_point(i); + assert!((orig - quant).abs() < 0.15); + } + } + + // test save-load + borrowed_storage + .load_quantization(&quantized_meta_path, &quantized_data_path) + .unwrap(); + + let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_quant = borrowed_storage.raw_scorer(query); + for i in 0..5 { + let orig = scorer_orig.score_point(i); + let quant = scorer_quant.score_point(i); + assert!((orig - quant).abs() < 0.15); + } + } } diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index 5b07f78083e..43ddf8797df 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -77,9 +77,12 @@ pub trait VectorStorage { /// Generate a `RawScorer` object which contains all required context for searching similar vector fn raw_scorer(&self, vector: Vec) -> Box; + // Generate RawScorer on quantized vectors if present fn quantized_raw_scorer(&self, vector: &[VectorElementType]) -> Option>; + // Generate quantized vectors and store them on disk fn quantize(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; + // Load quantized vectors from disk fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; fn score_points( diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index fca866bccd7..3ec5bebd32d 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -115,10 +115,7 @@ mod tests { sames += sames_count(&index_result, &plain_result); } let acc = 100.0 * sames as f64 / (attempts * top) as f64; - println!( - "sames = {:}, attempts = {:}, top = {:}, acc = {:}", - sames, attempts, top, acc - ); + println!("sames = {sames}, attempts = {attempts}, top = {top}, acc = {acc}"); assert!(acc > 40.0); } From f796067dcbe906b95df3c22e946839dbf63a8473 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 30 Jan 2023 10:04:02 +0000 Subject: [PATCH 34/55] add comments --- lib/api/src/grpc/proto/collections.proto | 2 +- lib/api/src/grpc/qdrant.rs | 1 + lib/segment/src/types.rs | 2 +- lib/storage/src/content_manager/collection_meta_ops.rs | 1 + 4 files changed, 4 insertions(+), 2 deletions(-) diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index 7c35f9db332..488ff2bd5b5 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -212,7 +212,7 @@ message CollectionConfig { HnswConfigDiff hnsw_config = 2; // Configuration of vector index OptimizersConfigDiff optimizer_config = 3; // Configuration of the optimizers WalConfigDiff wal_config = 4; // Configuration of the Write-Ahead-Log - optional QuantizationConfig quantization_config = 5; + optional QuantizationConfig quantization_config = 5; // Configuration of the vector quantization } enum TokenizerType { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 97c2312953d..c58eabecd91 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -294,6 +294,7 @@ pub struct CollectionConfig { /// Configuration of the Write-Ahead-Log #[prost(message, optional, tag = "4")] pub wal_config: ::core::option::Option, + /// Configuration of the vector quantization #[prost(message, optional, tag = "5")] pub quantization_config: ::core::option::Option, } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 526158f3e9a..b37f5c0a1a1 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -384,7 +384,7 @@ pub struct SegmentConfig { /// Defines payload storage type #[serde(default)] pub payload_storage_type: PayloadStorageType, - + /// Quantization parameters. If none - quantization is disabled. #[serde(default)] pub quantization_config: Option, } diff --git a/lib/storage/src/content_manager/collection_meta_ops.rs b/lib/storage/src/content_manager/collection_meta_ops.rs index 1bfa7cde30f..429928b5b2c 100644 --- a/lib/storage/src/content_manager/collection_meta_ops.rs +++ b/lib/storage/src/content_manager/collection_meta_ops.rs @@ -131,6 +131,7 @@ pub struct CreateCollection { /// Specify other collection to copy data from. #[serde(default)] pub init_from: Option, + /// Quantization parameters. If none - quantization is disabled. #[serde(default)] pub quantization_config: Option, } From 303cf8fd4357f5c4283071dd3ed1f1ea5b633a6b Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 8 Feb 2023 08:04:34 +0000 Subject: [PATCH 35/55] quantiles --- Cargo.lock | 142 +++++++++++++++--- lib/api/src/grpc/proto/collections.proto | 1 + lib/api/src/grpc/qdrant.rs | 2 + lib/collection/src/operations/conversions.rs | 16 +- .../segment_constructor/segment_builder.rs | 9 +- lib/segment/src/types.rs | 11 +- .../vector_storage/memmap_vector_storage.rs | 9 +- .../src/vector_storage/mmap_vectors.rs | 15 +- .../vector_storage/simple_vector_storage.rs | 23 ++- .../src/vector_storage/vector_storage_base.rs | 7 +- .../tests/hnsw_quantized_search_test.rs | 2 +- .../src/content_manager/conversions.rs | 7 +- 12 files changed, 193 insertions(+), 51 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cc1a702212f..0b85ffba420 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,7 +87,7 @@ dependencies = [ "mime", "percent-encoding", "pin-project-lite", - "rand", + "rand 0.8.5", "sha1", "smallvec", "tokio", @@ -242,7 +242,7 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fcb51a0695d8f838b1ee009b3fbf66bda078cd64590202a864a8f3e8c4315c47" dependencies = [ - "getrandom", + "getrandom 0.2.8", "once_cell", "version_check", ] @@ -254,7 +254,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf6ccdb167abbf410dcb915cabd428929d7f6a04980b54a11f26a39f1c7f7107" dependencies = [ "cfg-if", - "getrandom", + "getrandom 0.2.8", "once_cell", "version_check", ] @@ -313,7 +313,7 @@ dependencies = [ "log", "prost", "prost-types", - "rand", + "rand 0.8.5", "schemars", "segment", "serde", @@ -341,6 +341,15 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6" +[[package]] +name = "arrayvec" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd9fd44efafa8690358b7408d253adf110036b88f55672a933f01d616ad9b1b9" +dependencies = [ + "nodrop", +] + [[package]] name = "arrayvec" version = "0.7.2" @@ -561,6 +570,16 @@ dependencies = [ "wyz", ] +[[package]] +name = "blake2-rfc" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d6d530bdd2d52966a6d03b7a964add7ae1a288d25214066fd4b600f0f796400" +dependencies = [ + "arrayvec 0.4.12", + "constant_time_eq", +] + [[package]] name = "block-buffer" version = "0.10.3" @@ -818,7 +837,7 @@ dependencies = [ "ordered-float", "parking_lot", "pprof", - "rand", + "rand 0.8.5", "rmp-serde", "schemars", "segment", @@ -882,6 +901,12 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "constant_time_eq" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc" + [[package]] name = "convert_case" version = "0.4.0" @@ -1495,6 +1520,17 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "getrandom" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.9.0+wasi-snapshot-preview1", +] + [[package]] name = "getrandom" version = "0.2.8" @@ -2171,6 +2207,12 @@ dependencies = [ "libc", ] +[[package]] +name = "nodrop" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72ef4a56884ca558e5ddb05a1d1e7e1bfd9a68d9ed024c21704cc98872dae1bb" + [[package]] name = "nom" version = "7.1.1" @@ -2198,7 +2240,7 @@ version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a652d9771a63711fd3c3deb670acfbe5c30a4072e664d7a3bf5a9e1056ac72c3" dependencies = [ - "arrayvec", + "arrayvec 0.7.2", "itoa", ] @@ -2389,6 +2431,16 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "478c572c3d73181ff3c2539045f6eb99e5491218eae919370993b890cdbdd98e" +[[package]] +name = "permutation_iterator" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b55405179fe06e4e3820ddaf9f9b51cdff9e7496af9554acdb2b1921a86ca9cb" +dependencies = [ + "blake2-rfc", + "rand 0.7.3", +] + [[package]] name = "pest" version = "2.5.1" @@ -2601,8 +2653,8 @@ dependencies = [ "lazy_static", "num-traits", "quick-error 2.0.1", - "rand", - "rand_chacha", + "rand 0.8.5", + "rand_chacha 0.3.1", "rand_xorshift", "regex-syntax", "rusty-fork", @@ -2735,9 +2787,10 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#c34a173afd90d9daf6bd774298c8c19b1ae8583e" +source = "git+https://github.com/qdrant/quantization.git#1703f4d0374be995a399b3b0d049e6449828fde1" dependencies = [ "cc", + "permutation_iterator", "serde", "serde_json", ] @@ -2787,7 +2840,7 @@ dependencies = [ "getset", "protobuf", "raft-proto", - "rand", + "rand 0.8.5", "slog", "thiserror", ] @@ -2803,6 +2856,19 @@ dependencies = [ "protobuf-build", ] +[[package]] +name = "rand" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03" +dependencies = [ + "getrandom 0.1.16", + "libc", + "rand_chacha 0.2.2", + "rand_core 0.5.1", + "rand_hc", +] + [[package]] name = "rand" version = "0.8.5" @@ -2810,8 +2876,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402" +dependencies = [ + "ppv-lite86", + "rand_core 0.5.1", ] [[package]] @@ -2821,7 +2897,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19" +dependencies = [ + "getrandom 0.1.16", ] [[package]] @@ -2830,7 +2915,7 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom", + "getrandom 0.2.8", ] [[package]] @@ -2840,7 +2925,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" dependencies = [ "num-traits", - "rand", + "rand 0.8.5", +] + +[[package]] +name = "rand_hc" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c" +dependencies = [ + "rand_core 0.5.1", ] [[package]] @@ -2849,7 +2943,7 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d25bf25ec5ae4a3f1b92f929810509a2f53d7dca2f50b794ff57e3face536c8f" dependencies = [ - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -3259,7 +3353,7 @@ dependencies = [ "parking_lot", "pprof", "quantization", - "rand", + "rand 0.8.5", "rand_distr", "rayon", "rmp-serde", @@ -3486,7 +3580,7 @@ dependencies = [ "proptest", "prost", "raft", - "rand", + "rand 0.8.5", "reqwest", "schemars", "segment", @@ -3889,7 +3983,7 @@ dependencies = [ "indexmap", "pin-project", "pin-project-lite", - "rand", + "rand 0.8.5", "slab", "tokio", "tokio-util", @@ -4050,7 +4144,7 @@ version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "422ee0de9031b5b948b97a8fc04e3aa35230001a722ddd27943e0be31564ce4c" dependencies = [ - "getrandom", + "getrandom 0.2.8", "serde", ] @@ -4088,7 +4182,7 @@ dependencies = [ "fs2", "log", "memmap", - "rand", + "rand 0.8.5", "rand_distr", "serde", ] @@ -4114,6 +4208,12 @@ dependencies = [ "try-lock", ] +[[package]] +name = "wasi" +version = "0.9.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" + [[package]] name = "wasi" version = "0.10.0+wasi-snapshot-preview1" diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index 488ff2bd5b5..8e32bd79f85 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -156,6 +156,7 @@ message OptimizersConfigDiff { message QuantizationConfig { bool enable = 1; + optional float quantile = 2; } message CreateCollection { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index c58eabecd91..288cc6d63e4 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -174,6 +174,8 @@ pub struct OptimizersConfigDiff { pub struct QuantizationConfig { #[prost(bool, tag = "1")] pub enable: bool, + #[prost(float, optional, tag = "2")] + pub quantile: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/lib/collection/src/operations/conversions.rs b/lib/collection/src/operations/conversions.rs index 2b3193a4e55..dab57236357 100644 --- a/lib/collection/src/operations/conversions.rs +++ b/lib/collection/src/operations/conversions.rs @@ -211,9 +211,12 @@ impl From for api::grpc::qdrant::CollectionInfo { wal_capacity_mb: Some(config.wal_config.wal_capacity_mb as u64), wal_segments_ahead: Some(config.wal_config.wal_segments_ahead as u64), }), - quantization_config: config - .quantization_config - .map(|x| api::grpc::qdrant::QuantizationConfig { enable: x.enable }), + quantization_config: config.quantization_config.map(|x| { + api::grpc::qdrant::QuantizationConfig { + enable: x.enable, + quantile: x.quantile, + } + }), }), payload_schema: payload_schema .into_iter() @@ -374,9 +377,10 @@ impl TryFrom for CollectionConfig { None => return Err(Status::invalid_argument("Malformed WalConfig type")), Some(wal_config) => wal_config.into(), }, - quantization_config: config - .quantization_config - .map(|x| QuantizationConfig { enable: x.enable }), + quantization_config: config.quantization_config.map(|x| QuantizationConfig { + enable: x.enable, + quantile: x.quantile, + }), }) } } diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index d5e4ff4da8a..0837b41b9d4 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -219,10 +219,11 @@ impl SegmentBuilder { let vector_storage_path = get_vector_storage_path(segment_path, vector_name); let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); - vector_data - .vector_storage - .borrow_mut() - .quantize(&quantized_meta_path, &quantized_data_path)?; + vector_data.vector_storage.borrow_mut().quantize( + &quantized_meta_path, + &quantized_data_path, + quantization.quantile, + )?; } } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index b37f5c0a1a1..ffcb7370f9d 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -289,10 +289,19 @@ fn default_max_indexing_threads() -> usize { 0 } -#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] +#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq)] #[serde(rename_all = "snake_case")] pub struct QuantizationConfig { pub enable: bool, + pub quantile: Option, +} + +impl Eq for QuantizationConfig {} + +impl std::hash::Hash for QuantizationConfig { + fn hash(&self, state: &mut H) { + self.enable.hash(state); + } } pub const DEFAULT_HNSW_EF_CONSTRUCT: usize = 100; diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 8a8edc1bf2f..600a06ff88a 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -275,9 +275,14 @@ where } } - fn quantize(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { + fn quantize( + &mut self, + meta_path: &Path, + data_path: &Path, + quantile: Option, + ) -> OperationResult<()> { let mmap_store = self.mmap_store.as_mut().unwrap(); - mmap_store.quantize(TMetric::distance(), meta_path, data_path) + mmap_store.quantize(TMetric::distance(), meta_path, data_path, quantile) } fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 74716cb57a7..cea516c15a8 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use bitvec::vec::BitVec; use memmap2::{Mmap, MmapMut, MmapOptions}; use parking_lot::{RwLock, RwLockReadGuard}; +use quantization::encoder::EncodingParameters; use super::quantized_vector_storage::EncodedVectors; use crate::common::error_logging::LogError; @@ -105,6 +106,7 @@ impl MmapVectors { distance: Distance, meta_path: &Path, data_path: &Path, + quantile: Option, ) -> OperationResult<()> { self.enable_deleted_ram(); let quantized_vectors = EncodedVectors::encode( @@ -113,12 +115,15 @@ impl MmapVectors { self.raw_vector_offset(offset) }), Vec::new(), - match distance { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, + EncodingParameters { + distance_type: match distance { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + invert: distance == Distance::Euclid, + quantile, }, - distance == Distance::Euclid, ) .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; quantized_vectors.save(data_path, meta_path)?; diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 8b2122eff08..4c46a513c54 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -9,6 +9,7 @@ use atomic_refcell::AtomicRefCell; use bitvec::prelude::BitVec; use log::debug; use parking_lot::RwLock; +use quantization::encoder::EncodingParameters; use rocksdb::DB; use serde::{Deserialize, Serialize}; @@ -299,17 +300,25 @@ where } } - fn quantize(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { + fn quantize( + &mut self, + meta_path: &Path, + data_path: &Path, + quantile: Option, + ) -> OperationResult<()> { log::info!("Quantizing vectors..."); let quantized_vectors = EncodedVectors::encode( (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), Vec::new(), - match TMetric::distance() { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, + EncodingParameters { + distance_type: match TMetric::distance() { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + invert: TMetric::distance() == Distance::Euclid, + quantile, }, - TMetric::distance() == Distance::Euclid, ) .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; quantized_vectors.save(data_path, meta_path)?; @@ -479,7 +488,7 @@ mod tests { let quantized_data_path = dir.path().join("quantized.data"); borrowed_storage - .quantize(&quantized_meta_path, &quantized_data_path) + .quantize(&quantized_meta_path, &quantized_data_path, None) .unwrap(); let query = vec![0.0, 1.0, 1.1, 1.0]; diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index 43ddf8797df..1f9d9a8914c 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -81,7 +81,12 @@ pub trait VectorStorage { fn quantized_raw_scorer(&self, vector: &[VectorElementType]) -> Option>; // Generate quantized vectors and store them on disk - fn quantize(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; + fn quantize( + &mut self, + meta_path: &Path, + data_path: &Path, + quantile: Option, + ) -> OperationResult<()>; // Load quantized vectors from disk fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 3ec5bebd32d..07948d0648b 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -68,7 +68,7 @@ mod tests { vector_storage .vector_storage .borrow_mut() - .quantize(&quantized_meta_path, &quantized_data_path) + .quantize(&quantized_meta_path, &quantized_data_path, None) .unwrap() }); diff --git a/lib/storage/src/content_manager/conversions.rs b/lib/storage/src/content_manager/conversions.rs index a3925d422d5..1afc420c695 100644 --- a/lib/storage/src/content_manager/conversions.rs +++ b/lib/storage/src/content_manager/conversions.rs @@ -60,9 +60,10 @@ impl TryFrom for CollectionMetaOperations { init_from: value .init_from_collection .map(|v| InitFrom { collection: v }), - quantization_config: value - .quantization_config - .map(|v| QuantizationConfig { enable: v.enable }), + quantization_config: value.quantization_config.map(|v| QuantizationConfig { + enable: v.enable, + quantile: v.quantile, + }), }, ))) } From b668b96e9a6997d1fa9adef141b7b1a0a63a5ad5 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 8 Feb 2023 10:41:04 +0000 Subject: [PATCH 36/55] quantization mmap --- Cargo.lock | 2 +- lib/api/src/grpc/proto/collections.proto | 1 + lib/api/src/grpc/qdrant.rs | 2 + lib/collection/src/operations/conversions.rs | 2 + .../segment_constructor/segment_builder.rs | 2 +- .../segment_constructor_base.rs | 8 +- lib/segment/src/types.rs | 14 ++- .../vector_storage/memmap_vector_storage.rs | 27 +++--- .../src/vector_storage/mmap_vectors.rs | 88 ++++++++++++++----- lib/segment/src/vector_storage/mod.rs | 1 + .../vector_storage/quantized_mmap_storage.rs | 74 ++++++++++++++++ .../quantized_vector_storage.rs | 33 +++++-- .../vector_storage/simple_vector_storage.rs | 51 +++++++---- .../src/vector_storage/vector_storage_base.rs | 11 ++- .../tests/hnsw_quantized_search_test.rs | 14 ++- .../src/content_manager/conversions.rs | 1 + 16 files changed, 264 insertions(+), 67 deletions(-) create mode 100644 lib/segment/src/vector_storage/quantized_mmap_storage.rs diff --git a/Cargo.lock b/Cargo.lock index 0b85ffba420..5e6df3940ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2787,7 +2787,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#1703f4d0374be995a399b3b0d049e6449828fde1" +source = "git+https://github.com/qdrant/quantization.git#8f02fbb86b19718ce9584e7ee8a52719de307d45" dependencies = [ "cc", "permutation_iterator", diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index 8e32bd79f85..a290c9a4542 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -157,6 +157,7 @@ message OptimizersConfigDiff { message QuantizationConfig { bool enable = 1; optional float quantile = 2; + optional bool always_ram = 3; } message CreateCollection { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 288cc6d63e4..4908a325acd 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -176,6 +176,8 @@ pub struct QuantizationConfig { pub enable: bool, #[prost(float, optional, tag = "2")] pub quantile: ::core::option::Option, + #[prost(bool, optional, tag = "3")] + pub always_ram: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/lib/collection/src/operations/conversions.rs b/lib/collection/src/operations/conversions.rs index dab57236357..40e745b23e0 100644 --- a/lib/collection/src/operations/conversions.rs +++ b/lib/collection/src/operations/conversions.rs @@ -215,6 +215,7 @@ impl From for api::grpc::qdrant::CollectionInfo { api::grpc::qdrant::QuantizationConfig { enable: x.enable, quantile: x.quantile, + always_ram: x.always_ram, } }), }), @@ -380,6 +381,7 @@ impl TryFrom for CollectionConfig { quantization_config: config.quantization_config.map(|x| QuantizationConfig { enable: x.enable, quantile: x.quantile, + always_ram: x.always_ram, }), }) } diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 0837b41b9d4..0e4b32bbdd6 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -222,7 +222,7 @@ impl SegmentBuilder { vector_data.vector_storage.borrow_mut().quantize( &quantized_meta_path, &quantized_data_path, - quantization.quantile, + quantization, )?; } } diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 093ffcfcae3..555d6032bd4 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -115,9 +115,11 @@ fn create_segment( let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); if quantized_meta_path.exists() && quantized_data_path.exists() { - vector_storage - .borrow_mut() - .load_quantization(&quantized_meta_path, &quantized_data_path)?; + vector_storage.borrow_mut().load_quantization( + &quantized_meta_path, + &quantized_data_path, + quantization_config, + )?; } } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index ffcb7370f9d..6f5ac7eafcd 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -289,11 +289,23 @@ fn default_max_indexing_threads() -> usize { 0 } -#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq)] +#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone)] #[serde(rename_all = "snake_case")] pub struct QuantizationConfig { pub enable: bool, pub quantile: Option, + /// If true, quantized data is stored in RAM even if original data is memmapped. + /// If false, quantized data is stored like original data. + /// Default value is false. + pub always_ram: Option, +} + +impl PartialEq for QuantizationConfig { + fn eq(&self, other: &Self) -> bool { + self.enable == other.enable + && self.quantile == other.quantile + && self.always_ram == other.always_ram + } } impl Eq for QuantizationConfig {} diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 600a06ff88a..251b7df44ff 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -9,14 +9,13 @@ use std::sync::Arc; use atomic_refcell::AtomicRefCell; -use super::quantized_vector_storage::QuantizedRawScorer; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::{check_process_stopped, OperationResult}; use crate::spaces::metric::Metric; use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric}; use crate::spaces::tools::peek_top_largest_iterable; -use crate::types::{Distance, PointOffsetType, ScoreType}; +use crate::types::{Distance, PointOffsetType, QuantizationConfig, ScoreType}; use crate::vector_storage::mmap_vectors::MmapVectors; use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorage, VectorStorageSS}; @@ -262,11 +261,7 @@ where if let Some(quantized_data) = &mmap_store.quantized_vectors { if let Some(deleted_ram) = &mmap_store.deleted_ram { let query = TMetric::preprocess(vector).unwrap_or_else(|| vector.to_owned()); - Some(Box::new(QuantizedRawScorer { - query: quantized_data.encode_query(&query), - quantized_data, - deleted: deleted_ram, - })) + Some(quantized_data.raw_scorer(&query, deleted_ram)) } else { None } @@ -279,15 +274,25 @@ where &mut self, meta_path: &Path, data_path: &Path, - quantile: Option, + quantization_config: &QuantizationConfig, ) -> OperationResult<()> { let mmap_store = self.mmap_store.as_mut().unwrap(); - mmap_store.quantize(TMetric::distance(), meta_path, data_path, quantile) + mmap_store.quantize( + TMetric::distance(), + meta_path, + data_path, + quantization_config, + ) } - fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { + fn load_quantization( + &mut self, + meta_path: &Path, + data_path: &Path, + quantization_config: &QuantizationConfig, + ) -> OperationResult<()> { let mmap_store = self.mmap_store.as_mut().unwrap(); - mmap_store.load_quantization(meta_path, data_path) + mmap_store.load_quantization(meta_path, data_path, quantization_config) } fn score_points( diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index cea516c15a8..7fccdebc306 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -9,13 +9,14 @@ use memmap2::{Mmap, MmapMut, MmapOptions}; use parking_lot::{RwLock, RwLockReadGuard}; use quantization::encoder::EncodingParameters; -use super::quantized_vector_storage::EncodedVectors; +use super::quantized_mmap_storage::{QuantizedMmapStorage, QuantizedMmapStorageBuilder}; +use super::quantized_vector_storage::QuantizedVectorStorage; use crate::common::error_logging::LogError; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::{OperationError, OperationResult}; use crate::madvise; -use crate::types::{Distance, PointOffsetType}; +use crate::types::{Distance, PointOffsetType, QuantizationConfig}; const HEADER_SIZE: usize = 4; const DELETED_HEADER: &[u8; 4] = b"drop"; @@ -29,7 +30,7 @@ pub struct MmapVectors { deleted_mmap: Arc>, pub deleted_count: usize, pub deleted_ram: Option, - pub quantized_vectors: Option, + pub quantized_vectors: Option>, } fn open_read(path: &Path) -> OperationResult { @@ -101,39 +102,80 @@ impl MmapVectors { self.deleted_ram = Some(deleted); } + fn create_quantized_storage( + &self, + meta_path: &Path, + data_path: &Path, + encoding_parameters: EncodingParameters, + ) -> OperationResult> + where + TStorage: quantization::encoder::Storage, + TStorageBuilder: quantization::encoder::StorageBuilder, + { + let vector_data_iterator = (0..self.num_vectors as u32).map(|i| { + let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); + self.raw_vector_offset(offset) + }); + let quantized_vectors = quantization::encoder::EncodedVectors::>::encode( + vector_data_iterator, + Vec::new(), + encoding_parameters, + ) + .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; + quantized_vectors.save(data_path, meta_path)?; + Ok(Box::new(quantized_vectors)) + } + pub fn quantize( &mut self, distance: Distance, meta_path: &Path, data_path: &Path, - quantile: Option, + quantization_config: &QuantizationConfig, ) -> OperationResult<()> { self.enable_deleted_ram(); - let quantized_vectors = EncodedVectors::encode( - (0..self.num_vectors as u32).map(|i| { - let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); - self.raw_vector_offset(offset) - }), - Vec::new(), - EncodingParameters { - distance_type: match distance { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - invert: distance == Distance::Euclid, - quantile, + let encoding_parameters = EncodingParameters { + distance_type: match distance { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, }, - ) - .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; - quantized_vectors.save(data_path, meta_path)?; + invert: distance == Distance::Euclid, + quantile: quantization_config.quantile, + }; + let quantized_vectors = if quantization_config.always_ram == Some(true) { + self.create_quantized_storage::, Vec>( + meta_path, + data_path, + encoding_parameters, + )? + } else { + self.create_quantized_storage::( + meta_path, + data_path, + encoding_parameters, + )? + }; self.quantized_vectors = Some(quantized_vectors); Ok(()) } - pub fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { + pub fn load_quantization( + &mut self, + meta_path: &Path, + data_path: &Path, + quantization_config: &QuantizationConfig, + ) -> OperationResult<()> { self.enable_deleted_ram(); - self.quantized_vectors = Some(EncodedVectors::load(data_path, meta_path)?); + self.quantized_vectors = if quantization_config.always_ram == Some(true) { + Some(Box::new( + quantization::encoder::EncodedVectors::>::load(data_path, meta_path)?, + )) + } else { + Some(Box::new(quantization::encoder::EncodedVectors::< + QuantizedMmapStorage, + >::load(data_path, meta_path)?)) + }; Ok(()) } diff --git a/lib/segment/src/vector_storage/mod.rs b/lib/segment/src/vector_storage/mod.rs index e6179a69b61..fb41e84c1f9 100644 --- a/lib/segment/src/vector_storage/mod.rs +++ b/lib/segment/src/vector_storage/mod.rs @@ -1,6 +1,7 @@ pub mod chunked_vectors; pub mod memmap_vector_storage; mod mmap_vectors; +pub mod quantized_mmap_storage; mod quantized_vector_storage; pub mod simple_vector_storage; mod vector_storage_base; diff --git a/lib/segment/src/vector_storage/quantized_mmap_storage.rs b/lib/segment/src/vector_storage/quantized_mmap_storage.rs new file mode 100644 index 00000000000..9dda7b7465d --- /dev/null +++ b/lib/segment/src/vector_storage/quantized_mmap_storage.rs @@ -0,0 +1,74 @@ +use std::path::Path; + +use memmap2::{Mmap, MmapMut}; +use quantization::encoder::{EncodingParameters, Storage, StorageBuilder}; + +use crate::madvise; + +pub struct QuantizedMmapStorage { + mmap: Mmap, +} + +pub struct QuantizedMmapStorageBuilder { + mmap: MmapMut, + cursor_pos: usize, +} + +impl Storage for QuantizedMmapStorage { + fn ptr(&self) -> *const u8 { + self.mmap.as_ptr() + } + + fn from_file(path: &std::path::Path) -> std::io::Result { + let file = std::fs::OpenOptions::new() + .read(true) + .write(false) + .create(false) + .open(path)?; + let mmap = unsafe { Mmap::map(&file)? }; + madvise::madvise(&mmap, madvise::get_global())?; + Ok(Self { mmap }) + } + + fn save_to_file(&self, _path: &Path) -> std::io::Result<()> { + // do nothing because mmap is already saved + Ok(()) + } +} + +impl StorageBuilder for QuantizedMmapStorageBuilder { + fn build(self) -> QuantizedMmapStorage { + self.mmap.flush().unwrap(); + let mmap = self.mmap.make_read_only().unwrap(); // TODO: remove unwrap + QuantizedMmapStorage { mmap } + } + + fn extend_from_slice(&mut self, other: &[u8]) { + self.mmap[self.cursor_pos..other.len()].copy_from_slice(other); + self.cursor_pos += other.len(); + } +} + +impl QuantizedMmapStorageBuilder { + pub fn new( + path: &Path, + size: usize, + dim: usize, + encoding_parameters: &EncodingParameters, + ) -> std::io::Result { + let encoded_storage_size = encoding_parameters.estimate_encoded_storage_size(dim, size); + path.parent().map(std::fs::create_dir_all); + let file = std::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .open(path)?; + file.set_len(encoded_storage_size as u64)?; + let mmap = unsafe { MmapMut::map_mut(&file) }?; + madvise::madvise(&mmap, madvise::get_global())?; + Ok(Self { + mmap, + cursor_pos: 0, + }) + } +} diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs index 8efb135c422..1389b1d506e 100644 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -1,17 +1,40 @@ use bitvec::vec::BitVec; +use quantization::encoder::{EncodedQuery, EncodedVectors, Storage}; use super::{RawScorer, ScoredPointOffset}; use crate::types::{PointOffsetType, ScoreType}; -pub type EncodedVectors = quantization::encoder::EncodedVectors>; +pub trait QuantizedVectorStorage: Send + Sync { + fn raw_scorer<'a>(&'a self, query: &[f32], deleted: &'a BitVec) -> Box; +} + +impl QuantizedVectorStorage for EncodedVectors +where + TStorage: Storage + Send + Sync, +{ + fn raw_scorer<'a>(&'a self, query: &[f32], deleted: &'a BitVec) -> Box { + let query = self.encode_query(query); + Box::new(QuantizedRawScorer { + query, + deleted, + quantized_data: self, + }) + } +} -pub struct QuantizedRawScorer<'a> { - pub query: quantization::encoder::EncodedQuery, +pub struct QuantizedRawScorer<'a, TStorage> +where + TStorage: Storage, +{ + pub query: EncodedQuery, pub deleted: &'a BitVec, - pub quantized_data: &'a EncodedVectors, + pub quantized_data: &'a EncodedVectors, } -impl RawScorer for QuantizedRawScorer<'_> { +impl RawScorer for QuantizedRawScorer<'_, TStorage> +where + TStorage: Storage, +{ fn score_points(&self, points: &[PointOffsetType], scores: &mut [ScoredPointOffset]) -> usize { let mut size: usize = 0; for point_id in points.iter().copied() { diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 4c46a513c54..e0c471a762e 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -14,7 +14,7 @@ use rocksdb::DB; use serde::{Deserialize, Serialize}; use super::chunked_vectors::ChunkedVectors; -use super::quantized_vector_storage::{EncodedVectors, QuantizedRawScorer}; +use super::quantized_vector_storage::QuantizedVectorStorage; use super::vector_storage_base::VectorStorage; use crate::common::rocksdb_wrapper::DatabaseColumnWrapper; use crate::common::Flusher; @@ -23,7 +23,7 @@ use crate::entry::entry_point::{check_process_stopped, OperationError, Operation use crate::spaces::metric::Metric; use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric}; use crate::spaces::tools::peek_top_largest_iterable; -use crate::types::{Distance, PointOffsetType, ScoreType}; +use crate::types::{Distance, PointOffsetType, QuantizationConfig, ScoreType}; use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorageSS}; /// In-memory vector storage with on-update persistence using `store` @@ -33,7 +33,7 @@ pub struct SimpleVectorStorage { vectors: ChunkedVectors, deleted: BitVec, deleted_count: usize, - quantized_vectors: Option, + quantized_vectors: Option>, db_wrapper: DatabaseColumnWrapper, } @@ -290,11 +290,7 @@ where ) -> Option> { if let Some(quantized_data) = &self.quantized_vectors { let query = TMetric::preprocess(vector).unwrap_or_else(|| vector.to_owned()); - Some(Box::new(QuantizedRawScorer { - query: quantized_data.encode_query(&query), - quantized_data, - deleted: &self.deleted, - })) + Some(quantized_data.raw_scorer(&query, &self.deleted)) } else { None } @@ -304,10 +300,10 @@ where &mut self, meta_path: &Path, data_path: &Path, - quantile: Option, + quantization_config: &QuantizationConfig, ) -> OperationResult<()> { log::info!("Quantizing vectors..."); - let quantized_vectors = EncodedVectors::encode( + let quantized_vectors = quantization::encoder::EncodedVectors::>::encode( (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), Vec::new(), EncodingParameters { @@ -317,17 +313,24 @@ where Distance::Dot => quantization::encoder::SimilarityType::Dot, }, invert: TMetric::distance() == Distance::Euclid, - quantile, + quantile: quantization_config.quantile, }, ) .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; quantized_vectors.save(data_path, meta_path)?; - self.quantized_vectors = Some(quantized_vectors); + self.quantized_vectors = Some(Box::new(quantized_vectors)); Ok(()) } - fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { - self.quantized_vectors = Some(EncodedVectors::load(data_path, meta_path)?); + fn load_quantization( + &mut self, + meta_path: &Path, + data_path: &Path, + _quantization_config: &QuantizationConfig, + ) -> OperationResult<()> { + self.quantized_vectors = Some(Box::new( + quantization::encoder::EncodedVectors::>::load(data_path, meta_path)?, + )); Ok(()) } @@ -488,7 +491,15 @@ mod tests { let quantized_data_path = dir.path().join("quantized.data"); borrowed_storage - .quantize(&quantized_meta_path, &quantized_data_path, None) + .quantize( + &quantized_meta_path, + &quantized_data_path, + &QuantizationConfig { + enable: true, + quantile: None, + always_ram: None, + }, + ) .unwrap(); let query = vec![0.0, 1.0, 1.1, 1.0]; @@ -505,7 +516,15 @@ mod tests { // test save-load borrowed_storage - .load_quantization(&quantized_meta_path, &quantized_data_path) + .load_quantization( + &quantized_meta_path, + &quantized_data_path, + &QuantizationConfig { + enable: true, + quantile: None, + always_ram: None, + }, + ) .unwrap(); let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index 1f9d9a8914c..3cf22f75a1f 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -9,7 +9,7 @@ use rand::Rng; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::OperationResult; -use crate::types::{PointOffsetType, ScoreType}; +use crate::types::{PointOffsetType, QuantizationConfig, ScoreType}; #[derive(Copy, Clone, PartialEq, Debug, Default)] pub struct ScoredPointOffset { @@ -85,10 +85,15 @@ pub trait VectorStorage { &mut self, meta_path: &Path, data_path: &Path, - quantile: Option, + quantization_config: &QuantizationConfig, ) -> OperationResult<()>; // Load quantized vectors from disk - fn load_quantization(&mut self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; + fn load_quantization( + &mut self, + meta_path: &Path, + data_path: &Path, + quantization_config: &QuantizationConfig, + ) -> OperationResult<()>; fn score_points( &self, diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 07948d0648b..8c80f6950ca 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -12,8 +12,8 @@ mod tests { use segment::index::VectorIndex; use segment::segment_constructor::build_segment; use segment::types::{ - Distance, HnswConfig, Indexes, SearchParams, SegmentConfig, SeqNumberType, StorageType, - VectorDataConfig, + Distance, HnswConfig, Indexes, QuantizationConfig, SearchParams, SegmentConfig, + SeqNumberType, StorageType, VectorDataConfig, }; use segment::vector_storage::ScoredPointOffset; use tempfile::Builder; @@ -68,7 +68,15 @@ mod tests { vector_storage .vector_storage .borrow_mut() - .quantize(&quantized_meta_path, &quantized_data_path, None) + .quantize( + &quantized_meta_path, + &quantized_data_path, + &QuantizationConfig { + enable: true, + quantile: None, + always_ram: None, + }, + ) .unwrap() }); diff --git a/lib/storage/src/content_manager/conversions.rs b/lib/storage/src/content_manager/conversions.rs index 1afc420c695..a57a8659cdb 100644 --- a/lib/storage/src/content_manager/conversions.rs +++ b/lib/storage/src/content_manager/conversions.rs @@ -63,6 +63,7 @@ impl TryFrom for CollectionMetaOperations { quantization_config: value.quantization_config.map(|v| QuantizationConfig { enable: v.enable, quantile: v.quantile, + always_ram: v.always_ram, }), }, ))) From 970ac922f09d1710600a09751fb68a1531733e2b Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 8 Feb 2023 11:03:07 +0000 Subject: [PATCH 37/55] remove f32 --- .../src/vector_storage/quantized_vector_storage.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs index 1389b1d506e..ccef8ec233a 100644 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -2,17 +2,26 @@ use bitvec::vec::BitVec; use quantization::encoder::{EncodedQuery, EncodedVectors, Storage}; use super::{RawScorer, ScoredPointOffset}; +use crate::data_types::vectors::VectorElementType; use crate::types::{PointOffsetType, ScoreType}; pub trait QuantizedVectorStorage: Send + Sync { - fn raw_scorer<'a>(&'a self, query: &[f32], deleted: &'a BitVec) -> Box; + fn raw_scorer<'a>( + &'a self, + query: &[VectorElementType], + deleted: &'a BitVec, + ) -> Box; } impl QuantizedVectorStorage for EncodedVectors where TStorage: Storage + Send + Sync, { - fn raw_scorer<'a>(&'a self, query: &[f32], deleted: &'a BitVec) -> Box { + fn raw_scorer<'a>( + &'a self, + query: &[VectorElementType], + deleted: &'a BitVec, + ) -> Box { let query = self.encode_query(query); Box::new(QuantizedRawScorer { query, From 4cd30c335c8d3944f07bfb6c1e6796b75ae36060 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 07:58:04 +0000 Subject: [PATCH 38/55] mmap test --- .../vector_storage/memmap_vector_storage.rs | 78 +++++++++++++++++++ .../src/vector_storage/mmap_vectors.rs | 33 +++++--- .../vector_storage/quantized_mmap_storage.rs | 5 +- 3 files changed, 104 insertions(+), 12 deletions(-) diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 251b7df44ff..07dd7b552ae 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -501,4 +501,82 @@ mod tests { println!("slice[{idx}] = {element:?}"); } } + + #[test] + fn test_mmap_quantization() { + let dist = Distance::Dot; + let dir = Builder::new().prefix("storage_dir").tempdir().unwrap(); + let storage = open_memmap_vector_storage(dir.path(), 4, dist).unwrap(); + let mut borrowed_storage = storage.borrow_mut(); + + let vec1 = vec![1.0, 0.0, 1.0, 1.0]; + let vec2 = vec![1.0, 0.0, 1.0, 0.0]; + let vec3 = vec![1.0, 1.0, 1.0, 1.0]; + let vec4 = vec![1.0, 1.0, 0.0, 1.0]; + let vec5 = vec![1.0, 0.0, 0.0, 0.0]; + + { + let dir2 = Builder::new().prefix("db_dir").tempdir().unwrap(); + let db = open_db(dir2.path(), &[DB_VECTOR_CF]).unwrap(); + let storage2 = open_simple_vector_storage(db, DB_VECTOR_CF, 4, dist).unwrap(); + { + let mut borrowed_storage2 = storage2.borrow_mut(); + borrowed_storage2.put_vector(vec1).unwrap(); + borrowed_storage2.put_vector(vec2).unwrap(); + borrowed_storage2.put_vector(vec3).unwrap(); + borrowed_storage2.put_vector(vec4).unwrap(); + borrowed_storage2.put_vector(vec5).unwrap(); + } + borrowed_storage + .update_from(&*storage2.borrow(), &Default::default()) + .unwrap(); + } + + let quantized_data_path = dir.path().join("quantized.data"); + let quantized_meta_path = dir.path().join("quantized.meta"); + borrowed_storage + .quantize( + &quantized_data_path, + &quantized_meta_path, + &QuantizationConfig { + enable: true, + quantile: None, + always_ram: None, + }, + ) + .unwrap(); + + let query = vec![-1.0, -1.0, -1.0, -1.0]; + + { + let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_quant = borrowed_storage.raw_scorer(query.clone()); + for i in 0..5 { + let orig = scorer_orig.score_point(i); + let quant = scorer_quant.score_point(i); + assert!((orig - quant).abs() < 0.15); + } + } + + // test save-load + borrowed_storage + .load_quantization( + &quantized_meta_path, + &quantized_data_path, + &QuantizationConfig { + enable: true, + quantile: None, + always_ram: None, + }, + ) + .unwrap(); + + let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_quant = borrowed_storage.raw_scorer(query); + for i in 0..5 { + let orig = scorer_orig.score_point(i); + let quant = scorer_quant.score_point(i); + assert!((orig - quant).abs() < 0.15); + } + } } diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 7fccdebc306..6b6964d222c 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -102,28 +102,33 @@ impl MmapVectors { self.deleted_ram = Some(deleted); } - fn create_quantized_storage( + fn create_quantized_storage<'a, TStorage, TStorageBuilder>( &self, meta_path: &Path, data_path: &Path, encoding_parameters: EncodingParameters, - ) -> OperationResult> + storage_builder: TStorageBuilder, + ) -> OperationResult> where - TStorage: quantization::encoder::Storage, + TStorage: quantization::encoder::Storage + Send + Sync + 'a, TStorageBuilder: quantization::encoder::StorageBuilder, { let vector_data_iterator = (0..self.num_vectors as u32).map(|i| { let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); self.raw_vector_offset(offset) }); - let quantized_vectors = quantization::encoder::EncodedVectors::>::encode( - vector_data_iterator, - Vec::new(), - encoding_parameters, - ) - .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; + let quantized_vectors = Box::new( + quantization::encoder::EncodedVectors::encode( + vector_data_iterator, + storage_builder, + encoding_parameters, + ) + .map_err(|e| { + OperationError::service_error(format!("Cannot quantize vector data: {e}")) + })?, + ); quantized_vectors.save(data_path, meta_path)?; - Ok(Box::new(quantized_vectors)) + Ok(quantized_vectors) } pub fn quantize( @@ -148,12 +153,20 @@ impl MmapVectors { meta_path, data_path, encoding_parameters, + Vec::new(), )? } else { + let storage_builder = QuantizedMmapStorageBuilder::new( + data_path, + self.num_vectors, + self.dim, + &encoding_parameters, + )?; self.create_quantized_storage::( meta_path, data_path, encoding_parameters, + storage_builder, )? }; self.quantized_vectors = Some(quantized_vectors); diff --git a/lib/segment/src/vector_storage/quantized_mmap_storage.rs b/lib/segment/src/vector_storage/quantized_mmap_storage.rs index 9dda7b7465d..b9248795c2a 100644 --- a/lib/segment/src/vector_storage/quantized_mmap_storage.rs +++ b/lib/segment/src/vector_storage/quantized_mmap_storage.rs @@ -52,11 +52,12 @@ impl StorageBuilder for QuantizedMmapStorageBuilder { impl QuantizedMmapStorageBuilder { pub fn new( path: &Path, - size: usize, + vectors_count: usize, dim: usize, encoding_parameters: &EncodingParameters, ) -> std::io::Result { - let encoded_storage_size = encoding_parameters.estimate_encoded_storage_size(dim, size); + let encoded_storage_size = + encoding_parameters.estimate_encoded_storage_size(dim, vectors_count); path.parent().map(std::fs::create_dir_all); let file = std::fs::OpenOptions::new() .read(true) From f77fba8b6d4510c16b5cea1139571b0f518dccd2 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 08:07:47 +0000 Subject: [PATCH 39/55] fix mmap slice --- lib/segment/src/vector_storage/quantized_mmap_storage.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/segment/src/vector_storage/quantized_mmap_storage.rs b/lib/segment/src/vector_storage/quantized_mmap_storage.rs index b9248795c2a..971404a75af 100644 --- a/lib/segment/src/vector_storage/quantized_mmap_storage.rs +++ b/lib/segment/src/vector_storage/quantized_mmap_storage.rs @@ -44,7 +44,7 @@ impl StorageBuilder for QuantizedMmapStorageBuilder { } fn extend_from_slice(&mut self, other: &[u8]) { - self.mmap[self.cursor_pos..other.len()].copy_from_slice(other); + self.mmap[self.cursor_pos..self.cursor_pos + other.len()].copy_from_slice(other); self.cursor_pos += other.len(); } } From 265c3e2b3e12ed225d86e855697ee81c2b1415dd Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 08:51:50 +0000 Subject: [PATCH 40/55] fix mmap test --- lib/segment/src/vector_storage/memmap_vector_storage.rs | 4 ++-- lib/segment/src/vector_storage/simple_vector_storage.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 07dd7b552ae..26eba549bfc 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -536,8 +536,8 @@ mod tests { let quantized_meta_path = dir.path().join("quantized.meta"); borrowed_storage .quantize( - &quantized_data_path, &quantized_meta_path, + &quantized_data_path, &QuantizationConfig { enable: true, quantile: None, @@ -546,7 +546,7 @@ mod tests { ) .unwrap(); - let query = vec![-1.0, -1.0, -1.0, -1.0]; + let query = vec![0.5, 0.5, 0.5, 0.5]; { let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index e0c471a762e..4165899924f 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -502,7 +502,7 @@ mod tests { ) .unwrap(); - let query = vec![0.0, 1.0, 1.1, 1.0]; + let query = vec![0.5, 0.5, 0.5, 0.5]; { let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); From 6a04dffad7933226d6cad6bc05e85b5f858dfff8 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 13:03:49 +0000 Subject: [PATCH 41/55] use chunks for quantization storage --- Cargo.lock | 2 +- lib/segment/src/fixtures/index_fixtures.rs | 2 +- .../src/vector_storage/chunked_vectors.rs | 69 +++++++++++++++---- .../vector_storage/memmap_vector_storage.rs | 7 +- .../src/vector_storage/mmap_vectors.rs | 35 +++++++--- .../vector_storage/quantized_mmap_storage.rs | 15 ++-- .../vector_storage/simple_vector_storage.rs | 51 +++++++++----- 7 files changed, 129 insertions(+), 52 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5e6df3940ba..991b19ccff8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2787,7 +2787,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#8f02fbb86b19718ce9584e7ee8a52719de307d45" +source = "git+https://github.com/qdrant/quantization.git#5052f31ceb7b1a0146cfa5c8d57142f049d6ee90" dependencies = [ "cc", "permutation_iterator", diff --git a/lib/segment/src/fixtures/index_fixtures.rs b/lib/segment/src/fixtures/index_fixtures.rs index d1081c1f978..953b2119963 100644 --- a/lib/segment/src/fixtures/index_fixtures.rs +++ b/lib/segment/src/fixtures/index_fixtures.rs @@ -23,7 +23,7 @@ impl FilterContext for FakeFilterContext { } pub struct TestRawScorerProducer { - pub vectors: ChunkedVectors, + pub vectors: ChunkedVectors, pub deleted: BitVec, pub metric: PhantomData, } diff --git a/lib/segment/src/vector_storage/chunked_vectors.rs b/lib/segment/src/vector_storage/chunked_vectors.rs index 984ebbf5a56..237a68fb823 100644 --- a/lib/segment/src/vector_storage/chunked_vectors.rs +++ b/lib/segment/src/vector_storage/chunked_vectors.rs @@ -1,10 +1,12 @@ use std::cmp::max; +use std::fs::File; +use std::io::{Read, Write}; use std::mem; +use std::path::Path; -use crate::data_types::vectors::VectorElementType; -use crate::types::PointOffsetType; +use quantization::encoder::EncodingParameters; -type Chunk = Vec; +use crate::types::PointOffsetType; // chunk size in bytes const CHUNK_SIZE: usize = 32 * 1024 * 1024; @@ -12,19 +14,19 @@ const CHUNK_SIZE: usize = 32 * 1024 * 1024; // if dimension is too high, use this capacity const MIN_CHUNK_CAPACITY: usize = 16; -pub struct ChunkedVectors { +pub struct ChunkedVectors { dim: usize, len: usize, // amount of stored vectors chunk_capacity: usize, // max amount of vectors in each chunk - chunks: Vec, + chunks: Vec>, } -impl ChunkedVectors { - pub fn new(dim: usize) -> ChunkedVectors { +impl ChunkedVectors { + pub fn new(dim: usize) -> Self { assert_ne!(dim, 0, "The vector's dimension cannot be 0"); - let vector_size = dim * mem::size_of::(); + let vector_size = dim * mem::size_of::(); let chunk_capacity = max(MIN_CHUNK_CAPACITY, CHUNK_SIZE / vector_size); - ChunkedVectors { + Self { dim, len: 0, chunk_capacity, @@ -40,20 +42,23 @@ impl ChunkedVectors { self.len == 0 } - pub fn get(&self, key: PointOffsetType) -> &[VectorElementType] { - let key = key as usize; + pub fn get(&self, key: TKey) -> &[T] + where + TKey: num_traits::cast::AsPrimitive, + { + let key: usize = key.as_(); let chunk_data = &self.chunks[key / self.chunk_capacity]; let idx = (key % self.chunk_capacity) * self.dim; &chunk_data[idx..idx + self.dim] } - pub fn push(&mut self, vector: &[VectorElementType]) -> PointOffsetType { + pub fn push(&mut self, vector: &[T]) -> PointOffsetType { let new_id = self.len as PointOffsetType; self.insert(new_id, vector); new_id } - pub fn insert(&mut self, key: PointOffsetType, vector: &[VectorElementType]) { + pub fn insert(&mut self, key: PointOffsetType, vector: &[T]) { let key = key as usize; self.len = max(self.len, key + 1); while self.chunks.len() * self.chunk_capacity < self.len { @@ -63,9 +68,45 @@ impl ChunkedVectors { let chunk_data = &mut self.chunks[key / self.chunk_capacity]; let idx = (key % self.chunk_capacity) * self.dim; if chunk_data.len() < idx + self.dim { - chunk_data.resize(idx + self.dim, 0.); + chunk_data.resize(idx + self.dim, T::default()); } let data = &mut chunk_data[idx..idx + self.dim]; data.copy_from_slice(vector); } } + +impl quantization::encoder::Storage for ChunkedVectors { + fn get_vector_data(&self, index: usize, _vector_size: usize) -> &[u8] { + self.get(index) + } + + fn from_file(path: &Path, encoding_parameters: &EncodingParameters) -> std::io::Result { + let vector_data_size = encoding_parameters.get_vector_data_size(); + let mut vectors = Self::new(vector_data_size); + let mut file = File::open(path)?; + let mut buffer = vec![0u8; vector_data_size]; + while file.read_exact(&mut buffer).is_ok() { + vectors.push(&buffer); + } + Ok(vectors) + } + + fn save_to_file(&self, path: &Path) -> std::io::Result<()> { + let mut buffer = File::create(path)?; + for i in 0..self.len() { + buffer.write_all(self.get(i))?; + } + buffer.flush()?; + Ok(()) + } +} + +impl quantization::encoder::StorageBuilder> for ChunkedVectors { + fn build(self) -> ChunkedVectors { + self + } + + fn push_vector_data(&mut self, other: &[u8]) { + self.push(other); + } +} diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 26eba549bfc..a2136702196 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -292,7 +292,12 @@ where quantization_config: &QuantizationConfig, ) -> OperationResult<()> { let mmap_store = self.mmap_store.as_mut().unwrap(); - mmap_store.load_quantization(meta_path, data_path, quantization_config) + mmap_store.load_quantization( + meta_path, + data_path, + TMetric::distance(), + quantization_config, + ) } fn score_points( diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 6b6964d222c..5295f90366d 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -9,6 +9,7 @@ use memmap2::{Mmap, MmapMut, MmapOptions}; use parking_lot::{RwLock, RwLockReadGuard}; use quantization::encoder::EncodingParameters; +use super::chunked_vectors::ChunkedVectors; use super::quantized_mmap_storage::{QuantizedMmapStorage, QuantizedMmapStorageBuilder}; use super::quantized_vector_storage::QuantizedVectorStorage; use crate::common::error_logging::LogError; @@ -140,6 +141,7 @@ impl MmapVectors { ) -> OperationResult<()> { self.enable_deleted_ram(); let encoding_parameters = EncodingParameters { + dim: self.dim, distance_type: match distance { Distance::Cosine => quantization::encoder::SimilarityType::Dot, Distance::Euclid => quantization::encoder::SimilarityType::L2, @@ -148,18 +150,14 @@ impl MmapVectors { invert: distance == Distance::Euclid, quantile: quantization_config.quantile, }; + let quantized_vector_size = encoding_parameters.get_vector_data_size(); + let vectors = ChunkedVectors::::new(quantized_vector_size); let quantized_vectors = if quantization_config.always_ram == Some(true) { - self.create_quantized_storage::, Vec>( - meta_path, - data_path, - encoding_parameters, - Vec::new(), - )? + self.create_quantized_storage(meta_path, data_path, encoding_parameters, vectors)? } else { let storage_builder = QuantizedMmapStorageBuilder::new( data_path, self.num_vectors, - self.dim, &encoding_parameters, )?; self.create_quantized_storage::( @@ -177,17 +175,32 @@ impl MmapVectors { &mut self, meta_path: &Path, data_path: &Path, + distance: Distance, quantization_config: &QuantizationConfig, ) -> OperationResult<()> { self.enable_deleted_ram(); + let encoding_parameters = EncodingParameters { + dim: self.dim, + distance_type: match distance { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + invert: distance == Distance::Euclid, + quantile: quantization_config.quantile, + }; self.quantized_vectors = if quantization_config.always_ram == Some(true) { - Some(Box::new( - quantization::encoder::EncodedVectors::>::load(data_path, meta_path)?, - )) + Some(Box::new(quantization::encoder::EncodedVectors::< + ChunkedVectors, + >::load( + data_path, meta_path, &encoding_parameters + )?)) } else { Some(Box::new(quantization::encoder::EncodedVectors::< QuantizedMmapStorage, - >::load(data_path, meta_path)?)) + >::load( + data_path, meta_path, &encoding_parameters + )?)) }; Ok(()) } diff --git a/lib/segment/src/vector_storage/quantized_mmap_storage.rs b/lib/segment/src/vector_storage/quantized_mmap_storage.rs index 971404a75af..7e7ff3aa2fe 100644 --- a/lib/segment/src/vector_storage/quantized_mmap_storage.rs +++ b/lib/segment/src/vector_storage/quantized_mmap_storage.rs @@ -15,11 +15,14 @@ pub struct QuantizedMmapStorageBuilder { } impl Storage for QuantizedMmapStorage { - fn ptr(&self) -> *const u8 { - self.mmap.as_ptr() + fn get_vector_data(&self, index: usize, vector_size: usize) -> &[u8] { + &self.mmap[vector_size * index..vector_size * (index + 1)] } - fn from_file(path: &std::path::Path) -> std::io::Result { + fn from_file( + path: &Path, + _encoding_parameters: &EncodingParameters, + ) -> std::io::Result { let file = std::fs::OpenOptions::new() .read(true) .write(false) @@ -43,7 +46,7 @@ impl StorageBuilder for QuantizedMmapStorageBuilder { QuantizedMmapStorage { mmap } } - fn extend_from_slice(&mut self, other: &[u8]) { + fn push_vector_data(&mut self, other: &[u8]) { self.mmap[self.cursor_pos..self.cursor_pos + other.len()].copy_from_slice(other); self.cursor_pos += other.len(); } @@ -53,11 +56,9 @@ impl QuantizedMmapStorageBuilder { pub fn new( path: &Path, vectors_count: usize, - dim: usize, encoding_parameters: &EncodingParameters, ) -> std::io::Result { - let encoded_storage_size = - encoding_parameters.estimate_encoded_storage_size(dim, vectors_count); + let encoded_storage_size = encoding_parameters.get_vector_data_size() * vectors_count; path.parent().map(std::fs::create_dir_all); let file = std::fs::OpenOptions::new() .read(true) diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 4165899924f..1722eca524c 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -30,7 +30,7 @@ use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorageSS}; pub struct SimpleVectorStorage { dim: usize, metric: PhantomData, - vectors: ChunkedVectors, + vectors: ChunkedVectors, deleted: BitVec, deleted_count: usize, quantized_vectors: Option>, @@ -45,7 +45,7 @@ struct StoredRecord { pub struct SimpleRawScorer<'a, TMetric: Metric> { pub query: Vec, - pub vectors: &'a ChunkedVectors, + pub vectors: &'a ChunkedVectors, pub deleted: &'a BitVec, pub metric: PhantomData, } @@ -303,18 +303,22 @@ where quantization_config: &QuantizationConfig, ) -> OperationResult<()> { log::info!("Quantizing vectors..."); - let quantized_vectors = quantization::encoder::EncodedVectors::>::encode( - (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), - Vec::new(), - EncodingParameters { - distance_type: match TMetric::distance() { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - invert: TMetric::distance() == Distance::Euclid, - quantile: quantization_config.quantile, + let encoding_parameters = EncodingParameters { + dim: self.dim, + distance_type: match TMetric::distance() { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, }, + invert: TMetric::distance() == Distance::Euclid, + quantile: quantization_config.quantile, + }; + let quantized_vector_size = encoding_parameters.get_vector_data_size(); + let vectors = ChunkedVectors::::new(quantized_vector_size); + let quantized_vectors = quantization::encoder::EncodedVectors::encode( + (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), + vectors, + encoding_parameters, ) .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; quantized_vectors.save(data_path, meta_path)?; @@ -326,11 +330,24 @@ where &mut self, meta_path: &Path, data_path: &Path, - _quantization_config: &QuantizationConfig, + quantization_config: &QuantizationConfig, ) -> OperationResult<()> { - self.quantized_vectors = Some(Box::new( - quantization::encoder::EncodedVectors::>::load(data_path, meta_path)?, - )); + self.quantized_vectors = Some(Box::new(quantization::encoder::EncodedVectors::< + ChunkedVectors, + >::load( + data_path, + meta_path, + &EncodingParameters { + dim: self.dim, + distance_type: match TMetric::distance() { + Distance::Cosine => quantization::encoder::SimilarityType::Dot, + Distance::Euclid => quantization::encoder::SimilarityType::L2, + Distance::Dot => quantization::encoder::SimilarityType::Dot, + }, + invert: TMetric::distance() == Distance::Euclid, + quantile: quantization_config.quantile, + }, + )?)); Ok(()) } From f26c5d5644aae821e06802f80ee0d9df0bba0af4 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 13:07:55 +0000 Subject: [PATCH 42/55] fix build --- lib/collection/src/shards/replica_set.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/collection/src/shards/replica_set.rs b/lib/collection/src/shards/replica_set.rs index ed76b4daa30..cd462653ece 100644 --- a/lib/collection/src/shards/replica_set.rs +++ b/lib/collection/src/shards/replica_set.rs @@ -1476,6 +1476,7 @@ mod tests { optimizer_config: TEST_OPTIMIZERS_CONFIG.clone(), wal_config, hnsw_config: Default::default(), + quantization_config: None, }; let shared_config = Arc::new(RwLock::new(config.clone())); From b2a7bfcaa2c3caa201bcee35475b53ec570bda26 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 13:08:31 +0000 Subject: [PATCH 43/55] are you happy fmt --- lib/collection/src/optimizers_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/collection/src/optimizers_builder.rs b/lib/collection/src/optimizers_builder.rs index a34731b10fe..d0cac6a054e 100644 --- a/lib/collection/src/optimizers_builder.rs +++ b/lib/collection/src/optimizers_builder.rs @@ -2,8 +2,8 @@ use std::path::Path; use std::sync::Arc; use schemars::JsonSchema; -use segment::types::{HnswConfig, QuantizationConfig}; use segment::common::cpu::get_num_cpus; +use segment::types::{HnswConfig, QuantizationConfig}; use serde::{Deserialize, Serialize}; use crate::collection_manager::optimizers::indexing_optimizer::IndexingOptimizer; From 7079de77fb835a1b5936e3932dd759a6bcbc0fc3 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 18:16:50 +0000 Subject: [PATCH 44/55] update quantization library --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index c9d4733cd54..d921aa3da94 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2767,7 +2767,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#5052f31ceb7b1a0146cfa5c8d57142f049d6ee90" +source = "git+https://github.com/qdrant/quantization.git#08d3a2d68012fb6361c1015964837503d9dbf632" dependencies = [ "cc", "permutation_iterator", From 03f8dfcb54bb5e8401aec30fbd985ca195cbbf01 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 18:43:38 +0000 Subject: [PATCH 45/55] update quantization lib --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index d921aa3da94..5c2e1d10be5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2767,7 +2767,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#08d3a2d68012fb6361c1015964837503d9dbf632" +source = "git+https://github.com/qdrant/quantization.git#0637cdaa85a63ed0ce5ed693898cf1a8c1180412" dependencies = [ "cc", "permutation_iterator", From a6baf83aa9c6691b7fb5a5713c0c097149c02346 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 9 Feb 2023 19:23:51 +0000 Subject: [PATCH 46/55] update quantization lib --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 5c2e1d10be5..8ac546a01cc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2767,7 +2767,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#0637cdaa85a63ed0ce5ed693898cf1a8c1180412" +source = "git+https://github.com/qdrant/quantization.git#747a82bbed99d86db3d17b7975c07320600f804f" dependencies = [ "cc", "permutation_iterator", From e23bba554f69e528af85c4746d27f419b2c2394d Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 16 Feb 2023 13:37:11 +0300 Subject: [PATCH 47/55] integrate api changes --- Cargo.lock | 2 +- .../src/vector_storage/chunked_vectors.rs | 22 ++- .../src/vector_storage/mmap_vectors.rs | 112 +++-------- .../vector_storage/quantized_mmap_storage.rs | 23 ++- .../quantized_vector_storage.rs | 183 ++++++++++++++++-- .../vector_storage/simple_vector_storage.rs | 61 +++--- 6 files changed, 246 insertions(+), 157 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 721cf727dbd..5dd7a5f61cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2776,7 +2776,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#747a82bbed99d86db3d17b7975c07320600f804f" +source = "git+https://github.com/qdrant/quantization.git#7ce5517b1084eee54a47062a6247c97a4cf6af3f" dependencies = [ "cc", "permutation_iterator", diff --git a/lib/segment/src/vector_storage/chunked_vectors.rs b/lib/segment/src/vector_storage/chunked_vectors.rs index 237a68fb823..04aa57f8c4f 100644 --- a/lib/segment/src/vector_storage/chunked_vectors.rs +++ b/lib/segment/src/vector_storage/chunked_vectors.rs @@ -4,8 +4,6 @@ use std::io::{Read, Write}; use std::mem; use std::path::Path; -use quantization::encoder::EncodingParameters; - use crate::types::PointOffsetType; // chunk size in bytes @@ -75,20 +73,26 @@ impl ChunkedVectors { } } -impl quantization::encoder::Storage for ChunkedVectors { +impl quantization::EncodedStorage for ChunkedVectors { fn get_vector_data(&self, index: usize, _vector_size: usize) -> &[u8] { self.get(index) } - fn from_file(path: &Path, encoding_parameters: &EncodingParameters) -> std::io::Result { - let vector_data_size = encoding_parameters.get_vector_data_size(); - let mut vectors = Self::new(vector_data_size); + fn from_file(path: &Path, quantized_vector_size: usize, vectors_count: usize) -> std::io::Result { + let mut vectors = Self::new(quantized_vector_size); let mut file = File::open(path)?; - let mut buffer = vec![0u8; vector_data_size]; + let mut buffer = vec![0u8; quantized_vector_size]; while file.read_exact(&mut buffer).is_ok() { vectors.push(&buffer); } - Ok(vectors) + if vectors.len() == vectors_count { + Ok(vectors) + } else { + Err(std::io::Error::new( + std::io::ErrorKind::Other, + format!("Loaded vectors count {} is not equal to expected count {vectors_count}", vectors.len()), + )) + } } fn save_to_file(&self, path: &Path) -> std::io::Result<()> { @@ -101,7 +105,7 @@ impl quantization::encoder::Storage for ChunkedVectors { } } -impl quantization::encoder::StorageBuilder> for ChunkedVectors { +impl quantization::EncodedStorageBuilder> for ChunkedVectors { fn build(self) -> ChunkedVectors { self } diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 5295f90366d..fa1a1a673ab 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -7,15 +7,12 @@ use std::sync::Arc; use bitvec::vec::BitVec; use memmap2::{Mmap, MmapMut, MmapOptions}; use parking_lot::{RwLock, RwLockReadGuard}; -use quantization::encoder::EncodingParameters; -use super::chunked_vectors::ChunkedVectors; -use super::quantized_mmap_storage::{QuantizedMmapStorage, QuantizedMmapStorageBuilder}; -use super::quantized_vector_storage::QuantizedVectorStorage; +use super::quantized_vector_storage::{QuantizedVectors, create_quantized_vectors, load_quantized_vectors}; use crate::common::error_logging::LogError; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; -use crate::entry::entry_point::{OperationError, OperationResult}; +use crate::entry::entry_point::OperationResult; use crate::madvise; use crate::types::{Distance, PointOffsetType, QuantizationConfig}; @@ -31,7 +28,7 @@ pub struct MmapVectors { deleted_mmap: Arc>, pub deleted_count: usize, pub deleted_ram: Option, - pub quantized_vectors: Option>, + pub quantized_vectors: Option>, } fn open_read(path: &Path) -> OperationResult { @@ -103,35 +100,6 @@ impl MmapVectors { self.deleted_ram = Some(deleted); } - fn create_quantized_storage<'a, TStorage, TStorageBuilder>( - &self, - meta_path: &Path, - data_path: &Path, - encoding_parameters: EncodingParameters, - storage_builder: TStorageBuilder, - ) -> OperationResult> - where - TStorage: quantization::encoder::Storage + Send + Sync + 'a, - TStorageBuilder: quantization::encoder::StorageBuilder, - { - let vector_data_iterator = (0..self.num_vectors as u32).map(|i| { - let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); - self.raw_vector_offset(offset) - }); - let quantized_vectors = Box::new( - quantization::encoder::EncodedVectors::encode( - vector_data_iterator, - storage_builder, - encoding_parameters, - ) - .map_err(|e| { - OperationError::service_error(format!("Cannot quantize vector data: {e}")) - })?, - ); - quantized_vectors.save(data_path, meta_path)?; - Ok(quantized_vectors) - } - pub fn quantize( &mut self, distance: Distance, @@ -140,34 +108,20 @@ impl MmapVectors { quantization_config: &QuantizationConfig, ) -> OperationResult<()> { self.enable_deleted_ram(); - let encoding_parameters = EncodingParameters { - dim: self.dim, - distance_type: match distance { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - invert: distance == Distance::Euclid, - quantile: quantization_config.quantile, - }; - let quantized_vector_size = encoding_parameters.get_vector_data_size(); - let vectors = ChunkedVectors::::new(quantized_vector_size); - let quantized_vectors = if quantization_config.always_ram == Some(true) { - self.create_quantized_storage(meta_path, data_path, encoding_parameters, vectors)? - } else { - let storage_builder = QuantizedMmapStorageBuilder::new( - data_path, - self.num_vectors, - &encoding_parameters, - )?; - self.create_quantized_storage::( - meta_path, - data_path, - encoding_parameters, - storage_builder, - )? - }; - self.quantized_vectors = Some(quantized_vectors); + let vector_data_iterator = (0..self.num_vectors as u32).map(|i| { + let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); + self.raw_vector_offset(offset) + }); + self.quantized_vectors = Some(create_quantized_vectors( + vector_data_iterator, + quantization_config, + distance, + self.dim, + self.num_vectors, + meta_path, + data_path, + true, + )?); Ok(()) } @@ -179,29 +133,15 @@ impl MmapVectors { quantization_config: &QuantizationConfig, ) -> OperationResult<()> { self.enable_deleted_ram(); - let encoding_parameters = EncodingParameters { - dim: self.dim, - distance_type: match distance { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - invert: distance == Distance::Euclid, - quantile: quantization_config.quantile, - }; - self.quantized_vectors = if quantization_config.always_ram == Some(true) { - Some(Box::new(quantization::encoder::EncodedVectors::< - ChunkedVectors, - >::load( - data_path, meta_path, &encoding_parameters - )?)) - } else { - Some(Box::new(quantization::encoder::EncodedVectors::< - QuantizedMmapStorage, - >::load( - data_path, meta_path, &encoding_parameters - )?)) - }; + self.quantized_vectors = Some(load_quantized_vectors( + quantization_config, + distance, + self.dim, + self.num_vectors, + meta_path, + data_path, + true, + )?); Ok(()) } diff --git a/lib/segment/src/vector_storage/quantized_mmap_storage.rs b/lib/segment/src/vector_storage/quantized_mmap_storage.rs index 7e7ff3aa2fe..904de409f3b 100644 --- a/lib/segment/src/vector_storage/quantized_mmap_storage.rs +++ b/lib/segment/src/vector_storage/quantized_mmap_storage.rs @@ -1,7 +1,6 @@ use std::path::Path; use memmap2::{Mmap, MmapMut}; -use quantization::encoder::{EncodingParameters, Storage, StorageBuilder}; use crate::madvise; @@ -14,14 +13,15 @@ pub struct QuantizedMmapStorageBuilder { cursor_pos: usize, } -impl Storage for QuantizedMmapStorage { +impl quantization::EncodedStorage for QuantizedMmapStorage { fn get_vector_data(&self, index: usize, vector_size: usize) -> &[u8] { &self.mmap[vector_size * index..vector_size * (index + 1)] } fn from_file( path: &Path, - _encoding_parameters: &EncodingParameters, + quantized_vector_size: usize, + vectors_count: usize, ) -> std::io::Result { let file = std::fs::OpenOptions::new() .read(true) @@ -30,7 +30,16 @@ impl Storage for QuantizedMmapStorage { .open(path)?; let mmap = unsafe { Mmap::map(&file)? }; madvise::madvise(&mmap, madvise::get_global())?; - Ok(Self { mmap }) + + let expected_size = quantized_vector_size * vectors_count; + if mmap.len() == expected_size { + Ok(Self { mmap }) + } else { + Err(std::io::Error::new( + std::io::ErrorKind::Other, + format!("Loaded storage size {} is not equal to expected size {expected_size}", mmap.len()), + )) + } } fn save_to_file(&self, _path: &Path) -> std::io::Result<()> { @@ -39,7 +48,7 @@ impl Storage for QuantizedMmapStorage { } } -impl StorageBuilder for QuantizedMmapStorageBuilder { +impl quantization::EncodedStorageBuilder for QuantizedMmapStorageBuilder { fn build(self) -> QuantizedMmapStorage { self.mmap.flush().unwrap(); let mmap = self.mmap.make_read_only().unwrap(); // TODO: remove unwrap @@ -56,9 +65,9 @@ impl QuantizedMmapStorageBuilder { pub fn new( path: &Path, vectors_count: usize, - encoding_parameters: &EncodingParameters, + quantized_vector_size: usize, ) -> std::io::Result { - let encoded_storage_size = encoding_parameters.get_vector_data_size() * vectors_count; + let encoded_storage_size = quantized_vector_size * vectors_count; path.parent().map(std::fs::create_dir_all); let file = std::fs::OpenOptions::new() .read(true) diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs index ccef8ec233a..46dcab4cc9e 100644 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -1,21 +1,41 @@ +use std::path::Path; + use bitvec::vec::BitVec; -use quantization::encoder::{EncodedQuery, EncodedVectors, Storage}; +use super::chunked_vectors::ChunkedVectors; +use super::quantized_mmap_storage::{QuantizedMmapStorageBuilder, QuantizedMmapStorage}; use super::{RawScorer, ScoredPointOffset}; use crate::data_types::vectors::VectorElementType; -use crate::types::{PointOffsetType, ScoreType}; +use crate::entry::entry_point::{OperationResult, OperationError}; +use crate::types::{PointOffsetType, ScoreType, QuantizationConfig, Distance}; +use quantization::EncodedVectors; -pub trait QuantizedVectorStorage: Send + Sync { +pub trait QuantizedVectors: Send + Sync { fn raw_scorer<'a>( &'a self, query: &[VectorElementType], deleted: &'a BitVec, ) -> Box; + + fn save_to_file( + &self, + meta_path: &Path, + data_path: &Path, + ) -> OperationResult<()>; +} + +pub struct QuantizedRawScorer<'a, TEncodedQuery, TEncodedVectors> +where + TEncodedVectors: quantization::EncodedVectors, +{ + pub query: TEncodedQuery, + pub deleted: &'a BitVec, + pub quantized_data: &'a TEncodedVectors, } -impl QuantizedVectorStorage for EncodedVectors +impl QuantizedVectors for quantization::EncodedVectorsU8 where - TStorage: Storage + Send + Sync, + TStorage: quantization::EncodedStorage + Send + Sync, { fn raw_scorer<'a>( &'a self, @@ -29,20 +49,21 @@ where quantized_data: self, }) } -} -pub struct QuantizedRawScorer<'a, TStorage> -where - TStorage: Storage, -{ - pub query: EncodedQuery, - pub deleted: &'a BitVec, - pub quantized_data: &'a EncodedVectors, + fn save_to_file( + &self, + meta_path: &Path, + data_path: &Path, + ) -> OperationResult<()> { + self.save(data_path, meta_path)?; + Ok(()) + } } -impl RawScorer for QuantizedRawScorer<'_, TStorage> +impl RawScorer + for QuantizedRawScorer<'_, TEncodedQuery, TEncodedVectors> where - TStorage: Storage, + TEncodedVectors: quantization::EncodedVectors, { fn score_points(&self, points: &[PointOffsetType], scores: &mut [ScoredPointOffset]) -> usize { let mut size: usize = 0; @@ -74,3 +95,135 @@ where self.quantized_data.score_internal(point_a, point_b) } } + +#[allow(clippy::too_many_arguments)] +pub fn create_quantized_vectors<'a>( + vectors: impl IntoIterator + Clone, + quantization_config: &QuantizationConfig, + distance: Distance, + dim: usize, + count: usize, + meta_path: &Path, + data_path: &Path, + on_disk_vector_storage: bool, +) -> OperationResult> { + let vector_parameters = get_vector_parameters(distance, dim, count); + let is_ram = use_ram_quantization_storage(quantization_config, on_disk_vector_storage); + let quantized_vectors = if is_ram { + create_quantized_vectors_ram( + vectors, + quantization_config, + &vector_parameters, + )? + } else { + create_quantized_vectors_mmap( + vectors, + quantization_config, + &vector_parameters, + data_path, + )? + }; + quantized_vectors.save_to_file(meta_path, data_path)?; + Ok(quantized_vectors) +} + +fn use_ram_quantization_storage( + quantization_config: &QuantizationConfig, + on_disk_vector_storage: bool, +) -> bool { + !on_disk_vector_storage || quantization_config.always_ram == Some(true) +} + +fn get_vector_parameters( + distance: Distance, + dim: usize, + count: usize, +) -> quantization::VectorParameters { + quantization::VectorParameters { + dim, + count, + distance_type: match distance { + Distance::Cosine => quantization::DistanceType::Dot, + Distance::Euclid => quantization::DistanceType::L2, + Distance::Dot => quantization::DistanceType::Dot, + }, + invert: distance == Distance::Euclid, + } +} + +fn create_quantized_vectors_ram<'a>( + vectors: impl IntoIterator + Clone, + quantization_config: &QuantizationConfig, + vector_parameters: &quantization::VectorParameters, +) -> OperationResult> { + let quantized_vector_size = + quantization::EncodedVectorsU8::>::get_quantized_vector_size( + vector_parameters, + ); + let storage_builder = ChunkedVectors::::new(quantized_vector_size); + Ok(Box::new( + quantization::EncodedVectorsU8::encode( + vectors, + storage_builder, + vector_parameters, + quantization_config.quantile, + ) + .map_err(|e| { + OperationError::service_error(format!("Cannot quantize vector data: {e}")) + })?, + )) +} + +fn create_quantized_vectors_mmap<'a>( + vectors: impl IntoIterator + Clone, + quantization_config: &QuantizationConfig, + vector_parameters: &quantization::VectorParameters, + data_path: &Path, +) -> OperationResult> { + let quantized_vector_size = + quantization::EncodedVectorsU8::::get_quantized_vector_size( + vector_parameters, + ); + let storage_builder = QuantizedMmapStorageBuilder::new( + data_path, + vector_parameters.count, + quantized_vector_size, + )?; + Ok(Box::new( + quantization::EncodedVectorsU8::encode( + vectors, + storage_builder, + vector_parameters, + quantization_config.quantile, + ) + .map_err(|e| { + OperationError::service_error(format!("Cannot quantize vector data: {e}")) + })?, + )) +} + +pub fn load_quantized_vectors( + quantization_config: &QuantizationConfig, + distance: Distance, + dim: usize, + count: usize, + meta_path: &Path, + data_path: &Path, + on_disk_vector_storage: bool, +) -> OperationResult> { + let vector_parameters = get_vector_parameters(distance, dim, count); + let is_ram = use_ram_quantization_storage(quantization_config, on_disk_vector_storage); + if is_ram { + Ok(Box::new(quantization::EncodedVectorsU8::< + ChunkedVectors, + >::load( + data_path, meta_path, &vector_parameters + )?)) + } else { + Ok(Box::new(quantization::EncodedVectorsU8::< + QuantizedMmapStorage, + >::load( + data_path, meta_path, &vector_parameters + )?)) + } +} diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 1722eca524c..9bf9ef2d862 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -9,12 +9,11 @@ use atomic_refcell::AtomicRefCell; use bitvec::prelude::BitVec; use log::debug; use parking_lot::RwLock; -use quantization::encoder::EncodingParameters; use rocksdb::DB; use serde::{Deserialize, Serialize}; use super::chunked_vectors::ChunkedVectors; -use super::quantized_vector_storage::QuantizedVectorStorage; +use super::quantized_vector_storage::{QuantizedVectors, load_quantized_vectors}; use super::vector_storage_base::VectorStorage; use crate::common::rocksdb_wrapper::DatabaseColumnWrapper; use crate::common::Flusher; @@ -24,6 +23,7 @@ use crate::spaces::metric::Metric; use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric}; use crate::spaces::tools::peek_top_largest_iterable; use crate::types::{Distance, PointOffsetType, QuantizationConfig, ScoreType}; +use crate::vector_storage::quantized_vector_storage::create_quantized_vectors; use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorageSS}; /// In-memory vector storage with on-update persistence using `store` @@ -33,7 +33,7 @@ pub struct SimpleVectorStorage { vectors: ChunkedVectors, deleted: BitVec, deleted_count: usize, - quantized_vectors: Option>, + quantized_vectors: Option>, db_wrapper: DatabaseColumnWrapper, } @@ -302,27 +302,17 @@ where data_path: &Path, quantization_config: &QuantizationConfig, ) -> OperationResult<()> { - log::info!("Quantizing vectors..."); - let encoding_parameters = EncodingParameters { - dim: self.dim, - distance_type: match TMetric::distance() { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - invert: TMetric::distance() == Distance::Euclid, - quantile: quantization_config.quantile, - }; - let quantized_vector_size = encoding_parameters.get_vector_data_size(); - let vectors = ChunkedVectors::::new(quantized_vector_size); - let quantized_vectors = quantization::encoder::EncodedVectors::encode( - (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)), - vectors, - encoding_parameters, - ) - .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; - quantized_vectors.save(data_path, meta_path)?; - self.quantized_vectors = Some(Box::new(quantized_vectors)); + let vector_data_iterator = (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)); + self.quantized_vectors = Some(create_quantized_vectors( + vector_data_iterator, + quantization_config, + TMetric::distance(), + self.dim, + self.vectors.len(), + meta_path, + data_path, + false, + )?); Ok(()) } @@ -332,22 +322,15 @@ where data_path: &Path, quantization_config: &QuantizationConfig, ) -> OperationResult<()> { - self.quantized_vectors = Some(Box::new(quantization::encoder::EncodedVectors::< - ChunkedVectors, - >::load( - data_path, + self.quantized_vectors = Some(load_quantized_vectors( + quantization_config, + TMetric::distance(), + self.dim, + self.vectors.len(), meta_path, - &EncodingParameters { - dim: self.dim, - distance_type: match TMetric::distance() { - Distance::Cosine => quantization::encoder::SimilarityType::Dot, - Distance::Euclid => quantization::encoder::SimilarityType::L2, - Distance::Dot => quantization::encoder::SimilarityType::Dot, - }, - invert: TMetric::distance() == Distance::Euclid, - quantile: quantization_config.quantile, - }, - )?)); + data_path, + false, + )?); Ok(()) } From adf633d64aa76eb5751a465b9333fa853910abf8 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Thu, 16 Feb 2023 13:37:32 +0300 Subject: [PATCH 48/55] are you happy fmt --- .../src/vector_storage/chunked_vectors.rs | 11 +++- .../src/vector_storage/mmap_vectors.rs | 4 +- .../vector_storage/quantized_mmap_storage.rs | 5 +- .../quantized_vector_storage.rs | 53 ++++++------------- .../vector_storage/simple_vector_storage.rs | 2 +- 5 files changed, 33 insertions(+), 42 deletions(-) diff --git a/lib/segment/src/vector_storage/chunked_vectors.rs b/lib/segment/src/vector_storage/chunked_vectors.rs index 04aa57f8c4f..6be397b27d9 100644 --- a/lib/segment/src/vector_storage/chunked_vectors.rs +++ b/lib/segment/src/vector_storage/chunked_vectors.rs @@ -78,7 +78,11 @@ impl quantization::EncodedStorage for ChunkedVectors { self.get(index) } - fn from_file(path: &Path, quantized_vector_size: usize, vectors_count: usize) -> std::io::Result { + fn from_file( + path: &Path, + quantized_vector_size: usize, + vectors_count: usize, + ) -> std::io::Result { let mut vectors = Self::new(quantized_vector_size); let mut file = File::open(path)?; let mut buffer = vec![0u8; quantized_vector_size]; @@ -90,7 +94,10 @@ impl quantization::EncodedStorage for ChunkedVectors { } else { Err(std::io::Error::new( std::io::ErrorKind::Other, - format!("Loaded vectors count {} is not equal to expected count {vectors_count}", vectors.len()), + format!( + "Loaded vectors count {} is not equal to expected count {vectors_count}", + vectors.len() + ), )) } } diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index fa1a1a673ab..445fd8bb0c2 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -8,7 +8,9 @@ use bitvec::vec::BitVec; use memmap2::{Mmap, MmapMut, MmapOptions}; use parking_lot::{RwLock, RwLockReadGuard}; -use super::quantized_vector_storage::{QuantizedVectors, create_quantized_vectors, load_quantized_vectors}; +use super::quantized_vector_storage::{ + create_quantized_vectors, load_quantized_vectors, QuantizedVectors, +}; use crate::common::error_logging::LogError; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; diff --git a/lib/segment/src/vector_storage/quantized_mmap_storage.rs b/lib/segment/src/vector_storage/quantized_mmap_storage.rs index 904de409f3b..2a1fd4d3b9c 100644 --- a/lib/segment/src/vector_storage/quantized_mmap_storage.rs +++ b/lib/segment/src/vector_storage/quantized_mmap_storage.rs @@ -37,7 +37,10 @@ impl quantization::EncodedStorage for QuantizedMmapStorage { } else { Err(std::io::Error::new( std::io::ErrorKind::Other, - format!("Loaded storage size {} is not equal to expected size {expected_size}", mmap.len()), + format!( + "Loaded storage size {} is not equal to expected size {expected_size}", + mmap.len() + ), )) } } diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs index 46dcab4cc9e..fea52dac035 100644 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -1,14 +1,14 @@ use std::path::Path; use bitvec::vec::BitVec; +use quantization::EncodedVectors; use super::chunked_vectors::ChunkedVectors; -use super::quantized_mmap_storage::{QuantizedMmapStorageBuilder, QuantizedMmapStorage}; +use super::quantized_mmap_storage::{QuantizedMmapStorage, QuantizedMmapStorageBuilder}; use super::{RawScorer, ScoredPointOffset}; use crate::data_types::vectors::VectorElementType; -use crate::entry::entry_point::{OperationResult, OperationError}; -use crate::types::{PointOffsetType, ScoreType, QuantizationConfig, Distance}; -use quantization::EncodedVectors; +use crate::entry::entry_point::{OperationError, OperationResult}; +use crate::types::{Distance, PointOffsetType, QuantizationConfig, ScoreType}; pub trait QuantizedVectors: Send + Sync { fn raw_scorer<'a>( @@ -17,11 +17,7 @@ pub trait QuantizedVectors: Send + Sync { deleted: &'a BitVec, ) -> Box; - fn save_to_file( - &self, - meta_path: &Path, - data_path: &Path, - ) -> OperationResult<()>; + fn save_to_file(&self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; } pub struct QuantizedRawScorer<'a, TEncodedQuery, TEncodedVectors> @@ -50,11 +46,7 @@ where }) } - fn save_to_file( - &self, - meta_path: &Path, - data_path: &Path, - ) -> OperationResult<()> { + fn save_to_file(&self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { self.save(data_path, meta_path)?; Ok(()) } @@ -110,18 +102,9 @@ pub fn create_quantized_vectors<'a>( let vector_parameters = get_vector_parameters(distance, dim, count); let is_ram = use_ram_quantization_storage(quantization_config, on_disk_vector_storage); let quantized_vectors = if is_ram { - create_quantized_vectors_ram( - vectors, - quantization_config, - &vector_parameters, - )? + create_quantized_vectors_ram(vectors, quantization_config, &vector_parameters)? } else { - create_quantized_vectors_mmap( - vectors, - quantization_config, - &vector_parameters, - data_path, - )? + create_quantized_vectors_mmap(vectors, quantization_config, &vector_parameters, data_path)? }; quantized_vectors.save_to_file(meta_path, data_path)?; Ok(quantized_vectors) @@ -157,9 +140,9 @@ fn create_quantized_vectors_ram<'a>( vector_parameters: &quantization::VectorParameters, ) -> OperationResult> { let quantized_vector_size = - quantization::EncodedVectorsU8::>::get_quantized_vector_size( - vector_parameters, - ); + quantization::EncodedVectorsU8::>::get_quantized_vector_size( + vector_parameters, + ); let storage_builder = ChunkedVectors::::new(quantized_vector_size); Ok(Box::new( quantization::EncodedVectorsU8::encode( @@ -168,9 +151,7 @@ fn create_quantized_vectors_ram<'a>( vector_parameters, quantization_config.quantile, ) - .map_err(|e| { - OperationError::service_error(format!("Cannot quantize vector data: {e}")) - })?, + .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?, )) } @@ -181,9 +162,9 @@ fn create_quantized_vectors_mmap<'a>( data_path: &Path, ) -> OperationResult> { let quantized_vector_size = - quantization::EncodedVectorsU8::::get_quantized_vector_size( - vector_parameters, - ); + quantization::EncodedVectorsU8::::get_quantized_vector_size( + vector_parameters, + ); let storage_builder = QuantizedMmapStorageBuilder::new( data_path, vector_parameters.count, @@ -196,9 +177,7 @@ fn create_quantized_vectors_mmap<'a>( vector_parameters, quantization_config.quantile, ) - .map_err(|e| { - OperationError::service_error(format!("Cannot quantize vector data: {e}")) - })?, + .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?, )) } diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 9bf9ef2d862..c94263af0ab 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -13,7 +13,7 @@ use rocksdb::DB; use serde::{Deserialize, Serialize}; use super::chunked_vectors::ChunkedVectors; -use super::quantized_vector_storage::{QuantizedVectors, load_quantized_vectors}; +use super::quantized_vector_storage::{load_quantized_vectors, QuantizedVectors}; use super::vector_storage_base::VectorStorage; use crate::common::rocksdb_wrapper::DatabaseColumnWrapper; use crate::common::Flusher; From 62d93523a2a8d982f76f24abfcb0e88f9f0b2096 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Fri, 17 Feb 2023 13:00:37 +0300 Subject: [PATCH 49/55] change quantization api --- Cargo.lock | 2 +- lib/api/src/grpc/conversions.rs | 41 +++++++- lib/api/src/grpc/proto/collections.proto | 11 ++- lib/api/src/grpc/qdrant.rs | 25 ++++- lib/collection/src/operations/conversions.rs | 22 ++--- .../segment_constructor/segment_builder.rs | 24 +++-- .../segment_constructor_base.rs | 18 ++-- lib/segment/src/types.rs | 28 +++--- .../vector_storage/memmap_vector_storage.rs | 38 +++----- .../quantized_vector_storage.rs | 93 +++++++++++-------- .../vector_storage/simple_vector_storage.rs | 37 +++----- .../tests/hnsw_quantized_search_test.rs | 9 +- .../src/content_manager/conversions.rs | 13 +-- 13 files changed, 202 insertions(+), 159 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5dd7a5f61cf..f30b0e86794 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2776,7 +2776,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#7ce5517b1084eee54a47062a6247c97a4cf6af3f" +source = "git+https://github.com/qdrant/quantization.git#a2ec13ca47905924a782910a6d70fc2470f427e0" dependencies = [ "cc", "permutation_iterator", diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index a741138296d..4330580280a 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -21,9 +21,9 @@ use crate::grpc::qdrant::{ FieldCondition, Filter, GeoBoundingBox, GeoPoint, GeoRadius, HasIdCondition, HealthCheckReply, HnswConfigDiff, IsEmptyCondition, ListCollectionsResponse, ListValue, Match, NamedVectors, PayloadExcludeSelector, PayloadIncludeSelector, PayloadIndexParams, PayloadSchemaInfo, - PayloadSchemaType, PointId, Range, ScoredPoint, SearchParams, Struct, TextIndexParams, - TokenizerType, Value, ValuesCount, Vector, Vectors, VectorsSelector, WithPayloadSelector, - WithVectorsSelector, + PayloadSchemaType, PointId, QuantizationConfig, Range, ScalarU8Quantization, ScoredPoint, + SearchParams, Struct, TextIndexParams, TokenizerType, Value, ValuesCount, Vector, Vectors, + VectorsSelector, WithPayloadSelector, WithVectorsSelector, }; pub fn payload_to_proto(payload: segment::types::Payload) -> HashMap { @@ -475,6 +475,41 @@ impl TryFrom for segment::types::PointIdType { } } +impl From for QuantizationConfig { + fn from(value: segment::types::QuantizationConfig) -> Self { + match value { + segment::types::QuantizationConfig::ScalarU8(config) => Self { + quantization: Some(super::qdrant::quantization_config::Quantization::Scalar( + ScalarU8Quantization { + quantile: config.quantile, + always_ram: config.always_ram, + }, + )), + }, + } + } +} + +impl TryFrom for segment::types::QuantizationConfig { + type Error = Status; + + fn try_from(value: QuantizationConfig) -> Result { + let value = value + .quantization + .ok_or_else(|| Status::invalid_argument("Unable to convert quantization config"))?; + match value { + super::qdrant::quantization_config::Quantization::Scalar(config) => { + Ok(segment::types::QuantizationConfig::ScalarU8( + segment::types::ScalarU8Quantization { + quantile: config.quantile, + always_ram: config.always_ram, + }, + )) + } + } + } +} + fn conditions_helper_from_grpc( conditions: Vec, ) -> Result>, tonic::Status> { diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index a290c9a4542..503e33ff816 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -154,10 +154,15 @@ message OptimizersConfigDiff { optional uint64 max_optimization_threads = 8; } +message ScalarU8Quantization { + optional float quantile = 1; // Number of bits to use for quantization + optional bool always_ram = 2; // If true - quantized vectors always will be stored in RAM, ignoring the config of main storage +} + message QuantizationConfig { - bool enable = 1; - optional float quantile = 2; - optional bool always_ram = 3; + oneof quantization { + ScalarU8Quantization scalar = 1; + } } message CreateCollection { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index db0f59892cb..359bf44bcf4 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -171,16 +171,31 @@ pub struct OptimizersConfigDiff { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct QuantizationConfig { - #[prost(bool, tag = "1")] - pub enable: bool, - #[prost(float, optional, tag = "2")] +pub struct ScalarU8Quantization { + /// Number of bits to use for quantization + #[prost(float, optional, tag = "1")] pub quantile: ::core::option::Option, - #[prost(bool, optional, tag = "3")] + /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage + #[prost(bool, optional, tag = "2")] pub always_ram: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct QuantizationConfig { + #[prost(oneof = "quantization_config::Quantization", tags = "1")] + pub quantization: ::core::option::Option, +} +/// Nested message and enum types in `QuantizationConfig`. +pub mod quantization_config { + #[allow(clippy::derive_partial_eq_without_eq)] + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Quantization { + #[prost(message, tag = "1")] + Scalar(super::ScalarU8Quantization), + } +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct CreateCollection { /// Name of the collection #[prost(string, tag = "1")] diff --git a/lib/collection/src/operations/conversions.rs b/lib/collection/src/operations/conversions.rs index 2d8e1f511e6..b8997288578 100644 --- a/lib/collection/src/operations/conversions.rs +++ b/lib/collection/src/operations/conversions.rs @@ -4,7 +4,7 @@ use std::num::{NonZeroU32, NonZeroU64}; use api::grpc::conversions::{from_grpc_dist, payload_to_proto, proto_to_payloads}; use itertools::Itertools; use segment::data_types::vectors::{NamedVector, VectorStruct, DEFAULT_VECTOR_NAME}; -use segment::types::{Distance, QuantizationConfig}; +use segment::types::Distance; use tonic::Status; use super::config_diff::CollectionParamsDiff; @@ -239,13 +239,7 @@ impl From for api::grpc::qdrant::CollectionInfo { wal_capacity_mb: Some(config.wal_config.wal_capacity_mb as u64), wal_segments_ahead: Some(config.wal_config.wal_segments_ahead as u64), }), - quantization_config: config.quantization_config.map(|x| { - api::grpc::qdrant::QuantizationConfig { - enable: x.enable, - quantile: x.quantile, - always_ram: x.always_ram, - } - }), + quantization_config: config.quantization_config.map(|x| x.into()), }), payload_schema: payload_schema .into_iter() @@ -389,11 +383,13 @@ impl TryFrom for CollectionConfig { None => return Err(Status::invalid_argument("Malformed WalConfig type")), Some(wal_config) => wal_config.into(), }, - quantization_config: config.quantization_config.map(|x| QuantizationConfig { - enable: x.enable, - quantile: x.quantile, - always_ram: x.always_ram, - }), + quantization_config: { + if let Some(config) = config.quantization_config { + Some(config.try_into()?) + } else { + None + } + }, }) } } diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 0e4b32bbdd6..fbe7841c41f 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -211,20 +211,18 @@ impl SegmentBuilder { fn update_quantization(segment: &Segment, stopped: &AtomicBool) -> OperationResult<()> { if let Some(quantization) = &segment.config().quantization_config { - if quantization.enable { - let segment_path = segment.current_path.as_path(); - for (vector_name, vector_data) in &segment.vector_data { - check_process_stopped(stopped)?; + let segment_path = segment.current_path.as_path(); + for (vector_name, vector_data) in &segment.vector_data { + check_process_stopped(stopped)?; - let vector_storage_path = get_vector_storage_path(segment_path, vector_name); - let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); - let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); - vector_data.vector_storage.borrow_mut().quantize( - &quantized_meta_path, - &quantized_data_path, - quantization, - )?; - } + let vector_storage_path = get_vector_storage_path(segment_path, vector_name); + let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); + let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); + vector_data.vector_storage.borrow_mut().quantize( + &quantized_meta_path, + &quantized_data_path, + quantization, + )?; } } Ok(()) diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 555d6032bd4..5340cb6e56f 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -111,16 +111,14 @@ fn create_segment( }; if let Some(quantization_config) = &config.quantization_config { - if quantization_config.enable { - let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); - let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); - if quantized_meta_path.exists() && quantized_data_path.exists() { - vector_storage.borrow_mut().load_quantization( - &quantized_meta_path, - &quantized_data_path, - quantization_config, - )?; - } + let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); + let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); + if quantized_meta_path.exists() && quantized_data_path.exists() { + vector_storage.borrow_mut().load_quantization( + &quantized_meta_path, + &quantized_data_path, + quantization_config, + )?; } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 55a2df92399..75d15025048 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -315,31 +315,33 @@ fn default_max_indexing_threads() -> usize { #[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone)] #[serde(rename_all = "snake_case")] -pub struct QuantizationConfig { - pub enable: bool, +pub struct ScalarU8Quantization { + /// Number of bits to use for quantization pub quantile: Option, - /// If true, quantized data is stored in RAM even if original data is memmapped. - /// If false, quantized data is stored like original data. - /// Default value is false. + /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage pub always_ram: Option, } -impl PartialEq for QuantizationConfig { +impl PartialEq for ScalarU8Quantization { fn eq(&self, other: &Self) -> bool { - self.enable == other.enable - && self.quantile == other.quantile - && self.always_ram == other.always_ram + self.quantile == other.quantile && self.always_ram == other.always_ram } } -impl Eq for QuantizationConfig {} - -impl std::hash::Hash for QuantizationConfig { +impl std::hash::Hash for ScalarU8Quantization { fn hash(&self, state: &mut H) { - self.enable.hash(state); + self.always_ram.hash(state); } } +impl Eq for ScalarU8Quantization {} + +#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum QuantizationConfig { + ScalarU8(ScalarU8Quantization), +} + pub const DEFAULT_HNSW_EF_CONSTRUCT: usize = 100; impl Default for HnswConfig { diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index a2136702196..c9ed3539eb1 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -368,6 +368,7 @@ mod tests { use super::*; use crate::common::rocksdb_wrapper::{open_db, DB_VECTOR_CF}; + use crate::types::ScalarU8Quantization; use crate::vector_storage::simple_vector_storage::open_simple_vector_storage; #[test] @@ -537,50 +538,39 @@ mod tests { .unwrap(); } + let config = QuantizationConfig::ScalarU8(ScalarU8Quantization { + quantile: None, + always_ram: None, + }); + let quantized_data_path = dir.path().join("quantized.data"); let quantized_meta_path = dir.path().join("quantized.meta"); borrowed_storage - .quantize( - &quantized_meta_path, - &quantized_data_path, - &QuantizationConfig { - enable: true, - quantile: None, - always_ram: None, - }, - ) + .quantize(&quantized_meta_path, &quantized_data_path, &config) .unwrap(); let query = vec![0.5, 0.5, 0.5, 0.5]; { - let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); - let scorer_quant = borrowed_storage.raw_scorer(query.clone()); + let scorer_quant = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_orig = borrowed_storage.raw_scorer(query.clone()); for i in 0..5 { - let orig = scorer_orig.score_point(i); let quant = scorer_quant.score_point(i); + let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); } } // test save-load borrowed_storage - .load_quantization( - &quantized_meta_path, - &quantized_data_path, - &QuantizationConfig { - enable: true, - quantile: None, - always_ram: None, - }, - ) + .load_quantization(&quantized_meta_path, &quantized_data_path, &config) .unwrap(); - let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); - let scorer_quant = borrowed_storage.raw_scorer(query); + let scorer_quant = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_orig = borrowed_storage.raw_scorer(query); for i in 0..5 { - let orig = scorer_orig.score_point(i); let quant = scorer_quant.score_point(i); + let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); } } diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs index fea52dac035..64fd69d1be0 100644 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ b/lib/segment/src/vector_storage/quantized_vector_storage.rs @@ -8,7 +8,9 @@ use super::quantized_mmap_storage::{QuantizedMmapStorage, QuantizedMmapStorageBu use super::{RawScorer, ScoredPointOffset}; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::{OperationError, OperationResult}; -use crate::types::{Distance, PointOffsetType, QuantizationConfig, ScoreType}; +use crate::types::{ + Distance, PointOffsetType, QuantizationConfig, ScalarU8Quantization, ScoreType, +}; pub trait QuantizedVectors: Send + Sync { fn raw_scorer<'a>( @@ -100,21 +102,53 @@ pub fn create_quantized_vectors<'a>( on_disk_vector_storage: bool, ) -> OperationResult> { let vector_parameters = get_vector_parameters(distance, dim, count); - let is_ram = use_ram_quantization_storage(quantization_config, on_disk_vector_storage); - let quantized_vectors = if is_ram { - create_quantized_vectors_ram(vectors, quantization_config, &vector_parameters)? - } else { - create_quantized_vectors_mmap(vectors, quantization_config, &vector_parameters, data_path)? + let quantized_vectors = match quantization_config { + QuantizationConfig::ScalarU8(scalar_u8_config) => { + let is_ram = use_ram_quantization_storage(scalar_u8_config, on_disk_vector_storage); + if is_ram { + create_quantized_vectors_ram(vectors, scalar_u8_config, &vector_parameters)? + } else { + create_quantized_vectors_mmap( + vectors, + scalar_u8_config, + &vector_parameters, + data_path, + )? + } + } }; quantized_vectors.save_to_file(meta_path, data_path)?; Ok(quantized_vectors) } -fn use_ram_quantization_storage( +pub fn load_quantized_vectors( quantization_config: &QuantizationConfig, + distance: Distance, + dim: usize, + count: usize, + meta_path: &Path, + data_path: &Path, on_disk_vector_storage: bool, -) -> bool { - !on_disk_vector_storage || quantization_config.always_ram == Some(true) +) -> OperationResult> { + let vector_parameters = get_vector_parameters(distance, dim, count); + match quantization_config { + QuantizationConfig::ScalarU8(scalar_u8_config) => { + let is_ram = use_ram_quantization_storage(scalar_u8_config, on_disk_vector_storage); + if is_ram { + Ok(Box::new(quantization::EncodedVectorsU8::< + ChunkedVectors, + >::load( + data_path, meta_path, &vector_parameters + )?)) + } else { + Ok(Box::new(quantization::EncodedVectorsU8::< + QuantizedMmapStorage, + >::load( + data_path, meta_path, &vector_parameters + )?)) + } + } + } } fn get_vector_parameters( @@ -134,9 +168,16 @@ fn get_vector_parameters( } } +fn use_ram_quantization_storage( + config: &ScalarU8Quantization, + on_disk_vector_storage: bool, +) -> bool { + !on_disk_vector_storage || config.always_ram == Some(true) +} + fn create_quantized_vectors_ram<'a>( vectors: impl IntoIterator + Clone, - quantization_config: &QuantizationConfig, + config: &ScalarU8Quantization, vector_parameters: &quantization::VectorParameters, ) -> OperationResult> { let quantized_vector_size = @@ -149,7 +190,7 @@ fn create_quantized_vectors_ram<'a>( vectors, storage_builder, vector_parameters, - quantization_config.quantile, + config.quantile, ) .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?, )) @@ -157,7 +198,7 @@ fn create_quantized_vectors_ram<'a>( fn create_quantized_vectors_mmap<'a>( vectors: impl IntoIterator + Clone, - quantization_config: &QuantizationConfig, + config: &ScalarU8Quantization, vector_parameters: &quantization::VectorParameters, data_path: &Path, ) -> OperationResult> { @@ -175,34 +216,8 @@ fn create_quantized_vectors_mmap<'a>( vectors, storage_builder, vector_parameters, - quantization_config.quantile, + config.quantile, ) .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?, )) } - -pub fn load_quantized_vectors( - quantization_config: &QuantizationConfig, - distance: Distance, - dim: usize, - count: usize, - meta_path: &Path, - data_path: &Path, - on_disk_vector_storage: bool, -) -> OperationResult> { - let vector_parameters = get_vector_parameters(distance, dim, count); - let is_ram = use_ram_quantization_storage(quantization_config, on_disk_vector_storage); - if is_ram { - Ok(Box::new(quantization::EncodedVectorsU8::< - ChunkedVectors, - >::load( - data_path, meta_path, &vector_parameters - )?)) - } else { - Ok(Box::new(quantization::EncodedVectorsU8::< - QuantizedMmapStorage, - >::load( - data_path, meta_path, &vector_parameters - )?)) - } -} diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index c94263af0ab..73d72af2d3f 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -390,6 +390,7 @@ mod tests { use super::*; use crate::common::rocksdb_wrapper::{open_db, DB_VECTOR_CF}; + use crate::types::ScalarU8Quantization; #[test] fn test_score_points() { @@ -487,51 +488,39 @@ mod tests { borrowed_storage.put_vector(vec3).unwrap(); borrowed_storage.put_vector(vec4).unwrap(); + let config = QuantizationConfig::ScalarU8(ScalarU8Quantization { + quantile: None, + always_ram: None, + }); let quantized_meta_path = dir.path().join("quantized.meta"); let quantized_data_path = dir.path().join("quantized.data"); borrowed_storage - .quantize( - &quantized_meta_path, - &quantized_data_path, - &QuantizationConfig { - enable: true, - quantile: None, - always_ram: None, - }, - ) + .quantize(&quantized_meta_path, &quantized_data_path, &config) .unwrap(); let query = vec![0.5, 0.5, 0.5, 0.5]; { - let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); - let scorer_quant = borrowed_storage.raw_scorer(query.clone()); + let scorer_quant = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_orig = borrowed_storage.raw_scorer(query.clone()); for i in 0..5 { - let orig = scorer_orig.score_point(i); let quant = scorer_quant.score_point(i); + let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); } } // test save-load borrowed_storage - .load_quantization( - &quantized_meta_path, - &quantized_data_path, - &QuantizationConfig { - enable: true, - quantile: None, - always_ram: None, - }, - ) + .load_quantization(&quantized_meta_path, &quantized_data_path, &config) .unwrap(); - let scorer_orig = borrowed_storage.quantized_raw_scorer(&query).unwrap(); - let scorer_quant = borrowed_storage.raw_scorer(query); + let scorer_quant = borrowed_storage.quantized_raw_scorer(&query).unwrap(); + let scorer_orig = borrowed_storage.raw_scorer(query.clone()); for i in 0..5 { - let orig = scorer_orig.score_point(i); let quant = scorer_quant.score_point(i); + let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); } } diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 8c80f6950ca..5d370e77efd 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -12,8 +12,8 @@ mod tests { use segment::index::VectorIndex; use segment::segment_constructor::build_segment; use segment::types::{ - Distance, HnswConfig, Indexes, QuantizationConfig, SearchParams, SegmentConfig, - SeqNumberType, StorageType, VectorDataConfig, + Distance, HnswConfig, Indexes, QuantizationConfig, ScalarU8Quantization, SearchParams, + SegmentConfig, SeqNumberType, StorageType, VectorDataConfig, }; use segment::vector_storage::ScoredPointOffset; use tempfile::Builder; @@ -71,11 +71,10 @@ mod tests { .quantize( &quantized_meta_path, &quantized_data_path, - &QuantizationConfig { - enable: true, + &QuantizationConfig::ScalarU8(ScalarU8Quantization { quantile: None, always_ram: None, - }, + }), ) .unwrap() }); diff --git a/lib/storage/src/content_manager/conversions.rs b/lib/storage/src/content_manager/conversions.rs index a57a8659cdb..a0c2579fb90 100644 --- a/lib/storage/src/content_manager/conversions.rs +++ b/lib/storage/src/content_manager/conversions.rs @@ -1,7 +1,6 @@ use std::collections::BTreeMap; use collection::operations::types::VectorsConfig; -use segment::types::QuantizationConfig; use tonic::Status; use crate::content_manager::collection_meta_ops::{ @@ -60,11 +59,13 @@ impl TryFrom for CollectionMetaOperations { init_from: value .init_from_collection .map(|v| InitFrom { collection: v }), - quantization_config: value.quantization_config.map(|v| QuantizationConfig { - enable: v.enable, - quantile: v.quantile, - always_ram: v.always_ram, - }), + quantization_config: { + if let Some(config) = value.quantization_config { + Some(config.try_into()?) + } else { + None + } + }, }, ))) } From 63830a6ad1549d67faea7f7fd8d3b39ecd261bdf Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Sun, 19 Feb 2023 23:12:02 +0300 Subject: [PATCH 50/55] additional checks in tests --- lib/segment/src/vector_storage/memmap_vector_storage.rs | 8 ++++++++ lib/segment/src/vector_storage/simple_vector_storage.rs | 8 ++++++++ lib/segment/tests/hnsw_quantized_search_test.rs | 8 +++++++- 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index c9ed3539eb1..565e7b8be94 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -558,6 +558,10 @@ mod tests { let quant = scorer_quant.score_point(i); let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); + + let quant = scorer_quant.score_internal(0, i); + let orig = scorer_orig.score_internal(0, i); + assert!((orig - quant).abs() < 0.15); } } @@ -572,6 +576,10 @@ mod tests { let quant = scorer_quant.score_point(i); let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); + + let quant = scorer_quant.score_internal(0, i); + let orig = scorer_orig.score_internal(0, i); + assert!((orig - quant).abs() < 0.15); } } } diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index 73d72af2d3f..ecc5e5d96df 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -508,6 +508,10 @@ mod tests { let quant = scorer_quant.score_point(i); let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); + + let quant = scorer_quant.score_internal(0, i); + let orig = scorer_orig.score_internal(0, i); + assert!((orig - quant).abs() < 0.15); } } @@ -522,6 +526,10 @@ mod tests { let quant = scorer_quant.score_point(i); let orig = scorer_orig.score_point(i); assert!((orig - quant).abs() < 0.15); + + let quant = scorer_quant.score_internal(0, i); + let orig = scorer_orig.score_internal(0, i); + assert!((orig - quant).abs() < 0.15); } } } diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 5d370e77efd..8f1b3ee2429 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -65,6 +65,9 @@ mod tests { .unwrap(); } segment.vector_data.values_mut().for_each(|vector_storage| { + if quantized_data_path.exists() || quantized_meta_path.exists() { + panic!("quantization files shouldn't exists"); + } vector_storage .vector_storage .borrow_mut() @@ -76,7 +79,10 @@ mod tests { always_ram: None, }), ) - .unwrap() + .unwrap(); + if !quantized_data_path.exists() || !quantized_meta_path.exists() { + panic!("quantization was not saved"); + } }); let hnsw_config = HnswConfig { From 31ea0d801b96d4a3027064347a19fbc25529d2dd Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 27 Feb 2023 15:00:53 +0400 Subject: [PATCH 51/55] update quantization version --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index f30b0e86794..20bb893a579 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2776,7 +2776,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.1.0" -source = "git+https://github.com/qdrant/quantization.git#a2ec13ca47905924a782910a6d70fc2470f427e0" +source = "git+https://github.com/qdrant/quantization.git#7a4b773880c63c78e9ab9ce8bba613375ebd34f2" dependencies = [ "cc", "permutation_iterator", From ce4e3859e0540754e2af2a3316513909f6941f8e Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 27 Feb 2023 16:00:38 +0000 Subject: [PATCH 52/55] fix unit tests --- lib/segment/tests/payload_index_test.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/segment/tests/payload_index_test.rs b/lib/segment/tests/payload_index_test.rs index a5821127774..db89c0e6c59 100644 --- a/lib/segment/tests/payload_index_test.rs +++ b/lib/segment/tests/payload_index_test.rs @@ -148,7 +148,7 @@ mod tests { )]), index: Indexes::Plain {}, storage_type: StorageType::InMemory, - payload_storage_type: Default::default(), + ..Default::default() }; let mut plain_segment = build_segment(path_plain, &config).unwrap(); From f217f34de49263402e9088f95b3e11432990e651 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 28 Feb 2023 14:49:23 +0000 Subject: [PATCH 53/55] add quantization to storage config --- lib/storage/src/content_manager/toc.rs | 5 +++++ lib/storage/src/types.rs | 3 ++- lib/storage/tests/alias_tests.rs | 1 + 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/lib/storage/src/content_manager/toc.rs b/lib/storage/src/content_manager/toc.rs index 21ad59ee98b..7f323c7ba13 100644 --- a/lib/storage/src/content_manager/toc.rs +++ b/lib/storage/src/content_manager/toc.rs @@ -318,6 +318,11 @@ impl TableOfContent { Some(diff) => diff.update(&self.storage_config.hnsw_index)?, }; + let quantization_config = match quantization_config { + None => self.storage_config.quantization_config.clone(), + Some(diff) => Some(diff), + }; + let collection_config = CollectionConfig { wal_config, params: collection_params, diff --git a/lib/storage/src/types.rs b/lib/storage/src/types.rs index 03afdc0d971..0bdeba7a4c1 100644 --- a/lib/storage/src/types.rs +++ b/lib/storage/src/types.rs @@ -7,7 +7,7 @@ use collection::shards::shard::PeerId; use schemars::JsonSchema; use segment::common::anonymize::Anonymize; use segment::madvise; -use segment::types::HnswConfig; +use segment::types::{HnswConfig, QuantizationConfig}; use serde::{Deserialize, Serialize}; use tonic::transport::Uri; @@ -47,6 +47,7 @@ pub struct StorageConfig { pub wal: WalConfig, pub performance: PerformanceConfig, pub hnsw_index: HnswConfig, + pub quantization_config: Option, #[serde(default = "default_mmap_advice")] pub mmap_advice: madvise::Advice, #[serde(default)] diff --git a/lib/storage/tests/alias_tests.rs b/lib/storage/tests/alias_tests.rs index a24bf0ac0ee..42cc99053d7 100644 --- a/lib/storage/tests/alias_tests.rs +++ b/lib/storage/tests/alias_tests.rs @@ -49,6 +49,7 @@ mod tests { max_optimization_threads: 1, }, hnsw_index: Default::default(), + quantization_config: None, mmap_advice: madvise::Advice::Random, node_type: Default::default(), }; From 86f916fb9fb78b74dec3422a904ff879172b0e0c Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 1 Mar 2023 08:21:11 +0000 Subject: [PATCH 54/55] use quantization for all cardinality search cases --- lib/segment/src/index/hnsw_index/hnsw.rs | 59 ++++++++++++++----- .../vector_storage/memmap_vector_storage.rs | 27 +++++++++ .../vector_storage/simple_vector_storage.rs | 20 +++++++ .../src/vector_storage/vector_storage_base.rs | 7 +++ 4 files changed, 97 insertions(+), 16 deletions(-) diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index 35e1e43d5be..be085f36d1a 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -233,6 +233,32 @@ impl HNSWIndex { .map(|vector| self.search_with_graph(vector, filter, top, params)) .collect() } + + fn search_vectors_plain( + &self, + vectors: &[&[VectorElementType]], + filter: &Filter, + top: usize, + params: Option<&SearchParams>, + ) -> Vec> { + let payload_index = self.payload_index.borrow(); + let vector_storage = self.vector_storage.borrow(); + let mut filtered_iter = payload_index.query_points(filter); + let ignore_quantization = params.map(|p| p.ignore_quantization).unwrap_or(false); + if ignore_quantization { + vectors + .iter() + .map(|vector| vector_storage.score_points(vector, filtered_iter.as_mut(), top)) + .collect() + } else { + vectors + .iter() + .map(|vector| { + vector_storage.score_quantized_points(vector, filtered_iter.as_mut(), top) + }) + .collect() + } + } } impl VectorIndex for HNSWIndex { @@ -264,24 +290,24 @@ impl VectorIndex for HNSWIndex { // - to retrieve possible points and score them after // - to use HNSW index with filtering condition - let payload_index = self.payload_index.borrow(); - let vector_storage = self.vector_storage.borrow(); - - let plain_search = || -> Vec> { - let mut filtered_iter = payload_index.query_points(query_filter); - return vectors - .iter() - .map(|vector| { - vector_storage.score_points(vector, filtered_iter.as_mut(), top) - }) - .collect(); - }; - // if exact search is requested, we should not use HNSW index if exact { - return plain_search(); + let exact_params = params.map(|params| { + let mut params = *params; + params.ignore_quantization = true; // disable quantization for exact search + params + }); + let _timer = + ScopeDurationMeasurer::new(&self.searches_telemetry.exact_filtered); + return self.search_vectors_plain( + vectors, + query_filter, + top, + exact_params.as_ref(), + ); } + let payload_index = self.payload_index.borrow(); let query_cardinality = payload_index.estimate_cardinality(query_filter); // debug!("query_cardinality: {:#?}", query_cardinality); @@ -290,7 +316,7 @@ impl VectorIndex for HNSWIndex { // if cardinality is small - use plain index let _timer = ScopeDurationMeasurer::new(&self.searches_telemetry.small_cardinality); - return plain_search(); + return self.search_vectors_plain(vectors, query_filter, top, params); } if query_cardinality.min > self.config.indexing_threshold { @@ -304,6 +330,7 @@ impl VectorIndex for HNSWIndex { // Fast cardinality estimation is not enough, do sample estimation of cardinality + let vector_storage = self.vector_storage.borrow(); if sample_check_cardinality( vector_storage.sample_ids(), |idx| filter_context.check(idx), @@ -318,7 +345,7 @@ impl VectorIndex for HNSWIndex { // if cardinality is small - use plain index let _timer = ScopeDurationMeasurer::new(&self.searches_telemetry.small_cardinality); - plain_search() + self.search_vectors_plain(vectors, query_filter, top, params) } } } diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 565e7b8be94..2996a963fc0 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -270,6 +270,33 @@ where } } + fn score_quantized_points( + &self, + vector: &[VectorElementType], + points: &mut dyn Iterator, + top: usize, + ) -> Vec { + match self.quantized_raw_scorer(vector) { + Some(scorer) => { + let scores = points + .filter(|idx| { + !self + .mmap_store + .as_ref() + .unwrap() + .deleted(*idx) + .unwrap_or(true) + }) + .map(|idx| { + let score = scorer.score_point(idx); + ScoredPointOffset { idx, score } + }); + peek_top_largest_iterable(scores, top) + } + None => self.score_points(vector, points, top), + } + } + fn quantize( &mut self, meta_path: &Path, diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index ecc5e5d96df..f435f650da8 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -296,6 +296,26 @@ where } } + fn score_quantized_points( + &self, + vector: &[VectorElementType], + points: &mut dyn Iterator, + top: usize, + ) -> Vec { + match self.quantized_raw_scorer(vector) { + Some(scorer) => { + let scores = points + .filter(|idx| !self.deleted[*idx as usize]) + .map(|idx| { + let score = scorer.score_point(idx); + ScoredPointOffset { idx, score } + }); + peek_top_largest_iterable(scores, top) + } + None => self.score_points(vector, points, top), + } + } + fn quantize( &mut self, meta_path: &Path, diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index 3cf22f75a1f..99e1d357ad0 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -80,6 +80,13 @@ pub trait VectorStorage { // Generate RawScorer on quantized vectors if present fn quantized_raw_scorer(&self, vector: &[VectorElementType]) -> Option>; + // Try peek top nearest points from quantized vectors. If quantized vectors are not present, do it on raw vectors + fn score_quantized_points( + &self, + vector: &[VectorElementType], + points: &mut dyn Iterator, + top: usize, + ) -> Vec; // Generate quantized vectors and store them on disk fn quantize( &mut self, From a1c282e7d0128a1309dfe875ab1268fef5bde397 Mon Sep 17 00:00:00 2001 From: Andrey Vasnetsov Date: Fri, 3 Mar 2023 16:53:46 +0100 Subject: [PATCH 55/55] Integrate quantization suggestions 2 (#1520) * review api * wip: refactor quantization integrations * wip: refactor quantization integrations * wip: fmt * include quantization into snapshot * fmt --- docs/redoc/master/openapi.json | 100 +++++++- lib/api/src/grpc/conversions.rs | 48 ++-- lib/api/src/grpc/proto/collections.proto | 14 +- lib/api/src/grpc/qdrant.rs | 29 ++- .../optimizers/segment_optimizer.rs | 7 +- lib/collection/src/shards/local_shard.rs | 2 +- .../segment_constructor/segment_builder.rs | 13 +- .../segment_constructor_base.rs | 19 +- lib/segment/src/types.rs | 43 +++- .../vector_storage/memmap_vector_storage.rs | 52 ++-- .../src/vector_storage/mmap_vectors.rs | 60 ++--- lib/segment/src/vector_storage/mod.rs | 3 +- .../src/vector_storage/quantized/mod.rs | 4 + .../quantized/quantized_vectors_base.rs | 192 +++++++++++++++ .../quantized/scalar_quantized.rs | 97 ++++++++ .../scalar_quantized_mmap_storage.rs} | 50 ++++ .../quantized/scalar_quantized_ram_storage.rs | 47 ++++ .../quantized_vector_storage.rs | 223 ------------------ .../vector_storage/simple_vector_storage.rs | 58 ++--- .../src/vector_storage/vector_storage_base.rs | 11 +- .../tests/hnsw_quantized_search_test.rs | 22 +- .../content_manager/collection_meta_ops.rs | 1 + lib/storage/src/content_manager/toc.rs | 2 +- lib/storage/src/types.rs | 2 +- lib/storage/tests/alias_tests.rs | 2 +- src/actix/helpers.rs | 10 +- 26 files changed, 701 insertions(+), 410 deletions(-) create mode 100644 lib/segment/src/vector_storage/quantized/mod.rs create mode 100644 lib/segment/src/vector_storage/quantized/quantized_vectors_base.rs create mode 100644 lib/segment/src/vector_storage/quantized/scalar_quantized.rs rename lib/segment/src/vector_storage/{quantized_mmap_storage.rs => quantized/scalar_quantized_mmap_storage.rs} (59%) create mode 100644 lib/segment/src/vector_storage/quantized/scalar_quantized_ram_storage.rs delete mode 100644 lib/segment/src/vector_storage/quantized_vector_storage.rs diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json index a7500ebdc51..bf26cc4606d 100644 --- a/docs/redoc/master/openapi.json +++ b/docs/redoc/master/openapi.json @@ -3423,6 +3423,17 @@ }, "wal_config": { "$ref": "#/components/schemas/WalConfig" + }, + "quantization_config": { + "default": null, + "anyOf": [ + { + "$ref": "#/components/schemas/QuantizationConfig" + }, + { + "nullable": true + } + ] } } }, @@ -3638,6 +3649,52 @@ } } }, + "QuantizationConfig": { + "anyOf": [ + { + "$ref": "#/components/schemas/ScalarQuantization" + } + ] + }, + "ScalarQuantization": { + "type": "object", + "required": [ + "scalar" + ], + "properties": { + "scalar": { + "$ref": "#/components/schemas/ScalarQuantizationConfig" + } + } + }, + "ScalarQuantizationConfig": { + "type": "object", + "required": [ + "type" + ], + "properties": { + "type": { + "$ref": "#/components/schemas/ScalarType" + }, + "quantile": { + "description": "Quantile for quantization. Expected value range in (0, 1.0]. If not set - use the whole range of values", + "type": "number", + "format": "float", + "nullable": true + }, + "always_ram": { + "description": "If true - quantized vectors always will be stored in RAM, ignoring the config of main storage", + "type": "boolean", + "nullable": true + } + } + }, + "ScalarType": { + "type": "string", + "enum": [ + "int8" + ] + }, "PayloadIndexInfo": { "description": "Display payload field type & index information", "type": "object", @@ -4393,6 +4450,11 @@ "description": "Search without approximation. If set to true, search may run long but with exact results.", "default": false, "type": "boolean" + }, + "ignore_quantization": { + "description": "If set to true, search will ignore quantized vector data", + "default": false, + "type": "boolean" } } }, @@ -4768,6 +4830,18 @@ "nullable": true } ] + }, + "quantization_config": { + "description": "Quantization parameters. If none - quantization is disabled.", + "default": null, + "anyOf": [ + { + "$ref": "#/components/schemas/QuantizationConfig" + }, + { + "nullable": true + } + ] } } }, @@ -5663,7 +5737,8 @@ "Active", "Dead", "Partial", - "Initializing" + "Initializing", + "Listener" ] }, "RemoteShardInfo": { @@ -6076,6 +6151,18 @@ }, "payload_storage_type": { "$ref": "#/components/schemas/PayloadStorageType" + }, + "quantization_config": { + "description": "Quantization parameters. If none - quantization is disabled.", + "default": null, + "anyOf": [ + { + "$ref": "#/components/schemas/QuantizationConfig" + }, + { + "nullable": true + } + ] } } }, @@ -6721,7 +6808,16 @@ "format": "uri" }, "priority": { - "$ref": "#/components/schemas/SnapshotPriority" + "description": "Defines which data should be used as a source of truth if there are other replicas in the cluster. If set to `Snapshot`, the snapshot will be used as a source of truth, and the current state will be overwritten. If set to `Replica`, the current state will be used as a source of truth, and after recovery if will be synchronized with the snapshot.", + "default": null, + "anyOf": [ + { + "$ref": "#/components/schemas/SnapshotPriority" + }, + { + "nullable": true + } + ] } } }, diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index 4feb39b4698..b4b87b9f135 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -4,7 +4,6 @@ use std::time::Instant; use chrono::{NaiveDateTime, Timelike}; use segment::data_types::text_index::TextIndexType; use segment::data_types::vectors::VectorElementType; -use segment::types::{PayloadSelector, WithPayloadInterface}; use tonic::Status; use uuid::Uuid; @@ -22,7 +21,7 @@ use crate::grpc::qdrant::{ FieldCondition, Filter, GeoBoundingBox, GeoPoint, GeoRadius, HasIdCondition, HealthCheckReply, HnswConfigDiff, IsEmptyCondition, ListCollectionsResponse, ListValue, Match, NamedVectors, PayloadExcludeSelector, PayloadIncludeSelector, PayloadIndexParams, PayloadSchemaInfo, - PayloadSchemaType, PointId, QuantizationConfig, Range, ScalarU8Quantization, ScoredPoint, + PayloadSchemaType, PointId, QuantizationConfig, Range, ScalarQuantization, ScoredPoint, SearchParams, Struct, TextIndexParams, TokenizerType, Value, ValuesCount, Vector, Vectors, VectorsSelector, WithPayloadSelector, WithVectorsSelector, }; @@ -310,15 +309,15 @@ impl TryFrom for segment::types::WithPayloadInterface { impl From for WithPayloadSelector { fn from(value: segment::types::WithPayloadInterface) -> Self { let selector_options = match value { - WithPayloadInterface::Bool(flag) => SelectorOptions::Enable(flag), - WithPayloadInterface::Fields(fields) => { + segment::types::WithPayloadInterface::Bool(flag) => SelectorOptions::Enable(flag), + segment::types::WithPayloadInterface::Fields(fields) => { SelectorOptions::Include(PayloadIncludeSelector { fields }) } - WithPayloadInterface::Selector(selector) => match selector { - PayloadSelector::Include(s) => { + segment::types::WithPayloadInterface::Selector(selector) => match selector { + segment::types::PayloadSelector::Include(s) => { SelectorOptions::Include(PayloadIncludeSelector { fields: s.include }) } - PayloadSelector::Exclude(s) => { + segment::types::PayloadSelector::Exclude(s) => { SelectorOptions::Exclude(PayloadExcludeSelector { fields: s.exclude }) } }, @@ -479,9 +478,16 @@ impl TryFrom for segment::types::PointIdType { impl From for QuantizationConfig { fn from(value: segment::types::QuantizationConfig) -> Self { match value { - segment::types::QuantizationConfig::ScalarU8(config) => Self { + segment::types::QuantizationConfig::Scalar(segment::types::ScalarQuantization { + scalar: config, + }) => Self { quantization: Some(super::qdrant::quantization_config::Quantization::Scalar( - ScalarU8Quantization { + ScalarQuantization { + r#type: match config.r#type { + segment::types::ScalarType::Int8 => { + crate::grpc::qdrant::QuantizationType::Int8 as i32 + } + }, quantile: config.quantile, always_ram: config.always_ram, }, @@ -500,12 +506,26 @@ impl TryFrom for segment::types::QuantizationConfig { .ok_or_else(|| Status::invalid_argument("Unable to convert quantization config"))?; match value { super::qdrant::quantization_config::Quantization::Scalar(config) => { - Ok(segment::types::QuantizationConfig::ScalarU8( - segment::types::ScalarU8Quantization { - quantile: config.quantile, - always_ram: config.always_ram, + Ok(segment::types::ScalarQuantizationConfig { + r#type: match crate::grpc::qdrant::QuantizationType::from_i32(config.r#type) { + None => { + return Err(Status::invalid_argument( + "Error converting quantization type: None", + )); + } + Some(crate::grpc::qdrant::QuantizationType::UnknownQuantization) => { + return Err(Status::invalid_argument( + "Error converting quantization type: UnknownQuantization", + )); + } + Some(crate::grpc::qdrant::QuantizationType::Int8) => { + segment::types::ScalarType::Int8 + } }, - )) + quantile: config.quantile, + always_ram: config.always_ram, + } + .into()) } } } diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index 503e33ff816..2021fe67555 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -61,6 +61,11 @@ enum PayloadSchemaType { Text = 5; } +enum QuantizationType { + UnknownQuantization = 0; + Int8 = 1; +} + message OptimizerStatus { bool ok = 1; string error = 2; @@ -154,14 +159,15 @@ message OptimizersConfigDiff { optional uint64 max_optimization_threads = 8; } -message ScalarU8Quantization { - optional float quantile = 1; // Number of bits to use for quantization - optional bool always_ram = 2; // If true - quantized vectors always will be stored in RAM, ignoring the config of main storage +message ScalarQuantization { + QuantizationType type = 1; // Type of quantization + optional float quantile = 2; // Number of bits to use for quantization + optional bool always_ram = 3; // If true - quantized vectors always will be stored in RAM, ignoring the config of main storage } message QuantizationConfig { oneof quantization { - ScalarU8Quantization scalar = 1; + ScalarQuantization scalar = 1; } } diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 60b841bbc70..667a8545acf 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -171,12 +171,15 @@ pub struct OptimizersConfigDiff { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct ScalarU8Quantization { +pub struct ScalarQuantization { + /// Type of quantization + #[prost(enumeration = "QuantizationType", tag = "1")] + pub r#type: i32, /// Number of bits to use for quantization - #[prost(float, optional, tag = "1")] + #[prost(float, optional, tag = "2")] pub quantile: ::core::option::Option, /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage - #[prost(bool, optional, tag = "2")] + #[prost(bool, optional, tag = "3")] pub always_ram: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] @@ -191,7 +194,7 @@ pub mod quantization_config { #[derive(Clone, PartialEq, ::prost::Oneof)] pub enum Quantization { #[prost(message, tag = "1")] - Scalar(super::ScalarU8Quantization), + Scalar(super::ScalarQuantization), } } #[allow(clippy::derive_partial_eq_without_eq)] @@ -553,6 +556,24 @@ impl PayloadSchemaType { } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] +pub enum QuantizationType { + UnknownQuantization = 0, + Int8 = 1, +} +impl QuantizationType { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + QuantizationType::UnknownQuantization => "UnknownQuantization", + QuantizationType::Int8 => "Int8", + } + } +} +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] pub enum TokenizerType { Unknown = 0, Prefix = 1, diff --git a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs index 15ef6e6366b..07747793e82 100644 --- a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs @@ -136,7 +136,12 @@ pub trait SegmentOptimizer { true => PayloadStorageType::OnDisk, false => PayloadStorageType::InMemory, }, - quantization_config: self.quantization_config(), + quantization_config: if is_indexed { + // TODO: separate config for applying quantization + self.quantization_config() + } else { + Default::default() + }, }; Ok(SegmentBuilder::new( diff --git a/lib/collection/src/shards/local_shard.rs b/lib/collection/src/shards/local_shard.rs index afa467266f3..247057c523a 100644 --- a/lib/collection/src/shards/local_shard.rs +++ b/lib/collection/src/shards/local_shard.rs @@ -299,7 +299,7 @@ impl LocalShard { true => PayloadStorageType::OnDisk, false => PayloadStorageType::InMemory, }, - quantization_config: config.quantization_config.clone(), + quantization_config: Default::default(), }; let segment = thread::Builder::new() .name(format!("shard-build-{collection_id}-{id}")) diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index fbe7841c41f..9eb2ca08c1b 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; -use super::{get_vector_storage_path, QUANTIZED_DATA_PATH, QUANTIZED_META_PATH}; +use super::get_vector_storage_path; use crate::common::error_logging::LogError; use crate::entry::entry_point::{ check_process_stopped, OperationError, OperationResult, SegmentEntry, @@ -216,13 +216,10 @@ impl SegmentBuilder { check_process_stopped(stopped)?; let vector_storage_path = get_vector_storage_path(segment_path, vector_name); - let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); - let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); - vector_data.vector_storage.borrow_mut().quantize( - &quantized_meta_path, - &quantized_data_path, - quantization, - )?; + vector_data + .vector_storage + .borrow_mut() + .quantize(&vector_storage_path, quantization)?; } } Ok(()) diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 5340cb6e56f..91ac5fdf212 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -35,8 +35,6 @@ use crate::vector_storage::VectorStorageSS; pub const PAYLOAD_INDEX_PATH: &str = "payload_index"; pub const VECTOR_STORAGE_PATH: &str = "vector_storage"; pub const VECTOR_INDEX_PATH: &str = "vector_index"; -pub const QUANTIZED_DATA_PATH: &str = "quantized.data"; -pub const QUANTIZED_META_PATH: &str = "quantized.meta"; fn sp(t: T) -> Arc> { Arc::new(AtomicRefCell::new(t)) @@ -110,16 +108,13 @@ fn create_segment( )?, }; - if let Some(quantization_config) = &config.quantization_config { - let quantized_meta_path = vector_storage_path.join(QUANTIZED_META_PATH); - let quantized_data_path = vector_storage_path.join(QUANTIZED_DATA_PATH); - if quantized_meta_path.exists() && quantized_data_path.exists() { - vector_storage.borrow_mut().load_quantization( - &quantized_meta_path, - &quantized_data_path, - quantization_config, - )?; - } + if config.quantization_config.is_some() { + let quantized_data_path = vector_storage_path; + // Try to load quantization data from disk, if exists + // If not exists or it's a new segment, just ignore it + vector_storage + .borrow_mut() + .load_quantization(&quantized_data_path)?; } let vector_index: Arc> = match config.index { diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index c2a0a778935..a5ed2406da1 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -314,33 +314,60 @@ fn default_max_indexing_threads() -> usize { 0 } -#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone)] +#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] +#[serde(rename_all = "lowercase")] +pub enum ScalarType { + #[default] + Int8, +} + +#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)] #[serde(rename_all = "snake_case")] -pub struct ScalarU8Quantization { - /// Number of bits to use for quantization +pub struct ScalarQuantizationConfig { + /// Type of quantization to use + /// If `int8` - 8 bit quantization will be used + pub r#type: ScalarType, + /// Quantile for quantization. Expected value range in (0, 1.0]. If not set - use the whole range of values + #[serde(skip_serializing_if = "Option::is_none")] pub quantile: Option, /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage + #[serde(skip_serializing_if = "Option::is_none")] pub always_ram: Option, } -impl PartialEq for ScalarU8Quantization { +#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] +pub struct ScalarQuantization { + pub scalar: ScalarQuantizationConfig, +} + +impl PartialEq for ScalarQuantizationConfig { fn eq(&self, other: &Self) -> bool { - self.quantile == other.quantile && self.always_ram == other.always_ram + self.quantile == other.quantile + && self.always_ram == other.always_ram + && self.r#type == other.r#type } } -impl std::hash::Hash for ScalarU8Quantization { +impl std::hash::Hash for ScalarQuantizationConfig { fn hash(&self, state: &mut H) { self.always_ram.hash(state); + self.r#type.hash(state); } } -impl Eq for ScalarU8Quantization {} +impl Eq for ScalarQuantizationConfig {} #[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] +#[serde(untagged)] pub enum QuantizationConfig { - ScalarU8(ScalarU8Quantization), + Scalar(ScalarQuantization), +} + +impl From for QuantizationConfig { + fn from(config: ScalarQuantizationConfig) -> Self { + QuantizationConfig::Scalar(ScalarQuantization { scalar: config }) + } } pub const DEFAULT_HNSW_EF_CONSTRUCT: usize = 100; diff --git a/lib/segment/src/vector_storage/memmap_vector_storage.rs b/lib/segment/src/vector_storage/memmap_vector_storage.rs index 2996a963fc0..a138e9307a6 100644 --- a/lib/segment/src/vector_storage/memmap_vector_storage.rs +++ b/lib/segment/src/vector_storage/memmap_vector_storage.rs @@ -17,6 +17,7 @@ use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric}; use crate::spaces::tools::peek_top_largest_iterable; use crate::types::{Distance, PointOffsetType, QuantizationConfig, ScoreType}; use crate::vector_storage::mmap_vectors::MmapVectors; +use crate::vector_storage::quantized::quantized_vectors_base::QuantizedVectors; use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorage, VectorStorageSS}; fn vf_to_u8(v: &[T]) -> &[u8] { @@ -259,7 +260,7 @@ where ) -> Option> { let mmap_store = self.mmap_store.as_ref().unwrap(); if let Some(quantized_data) = &mmap_store.quantized_vectors { - if let Some(deleted_ram) = &mmap_store.deleted_ram { + if let Some(deleted_ram) = &mmap_store.deleted_flags { let query = TMetric::preprocess(vector).unwrap_or_else(|| vector.to_owned()); Some(quantized_data.raw_scorer(&query, deleted_ram)) } else { @@ -299,32 +300,16 @@ where fn quantize( &mut self, - meta_path: &Path, data_path: &Path, quantization_config: &QuantizationConfig, ) -> OperationResult<()> { let mmap_store = self.mmap_store.as_mut().unwrap(); - mmap_store.quantize( - TMetric::distance(), - meta_path, - data_path, - quantization_config, - ) + mmap_store.quantize(TMetric::distance(), data_path, quantization_config) } - fn load_quantization( - &mut self, - meta_path: &Path, - data_path: &Path, - quantization_config: &QuantizationConfig, - ) -> OperationResult<()> { + fn load_quantization(&mut self, data_path: &Path) -> OperationResult<()> { let mmap_store = self.mmap_store.as_mut().unwrap(); - mmap_store.load_quantization( - meta_path, - data_path, - TMetric::distance(), - quantization_config, - ) + mmap_store.load_quantization(data_path) } fn score_points( @@ -383,7 +368,13 @@ where } fn files(&self) -> Vec { - vec![self.vectors_path.clone(), self.deleted_path.clone()] + let mut files = vec![self.vectors_path.clone(), self.deleted_path.clone()]; + if let Some(Some(quantized_vectors)) = + &self.mmap_store.as_ref().map(|x| &x.quantized_vectors) + { + files.extend(quantized_vectors.files()) + } + files } } @@ -395,7 +386,7 @@ mod tests { use super::*; use crate::common::rocksdb_wrapper::{open_db, DB_VECTOR_CF}; - use crate::types::ScalarU8Quantization; + use crate::types::ScalarQuantizationConfig; use crate::vector_storage::simple_vector_storage::open_simple_vector_storage; #[test] @@ -565,16 +556,14 @@ mod tests { .unwrap(); } - let config = QuantizationConfig::ScalarU8(ScalarU8Quantization { + let config: QuantizationConfig = ScalarQuantizationConfig { + r#type: Default::default(), quantile: None, always_ram: None, - }); + } + .into(); - let quantized_data_path = dir.path().join("quantized.data"); - let quantized_meta_path = dir.path().join("quantized.meta"); - borrowed_storage - .quantize(&quantized_meta_path, &quantized_data_path, &config) - .unwrap(); + borrowed_storage.quantize(dir.path(), &config).unwrap(); let query = vec![0.5, 0.5, 0.5, 0.5]; @@ -593,12 +582,11 @@ mod tests { } // test save-load - borrowed_storage - .load_quantization(&quantized_meta_path, &quantized_data_path, &config) - .unwrap(); + borrowed_storage.load_quantization(dir.path()).unwrap(); let scorer_quant = borrowed_storage.quantized_raw_scorer(&query).unwrap(); let scorer_orig = borrowed_storage.raw_scorer(query); + for i in 0..5 { let quant = scorer_quant.score_point(i); let orig = scorer_orig.score_point(i); diff --git a/lib/segment/src/vector_storage/mmap_vectors.rs b/lib/segment/src/vector_storage/mmap_vectors.rs index 445fd8bb0c2..d9dcfe32559 100644 --- a/lib/segment/src/vector_storage/mmap_vectors.rs +++ b/lib/segment/src/vector_storage/mmap_vectors.rs @@ -8,15 +8,13 @@ use bitvec::vec::BitVec; use memmap2::{Mmap, MmapMut, MmapOptions}; use parking_lot::{RwLock, RwLockReadGuard}; -use super::quantized_vector_storage::{ - create_quantized_vectors, load_quantized_vectors, QuantizedVectors, -}; use crate::common::error_logging::LogError; use crate::common::Flusher; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::OperationResult; use crate::madvise; use crate::types::{Distance, PointOffsetType, QuantizationConfig}; +use crate::vector_storage::quantized::quantized_vectors_base::QuantizedVectorsStorage; const HEADER_SIZE: usize = 4; const DELETED_HEADER: &[u8; 4] = b"drop"; @@ -27,10 +25,10 @@ pub struct MmapVectors { pub dim: usize, pub num_vectors: usize, mmap: Mmap, - deleted_mmap: Arc>, + deleted_flags_mmap: Arc>, pub deleted_count: usize, - pub deleted_ram: Option, - pub quantized_vectors: Option>, + pub deleted_flags: Option, + pub quantized_vectors: Option, } fn open_read(path: &Path) -> OperationResult { @@ -86,64 +84,50 @@ impl MmapVectors { dim, num_vectors, mmap, - deleted_mmap: Arc::new(RwLock::new(deleted_mmap)), + deleted_flags_mmap: Arc::new(RwLock::new(deleted_mmap)), deleted_count, - deleted_ram: None, + deleted_flags: None, quantized_vectors: None, }) } - fn enable_deleted_ram(&mut self) { + fn init_deleted_flags(&mut self) { let mut deleted = BitVec::new(); deleted.resize(self.num_vectors, false); for i in 0..self.num_vectors { deleted.set(i, self.deleted(i as PointOffsetType).unwrap_or_default()); } - self.deleted_ram = Some(deleted); + self.deleted_flags = Some(deleted); } pub fn quantize( &mut self, distance: Distance, - meta_path: &Path, data_path: &Path, quantization_config: &QuantizationConfig, ) -> OperationResult<()> { - self.enable_deleted_ram(); + self.init_deleted_flags(); let vector_data_iterator = (0..self.num_vectors as u32).map(|i| { let offset = self.data_offset(i as PointOffsetType).unwrap_or_default(); self.raw_vector_offset(offset) }); - self.quantized_vectors = Some(create_quantized_vectors( + self.quantized_vectors = Some(QuantizedVectorsStorage::create( vector_data_iterator, quantization_config, distance, self.dim, self.num_vectors, - meta_path, data_path, true, )?); Ok(()) } - pub fn load_quantization( - &mut self, - meta_path: &Path, - data_path: &Path, - distance: Distance, - quantization_config: &QuantizationConfig, - ) -> OperationResult<()> { - self.enable_deleted_ram(); - self.quantized_vectors = Some(load_quantized_vectors( - quantization_config, - distance, - self.dim, - self.num_vectors, - meta_path, - data_path, - true, - )?); + pub fn load_quantization(&mut self, data_path: &Path) -> OperationResult<()> { + if QuantizedVectorsStorage::check_exists(data_path) { + self.init_deleted_flags(); + self.quantized_vectors = Some(QuantizedVectorsStorage::load(data_path, true)?); + } Ok(()) } @@ -176,14 +160,14 @@ impl MmapVectors { } pub fn read_deleted_map(&self) -> RwLockReadGuard { - self.deleted_mmap.read() + self.deleted_flags_mmap.read() } pub fn deleted(&self, key: PointOffsetType) -> Option { - if let Some(deleted_ram) = &self.deleted_ram { - deleted_ram.get(key as usize).map(|x| *x) + if let Some(deleted_flags) = &self.deleted_flags { + deleted_flags.get(key as usize).map(|x| *x) } else { - Self::check_deleted(&self.deleted_mmap.read(), key) + Self::check_deleted(&self.deleted_flags_mmap.read(), key) } } @@ -199,10 +183,10 @@ impl MmapVectors { pub fn delete(&mut self, key: PointOffsetType) -> OperationResult<()> { if key < (self.num_vectors as PointOffsetType) { - let mut deleted_mmap = self.deleted_mmap.write(); + let mut deleted_mmap = self.deleted_flags_mmap.write(); let flag = deleted_mmap.get_mut((key as usize) + HEADER_SIZE).unwrap(); - if let Some(deleted_ram) = &mut self.deleted_ram { + if let Some(deleted_ram) = &mut self.deleted_flags { deleted_ram.set(key as usize, true); } @@ -215,7 +199,7 @@ impl MmapVectors { } pub fn flusher(&self) -> Flusher { - let deleted_mmap = self.deleted_mmap.clone(); + let deleted_mmap = self.deleted_flags_mmap.clone(); Box::new(move || { deleted_mmap.read().flush()?; Ok(()) diff --git a/lib/segment/src/vector_storage/mod.rs b/lib/segment/src/vector_storage/mod.rs index fb41e84c1f9..a7d5981400c 100644 --- a/lib/segment/src/vector_storage/mod.rs +++ b/lib/segment/src/vector_storage/mod.rs @@ -1,8 +1,7 @@ pub mod chunked_vectors; pub mod memmap_vector_storage; mod mmap_vectors; -pub mod quantized_mmap_storage; -mod quantized_vector_storage; +mod quantized; pub mod simple_vector_storage; mod vector_storage_base; diff --git a/lib/segment/src/vector_storage/quantized/mod.rs b/lib/segment/src/vector_storage/quantized/mod.rs new file mode 100644 index 00000000000..65bc5309640 --- /dev/null +++ b/lib/segment/src/vector_storage/quantized/mod.rs @@ -0,0 +1,4 @@ +pub mod quantized_vectors_base; +mod scalar_quantized; +mod scalar_quantized_mmap_storage; +mod scalar_quantized_ram_storage; diff --git a/lib/segment/src/vector_storage/quantized/quantized_vectors_base.rs b/lib/segment/src/vector_storage/quantized/quantized_vectors_base.rs new file mode 100644 index 00000000000..1f58f81a286 --- /dev/null +++ b/lib/segment/src/vector_storage/quantized/quantized_vectors_base.rs @@ -0,0 +1,192 @@ +use std::path::{Path, PathBuf}; + +use bitvec::prelude::BitVec; +use serde::{Deserialize, Serialize}; + +use crate::common::file_operations::{atomic_save_json, read_json}; +use crate::data_types::vectors::VectorElementType; +use crate::entry::entry_point::OperationResult; +use crate::types::{Distance, QuantizationConfig, ScalarQuantization, ScalarQuantizationConfig}; +use crate::vector_storage::chunked_vectors::ChunkedVectors; +use crate::vector_storage::quantized::scalar_quantized::ScalarQuantizedVectors; +use crate::vector_storage::quantized::scalar_quantized_mmap_storage::{ + create_scalar_quantized_vectors_mmap, load_scalar_quantized_vectors_mmap, QuantizedMmapStorage, +}; +use crate::vector_storage::quantized::scalar_quantized_ram_storage::{ + create_scalar_quantized_vectors_ram, load_scalar_quantized_vectors_ram, +}; +use crate::vector_storage::RawScorer; + +pub const QUANTIZED_CONFIG_PATH: &str = "quantized.config.json"; + +#[derive(Deserialize, Serialize, Clone)] +pub struct QuantizedVectorsConfig { + pub quantization_config: QuantizationConfig, + pub vector_parameters: quantization::VectorParameters, +} + +pub enum QuantizedVectorStorageImpl { + ScalarRam(ScalarQuantizedVectors>), + ScalarMmap(ScalarQuantizedVectors), +} + +pub struct QuantizedVectorsStorage { + storage_impl: QuantizedVectorStorageImpl, + config: QuantizedVectorsConfig, + path: PathBuf, +} + +pub trait QuantizedVectors: Send + Sync { + fn raw_scorer<'a>( + &'a self, + query: &[VectorElementType], + deleted: &'a BitVec, + ) -> Box; + + fn save_to(&self, path: &Path) -> OperationResult<()>; + + /// List all files used by the quantized vectors storage + fn files(&self) -> Vec; +} + +impl QuantizedVectors for QuantizedVectorsStorage { + fn raw_scorer<'a>( + &'a self, + query: &[VectorElementType], + deleted: &'a BitVec, + ) -> Box { + match &self.storage_impl { + QuantizedVectorStorageImpl::ScalarRam(storage) => storage.raw_scorer(query, deleted), + QuantizedVectorStorageImpl::ScalarMmap(storage) => storage.raw_scorer(query, deleted), + } + } + + fn save_to(&self, path: &Path) -> OperationResult<()> { + match &self.storage_impl { + QuantizedVectorStorageImpl::ScalarRam(storage) => storage.save_to(path), + QuantizedVectorStorageImpl::ScalarMmap(storage) => storage.save_to(path), + } + } + + fn files(&self) -> Vec { + let mut result = vec![self.path.join(QUANTIZED_CONFIG_PATH)]; + let storage_files = match &self.storage_impl { + QuantizedVectorStorageImpl::ScalarRam(storage) => storage.files(), + QuantizedVectorStorageImpl::ScalarMmap(storage) => storage.files(), + }; + + result.extend(storage_files.into_iter().map(|file| self.path.join(file))); + result + } +} + +impl QuantizedVectorsStorage { + fn check_use_ram_quantization_storage( + config: &ScalarQuantizationConfig, + on_disk_vector_storage: bool, + ) -> bool { + !on_disk_vector_storage || config.always_ram == Some(true) + } + + fn construct_vector_parameters( + distance: Distance, + dim: usize, + count: usize, + ) -> quantization::VectorParameters { + quantization::VectorParameters { + dim, + count, + distance_type: match distance { + Distance::Cosine => quantization::DistanceType::Dot, + Distance::Euclid => quantization::DistanceType::L2, + Distance::Dot => quantization::DistanceType::Dot, + }, + invert: distance == Distance::Euclid, + } + } + + pub fn create<'a>( + vectors: impl IntoIterator + Clone, + quantization_config: &QuantizationConfig, + distance: Distance, + dim: usize, + count: usize, + path: &Path, + on_disk_vector_storage: bool, + ) -> OperationResult { + let vector_parameters = Self::construct_vector_parameters(distance, dim, count); + + let quantized_storage = match quantization_config { + QuantizationConfig::Scalar(ScalarQuantization { + scalar: scalar_config, + }) => { + let in_ram = + Self::check_use_ram_quantization_storage(scalar_config, on_disk_vector_storage); + if in_ram { + let storage = create_scalar_quantized_vectors_ram( + vectors, + scalar_config, + &vector_parameters, + )?; + QuantizedVectorStorageImpl::ScalarRam(storage) + } else { + let storage = create_scalar_quantized_vectors_mmap( + vectors, + scalar_config, + &vector_parameters, + path, + )?; + QuantizedVectorStorageImpl::ScalarMmap(storage) + } + } + }; + + let quantized_vectors_config = QuantizedVectorsConfig { + quantization_config: quantization_config.clone(), + vector_parameters, + }; + + let quantized_vectors = QuantizedVectorsStorage { + storage_impl: quantized_storage, + config: quantized_vectors_config, + path: path.to_path_buf(), + }; + + quantized_vectors.save_to(path)?; + atomic_save_json(&path.join(QUANTIZED_CONFIG_PATH), &quantized_vectors.config)?; + Ok(quantized_vectors) + } + + pub fn check_exists(path: &Path) -> bool { + path.join(QUANTIZED_CONFIG_PATH).exists() + } + + pub fn load(data_path: &Path, on_disk_vector_storage: bool) -> OperationResult { + let config: QuantizedVectorsConfig = read_json(&data_path.join(QUANTIZED_CONFIG_PATH))?; + let quantized_store = match &config.quantization_config { + QuantizationConfig::Scalar(ScalarQuantization { + scalar: scalar_u8_config, + }) => { + let is_ram = Self::check_use_ram_quantization_storage( + scalar_u8_config, + on_disk_vector_storage, + ); + if is_ram { + let storage = + load_scalar_quantized_vectors_ram(data_path, &config.vector_parameters)?; + QuantizedVectorStorageImpl::ScalarRam(storage) + } else { + let storage = + load_scalar_quantized_vectors_mmap(data_path, &config.vector_parameters)?; + QuantizedVectorStorageImpl::ScalarMmap(storage) + } + } + }; + + Ok(QuantizedVectorsStorage { + storage_impl: quantized_store, + config, + path: data_path.to_path_buf(), + }) + } +} diff --git a/lib/segment/src/vector_storage/quantized/scalar_quantized.rs b/lib/segment/src/vector_storage/quantized/scalar_quantized.rs new file mode 100644 index 00000000000..25e464a296c --- /dev/null +++ b/lib/segment/src/vector_storage/quantized/scalar_quantized.rs @@ -0,0 +1,97 @@ +use std::path::{Path, PathBuf}; + +use bitvec::prelude::BitVec; +use quantization::EncodedVectors; + +use crate::data_types::vectors::VectorElementType; +use crate::entry::entry_point::OperationResult; +use crate::types::{PointOffsetType, ScoreType}; +use crate::vector_storage::quantized::quantized_vectors_base::QuantizedVectors; +use crate::vector_storage::{RawScorer, ScoredPointOffset}; + +pub const QUANTIZED_DATA_PATH: &str = "quantized.data"; +pub const QUANTIZED_META_PATH: &str = "quantized.meta.json"; + +pub struct ScalarQuantizedRawScorer<'a, TEncodedQuery, TEncodedVectors> +where + TEncodedVectors: quantization::EncodedVectors, +{ + pub query: TEncodedQuery, + pub deleted: &'a BitVec, + pub quantized_data: &'a TEncodedVectors, +} + +impl RawScorer + for ScalarQuantizedRawScorer<'_, TEncodedQuery, TEncodedVectors> +where + TEncodedVectors: quantization::EncodedVectors, +{ + fn score_points(&self, points: &[PointOffsetType], scores: &mut [ScoredPointOffset]) -> usize { + let mut size: usize = 0; + for point_id in points.iter().copied() { + if self.deleted[point_id as usize] { + continue; + } + scores[size] = ScoredPointOffset { + idx: point_id, + score: self.quantized_data.score_point(&self.query, point_id), + }; + size += 1; + if size == scores.len() { + return size; + } + } + size + } + + fn check_point(&self, point: PointOffsetType) -> bool { + (point as usize) < self.deleted.len() && !self.deleted[point as usize] + } + + fn score_point(&self, point: PointOffsetType) -> ScoreType { + self.quantized_data.score_point(&self.query, point) + } + + fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType { + self.quantized_data.score_internal(point_a, point_b) + } +} + +pub struct ScalarQuantizedVectors { + storage: quantization::EncodedVectorsU8, +} + +impl ScalarQuantizedVectors { + pub fn new(storage: quantization::EncodedVectorsU8) -> Self { + Self { storage } + } +} + +impl QuantizedVectors for ScalarQuantizedVectors +where + TStorage: quantization::EncodedStorage + Send + Sync, +{ + fn raw_scorer<'a>( + &'a self, + query: &[VectorElementType], + deleted: &'a BitVec, + ) -> Box { + let query = self.storage.encode_query(query); + Box::new(ScalarQuantizedRawScorer { + query, + deleted, + quantized_data: &self.storage, + }) + } + + fn save_to(&self, path: &Path) -> OperationResult<()> { + let data_path = path.join(QUANTIZED_DATA_PATH); + let meta_path = path.join(QUANTIZED_META_PATH); + self.storage.save(&data_path, &meta_path)?; + Ok(()) + } + + fn files(&self) -> Vec { + vec![QUANTIZED_DATA_PATH.into(), QUANTIZED_META_PATH.into()] + } +} diff --git a/lib/segment/src/vector_storage/quantized_mmap_storage.rs b/lib/segment/src/vector_storage/quantized/scalar_quantized_mmap_storage.rs similarity index 59% rename from lib/segment/src/vector_storage/quantized_mmap_storage.rs rename to lib/segment/src/vector_storage/quantized/scalar_quantized_mmap_storage.rs index 2a1fd4d3b9c..4e445a39e6f 100644 --- a/lib/segment/src/vector_storage/quantized_mmap_storage.rs +++ b/lib/segment/src/vector_storage/quantized/scalar_quantized_mmap_storage.rs @@ -1,8 +1,14 @@ use std::path::Path; use memmap2::{Mmap, MmapMut}; +use quantization::EncodedVectors; +use crate::entry::entry_point::{OperationError, OperationResult}; use crate::madvise; +use crate::types::ScalarQuantizationConfig; +use crate::vector_storage::quantized::scalar_quantized::{ + ScalarQuantizedVectors, QUANTIZED_DATA_PATH, QUANTIZED_META_PATH, +}; pub struct QuantizedMmapStorage { mmap: Mmap, @@ -86,3 +92,47 @@ impl QuantizedMmapStorageBuilder { }) } } + +pub fn create_scalar_quantized_vectors_mmap<'a>( + vectors: impl IntoIterator + Clone, + config: &ScalarQuantizationConfig, + vector_parameters: &quantization::VectorParameters, + data_path: &Path, +) -> OperationResult> { + let quantized_vector_size = + quantization::EncodedVectorsU8::::get_quantized_vector_size( + vector_parameters, + ); + let mmap_data_path = data_path.join(QUANTIZED_DATA_PATH); + + let storage_builder = QuantizedMmapStorageBuilder::new( + mmap_data_path.as_path(), + vector_parameters.count, + quantized_vector_size, + )?; + let quantized_vectors = quantization::EncodedVectorsU8::encode( + vectors, + storage_builder, + vector_parameters, + config.quantile, + ) + .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; + + Ok(ScalarQuantizedVectors::new(quantized_vectors)) +} + +pub fn load_scalar_quantized_vectors_mmap( + path: &Path, + vector_parameters: &quantization::VectorParameters, +) -> OperationResult> { + let data_path = path.join(QUANTIZED_DATA_PATH); + let meta_path = path.join(QUANTIZED_META_PATH); + + let storage = quantization::EncodedVectorsU8::::load( + &data_path, + &meta_path, + vector_parameters, + )?; + + Ok(ScalarQuantizedVectors::new(storage)) +} diff --git a/lib/segment/src/vector_storage/quantized/scalar_quantized_ram_storage.rs b/lib/segment/src/vector_storage/quantized/scalar_quantized_ram_storage.rs new file mode 100644 index 00000000000..17e594bbe0d --- /dev/null +++ b/lib/segment/src/vector_storage/quantized/scalar_quantized_ram_storage.rs @@ -0,0 +1,47 @@ +use std::path::Path; + +use quantization::EncodedVectors; + +use crate::entry::entry_point::{OperationError, OperationResult}; +use crate::types::ScalarQuantizationConfig; +use crate::vector_storage::chunked_vectors::ChunkedVectors; +use crate::vector_storage::quantized::scalar_quantized::{ + ScalarQuantizedVectors, QUANTIZED_DATA_PATH, QUANTIZED_META_PATH, +}; + +pub fn create_scalar_quantized_vectors_ram<'a>( + vectors: impl IntoIterator + Clone, + config: &ScalarQuantizationConfig, + vector_parameters: &quantization::VectorParameters, +) -> OperationResult>> { + let quantized_vector_size = + quantization::EncodedVectorsU8::>::get_quantized_vector_size( + vector_parameters, + ); + let storage_builder = ChunkedVectors::::new(quantized_vector_size); + let quantized_vectors = quantization::EncodedVectorsU8::encode( + vectors, + storage_builder, + vector_parameters, + config.quantile, + ) + .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?; + + Ok(ScalarQuantizedVectors::new(quantized_vectors)) +} + +pub fn load_scalar_quantized_vectors_ram( + path: &Path, + vector_parameters: &quantization::VectorParameters, +) -> OperationResult>> { + let data_path = path.join(QUANTIZED_DATA_PATH); + let meta_path = path.join(QUANTIZED_META_PATH); + + let storage = quantization::EncodedVectorsU8::>::load( + &data_path, + &meta_path, + vector_parameters, + )?; + + Ok(ScalarQuantizedVectors::new(storage)) +} diff --git a/lib/segment/src/vector_storage/quantized_vector_storage.rs b/lib/segment/src/vector_storage/quantized_vector_storage.rs deleted file mode 100644 index 64fd69d1be0..00000000000 --- a/lib/segment/src/vector_storage/quantized_vector_storage.rs +++ /dev/null @@ -1,223 +0,0 @@ -use std::path::Path; - -use bitvec::vec::BitVec; -use quantization::EncodedVectors; - -use super::chunked_vectors::ChunkedVectors; -use super::quantized_mmap_storage::{QuantizedMmapStorage, QuantizedMmapStorageBuilder}; -use super::{RawScorer, ScoredPointOffset}; -use crate::data_types::vectors::VectorElementType; -use crate::entry::entry_point::{OperationError, OperationResult}; -use crate::types::{ - Distance, PointOffsetType, QuantizationConfig, ScalarU8Quantization, ScoreType, -}; - -pub trait QuantizedVectors: Send + Sync { - fn raw_scorer<'a>( - &'a self, - query: &[VectorElementType], - deleted: &'a BitVec, - ) -> Box; - - fn save_to_file(&self, meta_path: &Path, data_path: &Path) -> OperationResult<()>; -} - -pub struct QuantizedRawScorer<'a, TEncodedQuery, TEncodedVectors> -where - TEncodedVectors: quantization::EncodedVectors, -{ - pub query: TEncodedQuery, - pub deleted: &'a BitVec, - pub quantized_data: &'a TEncodedVectors, -} - -impl QuantizedVectors for quantization::EncodedVectorsU8 -where - TStorage: quantization::EncodedStorage + Send + Sync, -{ - fn raw_scorer<'a>( - &'a self, - query: &[VectorElementType], - deleted: &'a BitVec, - ) -> Box { - let query = self.encode_query(query); - Box::new(QuantizedRawScorer { - query, - deleted, - quantized_data: self, - }) - } - - fn save_to_file(&self, meta_path: &Path, data_path: &Path) -> OperationResult<()> { - self.save(data_path, meta_path)?; - Ok(()) - } -} - -impl RawScorer - for QuantizedRawScorer<'_, TEncodedQuery, TEncodedVectors> -where - TEncodedVectors: quantization::EncodedVectors, -{ - fn score_points(&self, points: &[PointOffsetType], scores: &mut [ScoredPointOffset]) -> usize { - let mut size: usize = 0; - for point_id in points.iter().copied() { - if self.deleted[point_id as usize] { - continue; - } - scores[size] = ScoredPointOffset { - idx: point_id, - score: self.quantized_data.score_point(&self.query, point_id), - }; - size += 1; - if size == scores.len() { - return size; - } - } - size - } - - fn check_point(&self, point: PointOffsetType) -> bool { - (point as usize) < self.deleted.len() && !self.deleted[point as usize] - } - - fn score_point(&self, point: PointOffsetType) -> ScoreType { - self.quantized_data.score_point(&self.query, point) - } - - fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType { - self.quantized_data.score_internal(point_a, point_b) - } -} - -#[allow(clippy::too_many_arguments)] -pub fn create_quantized_vectors<'a>( - vectors: impl IntoIterator + Clone, - quantization_config: &QuantizationConfig, - distance: Distance, - dim: usize, - count: usize, - meta_path: &Path, - data_path: &Path, - on_disk_vector_storage: bool, -) -> OperationResult> { - let vector_parameters = get_vector_parameters(distance, dim, count); - let quantized_vectors = match quantization_config { - QuantizationConfig::ScalarU8(scalar_u8_config) => { - let is_ram = use_ram_quantization_storage(scalar_u8_config, on_disk_vector_storage); - if is_ram { - create_quantized_vectors_ram(vectors, scalar_u8_config, &vector_parameters)? - } else { - create_quantized_vectors_mmap( - vectors, - scalar_u8_config, - &vector_parameters, - data_path, - )? - } - } - }; - quantized_vectors.save_to_file(meta_path, data_path)?; - Ok(quantized_vectors) -} - -pub fn load_quantized_vectors( - quantization_config: &QuantizationConfig, - distance: Distance, - dim: usize, - count: usize, - meta_path: &Path, - data_path: &Path, - on_disk_vector_storage: bool, -) -> OperationResult> { - let vector_parameters = get_vector_parameters(distance, dim, count); - match quantization_config { - QuantizationConfig::ScalarU8(scalar_u8_config) => { - let is_ram = use_ram_quantization_storage(scalar_u8_config, on_disk_vector_storage); - if is_ram { - Ok(Box::new(quantization::EncodedVectorsU8::< - ChunkedVectors, - >::load( - data_path, meta_path, &vector_parameters - )?)) - } else { - Ok(Box::new(quantization::EncodedVectorsU8::< - QuantizedMmapStorage, - >::load( - data_path, meta_path, &vector_parameters - )?)) - } - } - } -} - -fn get_vector_parameters( - distance: Distance, - dim: usize, - count: usize, -) -> quantization::VectorParameters { - quantization::VectorParameters { - dim, - count, - distance_type: match distance { - Distance::Cosine => quantization::DistanceType::Dot, - Distance::Euclid => quantization::DistanceType::L2, - Distance::Dot => quantization::DistanceType::Dot, - }, - invert: distance == Distance::Euclid, - } -} - -fn use_ram_quantization_storage( - config: &ScalarU8Quantization, - on_disk_vector_storage: bool, -) -> bool { - !on_disk_vector_storage || config.always_ram == Some(true) -} - -fn create_quantized_vectors_ram<'a>( - vectors: impl IntoIterator + Clone, - config: &ScalarU8Quantization, - vector_parameters: &quantization::VectorParameters, -) -> OperationResult> { - let quantized_vector_size = - quantization::EncodedVectorsU8::>::get_quantized_vector_size( - vector_parameters, - ); - let storage_builder = ChunkedVectors::::new(quantized_vector_size); - Ok(Box::new( - quantization::EncodedVectorsU8::encode( - vectors, - storage_builder, - vector_parameters, - config.quantile, - ) - .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?, - )) -} - -fn create_quantized_vectors_mmap<'a>( - vectors: impl IntoIterator + Clone, - config: &ScalarU8Quantization, - vector_parameters: &quantization::VectorParameters, - data_path: &Path, -) -> OperationResult> { - let quantized_vector_size = - quantization::EncodedVectorsU8::::get_quantized_vector_size( - vector_parameters, - ); - let storage_builder = QuantizedMmapStorageBuilder::new( - data_path, - vector_parameters.count, - quantized_vector_size, - )?; - Ok(Box::new( - quantization::EncodedVectorsU8::encode( - vectors, - storage_builder, - vector_parameters, - config.quantile, - ) - .map_err(|e| OperationError::service_error(format!("Cannot quantize vector data: {e}")))?, - )) -} diff --git a/lib/segment/src/vector_storage/simple_vector_storage.rs b/lib/segment/src/vector_storage/simple_vector_storage.rs index f435f650da8..f3eaf5250e7 100644 --- a/lib/segment/src/vector_storage/simple_vector_storage.rs +++ b/lib/segment/src/vector_storage/simple_vector_storage.rs @@ -13,7 +13,6 @@ use rocksdb::DB; use serde::{Deserialize, Serialize}; use super::chunked_vectors::ChunkedVectors; -use super::quantized_vector_storage::{load_quantized_vectors, QuantizedVectors}; use super::vector_storage_base::VectorStorage; use crate::common::rocksdb_wrapper::DatabaseColumnWrapper; use crate::common::Flusher; @@ -23,7 +22,9 @@ use crate::spaces::metric::Metric; use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric}; use crate::spaces::tools::peek_top_largest_iterable; use crate::types::{Distance, PointOffsetType, QuantizationConfig, ScoreType}; -use crate::vector_storage::quantized_vector_storage::create_quantized_vectors; +use crate::vector_storage::quantized::quantized_vectors_base::{ + QuantizedVectors, QuantizedVectorsStorage, +}; use crate::vector_storage::{RawScorer, ScoredPointOffset, VectorStorageSS}; /// In-memory vector storage with on-update persistence using `store` @@ -33,7 +34,7 @@ pub struct SimpleVectorStorage { vectors: ChunkedVectors, deleted: BitVec, deleted_count: usize, - quantized_vectors: Option>, + quantized_vectors: Option, db_wrapper: DatabaseColumnWrapper, } @@ -318,39 +319,26 @@ where fn quantize( &mut self, - meta_path: &Path, - data_path: &Path, + path: &Path, quantization_config: &QuantizationConfig, ) -> OperationResult<()> { let vector_data_iterator = (0..self.vectors.len() as u32).map(|i| self.vectors.get(i)); - self.quantized_vectors = Some(create_quantized_vectors( + self.quantized_vectors = Some(QuantizedVectorsStorage::create( vector_data_iterator, quantization_config, TMetric::distance(), self.dim, self.vectors.len(), - meta_path, - data_path, + path, false, )?); Ok(()) } - fn load_quantization( - &mut self, - meta_path: &Path, - data_path: &Path, - quantization_config: &QuantizationConfig, - ) -> OperationResult<()> { - self.quantized_vectors = Some(load_quantized_vectors( - quantization_config, - TMetric::distance(), - self.dim, - self.vectors.len(), - meta_path, - data_path, - false, - )?); + fn load_quantization(&mut self, path: &Path) -> OperationResult<()> { + if QuantizedVectorsStorage::check_exists(path) { + self.quantized_vectors = Some(QuantizedVectorsStorage::load(path, false)?); + } Ok(()) } @@ -400,7 +388,11 @@ where } fn files(&self) -> Vec { - vec![] + if let Some(quantized_vectors) = &self.quantized_vectors { + quantized_vectors.files() + } else { + vec![] + } } } @@ -410,7 +402,7 @@ mod tests { use super::*; use crate::common::rocksdb_wrapper::{open_db, DB_VECTOR_CF}; - use crate::types::ScalarU8Quantization; + use crate::types::ScalarQuantizationConfig; #[test] fn test_score_points() { @@ -508,16 +500,14 @@ mod tests { borrowed_storage.put_vector(vec3).unwrap(); borrowed_storage.put_vector(vec4).unwrap(); - let config = QuantizationConfig::ScalarU8(ScalarU8Quantization { + let config: QuantizationConfig = ScalarQuantizationConfig { + r#type: Default::default(), quantile: None, always_ram: None, - }); - let quantized_meta_path = dir.path().join("quantized.meta"); - let quantized_data_path = dir.path().join("quantized.data"); + } + .into(); - borrowed_storage - .quantize(&quantized_meta_path, &quantized_data_path, &config) - .unwrap(); + borrowed_storage.quantize(dir.path(), &config).unwrap(); let query = vec![0.5, 0.5, 0.5, 0.5]; @@ -536,9 +526,7 @@ mod tests { } // test save-load - borrowed_storage - .load_quantization(&quantized_meta_path, &quantized_data_path, &config) - .unwrap(); + borrowed_storage.load_quantization(dir.path()).unwrap(); let scorer_quant = borrowed_storage.quantized_raw_scorer(&query).unwrap(); let scorer_orig = borrowed_storage.raw_scorer(query.clone()); diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index 99e1d357ad0..00bd7f8c9d9 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -80,6 +80,7 @@ pub trait VectorStorage { // Generate RawScorer on quantized vectors if present fn quantized_raw_scorer(&self, vector: &[VectorElementType]) -> Option>; + // Try peek top nearest points from quantized vectors. If quantized vectors are not present, do it on raw vectors fn score_quantized_points( &self, @@ -87,20 +88,16 @@ pub trait VectorStorage { points: &mut dyn Iterator, top: usize, ) -> Vec; + // Generate quantized vectors and store them on disk fn quantize( &mut self, - meta_path: &Path, data_path: &Path, quantization_config: &QuantizationConfig, ) -> OperationResult<()>; + // Load quantized vectors from disk - fn load_quantization( - &mut self, - meta_path: &Path, - data_path: &Path, - quantization_config: &QuantizationConfig, - ) -> OperationResult<()>; + fn load_quantization(&mut self, data_path: &Path) -> OperationResult<()>; fn score_points( &self, diff --git a/lib/segment/tests/hnsw_quantized_search_test.rs b/lib/segment/tests/hnsw_quantized_search_test.rs index 8f1b3ee2429..e126c98e103 100644 --- a/lib/segment/tests/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/hnsw_quantized_search_test.rs @@ -12,8 +12,8 @@ mod tests { use segment::index::VectorIndex; use segment::segment_constructor::build_segment; use segment::types::{ - Distance, HnswConfig, Indexes, QuantizationConfig, ScalarU8Quantization, SearchParams, - SegmentConfig, SeqNumberType, StorageType, VectorDataConfig, + Distance, HnswConfig, Indexes, ScalarQuantizationConfig, SearchParams, SegmentConfig, + SeqNumberType, StorageType, VectorDataConfig, }; use segment::vector_storage::ScoredPointOffset; use tempfile::Builder; @@ -29,8 +29,7 @@ mod tests { fn hnsw_quantized_search_test(distance: Distance) { let stopped = AtomicBool::new(false); let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); - let quantized_meta_path = dir.path().join("quantized.meta"); - let quantized_data_path = dir.path().join("quantized.data"); + let quantized_data_path = dir.path(); let dim = 128; let m = 16; @@ -65,24 +64,19 @@ mod tests { .unwrap(); } segment.vector_data.values_mut().for_each(|vector_storage| { - if quantized_data_path.exists() || quantized_meta_path.exists() { - panic!("quantization files shouldn't exists"); - } vector_storage .vector_storage .borrow_mut() .quantize( - &quantized_meta_path, - &quantized_data_path, - &QuantizationConfig::ScalarU8(ScalarU8Quantization { + quantized_data_path, + &ScalarQuantizationConfig { + r#type: Default::default(), quantile: None, always_ram: None, - }), + } + .into(), ) .unwrap(); - if !quantized_data_path.exists() || !quantized_meta_path.exists() { - panic!("quantization was not saved"); - } }); let hnsw_config = HnswConfig { diff --git a/lib/storage/src/content_manager/collection_meta_ops.rs b/lib/storage/src/content_manager/collection_meta_ops.rs index ab0f8ef5938..95a812cecd1 100644 --- a/lib/storage/src/content_manager/collection_meta_ops.rs +++ b/lib/storage/src/content_manager/collection_meta_ops.rs @@ -133,6 +133,7 @@ pub struct CreateCollection { pub init_from: Option, /// Quantization parameters. If none - quantization is disabled. #[serde(default)] + #[serde(alias = "quantization")] pub quantization_config: Option, } diff --git a/lib/storage/src/content_manager/toc.rs b/lib/storage/src/content_manager/toc.rs index 7f323c7ba13..71fba0cbddf 100644 --- a/lib/storage/src/content_manager/toc.rs +++ b/lib/storage/src/content_manager/toc.rs @@ -319,7 +319,7 @@ impl TableOfContent { }; let quantization_config = match quantization_config { - None => self.storage_config.quantization_config.clone(), + None => self.storage_config.quantization.clone(), Some(diff) => Some(diff), }; diff --git a/lib/storage/src/types.rs b/lib/storage/src/types.rs index 0bdeba7a4c1..2a4725d6073 100644 --- a/lib/storage/src/types.rs +++ b/lib/storage/src/types.rs @@ -47,7 +47,7 @@ pub struct StorageConfig { pub wal: WalConfig, pub performance: PerformanceConfig, pub hnsw_index: HnswConfig, - pub quantization_config: Option, + pub quantization: Option, #[serde(default = "default_mmap_advice")] pub mmap_advice: madvise::Advice, #[serde(default)] diff --git a/lib/storage/tests/alias_tests.rs b/lib/storage/tests/alias_tests.rs index 42cc99053d7..f6db84038f4 100644 --- a/lib/storage/tests/alias_tests.rs +++ b/lib/storage/tests/alias_tests.rs @@ -49,7 +49,7 @@ mod tests { max_optimization_threads: 1, }, hnsw_index: Default::default(), - quantization_config: None, + quantization: None, mmap_advice: madvise::Advice::Random, node_type: Default::default(), }; diff --git a/src/actix/helpers.rs b/src/actix/helpers.rs index 226e13e5036..76305537bd8 100644 --- a/src/actix/helpers.rs +++ b/src/actix/helpers.rs @@ -38,8 +38,14 @@ where let mut resp = match err { StorageError::BadInput { .. } => HttpResponse::BadRequest(), StorageError::NotFound { .. } => HttpResponse::NotFound(), - StorageError::ServiceError { .. } => { - log::warn!("error processing request: {}", err); + StorageError::ServiceError { + description, + backtrace, + } => { + log::warn!("error processing request: {}", description); + if let Some(backtrace) = backtrace { + log::warn!("backtrace: {}", backtrace); + } HttpResponse::InternalServerError() } StorageError::BadRequest { .. } => HttpResponse::BadRequest(),