diff --git a/c_api/IndexIVF_c_ex.cpp b/c_api/IndexIVF_c_ex.cpp index 989264a471..6f8d8a498d 100644 --- a/c_api/IndexIVF_c_ex.cpp +++ b/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(index)->compute_distance_to_codes_for_list( + list_no, x, n, codes, dists); + return 0; + } + CATCH_AND_HANDLE +} diff --git a/c_api/IndexIVF_c_ex.h b/c_api/IndexIVF_c_ex.h index e2effd3525..e62fbd1964 100644 --- a/c_api/IndexIVF_c_ex.h +++ b/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 \ No newline at end of file +#endif diff --git a/faiss/IndexIVF.h b/faiss/IndexIVF.h index 45c65ef839..b7fe7e5777 100644 --- a/faiss/IndexIVF.h +++ b/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(); }; diff --git a/faiss/IndexScalarQuantizer.cpp b/faiss/IndexScalarQuantizer.cpp index e5ed60ee33..ffba2c6405 100644 --- a/faiss/IndexScalarQuantizer.cpp +++ b/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 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 diff --git a/faiss/IndexScalarQuantizer.h b/faiss/IndexScalarQuantizer.h index 27332500c1..c5dd8aff73 100644 --- a/faiss/IndexScalarQuantizer.h +++ b/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 diff --git a/faiss/impl/ScalarQuantizer.cpp b/faiss/impl/ScalarQuantizer.cpp index d013e09035..f9c4f1901c 100644 --- a/faiss/impl/ScalarQuantizer.cpp +++ b/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 * diff --git a/faiss/impl/ScalarQuantizer.h b/faiss/impl/ScalarQuantizer.h index 550a979092..15aa7af09a 100644 --- a/faiss/impl/ScalarQuantizer.h +++ b/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(