Implement distance_to_code for IVFRaBitQFastScanScanner (#4822)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4822

The IVFRaBitQFastScanScanner previously threw an error for distance_to_code, making it incompatible with use cases that need per-code distance evaluation (e.g., custom inverted list scanning, re-ranking).

This implements distance_to_code by delegating to the existing RaBitQDistanceComputer, assuming the caller provides standard-format RaBitQ codes (as produced by encode_vectors), not FastScan-packed codes. This mirrors the approach used by IndexIVFRaBitQ's scanner.

The scanner's set_query and set_list methods now also initialize a FlatCodesDistanceComputer with the appropriate centroid, qb, and centered parameters so that distance_to_code can be called after both are set.

Reviewed By: mdouze

Differential Revision: D93804858

fbshipit-source-id: b0171ba23a97f7f7a744b6bfcb360ec3ff17ae09

faiss/IndexIVFRaBitQFastScan.cpp

@@ -807,7 +807,6 @@ float IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::
 namespace {
 
 /// Provides IVF scanner interface using FastScan's SIMD batch processing.
-/// Note: distance_to_code() not supported due to packed code format.
 struct IVFRaBitQFastScanScanner : InvertedListScanner {
     static constexpr int impl = 10;
     static constexpr size_t nq = 1;
@@ -824,6 +823,9 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
     QueryFactorsData query_factors;
     FastScanDistancePostProcessing context;
 
+    std::unique_ptr<FlatCodesDistanceComputer> dc;
+    std::vector<float> centroid;
+
     IVFRaBitQFastScanScanner(
             const IndexIVFRaBitQFastScan& index,
             bool store_pairs,
@@ -852,12 +854,23 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
 
         index.compute_LUT_uint8(
                 1, xi, cq, dis_tables, biases, &normalizers[0], context);
+
+        // Set up distance computer for distance_to_code
+        centroid.resize(index.d);
+        index.quantizer->reconstruct(list_no, centroid.data());
+        dc.reset(index.rabitq.get_distance_computer(
+                index.qb, centroid.data(), index.centered));
+        dc->set_query(xi);
     }
 
-    float distance_to_code(const uint8_t* /* code */) const override {
-        FAISS_THROW_MSG("distance_to_code not implemented for FastScan");
+    float distance_to_code(const uint8_t* code) const override {
+        FAISS_THROW_IF_NOT_MSG(
+                dc,
+                "set_query and set_list must be called before distance_to_code");
+        return dc->distance_to_code(code);
     }
 
+   public:
     size_t scan_codes(
             size_t ntotal,
             const uint8_t* codes,

tests/test_fast_scan_distance_to_code.cpp

@@ -182,3 +182,27 @@ TEST(IVFPQFastScan, SearchAndEncodeIVFPQ_IP) {
 TEST(IVFPQFastScan, SearchAndEncodeIVFPQFastScan_IP) {
     test_search_and_encode("IVF32,PQ16x4fs", METRIC_INNER_PRODUCT);
 }
+
+TEST(IVFRaBitQFastScan, SearchAndEncodeIVFRaBitQFastScan_L2) {
+    test_search_and_encode("IVF32,RaBitQfs", METRIC_L2);
+}
+
+TEST(IVFRaBitQFastScan, SearchAndEncodeIVFRaBitQFastScan_IP) {
+    test_search_and_encode("IVF32,RaBitQfs", METRIC_INNER_PRODUCT);
+}
+
+TEST(IVFRaBitQFastScan, SearchAndEncodeIVFRaBitQFastScan2_L2) {
+    test_search_and_encode("IVF32,RaBitQfs2", METRIC_L2);
+}
+
+TEST(IVFRaBitQFastScan, SearchAndEncodeIVFRaBitQFastScan2_IP) {
+    test_search_and_encode("IVF32,RaBitQfs2", METRIC_INNER_PRODUCT);
+}
+
+TEST(IVFRaBitQFastScan, SearchAndEncodeIVFRaBitQFastScan4_L2) {
+    test_search_and_encode("IVF32,RaBitQfs4", METRIC_L2);
+}
+
+TEST(IVFRaBitQFastScan, SearchAndEncodeIVFRaBitQFastScan4_IP) {
+    test_search_and_encode("IVF32,RaBitQfs4", METRIC_INNER_PRODUCT);
+}