fix: rabitq (#556)

Signed-off-by: usamoi <usamoi@outlook.com>

Author: usamoi

Cargo.lock

@@ -118,14 +118,6 @@ impl IndexOptions {
                 }
                 Ok(())
             }
-            QuantizationOptions::RaBitQ(_) => {
-                if !matches!(self.vector.v, VectorKind::Vecf32) {
-                    return Err(ValidationError::new(
-                        "scalar quantization or product quantization is not support for `vector`",
-                    ));
-                }
-                Ok(())
-            }
         }
     }
     fn validate_self(&self) -> Result<(), ValidationError> {
@@ -156,6 +148,18 @@ impl IndexOptions {
                     ));
                 }
             }
+            IndexingOptions::Rabitq(_) => {
+                if !matches!(self.vector.d, DistanceKind::L2) {
+                    return Err(ValidationError::new(
+                        "inverted_index is not support for distance that is not l2",
+                    ));
+                }
+                if !matches!(self.vector.v, VectorKind::Vecf32) {
+                    return Err(ValidationError::new(
+                        "inverted_index is not support for vectors that are not vector",
+                    ));
+                }
+            }
         }
         Ok(())
     }
@@ -289,6 +293,7 @@ pub enum IndexingOptions {
     Ivf(IvfIndexingOptions),
     Hnsw(HnswIndexingOptions),
     InvertedIndex(InvertedIndexingOptions),
+    Rabitq(RabitqIndexingOptions),
 }
 
 impl IndexingOptions {
@@ -310,6 +315,12 @@ impl IndexingOptions {
         };
         x
     }
+    pub fn unwrap_rabitq(self) -> RabitqIndexingOptions {
+        let IndexingOptions::Rabitq(x) = self else {
+            unreachable!()
+        };
+        x
+    }
 }
 
 impl Default for IndexingOptions {
@@ -324,7 +335,8 @@ impl Validate for IndexingOptions {
             Self::Flat(x) => x.validate(),
             Self::Ivf(x) => x.validate(),
             Self::Hnsw(x) => x.validate(),
-            Self::InvertedIndex(_) => Ok(()),
+            Self::InvertedIndex(x) => x.validate(),
+            Self::Rabitq(x) => x.validate(),
         }
     }
 }
