Byte storage integration into segment (qdrant#4049)

* byte storage with quantization raw scorer integration config and test are you happy fmt fn renamings cow refactor use quantization branch quantization update * are you happy clippy * don't use distance in quantized scorers * fix build * add fn quantization_preprocess * apply preprocessing for only cosine float metric * fix sparse vectors tests * update openapi * more complicated integration test * update openapi comment * mmap byte storages support * fix async test * move .unwrap closer to the actual check of the vector presence * fmt * remove distance similarity function * avoid copying data while working with cow --------- Co-authored-by: generall <andrey@vasnetsov.com>
AKJUS · Apr 22, 2024 · 19cda34 · 19cda34
1 parent ef5b606
commit 19cda34
Show file tree

Hide file tree

Showing 47 changed files with 1,039 additions and 217 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json
@@ -9153,6 +9153,17 @@
                 "nullable": true
               }
             ]
+          },
+          "datatype": {
+            "description": "Vector specific configuration to set specific storage element type",
+            "anyOf": [
+              {
+                "$ref": "#/components/schemas/VectorStorageDatatype"
+              },
+              {
+                "nullable": true
+              }
+            ]
           }
         }
       },
@@ -9244,6 +9255,25 @@
       "MaxSimConfig": {
         "type": "object"
       },
+      "VectorStorageDatatype": {
+        "description": "Storage types for vectors",
+        "oneOf": [
+          {
+            "description": "Single-precsion floating point",
+            "type": "string",
+            "enum": [
+              "Float"
+            ]
+          },
+          {
+            "description": "Unsigned 8-bit integer",
+            "type": "string",
+            "enum": [
+              "Uint8"
+            ]
+          }
+        ]
+      },
       "SparseVectorDataConfig": {
         "description": "Config of single sparse vector data storage",
         "type": "object",

diff --git a/lib/collection/src/collection_manager/holders/proxy_segment.rs b/lib/collection/src/collection_manager/holders/proxy_segment.rs
@@ -1400,6 +1400,7 @@ mod tests {
                         index: Indexes::Plain {},
                         quantization_config: None,
                         multi_vec_config: None,
+                        datatype: None,
                     },
                 ),
                 (
@@ -1411,6 +1412,7 @@ mod tests {
                         index: Indexes::Plain {},
                         quantization_config: None,
                         multi_vec_config: None,
+                        datatype: None,
                     },
                 ),
             ]),

diff --git a/lib/collection/src/config.rs b/lib/collection/src/config.rs
@@ -349,6 +349,8 @@ impl CollectionParams {
                         },
                         // TODO(colbert) add `multivec` to `VectorParams`
                         multi_vec_config: None,
+                        // TODO(byte_storage)
+                        datatype: None,
                     },
                 )
             })

diff --git a/lib/segment/benches/multi_vector_search.rs b/lib/segment/benches/multi_vector_search.rs
@@ -47,6 +47,7 @@ fn multi_vector_search_benchmark(c: &mut Criterion) {
                 index: Indexes::Plain {},
                 quantization_config: None,
                 multi_vec_config: Some(MultiVectorConfig::default()), // uses multivec config
+                datatype: None,
             },
         )]),
         sparse_vector_data: Default::default(),

diff --git a/lib/segment/src/compat.rs b/lib/segment/src/compat.rs
@@ -53,6 +53,7 @@ impl From<SegmentConfigV5> for SegmentConfig {
                         .then_some(VectorStorageType::Mmap)
                         .unwrap_or_else(|| old_segment.storage_type.into()),
                     multi_vec_config: None,
+                    datatype: None,
                 };
 
                 (vector_name, new_data)

diff --git a/lib/segment/src/data_types/named_vectors.rs b/lib/segment/src/data_types/named_vectors.rs
@@ -4,9 +4,13 @@ use std::collections::HashMap;
 use sparse::common::sparse_vector::SparseVector;
 
 use super::tiny_map;
-use super::vectors::{DenseVector, MultiDenseVector, Vector, VectorElementType, VectorRef};
+use super::vectors::{
+    DenseVector, MultiDenseVector, Vector, VectorElementType, VectorElementTypeByte, VectorRef,
+};
 use crate::common::operation_error::OperationError;
