refactor to extra 2 fp32 values per vector, add unit tests

Signed-off-by: Alexandr Guzhva <[email protected]>

Author: alexanderguzhva

faiss/IndexIVF.h

@@ -313,7 +313,8 @@ struct IndexIVF : Index, IndexIVFInterface {
     /** Get a scanner for this index (store_pairs means ignore labels)
      *
      * The default search implementation uses this to compute the distances.
-     * Use sel instead of params->sel, because sel can get overriden.
+     * Use sel instead of params->sel, because sel is initialized with
+     * params->sel, but may get overriden by IndexIVF's internal logic.
      */
     virtual InvertedListScanner* get_InvertedListScanner(
             bool store_pairs = false,

faiss/IndexIVFRaBitQ.cpp

@@ -17,7 +17,7 @@ IndexIVFRaBitQ::IndexIVFRaBitQ(
         const size_t d,
         const size_t nlist,
         MetricType metric)
-        : IndexIVF(quantizer, d, nlist, 0, metric), rabitq(d) {
+        : IndexIVF(quantizer, d, nlist, 0, metric), rabitq(d, metric) {
     code_size = rabitq.code_size;
     invlists->code_size = code_size;
     is_trained = false;
@@ -120,7 +120,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     std::vector<float> reconstructed_centroid;
     std::vector<float> query_vector;
 
-    std::unique_ptr<RaBitQuantizer::RaBitDistanceComputer> dc;
+    std::unique_ptr<FlatCodesDistanceComputer> dc;
 
     uint8_t qb = 0;
 
@@ -164,7 +164,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
             // both query_vector and centroid are available!
             // set up DistanceComputer
             dc.reset(ivf_rabitq.rabitq.get_distance_computer(
-                    qb, ivf_rabitq.metric_type, reconstructed_centroid.data()));
+                    qb, reconstructed_centroid.data()));
 
             dc->set_query(query_vector.data());
         }
@@ -249,9 +249,8 @@ float IVFRaBitDistanceComputer::operator()(idx_t i) {
     // compute the distance
     float distance = 0;
 
-    std::unique_ptr<RaBitQuantizer::RaBitDistanceComputer> dc(
-            parent->rabitq.get_distance_computer(
-                    parent->qb, parent->metric_type, centroid.data()));
+    std::unique_ptr<FlatCodesDistanceComputer> dc(
+            parent->rabitq.get_distance_computer(parent->qb, centroid.data()));
     dc->set_query(q);
     distance = dc->distance_to_code(code);
 

faiss/IndexRaBitQ.cpp

@@ -8,7 +8,7 @@ namespace faiss {
 IndexRaBitQ::IndexRaBitQ() = default;
 
 IndexRaBitQ::IndexRaBitQ(idx_t d, MetricType metric)
-        : IndexFlatCodes(0, d, metric), rabitq(d) {
+        : IndexFlatCodes(0, d, metric), rabitq(d, metric) {
     code_size = rabitq.code_size;
 
     is_trained = false;
@@ -47,17 +47,17 @@ void IndexRaBitQ::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
 }
 
 FlatCodesDistanceComputer* IndexRaBitQ::get_FlatCodesDistanceComputer() const {
-    RaBitQuantizer::RaBitDistanceComputer* dc =
-            rabitq.get_distance_computer(qb, metric_type, center.data());
+    FlatCodesDistanceComputer* dc =
+            rabitq.get_distance_computer(qb, center.data());
     dc->code_size = rabitq.code_size;
     dc->codes = codes.data();
     return dc;
 }
 
 FlatCodesDistanceComputer* IndexRaBitQ::get_quantized_distance_computer(
         const uint8_t qb) const {
-    RaBitQuantizer::RaBitDistanceComputer* dc =
-            rabitq.get_distance_computer(qb, metric_type, center.data());
+    FlatCodesDistanceComputer* dc =
+            rabitq.get_distance_computer(qb, center.data());
     dc->code_size = rabitq.code_size;
     dc->codes = codes.data();
     return dc;