@@ -428,15 +440,35 @@ impl Default for HnswIndexingOptions {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
+#[serde(deny_unknown_fields)]
+pub struct RabitqIndexingOptions {
+    #[serde(default = "RabitqIndexingOptions::default_nlist")]
+    #[validate(range(min = 1, max = 1_000_000))]
+    pub nlist: u32,
+}
+
+impl RabitqIndexingOptions {
+    fn default_nlist() -> u32 {
+        1000
+    }
+}
+
+impl Default for RabitqIndexingOptions {
+    fn default() -> Self {
+        Self {
+            nlist: Self::default_nlist(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 #[serde(rename_all = "snake_case")]
 pub enum QuantizationOptions {
     Trivial(TrivialQuantizationOptions),
     Scalar(ScalarQuantizationOptions),
     Product(ProductQuantizationOptions),
-    #[serde(rename = "rabitq")]
-    RaBitQ(RaBitQuantizationOptions),
 }
 
 impl Validate for QuantizationOptions {
@@ -445,7 +477,6 @@ impl Validate for QuantizationOptions {
             Self::Trivial(x) => x.validate(),
             Self::Scalar(x) => x.validate(),
             Self::Product(x) => x.validate(),
-            Self::RaBitQ(x) => x.validate(),
         }
     }
 }
@@ -466,16 +497,6 @@ impl Default for TrivialQuantizationOptions {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
-#[serde(deny_unknown_fields)]
-pub struct RaBitQuantizationOptions {}
-
-impl Default for RaBitQuantizationOptions {
-    fn default() -> Self {
-        Self {}
-    }
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize, Validate)]
 #[serde(deny_unknown_fields)]
 #[validate(schema(function = "Self::validate_self"))]
@@ -558,6 +579,9 @@ pub struct SearchOptions {
     #[validate(range(min = 1, max = 65535))]
     pub hnsw_ef_search: u32,
     #[validate(range(min = 1, max = 65535))]
+    pub rabitq_nprobe: u32,
+    pub rabitq_fast_scan: bool,
+    #[validate(range(min = 1, max = 65535))]
     pub diskann_ef_search: u32,
 }
 

crates/base/src/index.rs

@@ -32,15 +32,18 @@ pub trait VectorOwned: Clone + Serialize + for<'a> Deserialize<'a> + 'static {
 }
 
 pub trait VectorBorrowed: Copy + PartialEq + PartialOrd {
+    // it will be depcrated
     type Scalar: ScalarLike;
+
+    // it will be depcrated
+    fn to_vec(&self) -> Vec<Self::Scalar>;
+
     type Owned: VectorOwned<Scalar = Self::Scalar>;
 
     fn own(&self) -> Self::Owned;
 
     fn dims(&self) -> u32;
 
-    fn to_vec(&self) -> Vec<Self::Scalar>;
-
     fn norm(&self) -> F32;
 
     fn operator_dot(self, rhs: Self) -> F32;

crates/base/src/vector/mod.rs

@@ -358,9 +358,12 @@ impl<'a> VectorBorrowed for SVecf32Borrowed<'a> {
         let dims = end - start;
         let s = self.indexes.partition_point(|&x| x < start);
         let e = self.indexes.partition_point(|&x| x < end);
-        let indexes = self.indexes[s..e].iter().map(|x| x - start);
-        let values = &self.values[s..e];
-        Self::Owned::new_checked(dims, indexes.collect::<Vec<_>>(), values.to_vec())
+        let indexes = self.indexes[s..e]
+            .iter()
+            .map(|x| x - start)
+            .collect::<Vec<_>>();
+        let values = self.values[s..e].to_vec();
+        Self::Owned::new_checked(dims, indexes, values)
     }
 }
 

crates/base/src/vector/svecf32.rs

@@ -142,6 +142,8 @@ impl QueryArguments {
             flat_pq_fast_scan: false,
             ivf_sq_fast_scan: false,
             ivf_pq_fast_scan: false,
+            rabitq_fast_scan: true,
+            rabitq_nprobe: self.probe,
         }
     }
 }

crates/cli/src/args.rs

@@ -1,5 +1,7 @@
 use crate::vec2::Vec2;
 use base::operator::{Borrowed, Operator, Owned, Scalar};
+use base::scalar::ScalarLike;
+use base::scalar::F32;
 use base::search::Vectors;
 use base::vector::VectorBorrowed;
 use base::vector::VectorOwned;
@@ -18,6 +20,23 @@ pub fn sample<O: Operator>(vectors: &impl Vectors<O>) -> Vec2<Scalar<O>> {
     samples
 }
 
+pub fn sample_cast<O: Operator>(vectors: &impl Vectors<O>) -> Vec2<F32> {
+    let n = vectors.len();
+    let m = std::cmp::min(SAMPLES as u32, n);
+    let f = super::rand::sample_u32(&mut rand::thread_rng(), n, m);
+    let mut samples = Vec2::zeros((m as usize, vectors.dims() as usize));
+    for i in 0..m {
+        let v = vectors
+            .vector(f[i as usize] as u32)
+            .to_vec()
+            .into_iter()
+            .map(|x| x.to_f())
+            .collect::<Vec<_>>();
+        samples[(i as usize,)].copy_from_slice(&v);
+    }
+    samples
+}
+
 pub fn sample_subvector_transform<O: Operator>(
     vectors: &impl Vectors<O>,
     s: usize,

crates/common/src/sample.rs

@@ -30,6 +30,7 @@ flat = { path = "../flat" }
 hnsw = { path = "../hnsw" }
 inverted = { path = "../inverted" }
 ivf = { path = "../ivf" }
+rabitq = { path = "../rabitq" }
 
 [lints]
 workspace = true

crates/index/Cargo.toml

@@ -6,13 +6,15 @@ use flat::Flat;
 use hnsw::Hnsw;
 use inverted::InvertedIndex;
 use ivf::Ivf;
+use rabitq::Rabitq;
 use std::path::Path;
 
 pub enum SealedIndexing<O: Op> {
     Flat(Flat<O>),
     Ivf(Ivf<O>),
     Hnsw(Hnsw<O>),
     InvertedIndex(InvertedIndex<O>),
+    Rabitq(Rabitq<O>),
 }
 
 impl<O: Op> SealedIndexing<O> {
@@ -28,6 +30,7 @@ impl<O: Op> SealedIndexing<O> {
             IndexingOptions::InvertedIndex(_) => {
                 Self::InvertedIndex(InvertedIndex::create(path, options, source))
             }
+            IndexingOptions::Rabitq(_) => Self::Rabitq(Rabitq::create(path, options, source)),
         }
     }
 
@@ -37,6 +40,7 @@ impl<O: Op> SealedIndexing<O> {
             IndexingOptions::Ivf(_) => Self::Ivf(Ivf::open(path)),
             IndexingOptions::Hnsw(_) => Self::Hnsw(Hnsw::open(path)),
             IndexingOptions::InvertedIndex(_) => Self::InvertedIndex(InvertedIndex::open(path)),
+            IndexingOptions::Rabitq(_) => Self::Rabitq(Rabitq::open(path)),
         }
     }
 
@@ -50,6 +54,7 @@ impl<O: Op> SealedIndexing<O> {
             SealedIndexing::Ivf(x) => x.vbase(vector, opts),
             SealedIndexing::Hnsw(x) => x.vbase(vector, opts),
             SealedIndexing::InvertedIndex(x) => x.vbase(vector, opts),
+            SealedIndexing::Rabitq(x) => x.vbase(vector, opts),
         }
     }
 
