feat: add metrics dot and cos

Signed-off-by: cutecutecat <junyuchen@tensorchord.ai>

Author: cutecutecat

crates/base/src/index.rs

@@ -149,9 +149,9 @@ impl IndexOptions {
                 }
             }
             IndexingOptions::Rabitq(_) => {
-                if !matches!(self.vector.d, DistanceKind::L2) {
+                if !matches!(self.vector.d, DistanceKind::L2 | DistanceKind::Dot) {
                     return Err(ValidationError::new(
-                        "rabitq is not support for distance that is not l2",
+                        "rabitq is not support for distance that is not l2 or dot",
                     ));
                 }
                 if !matches!(self.vector.v, VectorKind::Vecf32) {

crates/k_means/src/lib.rs

@@ -71,6 +71,27 @@ pub fn k_means_lookup<S: ScalarLike>(vector: &[S], centroids: &Vec2<S>) -> usize
     result.1
 }
 
+/// returns (centroid_dot_dis, index)
+pub fn k_means_lookup_by_dot<S: ScalarLike>(
+    vector: &[S],
+    centroids: &Vec2<S>,
+    centroids_square: &[f32],
+) -> (f32, usize) {
+    assert_ne!(centroids.shape_0(), 0);
+    let vector_square = S::reduce_sum_of_x2(vector);
+    let mut result = (f32::INFINITY, f32::INFINITY, 0);
+
+    for i in 0..centroids.shape_0() {
+        let centroid_square = centroids_square[i];
+        let dot = S::reduce_sum_of_xy(vector, &centroids[(i,)]);
+        let l2_dis = vector_square + centroid_square - 2.0 * dot;
+        if l2_dis <= result.0 {
+            result = (l2_dis, -dot, i);
+        }
+    }
+    (result.1, result.2)
+}
+
 pub fn k_means_lookup_many<S: ScalarLike>(vector: &[S], centroids: &Vec2<S>) -> Vec<(f32, usize)> {
     assert_ne!(centroids.shape_0(), 0);
     let mut seq = Vec::new();
@@ -80,3 +101,33 @@ pub fn k_means_lookup_many<S: ScalarLike>(vector: &[S], centroids: &Vec2<S>) ->
     }
     seq
 }
+
+/// returns Vec of <l2_dis, (centroid_dot_dis, vector_l2_norm, centroids_l2_norm, index)>
+pub fn k_means_lookup_many_by_dot<S: ScalarLike>(
+    vector: &[S],
+    centroids: &Vec2<S>,
+    centroids_square: &[f32],
+) -> Vec<(f32, (f32, f32, f32, usize))> {
+    assert_ne!(centroids.shape_0(), 0);
+    let vector_square = S::reduce_sum_of_x2(vector);
+    let mut seq = Vec::new();
+
+    for i in 0..centroids.shape_0() {
+        let centroid_square = centroids_square[i];
+        let dot = S::reduce_sum_of_xy(vector, &centroids[(i,)]);
+        let l2_dis = vector_square + centroid_square - 2.0 * dot;
+        seq.push((l2_dis, (-dot, vector_square, centroid_square, i)));
+    }
+    seq
+}
+
+pub fn centroids_square<S: ScalarLike>(centroids: &Vec2<S>) -> Vec<f32> {
+    assert_ne!(centroids.shape_0(), 0);
+    let mut seq = Vec::new();
+
+    for i in 0..centroids.shape_0() {
+        let centroids_square = S::reduce_sum_of_x2(&centroids[(i,)]);
+        seq.push(centroids_square);
+    }
+    seq
+}

crates/quantization/src/quantize.rs

@@ -152,9 +152,13 @@ mod mul_add_round {
 #[inline(always)]
 pub fn quantize<const N: u8>(lut: &[f32]) -> (f32, f32, Vec<u8>) {
     let (min, max) = f32::reduce_min_max_of_x(lut);
-    let k = 0.0f32.max((max - min) / (N as f32));
-    let b = min;
-    (k, b, mul_add_round::mul_add_round(lut, 1.0 / k, -b / k))
+    let delta = 0.0f32.max((max - min) / (N as f32));
+    let lower_bound = min;
+    (
+        delta,
+        lower_bound,
+        mul_add_round::mul_add_round(lut, 1.0 / delta, -lower_bound / delta),
+    )
 }
 
 #[inline(always)]

crates/rabitq/src/lib.rs

@@ -17,7 +17,9 @@ use common::json::Json;
 use common::mmap_array::MmapArray;
 use common::remap::RemappedCollection;
 use common::vec2::Vec2;
-use k_means::{k_means, k_means_lookup, k_means_lookup_many};
+use k_means::{
+    centroids_square, k_means, k_means_lookup, k_means_lookup_by_dot, k_means_lookup_many_by_dot,
+};
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use std::fs::create_dir;
 use std::path::Path;
@@ -30,6 +32,7 @@ pub struct Rabitq<O: Op> {
     payloads: MmapArray<Payload>,
     offsets: Json<Vec<u32>>,
     projected_centroids: Json<Vec2<f32>>,
+    centroids_square: Json<Vec<f32>>,
     projection: Json<Vec<Vec<f32>>>,
 }
 
@@ -70,15 +73,18 @@ impl<O: Op> Rabitq<O> {
     ) -> Box<dyn Iterator<Item = Element> + 'a> {
         let projected_query = O::proj(&self.projection, O::cast(vector));
         let lists = select(
-            k_means_lookup_many(&projected_query, &self.projected_centroids),
+            k_means_lookup_many_by_dot(&projected_query, &self.centroids, &self.centroids_square),
             opts.rabitq_nprobe as usize,
         );
         let mut heap = Vec::new();
-        for &(_, i) in lists.iter() {
-            let preprocessed = self.quantization.preprocess(&O::residual(
+        for &(_, payload) in lists.iter() {
+            let (centroid_dot_dis, original_square, centroids_square, i) = payload;
+            let preprocessed = self.quantization.preprocess(
                 &projected_query,
-                &self.projected_centroids[(i,)],
-            ));
+                centroid_dot_dis,
+                original_square,
+                centroids_square,
+            );
             let start = self.offsets[i];
             let end = self.offsets[i + 1];
             self.quantization.push_batch(
@@ -135,6 +141,7 @@ fn from_nothing<O: Op>(
     let samples = O::sample(collection, nlist);
     rayon::check();
     let centroids: Vec2<f32> = k_means(nlist as usize, samples, true, spherical_centroids, false);
+    let centroids_squares = centroids_square(&centroids);
     rayon::check();
     let ls = (0..collection.len())
         .into_par_iter()
@@ -174,8 +181,12 @@ fn from_nothing<O: Op>(
         collection.len(),
         |vector| {
             let vector = O::cast(collection.vector(vector));
-            let target = k_means_lookup(vector, &centroids);
-            O::proj(&projection, &O::residual(vector, &centroids[(target,)]))
+            let (centroid_dot_dis, target) =
+                k_means_lookup_by_dot(vector, &centroids, &centroids_squares);
+            (
+                O::proj(&projection, &O::residual(vector, &centroids[(target,)])),
+                centroid_dot_dis,
+            )
         },
     );
     let projected_centroids = Vec2::from_vec(
@@ -184,6 +195,7 @@ fn from_nothing<O: Op>(
             .flat_map(|x| O::proj(&projection, &centroids[(x,)]))
             .collect(),
     );
+    let projected_centroids_square = centroids_square(&projected_centroids);
     let payloads = MmapArray::create(
         path.as_ref().join("payloads"),
         (0..collection.len()).map(|i| collection.payload(i)),
@@ -193,12 +205,17 @@ fn from_nothing<O: Op>(
         path.as_ref().join("projected_centroids"),
         projected_centroids,
     );
+    let centroids_square = Json::create(
+        path.as_ref().join("centroids_square"),
+        projected_centroids_square,
+    );
     let projection = Json::create(path.as_ref().join("projection"), projection);
     Rabitq {
         storage,
         payloads,
         offsets,
         projected_centroids,
+        centroids_square,
         quantization,
         projection,
     }
@@ -210,18 +227,20 @@ fn open<O: Op>(path: impl AsRef<Path>) -> Rabitq<O> {
     let payloads = MmapArray::open(path.as_ref().join("payloads"));
     let offsets = Json::open(path.as_ref().join("offsets"));
     let projected_centroids = Json::open(path.as_ref().join("projected_centroids"));
+    let centroids_square = Json::open(path.as_ref().join("centroids_square"));
     let projection = Json::open(path.as_ref().join("projection"));
     Rabitq {
         storage,
         quantization,
         payloads,
         offsets,
         projected_centroids,
+        centroids_square,
         projection,
     }
 }
 
-fn select(mut lists: Vec<(f32, usize)>, n: usize) -> Vec<(f32, usize)> {
+fn select<T>(mut lists: Vec<(f32, T)>, n: usize) -> Vec<(f32, T)> {
     if lists.is_empty() || n == 0 {
         return Vec::new();
     }