diff --git a/docs/grpc/docs.md b/docs/grpc/docs.md index aad362c9ab..6bcf9e14dc 100644 --- a/docs/grpc/docs.md +++ b/docs/grpc/docs.md @@ -93,6 +93,7 @@ - [VectorsConfigDiff](#qdrant-VectorsConfigDiff) - [WalConfigDiff](#qdrant-WalConfigDiff) + - [BinaryQuantizationEncoding](#qdrant-BinaryQuantizationEncoding) - [CollectionStatus](#qdrant-CollectionStatus) - [CompressionRatio](#qdrant-CompressionRatio) - [Datatype](#qdrant-Datatype) @@ -395,6 +396,7 @@ | Field | Type | Label | Description | | ----- | ---- | ----- | ----------- | | always_ram | [bool](#bool) | optional | If true - quantized vectors always will be stored in RAM, ignoring the config of main storage | +| encoding | [BinaryQuantizationEncoding](#qdrant-BinaryQuantizationEncoding) | optional | Binary quantization encoding method | @@ -1848,6 +1850,19 @@ Note: 1kB = 1 vector of size 256. | + + +### BinaryQuantizationEncoding + + +| Name | Number | Description | +| ---- | ------ | ----------- | +| OneBit | 0 | | +| TwoBits | 1 | | +| OneAndHalfBits | 2 | | + + + ### CollectionStatus diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json index da2bc3bea1..e57deab16b 100644 --- a/docs/redoc/master/openapi.json +++ b/docs/redoc/master/openapi.json @@ -7016,9 +7016,27 @@ "always_ram": { "type": "boolean", "nullable": true + }, + "encoding": { + "anyOf": [ + { + "$ref": "#/components/schemas/BinaryQuantizationEncoding" + }, + { + "nullable": true + } + ] } } }, + "BinaryQuantizationEncoding": { + "type": "string", + "enum": [ + "one_bit", + "two_bits", + "one_and_half_bits" + ] + }, "Datatype": { "type": "string", "enum": [ diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index 2b64e87108..0e60e1e02e 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -35,7 +35,6 @@ use super::qdrant::{ }; use super::{Expression, Formula, RecoQuery, Usage}; use crate::conversions::json::{self, json_to_proto}; -use crate::grpc; use crate::grpc::qdrant::condition::ConditionOneOf; use crate::grpc::qdrant::r#match::MatchValue; use crate::grpc::qdrant::payload_index_params::IndexParams; @@ -55,7 +54,8 @@ use crate::grpc::qdrant::{ with_vectors_selector, }; use crate::grpc::{ - DecayParamsExpression, DivExpression, GeoDistance, MultExpression, PowExpression, SumExpression, + self, BinaryQuantizationEncoding, DecayParamsExpression, DivExpression, GeoDistance, + MultExpression, PowExpression, SumExpression, }; use crate::rest::models::{CollectionsResponse, VersionInfo}; use crate::rest::schema as rest; @@ -1080,11 +1080,50 @@ impl TryFrom for segment::types::ProductQuantization { } } +impl From for BinaryQuantizationEncoding { + fn from(value: segment::types::BinaryQuantizationEncoding) -> Self { + match value { + segment::types::BinaryQuantizationEncoding::OneBit => { + BinaryQuantizationEncoding::OneBit + } + segment::types::BinaryQuantizationEncoding::TwoBits => { + BinaryQuantizationEncoding::TwoBits + } + segment::types::BinaryQuantizationEncoding::OneAndHalfBits => { + BinaryQuantizationEncoding::OneAndHalfBits + } + } + } +} + +impl From for segment::types::BinaryQuantizationEncoding { + fn from(value: BinaryQuantizationEncoding) -> Self { + match value { + BinaryQuantizationEncoding::OneBit => { + segment::types::BinaryQuantizationEncoding::OneBit + } + BinaryQuantizationEncoding::TwoBits => { + segment::types::BinaryQuantizationEncoding::TwoBits + } + BinaryQuantizationEncoding::OneAndHalfBits => { + segment::types::BinaryQuantizationEncoding::OneAndHalfBits + } + } + } +} + impl From for BinaryQuantization { fn from(value: segment::types::BinaryQuantization) -> Self { let segment::types::BinaryQuantization { binary } = value; - let segment::types::BinaryQuantizationConfig { always_ram } = binary; - BinaryQuantization { always_ram } + let segment::types::BinaryQuantizationConfig { + always_ram, + encoding, + } = binary; + BinaryQuantization { + always_ram, + encoding: encoding + .map(|encoding| i32::from(BinaryQuantizationEncoding::from(encoding))), + } } } @@ -1092,9 +1131,19 @@ impl TryFrom for segment::types::BinaryQuantization { type Error = Status; fn try_from(value: BinaryQuantization) -> Result { - let BinaryQuantization { always_ram } = value; + let BinaryQuantization { + always_ram, + encoding, + } = value; + let encoding = encoding + .map(BinaryQuantizationEncoding::try_from) + .transpose() + .map_err(|_| Status::invalid_argument("Unknown binary quantization encoding"))?; Ok(segment::types::BinaryQuantization { - binary: segment::types::BinaryQuantizationConfig { always_ram }, + binary: segment::types::BinaryQuantizationConfig { + always_ram, + encoding: encoding.map(segment::types::BinaryQuantizationEncoding::from), + }, }) } } diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index 4e91f40a73..963457624a 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -295,8 +295,15 @@ message ProductQuantization { optional bool always_ram = 2; // If true - quantized vectors always will be stored in RAM, ignoring the config of main storage } +enum BinaryQuantizationEncoding { + OneBit = 0; + TwoBits = 1; + OneAndHalfBits = 2; +} + message BinaryQuantization { optional bool always_ram = 1; // If true - quantized vectors always will be stored in RAM, ignoring the config of main storage + optional BinaryQuantizationEncoding encoding = 2; // Binary quantization encoding method } message QuantizationConfig { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index c759b06a83..94a6c42bc3 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -430,6 +430,9 @@ pub struct BinaryQuantization { /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage #[prost(bool, optional, tag = "1")] pub always_ram: ::core::option::Option, + /// Binary quantization encoding method + #[prost(enumeration = "BinaryQuantizationEncoding", optional, tag = "2")] + pub encoding: ::core::option::Option, } #[derive(validator::Validate)] #[derive(serde::Serialize)] @@ -1650,6 +1653,36 @@ impl CompressionRatio { #[derive(serde::Serialize)] #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] +pub enum BinaryQuantizationEncoding { + OneBit = 0, + TwoBits = 1, + OneAndHalfBits = 2, +} +impl BinaryQuantizationEncoding { + /// 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 { + BinaryQuantizationEncoding::OneBit => "OneBit", + BinaryQuantizationEncoding::TwoBits => "TwoBits", + BinaryQuantizationEncoding::OneAndHalfBits => "OneAndHalfBits", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "OneBit" => Some(Self::OneBit), + "TwoBits" => Some(Self::TwoBits), + "OneAndHalfBits" => Some(Self::OneAndHalfBits), + _ => None, + } + } +} +#[derive(serde::Serialize)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] pub enum ShardingMethod { /// Auto-sharding based on record ids Auto = 0, diff --git a/lib/quantization/benches/binary.rs b/lib/quantization/benches/binary.rs index 044cc6ceb9..126102d3bd 100644 --- a/lib/quantization/benches/binary.rs +++ b/lib/quantization/benches/binary.rs @@ -4,7 +4,7 @@ use common::counter::hardware_counter::HardwareCounterCell; use criterion::{Criterion, criterion_group, criterion_main}; use permutation_iterator::Permutor; use quantization::encoded_vectors::{DistanceType, EncodedVectors, VectorParameters}; -use quantization::encoded_vectors_binary::EncodedVectorsBin; +use quantization::encoded_vectors_binary::{EncodedVectorsBin, Encoding}; use rand::{Rng, SeedableRng}; fn generate_number(rng: &mut rand::rngs::StdRng) -> f32 { @@ -39,6 +39,7 @@ fn binary_bench(c: &mut Criterion) { distance_type: DistanceType::Dot, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -78,6 +79,7 @@ fn binary_bench(c: &mut Criterion) { distance_type: DistanceType::Dot, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); diff --git a/lib/quantization/src/encoded_vectors_binary.rs b/lib/quantization/src/encoded_vectors_binary.rs index 53b18900ba..fb246d6480 100644 --- a/lib/quantization/src/encoded_vectors_binary.rs +++ b/lib/quantization/src/encoded_vectors_binary.rs @@ -1,5 +1,5 @@ use std::marker::PhantomData; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; use common::counter::hardware_counter::HardwareCounterCell; @@ -8,6 +8,7 @@ use memory::mmap_ops::{transmute_from_u8_to_slice, transmute_to_u8_slice}; use serde::{Deserialize, Serialize}; use crate::encoded_vectors::validate_vector_parameters; +use crate::vector_stats::VectorStats; use crate::{ DistanceType, EncodedStorage, EncodedStorageBuilder, EncodedVectors, EncodingError, VectorParameters, @@ -16,9 +17,24 @@ use crate::{ pub struct EncodedVectorsBin { encoded_vectors: TStorage, metadata: Metadata, + vector_stats: Option, bits_store_type: PhantomData, } +#[derive(Clone, Copy, Eq, PartialEq, Debug, Serialize, Deserialize, Default)] +pub enum Encoding { + #[default] + OneBit, + TwoBits, + OneAndHalfBits, +} + +impl Encoding { + pub fn is_one(&self) -> bool { + matches!(self, Encoding::OneBit) + } +} + pub struct EncodedBinVector { encoded_vector: Vec, } @@ -26,6 +42,9 @@ pub struct EncodedBinVector { #[derive(Serialize, Deserialize)] struct Metadata { vector_parameters: VectorParameters, + #[serde(default)] + #[serde(skip_serializing_if = "Encoding::is_one")] + encoding: Encoding, } pub trait BitsStoreType: @@ -155,7 +174,7 @@ impl BitsStoreType for u128 { } fn get_storage_size(size: usize) -> usize { - let bits_count = 8 * std::mem::size_of::(); + let bits_count = u8::BITS as usize * std::mem::size_of::(); let mut result = size / bits_count; if size % bits_count != 0 { result += 1; @@ -175,16 +194,24 @@ impl orig_data: impl Iterator + 'a> + Clone, mut storage_builder: impl EncodedStorageBuilder, vector_parameters: &VectorParameters, + encoding: Encoding, stopped: &AtomicBool, ) -> Result { debug_assert!(validate_vector_parameters(orig_data.clone(), vector_parameters).is_ok()); + let vector_stats = match encoding { + Encoding::OneBit => None, + Encoding::TwoBits | Encoding::OneAndHalfBits => { + Some(VectorStats::build(orig_data.clone(), vector_parameters)) + } + }; + for vector in orig_data { if stopped.load(Ordering::Relaxed) { return Err(EncodingError::Stopped); } - let encoded_vector = Self::encode_vector(vector.as_ref()); + let encoded_vector = Self::encode_vector(vector.as_ref(), &vector_stats, encoding); let encoded_vector_slice = encoded_vector.encoded_vector.as_slice(); let bytes = transmute_to_u8_slice(encoded_vector_slice); storage_builder.push_vector_data(bytes); @@ -194,15 +221,37 @@ impl encoded_vectors: storage_builder.build(), metadata: Metadata { vector_parameters: vector_parameters.clone(), + encoding, }, + vector_stats, bits_store_type: PhantomData, }) } - fn encode_vector(vector: &[f32]) -> EncodedBinVector { - let mut encoded_vector = - vec![Default::default(); TBitsStoreType::get_storage_size(vector.len())]; + fn encode_vector( + vector: &[f32], + vector_stats: &Option, + encoding: Encoding, + ) -> EncodedBinVector { + let encoded_vector_size = + Self::get_quantized_vector_size_from_params(vector.len(), encoding) + / std::mem::size_of::(); + let mut encoded_vector = vec![Default::default(); encoded_vector_size]; + match encoding { + Encoding::OneBit => Self::encode_one_bit_vector(vector, &mut encoded_vector), + Encoding::TwoBits => { + Self::encode_two_bits_vector(vector, &mut encoded_vector, vector_stats) + } + Encoding::OneAndHalfBits => { + Self::encode_one_and_half_bits_vector(vector, &mut encoded_vector, vector_stats) + } + } + + EncodedBinVector { encoded_vector } + } + + fn encode_one_bit_vector(vector: &[f32], encoded_vector: &mut [TBitsStoreType]) { let bits_count = u8::BITS as usize * std::mem::size_of::(); let one = TBitsStoreType::one(); for (i, &v) in vector.iter().enumerate() { @@ -212,17 +261,137 @@ impl encoded_vector[i / bits_count] |= one << (i % bits_count); } } - - EncodedBinVector { encoded_vector } } - pub fn get_quantized_vector_size_from_params(vector_parameters: &VectorParameters) -> usize { - TBitsStoreType::get_storage_size(vector_parameters.dim) + fn encode_two_bits_vector( + vector: &[f32], + encoded_vector: &mut [TBitsStoreType], + vector_stats: &Option, + ) { + let Some(vector_stats) = vector_stats else { + debug_assert!(false, "Vector stats must be provided for two bits encoding"); + // If vector stats are not provided, we cannot encode two bits + // So we fall back to one bit encoding + return Self::encode_one_bit_vector(vector, encoded_vector); + }; + + let bits_count = u8::BITS as usize * std::mem::size_of::(); + let one = TBitsStoreType::one(); + for (i, &v) in vector.iter().enumerate() { + let mean = vector_stats.elements_stats[i].mean; + let sd = vector_stats.elements_stats[i].stddev; + + let (b1, b2) = Self::encode_two_bits_value(v, mean, sd); + + if b1 { + encoded_vector[i / bits_count] |= one << (i % bits_count); + } + if b2 { + let j = vector.len() + i; + encoded_vector[j / bits_count] |= one << (j % bits_count); + } + } + } + + fn encode_one_and_half_bits_vector( + vector: &[f32], + encoded_vector: &mut [TBitsStoreType], + vector_stats: &Option, + ) { + let Some(vector_stats) = vector_stats else { + debug_assert!(false, "Vector stats must be provided for two bits encoding"); + // If vector stats are not provided, we cannot encode one and half bits + // So we fall back to one bit encoding + return Self::encode_one_bit_vector(vector, encoded_vector); + }; + + // One and half bit encoding is a 2bit quantization but first bit, + // which describes that value is less that sigma, + // is united with the bit from the next value using OR operand. + // Scoring for 1.5bit quantization is the same as for 2bit and 1bit quantization. + // + // Example 1: + // `Value1` has `[1,0]` 2bits encoding, value `Value2` has `[1,1]` 2bits encoding. + // The resulting 1.5bit encoding will be `[value1[0], value2[0], value1[1] | value2[1]] = [1, 1, 1]`. + // + // Example 2: + // `Value1` has `[0,0]` 2bits encoding, value `Value2` has `[1,0]` 2bits encoding. + // The resulting 1.5bit encoding will be `[value1[0], value2[0], value1[1] | value2[1]] = [0, 1, 0]`. + let bits_count = u8::BITS as usize * std::mem::size_of::(); + let one = TBitsStoreType::one(); + for (i, &v) in vector.iter().enumerate() { + let mean = vector_stats.elements_stats[i].mean; + let sd = vector_stats.elements_stats[i].stddev; + let (b1, b2) = Self::encode_two_bits_value(v, mean, sd); + + if b1 { + encoded_vector[i / bits_count] |= one << (i % bits_count); + } + if b2 { + let j = vector.len() + i / 2; + encoded_vector[j / bits_count] |= one << (j % bits_count); + } + } + } + + fn encode_two_bits_value(value: f32, mean: f32, sd: f32) -> (bool, bool) { + // Two bit encoding is a regular BQ with "zero". + // It uses 2 bits per value and encodes values in the following way: + // 00 - if the value is in the range [-2*sigma; -sigma); + // 10 - if the value is in the range [-sigma; sigma); + // 11 - if the value is in the range [sigma; 2*sigma]; + // where sigma is the standard deviation of the value. + // + // Scoring for 2bit quantization is the same as for 1bit quantization. + + if sd < f32::EPSILON { + // If standard deviation is zero, + // we cannot calculate z-score count so use regular BQ with zero-comparison. + return (value > 0.0, false); + } + + // How many ranges we want to divide the values into, it's 3: + // [-2*sigma; -sigma), [-sigma; sigma), [sigma; 2*sigma) + let ranges = 3; + + // Calculate z-score for the value + let v_z = (value - mean) / sd; + + let min_border = -2.0; // -2*sigma + let max_border = 2.0; // 2*sigma + + // Normalize z-score to the range [-2*sigma; 2*sigma] + let normalized_z = (v_z - min_border) / (max_border - min_border); + + // Calculate index in the ranges list: [-2*sigma; -sigma), [-sigma; sigma), [sigma; 2*sigma) + let index = normalized_z * (ranges as f32); + + // Index 0 and less is [0, 0] encoding + // Index 1 is [1, 0] encoding + // Index 2 and more is [1, 1] encoding + if index >= 1.0 { + let count_ones = (index.floor() as usize).min(2); + (count_ones > 0, count_ones > 1) + } else { + (false, false) + } + } + + pub fn get_quantized_vector_size_from_params(dim: usize, encoding: Encoding) -> usize { + let extended_dim = match encoding { + Encoding::OneBit => dim, + Encoding::TwoBits => dim * 2, + Encoding::OneAndHalfBits => (dim * 3).div_ceil(2), // ceil(dim * 1.5) + }; + TBitsStoreType::get_storage_size(extended_dim.max(1)) * std::mem::size_of::() } fn get_quantized_vector_size(&self) -> usize { - Self::get_quantized_vector_size_from_params(&self.metadata.vector_parameters) + Self::get_quantized_vector_size_from_params( + self.metadata.vector_parameters.dim, + self.metadata.encoding, + ) } fn calculate_metric(&self, v1: &[TBitsStoreType], v2: &[TBitsStoreType]) -> f32 { @@ -273,6 +442,12 @@ impl pub fn vectors_count(&self) -> usize { self.metadata.vector_parameters.count } + + fn get_vector_stats_path_from_meta_path(meta_path: &Path) -> Option { + let mut vector_stats_path = meta_path.parent()?.to_owned(); + vector_stats_path.push("vector_stats.json"); + Some(vector_stats_path) + } } impl EncodedVectors @@ -286,6 +461,21 @@ impl EncodedVectors data_path.parent().map(std::fs::create_dir_all); self.encoded_vectors.save_to_file(data_path)?; + + let vector_stats_path = + Self::get_vector_stats_path_from_meta_path(meta_path).ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Failed to get vector stats path", + ) + })?; + if let Some(vector_stats) = &self.vector_stats { + vector_stats_path.parent().map(std::fs::create_dir_all); + atomic_save_json(&vector_stats_path, &vector_stats)?; + } else if let Ok(true) = std::fs::exists(&vector_stats_path) { + std::fs::remove_file(&vector_stats_path)?; + } + Ok(()) } @@ -296,12 +486,31 @@ impl EncodedVectors ) -> std::io::Result { let contents = std::fs::read_to_string(meta_path)?; let metadata: Metadata = serde_json::from_str(&contents)?; - let quantized_vector_size = Self::get_quantized_vector_size_from_params(vector_parameters); + let quantized_vector_size = + Self::get_quantized_vector_size_from_params(vector_parameters.dim, metadata.encoding); let encoded_vectors = TStorage::from_file(data_path, quantized_vector_size, vector_parameters.count)?; + + let vector_stats_path = + Self::get_vector_stats_path_from_meta_path(meta_path).ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Failed to get vector stats path", + ) + })?; + let vector_stats = match metadata.encoding { + Encoding::OneBit => None, + Encoding::TwoBits | Encoding::OneAndHalfBits => { + let vector_stats_contents = std::fs::read_to_string(&vector_stats_path)?; + let vector_stats: VectorStats = serde_json::from_str(&vector_stats_contents)?; + Some(vector_stats) + } + }; + let result = Self { metadata, encoded_vectors, + vector_stats, bits_store_type: PhantomData, }; Ok(result) @@ -313,7 +522,7 @@ impl EncodedVectors fn encode_query(&self, query: &[f32]) -> EncodedBinVector { debug_assert!(query.len() == self.metadata.vector_parameters.dim); - Self::encode_vector(query) + Self::encode_vector(query, &self.vector_stats, self.metadata.encoding) } fn score_point( diff --git a/lib/quantization/src/lib.rs b/lib/quantization/src/lib.rs index 655226eee6..3449dec9da 100644 --- a/lib/quantization/src/lib.rs +++ b/lib/quantization/src/lib.rs @@ -5,6 +5,7 @@ pub mod encoded_vectors_pq; pub mod encoded_vectors_u8; pub mod kmeans; pub mod quantile; +pub mod vector_stats; use std::fmt::Display; use std::sync::{Arc, Condvar, Mutex}; diff --git a/lib/quantization/src/vector_stats.rs b/lib/quantization/src/vector_stats.rs new file mode 100644 index 0000000000..c2ae91d31b --- /dev/null +++ b/lib/quantization/src/vector_stats.rs @@ -0,0 +1,97 @@ +use serde::{Deserialize, Serialize}; + +use crate::VectorParameters; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct VectorStats { + pub elements_stats: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct VectorElementStats { + pub min: f32, + pub max: f32, + pub mean: f32, + pub stddev: f32, +} + +impl Default for VectorElementStats { + fn default() -> Self { + VectorElementStats { + min: f32::MAX, + max: f32::MIN, + mean: 0.0, + stddev: 0.0, + } + } +} + +impl VectorStats { + pub fn build<'a>( + data: impl Iterator + 'a>, + vector_params: &VectorParameters, + ) -> Self { + // The Welford's Algorithm. + let mut stats = VectorStats { + elements_stats: vec![VectorElementStats::default(); vector_params.dim], + }; + + // For internal calculations use higher precision. + let mut m2 = vec![0.0f64; vector_params.dim]; + let mut means = vec![0.0f64; vector_params.dim]; + + let mut count = 0; + for vector in data { + let vector = vector.as_ref(); + count += 1; + + debug_assert_eq!( + vector.len(), + vector_params.dim, + "Vector length does not match the expected dimension" + ); + + for (((&value, element_stats), mean), m2) in vector + .iter() + .zip(stats.elements_stats.iter_mut()) + .zip(means.iter_mut()) + .zip(m2.iter_mut()) + { + element_stats.min = if value < element_stats.min { + value + } else { + element_stats.min + }; + element_stats.max = if value > element_stats.max { + value + } else { + element_stats.max + }; + + let delta = f64::from(value) - *mean; + *mean += delta / count as f64; + *m2 += delta * (f64::from(value) - *mean); + } + } + + debug_assert_eq!( + count, vector_params.count, + "Count of vectors processed does not match the expected count in vector parameters" + ); + for ((element_stats, means), m2) in stats + .elements_stats + .iter_mut() + .zip(means.iter()) + .zip(m2.iter()) + { + element_stats.stddev = if count > 1 { + (*m2 / (count - 1) as f64).sqrt() as f32 + } else { + 0.0 + }; + element_stats.mean = *means as f32; + } + + stats + } +} diff --git a/lib/quantization/tests/integration/empty_storage.rs b/lib/quantization/tests/integration/empty_storage.rs index f44e8e5e63..e1f4285f5f 100644 --- a/lib/quantization/tests/integration/empty_storage.rs +++ b/lib/quantization/tests/integration/empty_storage.rs @@ -101,6 +101,7 @@ mod tests { vector_data.iter(), Vec::::new(), &vector_parameters, + quantization::encoded_vectors_binary::Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); diff --git a/lib/quantization/tests/integration/main.rs b/lib/quantization/tests/integration/main.rs index b0766e443a..bf3cbb72f0 100644 --- a/lib/quantization/tests/integration/main.rs +++ b/lib/quantization/tests/integration/main.rs @@ -9,6 +9,8 @@ pub mod test_avx2; #[cfg(test)] pub mod test_binary; #[cfg(test)] +pub mod test_binary_encodings; +#[cfg(test)] pub mod test_neon; #[cfg(test)] pub mod test_pq; diff --git a/lib/quantization/tests/integration/test_binary.rs b/lib/quantization/tests/integration/test_binary.rs index b2ffbb563c..0f69985754 100644 --- a/lib/quantization/tests/integration/test_binary.rs +++ b/lib/quantization/tests/integration/test_binary.rs @@ -4,7 +4,7 @@ mod tests { use common::counter::hardware_counter::HardwareCounterCell; use quantization::encoded_vectors::{DistanceType, EncodedVectors, VectorParameters}; - use quantization::encoded_vectors_binary::{BitsStoreType, EncodedVectorsBin}; + use quantization::encoded_vectors_binary::{BitsStoreType, EncodedVectorsBin, Encoding}; use rand::{Rng, SeedableRng}; use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; @@ -49,6 +49,7 @@ mod tests { distance_type: DistanceType::Dot, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -95,6 +96,7 @@ mod tests { distance_type: DistanceType::Dot, invert: true, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -141,6 +143,7 @@ mod tests { distance_type: DistanceType::Dot, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -184,6 +187,7 @@ mod tests { distance_type: DistanceType::Dot, invert: true, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -226,6 +230,7 @@ mod tests { distance_type: DistanceType::L1, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -287,6 +292,7 @@ mod tests { distance_type: DistanceType::L1, invert: true, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -348,6 +354,7 @@ mod tests { distance_type: DistanceType::L1, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -406,6 +413,7 @@ mod tests { distance_type: DistanceType::L1, invert: true, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -464,6 +472,7 @@ mod tests { distance_type: DistanceType::L2, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -525,6 +534,7 @@ mod tests { distance_type: DistanceType::L2, invert: true, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -586,6 +596,7 @@ mod tests { distance_type: DistanceType::L2, invert: false, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); @@ -644,6 +655,7 @@ mod tests { distance_type: DistanceType::L2, invert: true, }, + Encoding::OneBit, &AtomicBool::new(false), ) .unwrap(); diff --git a/lib/quantization/tests/integration/test_binary_encodings.rs b/lib/quantization/tests/integration/test_binary_encodings.rs new file mode 100644 index 0000000000..ceafb450e6 --- /dev/null +++ b/lib/quantization/tests/integration/test_binary_encodings.rs @@ -0,0 +1,120 @@ +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicBool; + + use common::counter::hardware_counter::HardwareCounterCell; + use quantization::encoded_vectors::{DistanceType, EncodedVectors, VectorParameters}; + use quantization::encoded_vectors_binary::{BitsStoreType, EncodedVectorsBin, Encoding}; + use rand::{Rng, SeedableRng}; + + use crate::metrics::dot_similarity; + + fn generate_number(rng: &mut rand::rngs::StdRng) -> f32 { + rng.random_range(-1.0..1.0) + } + + fn generate_vector(dim: usize, rng: &mut rand::rngs::StdRng) -> Vec { + (0..dim) + .map(|_| generate_number(rng) / (dim as f32).sqrt()) + .collect() + } + + fn get_top(scores: &[f32], count: usize, invert: bool) -> Vec { + let mut indices: Vec = (0..scores.len()).collect(); + indices.sort_by(|&a, &b| scores[b].partial_cmp(&scores[a]).unwrap()); + if invert { + indices.reverse(); + } + indices.into_iter().take(count).collect() + } + + fn match_count(ids1: &[usize], ids2: &[usize]) -> usize { + ids1.iter().filter(|&&id| ids2.contains(&id)).count() + } + + #[test] + fn test_binary_dot() { + test_binary_dot_impl::(0, false); + test_binary_dot_impl::(600, false); + test_binary_dot_impl::(601, false); + test_binary_dot_impl::(600, false); + } + + #[test] + fn test_binary_dot_inverted() { + test_binary_dot_impl::(700, true); + test_binary_dot_impl::(700, true); + } + + fn test_binary_dot_impl(vector_dim: usize, invert: bool) { + let vectors_count = 1000; + + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encodings = [ + Encoding::OneBit, + Encoding::OneAndHalfBits, + Encoding::TwoBits, + ]; + + let encoded: Vec<_> = encodings + .iter() + .map(|&encoding| { + EncodedVectorsBin::::encode( + vector_data.iter(), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::Dot, + invert, + }, + encoding, + &AtomicBool::new(false), + ) + .unwrap() + }) + .collect(); + + let top = 10; + let query: Vec = generate_vector(vector_dim, &mut rng); + + let orig_scores: Vec = vector_data + .iter() + .map(|vector| dot_similarity(&query, vector)) + .collect(); + let original_top = get_top(&orig_scores, top, invert); + + let tops = encoded + .iter() + .map(|encoded| { + let query_encoded = encoded.encode_query(&query); + let scores: Vec = (0..vector_data.len()) + .map(|index| { + encoded.score_point( + &query_encoded, + index as u32, + &HardwareCounterCell::new(), + ) + }) + .collect(); + let tops = get_top(&scores, top, false); + match_count(&original_top, &tops) + }) + .collect::>(); + + // Check if encoding has more accuracy than previous one + for i in 1..tops.len() { + assert!( + tops[i] >= tops[i - 1], + "Encoding {} has less accuracy than encoding {}", + i, + i - 1 + ); + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_vector_storage/tests.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_vector_storage/tests.rs index 50ebfd6913..32f9025656 100644 --- a/lib/segment/src/index/hnsw_index/gpu/gpu_vector_storage/tests.rs +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_vector_storage/tests.rs @@ -18,8 +18,9 @@ use crate::index::hnsw_index::gpu::shader_builder::ShaderBuilder; use crate::spaces::metric::Metric; use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric, ManhattanMetric}; use crate::types::{ - BinaryQuantization, BinaryQuantizationConfig, Distance, ProductQuantization, - ProductQuantizationConfig, QuantizationConfig, ScalarQuantization, ScalarQuantizationConfig, + BinaryQuantization, BinaryQuantizationConfig, BinaryQuantizationEncoding, Distance, + ProductQuantization, ProductQuantizationConfig, QuantizationConfig, ScalarQuantization, + ScalarQuantizationConfig, }; use crate::vector_storage::dense::simple_dense_vector_storage::{ open_simple_dense_byte_vector_storage, open_simple_dense_full_vector_storage, @@ -144,65 +145,75 @@ fn test_gpu_vector_storage_sq( } #[rstest] -#[case::cosine_f32( +#[case::cosine_f32_one_bit( Distance::Cosine, TestStorageType::Dense(TestElementType::Float32), 273, - 2057 + 2057, + BinaryQuantizationEncoding::OneBit )] -#[case::dot_f32( +#[case::dot_f32_one_and_half_bits( Distance::Dot, TestStorageType::Dense(TestElementType::Float32), 256, - 512 + 512, + BinaryQuantizationEncoding::OneAndHalfBits )] #[case::euclid_f32( Distance::Euclid, TestStorageType::Dense(TestElementType::Float32), 273, - 2057 + 2057, + BinaryQuantizationEncoding::OneBit )] -#[case::manhattan_f32( +#[case::manhattan_f32_two_bits( Distance::Manhattan, TestStorageType::Dense(TestElementType::Float32), 273, - 2057 + 2057, + BinaryQuantizationEncoding::TwoBits )] #[case::small_dimension( Distance::Cosine, TestStorageType::Dense(TestElementType::Float32), 17, - 2057 + 2057, + BinaryQuantizationEncoding::OneBit )] #[case::cosine_f16( Distance::Cosine, TestStorageType::Dense(TestElementType::Float16), 273, - 2057 + 2057, + BinaryQuantizationEncoding::OneBit )] #[case::cosine_u8( Distance::Cosine, TestStorageType::Dense(TestElementType::Uint8), 273, - 2057 + 2057, + BinaryQuantizationEncoding::OneBit )] #[case::cosine_multi_f32( Distance::Cosine, TestStorageType::Multi(TestElementType::Float32), 67, - 2057 + 2057, + BinaryQuantizationEncoding::OneBit )] #[case::cosine_multi_u8( Distance::Cosine, TestStorageType::Multi(TestElementType::Uint8), 273, - 2057 + 2057, + BinaryQuantizationEncoding::OneBit )] fn test_gpu_vector_storage_bq( #[case] distance: Distance, #[case] storage_type: TestStorageType, #[case] dim: usize, #[case] num_vectors: usize, + #[case] encoding: BinaryQuantizationEncoding, ) { let _ = env_logger::builder() .is_test(true) @@ -212,6 +223,7 @@ fn test_gpu_vector_storage_bq( let quantization_config = QuantizationConfig::Binary(BinaryQuantization { binary: BinaryQuantizationConfig { always_ram: Some(true), + encoding: Some(encoding), }, }); diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 0223f79a84..87d7b5b926 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -660,11 +660,29 @@ impl Hash for ScalarQuantizationConfig { impl Eq for ScalarQuantizationConfig {} +#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash, Default)] +#[serde(rename_all = "snake_case")] +pub enum BinaryQuantizationEncoding { + #[default] + OneBit, + TwoBits, + OneAndHalfBits, +} + +impl BinaryQuantizationEncoding { + pub fn is_one_bit(&self) -> bool { + matches!(self, BinaryQuantizationEncoding::OneBit) + } +} + #[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] pub struct BinaryQuantizationConfig { #[serde(skip_serializing_if = "Option::is_none")] pub always_ram: Option, + #[serde(default)] + #[serde(skip_serializing_if = "Option::is_none")] + pub encoding: Option, } #[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Eq, Hash)] diff --git a/lib/segment/src/vector_storage/quantized/quantized_vectors.rs b/lib/segment/src/vector_storage/quantized/quantized_vectors.rs index f267dcf091..7ceff8abc6 100644 --- a/lib/segment/src/vector_storage/quantized/quantized_vectors.rs +++ b/lib/segment/src/vector_storage/quantized/quantized_vectors.rs @@ -20,9 +20,9 @@ use crate::common::vector_utils::TrySetCapacityExact; use crate::data_types::primitive::PrimitiveVectorElement; use crate::data_types::vectors::{QueryVector, VectorElementType}; use crate::types::{ - BinaryQuantization, BinaryQuantizationConfig, CompressionRatio, Distance, MultiVectorConfig, - ProductQuantization, ProductQuantizationConfig, QuantizationConfig, ScalarQuantization, - ScalarQuantizationConfig, VectorStorageDatatype, + BinaryQuantization, BinaryQuantizationConfig, BinaryQuantizationEncoding, CompressionRatio, + Distance, MultiVectorConfig, ProductQuantization, ProductQuantizationConfig, + QuantizationConfig, ScalarQuantization, ScalarQuantizationConfig, VectorStorageDatatype, }; use crate::vector_storage::chunked_vectors::ChunkedVectors; use crate::vector_storage::quantized::quantized_mmap_storage::{ @@ -862,16 +862,35 @@ impl QuantizedVectors { on_disk_vector_storage: bool, stopped: &AtomicBool, ) -> OperationResult { + let encoding = match binary_config.encoding { + Some(BinaryQuantizationEncoding::OneBit) => { + quantization::encoded_vectors_binary::Encoding::OneBit + } + Some(BinaryQuantizationEncoding::TwoBits) => { + quantization::encoded_vectors_binary::Encoding::TwoBits + } + Some(BinaryQuantizationEncoding::OneAndHalfBits) => { + quantization::encoded_vectors_binary::Encoding::OneAndHalfBits + } + None => quantization::encoded_vectors_binary::Encoding::OneBit, + }; let quantized_vector_size = EncodedVectorsBin::::get_quantized_vector_size_from_params( - vector_parameters, + vector_parameters.dim, + encoding, ); let in_ram = Self::is_ram(binary_config.always_ram, on_disk_vector_storage); if in_ram { let mut storage_builder = ChunkedVectors::::new(quantized_vector_size); storage_builder.try_set_capacity_exact(vector_parameters.count)?; Ok(QuantizedVectorStorage::BinaryRam( - EncodedVectorsBin::encode(vectors, storage_builder, vector_parameters, stopped)?, + EncodedVectorsBin::encode( + vectors, + storage_builder, + vector_parameters, + encoding, + stopped, + )?, )) } else { let mmap_data_path = path.join(QUANTIZED_DATA_PATH); @@ -881,7 +900,13 @@ impl QuantizedVectors { quantized_vector_size, )?; Ok(QuantizedVectorStorage::BinaryMmap( - EncodedVectorsBin::encode(vectors, storage_builder, vector_parameters, stopped)?, + EncodedVectorsBin::encode( + vectors, + storage_builder, + vector_parameters, + encoding, + stopped, + )?, )) } } @@ -897,16 +922,34 @@ impl QuantizedVectors { on_disk_vector_storage: bool, stopped: &AtomicBool, ) -> OperationResult { + let encoding = match binary_config.encoding { + Some(BinaryQuantizationEncoding::OneBit) => { + quantization::encoded_vectors_binary::Encoding::OneBit + } + Some(BinaryQuantizationEncoding::TwoBits) => { + quantization::encoded_vectors_binary::Encoding::TwoBits + } + Some(BinaryQuantizationEncoding::OneAndHalfBits) => { + quantization::encoded_vectors_binary::Encoding::OneAndHalfBits + } + None => quantization::encoded_vectors_binary::Encoding::OneBit, + }; let quantized_vector_size = EncodedVectorsBin::::get_quantized_vector_size_from_params( - vector_parameters, + vector_parameters.dim, + encoding, ); let in_ram = Self::is_ram(binary_config.always_ram, on_disk_vector_storage); if in_ram { let mut storage_builder = ChunkedVectors::::new(quantized_vector_size); storage_builder.try_set_capacity_exact(vector_parameters.count)?; - let quantized_storage = - EncodedVectorsBin::encode(vectors, storage_builder, vector_parameters, stopped)?; + let quantized_storage = EncodedVectorsBin::encode( + vectors, + storage_builder, + vector_parameters, + encoding, + stopped, + )?; Ok(QuantizedVectorStorage::BinaryRamMulti( QuantizedMultivectorStorage::new( vector_parameters.dim, @@ -922,8 +965,13 @@ impl QuantizedVectors { vector_parameters.count, quantized_vector_size, )?; - let quantized_storage = - EncodedVectorsBin::encode(vectors, storage_builder, vector_parameters, stopped)?; + let quantized_storage = EncodedVectorsBin::encode( + vectors, + storage_builder, + vector_parameters, + encoding, + stopped, + )?; let offsets_path = path.join(QUANTIZED_OFFSETS_PATH); create_offsets_file_from_iter(&offsets_path, vector_parameters.count, offsets)?; Ok(QuantizedVectorStorage::BinaryMmapMulti( diff --git a/lib/segment/src/vector_storage/tests/custom_query_scorer_equivalency.rs b/lib/segment/src/vector_storage/tests/custom_query_scorer_equivalency.rs index dd1aec626c..ab9e800b34 100644 --- a/lib/segment/src/vector_storage/tests/custom_query_scorer_equivalency.rs +++ b/lib/segment/src/vector_storage/tests/custom_query_scorer_equivalency.rs @@ -93,6 +93,7 @@ fn product_x4() -> WithQuantization { fn binary() -> WithQuantization { let config = BinaryQuantizationConfig { always_ram: Some(true), + encoding: None, } .into(); diff --git a/lib/segment/tests/integration/byte_storage_quantization_test.rs b/lib/segment/tests/integration/byte_storage_quantization_test.rs index 9928ef76a1..3694aa732f 100644 --- a/lib/segment/tests/integration/byte_storage_quantization_test.rs +++ b/lib/segment/tests/integration/byte_storage_quantization_test.rs @@ -297,7 +297,11 @@ fn test_byte_storage_binary_quantization_hnsw( always_ram: None, } .into(), - QuantizationVariant::Binary => BinaryQuantizationConfig { always_ram: None }.into(), + QuantizationVariant::Binary => BinaryQuantizationConfig { + always_ram: None, + encoding: None, + } + .into(), }; segment_byte diff --git a/lib/segment/tests/integration/multivector_quantization_test.rs b/lib/segment/tests/integration/multivector_quantization_test.rs index afaff14c5b..90bf89d168 100644 --- a/lib/segment/tests/integration/multivector_quantization_test.rs +++ b/lib/segment/tests/integration/multivector_quantization_test.rs @@ -273,6 +273,7 @@ fn test_multivector_quantization_hnsw( .into(), QuantizationVariant::Binary => BinaryQuantizationConfig { always_ram: Some(false), + encoding: None, } .into(), };