@@ -59,6 +64,7 @@ impl<O: Op> SealedIndexing<O> {
             SealedIndexing::Ivf(x) => x.len(),
             SealedIndexing::Hnsw(x) => x.len(),
             SealedIndexing::InvertedIndex(x) => x.len(),
+            SealedIndexing::Rabitq(x) => x.len(),
         }
     }
 
@@ -68,6 +74,7 @@ impl<O: Op> SealedIndexing<O> {
             SealedIndexing::Ivf(x) => x.vector(i),
             SealedIndexing::Hnsw(x) => x.vector(i),
             SealedIndexing::InvertedIndex(x) => x.vector(i),
+            SealedIndexing::Rabitq(x) => x.vector(i),
         }
     }
 
@@ -77,6 +84,7 @@ impl<O: Op> SealedIndexing<O> {
             SealedIndexing::Ivf(x) => x.payload(i),
             SealedIndexing::Hnsw(x) => x.payload(i),
             SealedIndexing::InvertedIndex(x) => x.payload(i),
+            SealedIndexing::Rabitq(x) => x.payload(i),
         }
     }
 }

crates/index/src/indexing/sealed.rs

@@ -29,6 +29,7 @@ use inverted::operator::OperatorInvertedIndex;
 use ivf::operator::OperatorIvf;
 use parking_lot::Mutex;
 use quantization::operator::OperatorQuantization;
+use rabitq::operator::OperatorRabitq;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::collections::HashSet;
@@ -43,12 +44,22 @@ use thiserror::Error;
 use validator::Validate;
 
 pub trait Op:
-    Operator + OperatorQuantization + OperatorStorage + OperatorIvf + OperatorInvertedIndex
+    Operator
+    + OperatorQuantization
+    + OperatorStorage
+    + OperatorIvf
+    + OperatorInvertedIndex
+    + OperatorRabitq
 {
 }
 
 impl<
-        T: Operator + OperatorQuantization + OperatorStorage + OperatorIvf + OperatorInvertedIndex,
+        T: Operator
+            + OperatorQuantization
+            + OperatorStorage
+            + OperatorIvf
+            + OperatorInvertedIndex
+            + OperatorRabitq,
     > Op for T
 {
 }

crates/index/src/lib.rs

@@ -122,6 +122,7 @@ impl<O: Op> SealedSegment<O> {
             SealedIndexing::Ivf(x) => x,
             SealedIndexing::Hnsw(x) => x,
             SealedIndexing::InvertedIndex(x) => x,
+            SealedIndexing::Rabitq(x) => x,
         }
     }
 }