MB-61093 Add method to compute distance from codes for IVF index

c_api/IndexIVF_c_ex.cpp

@@ -70,3 +70,18 @@ int faiss_IndexIVF_search_preassigned_with_params(
     }
     CATCH_AND_HANDLE
 }
+
+int faiss_IndexIVF_compute_distance_to_codes_for_list(
+        FaissIndexIVF* index,
+        idx_t list_no,
+        const float* x,
+        idx_t n,
+        const uint8_t* codes,
+        float* dists) {
+    try {
+        reinterpret_cast<IndexIVF*>(index)->compute_distance_to_codes_for_list(
+               list_no, x, n, codes, dists);
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}

c_api/IndexIVF_c_ex.h

@@ -68,9 +68,31 @@ int faiss_IndexIVF_search_preassigned_with_params(
         int store_pairs,
         const FaissSearchParametersIVF* params);
 
+/*
+    Given a query vector x, compute distance to provided codes
+    for the input list_no. This is a special purpose method
+    to be used as a flat distance computer for an inverted
+    list where codes are provided externally. This allows to
+    use the quantizer independently while computing distance
+    for the quantized codes.
+
+    @param list_no list number for inverted list
+    @param x - input query vector
+    @param n - number of codes
+    @param codes - input codes
+    @param dists - output computed distances
+*/
+
+int faiss_IndexIVF_compute_distance_to_codes_for_list(
+        FaissIndexIVF* index,
+        idx_t list_no,
+        const float* x,
+        idx_t n,
+        const uint8_t* codes,
+        float* dists);
 
 #ifdef __cplusplus
 }
 #endif
 
-#endif
+#endif

faiss/IndexIVF.h

@@ -435,6 +435,29 @@ struct IndexIVF : Index, IndexIVFInterface {
     size_t sa_code_size() const override;
     void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
 
+
+    /** Given a query vector x, compute distance to provided codes
+     * for the input list_no. This is a special purpose method
+     * to be used as a flat distance computer for an inverted
+     * list where codes are provided externally. This allows to
+     * use the quantizer independently while computing distance
+     * for the quantized codes.
+     *
+     * @param list_no list number for inverted list
+     * @param x - input query vector
+     * @param n - number of codes
+     * @param codes - input codes
+     * @param dists - output computed distances
+     */
+
+    virtual void compute_distance_to_codes_for_list(
+        const idx_t list_no,
+        const float* x,
+        idx_t n,
+        const uint8_t* codes,
+        float* dists) const {};
+
+
     IndexIVF();
 };
 

faiss/IndexScalarQuantizer.cpp

@@ -282,4 +282,30 @@ void IndexIVFScalarQuantizer::reconstruct_from_offset(
     }
 }
 
+void IndexIVFScalarQuantizer::compute_distance_to_codes_for_list(
+        const idx_t list_no,
+        const float* x,
+        idx_t n,
+        const uint8_t* codes,
+        float* dists) const {
+
+    ScalarQuantizer::SQDistanceComputer* dc =
+            sq.get_distance_computer(metric_type);
+    dc->code_size = sq.code_size;
+
+    if (by_residual) {
+        // shift of x_in wrt centroid
+        std::vector<float> tmp(d);
+        quantizer->compute_residual(x, tmp.data(), list_no);
+        dc->set_query(tmp.data());
+    } else {
+        dc->set_query(x);
+    }
+
+    dc->distance_to_codes(n, codes, dists);
+
+    return;
+}
+
+
 } // namespace faiss

faiss/IndexScalarQuantizer.h

@@ -103,6 +103,15 @@ struct IndexIVFScalarQuantizer : IndexIVF {
 
     /* standalone codec interface */
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+
+    void compute_distance_to_codes_for_list(
+        const idx_t list_no,
+        const float* x,
+        idx_t n,
+        const uint8_t* codes,
+        float* dists) const override;
+
 };
 
 } // namespace faiss

faiss/impl/ScalarQuantizer.cpp

@@ -1543,6 +1543,14 @@ SQDistanceComputer* ScalarQuantizer::get_distance_computer(
     }
 }
 
+void SQDistanceComputer::distance_to_codes(idx_t n, const uint8_t* codes, float* dists) {
+    for (idx_t i = 0; i < n; i++) {
+        const uint8_t* code = codes + i * code_size;
+        dists[i] = query_to_code(code);
+    }
+    return;
+}
+
 /*******************************************************************
  * IndexScalarQuantizer/IndexIVFScalarQuantizer scanner object
  *

faiss/impl/ScalarQuantizer.h

@@ -104,6 +104,8 @@ struct ScalarQuantizer : Quantizer {
         float distance_to_code(const uint8_t* code) final {
             return query_to_code(code);
         }
+
+	void distance_to_codes(idx_t n, const uint8_t* codes, float* dists);
     };
 
     SQDistanceComputer* get_distance_computer(