diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp index cdb2a7c1f..de5044728 100644 --- a/faiss/IndexIVFRaBitQFastScan.cpp +++ b/faiss/IndexIVFRaBitQFastScan.cpp @@ -807,7 +807,6 @@ float IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler:: 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 dc; + std::vector 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, diff --git a/tests/test_fast_scan_distance_to_code.cpp b/tests/test_fast_scan_distance_to_code.cpp index 1af25d577..258258f64 100644 --- a/tests/test_fast_scan_distance_to_code.cpp +++ b/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); +}