-use crate::types::Distance;
+use crate::spaces::metric::Metric;
+use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric, ManhattanMetric};
+use crate::types::{Distance, VectorDataConfig, VectorStorageDatatype};
 
 type CowKey<'a> = Cow<'a, str>;
 
@@ -48,6 +52,15 @@ impl<'a> CowVector<'a> {
     }
 }
 
+impl<'a> From<Cow<'a, [VectorElementType]>> for CowVector<'a> {
+    fn from(v: Cow<'a, [VectorElementType]>) -> Self {
+        match v {
+            Cow::Borrowed(v) => CowVector::Dense(Cow::Borrowed(v)),
+            Cow::Owned(v) => CowVector::Dense(Cow::Owned(v)),
+        }
+    }
+}
+
 impl<'a> From<Vector> for CowVector<'a> {
     fn from(v: Vector) -> Self {
         match v {
@@ -118,6 +131,18 @@ impl<'a> TryFrom<CowVector<'a>> for DenseVector {
     }
 }
 
+impl<'a> TryFrom<CowVector<'a>> for Cow<'a, [VectorElementType]> {
+    type Error = OperationError;
+
+    fn try_from(value: CowVector<'a>) -> Result<Self, Self::Error> {
+        match value {
+            CowVector::Dense(v) => Ok(v),
+            CowVector::Sparse(_) => Err(OperationError::WrongSparse),
+            CowVector::MultiDense(_) => Err(OperationError::WrongMulti),
+        }
+    }
+}
+
 impl<'a> From<VectorRef<'a>> for CowVector<'a> {
     fn from(v: VectorRef<'a>) -> Self {
         match v {
@@ -203,31 +228,66 @@ impl<'a> NamedVectors<'a> {
         self.map.get(key).map(|v| v.as_vec_ref())
     }
 
-    pub fn preprocess<F>(&mut self, distance_map: F)
-    where
-        F: Fn(&str) -> Distance,
-    {
+    pub fn preprocess<'b>(&'b mut self, get_vector_data: impl Fn(&str) -> &'b VectorDataConfig) {
         for (name, vector) in self.map.iter_mut() {
-            let distance = distance_map(name);
             match vector {
                 CowVector::Dense(v) => {
-                    let preprocessed_vector = distance.preprocess_vector(v.to_vec());
+                    let config = get_vector_data(name.as_ref());
+                    let preprocessed_vector = Self::preprocess_dense_vector(v.to_vec(), config);
                     *vector = CowVector::Dense(Cow::Owned(preprocessed_vector))
                 }
                 CowVector::Sparse(v) => {
                     // sort by indices to enable faster dot product and overlap checks
                     v.to_mut().sort_by_indices();
                 }
                 CowVector::MultiDense(v) => {
+                    let config = get_vector_data(name.as_ref());
                     for dense_vector in v.to_mut().multi_vectors_mut() {
-                        let preprocessed_vector = distance.preprocess_vector(dense_vector.to_vec());
+                        let preprocessed_vector =
+                            Self::preprocess_dense_vector(dense_vector.to_vec(), config);
                         // replace dense vector with preprocessed vector
                         dense_vector.copy_from_slice(&preprocessed_vector);
                     }
                 }
             }
         }
     }
+
+    fn preprocess_dense_vector(
+        dense_vector: DenseVector,
+        config: &VectorDataConfig,
+    ) -> DenseVector {
+        match config.datatype {
+            Some(VectorStorageDatatype::Float) | None => match config.distance {
+                Distance::Cosine => {
+                    <CosineMetric as Metric<VectorElementType>>::preprocess(dense_vector)
+                }
+                Distance::Euclid => {
+                    <EuclidMetric as Metric<VectorElementType>>::preprocess(dense_vector)
+                }
+                Distance::Dot => {
+                    <DotProductMetric as Metric<VectorElementType>>::preprocess(dense_vector)
+                }
+                Distance::Manhattan => {
+                    <ManhattanMetric as Metric<VectorElementType>>::preprocess(dense_vector)
+                }
+            },
+            Some(VectorStorageDatatype::Uint8) => match config.distance {
+                Distance::Cosine => {
+                    <CosineMetric as Metric<VectorElementTypeByte>>::preprocess(dense_vector)
+                }
+                Distance::Euclid => {
+                    <EuclidMetric as Metric<VectorElementTypeByte>>::preprocess(dense_vector)
+                }
+                Distance::Dot => {
+                    <DotProductMetric as Metric<VectorElementTypeByte>>::preprocess(dense_vector)
+                }
+                Distance::Manhattan => {
+                    <ManhattanMetric as Metric<VectorElementTypeByte>>::preprocess(dense_vector)
+                }
+            },
+        }
+    }
 }
 
 impl<'a> IntoIterator for NamedVectors<'a> {

diff --git a/lib/segment/src/data_types/primitive.rs b/lib/segment/src/data_types/primitive.rs
@@ -1,41 +1,63 @@
 use std::borrow::Cow;
 
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
-use crate::common::operation_error::OperationResult;
-use crate::data_types::named_vectors::CowVector;
-use crate::data_types::vectors::{VectorElementType, VectorElementTypeByte, VectorRef};
+use crate::data_types::vectors::{VectorElementType, VectorElementTypeByte};
+use crate::types::QuantizationConfig;
 
 pub trait PrimitiveVectorElement:
     Copy + Clone + Default + Serialize + for<'a> Deserialize<'a>
 {
-    fn from_vector_ref(vector: VectorRef) -> OperationResult<Cow<[Self]>>;
+    fn slice_from_float_cow(vector: Cow<[VectorElementType]>) -> Cow<[Self]>;
 
-    fn vector_to_cow(vector: &[Self]) -> CowVector;
+    fn slice_to_float_cow(vector: Cow<[Self]>) -> Cow<[VectorElementType]>;
+
+    fn quantization_preprocess<'a>(
+        quantization_config: &QuantizationConfig,
+        vector: &'a [Self],
+    ) -> Cow<'a, [f32]>;
 }
+
 impl PrimitiveVectorElement for VectorElementType {
-    fn from_vector_ref(vector: VectorRef) -> OperationResult<Cow<[Self]>> {
-        let vector_ref: &[Self] = vector.try_into()?;
-        Ok(Cow::from(vector_ref))
+    fn slice_from_float_cow(vector: Cow<[VectorElementType]>) -> Cow<[Self]> {
+        vector
     }
 
-    fn vector_to_cow(vector: &[Self]) -> CowVector {
-        vector.into()
+    fn slice_to_float_cow(vector: Cow<[Self]>) -> Cow<[VectorElementType]> {
+        vector
+    }
+
+    fn quantization_preprocess<'a>(
+        _quantization_config: &QuantizationConfig,
+        vector: &'a [Self],
+    ) -> Cow<'a, [f32]> {
+        Cow::Borrowed(vector)
     }
 }
 
 impl PrimitiveVectorElement for VectorElementTypeByte {
-    fn from_vector_ref(vector: VectorRef) -> OperationResult<Cow<[Self]>> {
-        let vector_ref: &[VectorElementType] = vector.try_into()?;
-        let byte_vector = vector_ref.iter().map(|&x| x as u8).collect::<Vec<u8>>();
-        Ok(Cow::from(byte_vector))
+    fn slice_from_float_cow(vector: Cow<[VectorElementType]>) -> Cow<[Self]> {
+        Cow::Owned(vector.iter().map(|&x| x as u8).collect())
     }
 
-    fn vector_to_cow(vector: &[Self]) -> CowVector {
-        vector
-            .iter()
-            .map(|&x| x as VectorElementType)
-            .collect::<Vec<VectorElementType>>()
-            .into()
+    fn slice_to_float_cow(vector: Cow<[Self]>) -> Cow<[VectorElementType]> {
+        Cow::Owned(vector.iter().map(|&x| x as VectorElementType).collect_vec())
+    }
+
+    fn quantization_preprocess<'a>(
+        quantization_config: &QuantizationConfig,
+        vector: &'a [Self],
+    ) -> Cow<'a, [f32]> {
+        if let QuantizationConfig::Binary(_) = quantization_config {
+            Cow::from(
+                vector
+                    .iter()
+                    .map(|&x| (x as VectorElementType) - 127.0)
+                    .collect_vec(),
+            )
+        } else {
+            Cow::from(vector.iter().map(|&x| x as VectorElementType).collect_vec())
+        }
     }
 }
diff --git a/lib/segment/src/fixtures/payload_fixtures.rs b/lib/segment/src/fixtures/payload_fixtures.rs
@@ -8,7 +8,7 @@ use rand::seq::SliceRandom;
 use rand::Rng;
 use serde_json::{json, Value};
 
-use crate::data_types::vectors::{DenseVector, MultiDenseVector};
+use crate::data_types::vectors::{DenseVector, MultiDenseVector, VectorElementType};
 use crate::types::{
     AnyVariants, Condition, ExtendedPointId, FieldCondition, Filter, HasIdCondition,
     IsEmptyCondition, Match, MatchAny, Payload, PayloadField, Range as RangeCondition, ValuesCount,
@@ -121,6 +121,16 @@ pub fn random_vector<R: Rng + ?Sized>(rnd_gen: &mut R, size: usize) -> DenseVect
     (0..size).map(|_| rnd_gen.gen()).collect()
 }
 
+pub fn random_dense_byte_vector<R: Rng + ?Sized>(rnd_gen: &mut R, size: usize) -> DenseVector {
+    (0..size)
+        .map(|_| {
+            rnd_gen
+                .gen_range::<VectorElementType, _>(0.0..=255.0)
+                .round()
+        })
+        .collect()
+}
+
 pub fn random_multi_vector<R: Rng + ?Sized>(
     rnd_gen: &mut R,
     vector_size: usize,

diff --git a/lib/segment/src/index/hnsw_index/graph_layers.rs b/lib/segment/src/index/hnsw_index/graph_layers.rs
@@ -451,7 +451,7 @@ mod tests {
 
         let top = 5;
         let query = random_vector(&mut rng, dim);
-        let processed_query = M::preprocess(query.clone());
+        let processed_query = <M as Metric<VectorElementType>>::preprocess(query.clone());
         let mut reference_top = FixedLengthPriorityQueue::new(top);
         for idx in 0..vector_holder.vectors.len() as PointOffsetType {
             let vec = &vector_holder.vectors.get(idx);

diff --git a/lib/segment/src/index/hnsw_index/graph_layers_builder.rs b/lib/segment/src/index/hnsw_index/graph_layers_builder.rs
@@ -632,7 +632,7 @@ mod tests {
 
         let top = 5;
         let query = random_vector(&mut rng, dim);
-        let processed_query = M::preprocess(query.clone());
+        let processed_query = <M as Metric<VectorElementType>>::preprocess(query.clone());
         let mut reference_top = FixedLengthPriorityQueue::new(top);
         for idx in 0..vector_holder.vectors.len() as PointOffsetType {
             let vec = &vector_holder.vectors.get(idx);
@@ -715,7 +715,7 @@ mod tests {
 
         let top = 5;
         let query = random_vector(&mut rng, dim);
-        let processed_query = M::preprocess(query.clone());
+        let processed_query = <M as Metric<VectorElementType>>::preprocess(query.clone());
         let mut reference_top = FixedLengthPriorityQueue::new(top);
         for idx in 0..vector_holder.vectors.len() as PointOffsetType {
             let vec = &vector_holder.vectors.get(idx);

diff --git a/lib/segment/src/index/hnsw_index/tests/test_graph_connectivity.rs b/lib/segment/src/index/hnsw_index/tests/test_graph_connectivity.rs
@@ -46,6 +46,7 @@ fn test_graph_connectivity() {
                 index: Indexes::Plain {},
                 quantization_config: None,
                 multi_vec_config: None,
+                datatype: None,
             },
         )]),
         payload_storage_type: Default::default(),