diff --git a/c_api/IndexIVF_c_ex.cpp b/c_api/IndexIVF_c_ex.cpp index 5b72e15c94..39041a5265 100644 --- a/c_api/IndexIVF_c_ex.cpp +++ b/c_api/IndexIVF_c_ex.cpp @@ -31,4 +31,39 @@ int faiss_SearchParametersIVF_new_with_sel( *p_sp = reinterpret_cast(sp); } CATCH_AND_HANDLE +} + +int faiss_Search_closest_eligible_centroids( + FaissIndex* index, + int n, + float* query, + int k, + float* centroid_distances, + idx_t* centroid_ids) { + try { + faiss::IndexIVF* index_ivf = reinterpret_cast(index); + assert(index_ivf); + + index_ivf->quantizer->search(n, query, k, centroid_distances, centroid_ids); + } + CATCH_AND_HANDLE +} + +int faiss_IndexIVF_search_preassigned_with_params( + const FaissIndexIVF* index, + idx_t n, + const float* x, + idx_t k, + const idx_t* assign, + const float* centroid_dis, + float* distances, + idx_t* labels, + int store_pairs, + const FaissSearchParametersIVF* params) { + try { + reinterpret_cast(index)->search_preassigned( + n, x, k, assign, centroid_dis, distances, labels, store_pairs, + reinterpret_cast(params)); + } + CATCH_AND_HANDLE } \ No newline at end of file diff --git a/c_api/IndexIVF_c_ex.h b/c_api/IndexIVF_c_ex.h index e82c5c106f..1aa9378b66 100644 --- a/c_api/IndexIVF_c_ex.h +++ b/c_api/IndexIVF_c_ex.h @@ -26,6 +26,47 @@ int faiss_IndexIVF_set_direct_map( int faiss_SearchParametersIVF_new_with_sel( FaissSearchParametersIVF** p_sp, FaissIDSelector* sel); +/* + Return 'k' centroids in the index closest to the query vector. + + @param n: number of queries. + @param query: query vector. + @param k: count of closest number of vectors. + @param centroid_distances: output distances, size n * k. + @param centroid_ids: output centroid IDs, size n * k. +*/ +int faiss_Search_closest_eligible_centroids( + FaissIndex* index, + int n, + float* query, + int k, + float* centroid_distances, + idx_t* centroid_ids +); + +/* + Search the clusters whose IDs are in 'assign' and + return the 'k' nearest neighbours from among them. + + @param n: number of queries. + @param x: query vector, size n * d. + @param k: count of nearest neighbours to be returned for each query. + @param centroid_ids: output centroid IDs, size n * k. + @param distance: output distances, size n * k + @param labels: output labels, size n * k +*/ +int faiss_IndexIVF_search_preassigned_with_params( + const FaissIndexIVF* index, + idx_t n, + const float* x, + idx_t k, + const idx_t* assign, + const float* centroid_dis, + float* distances, + idx_t* labels, + int store_pairs, + const FaissSearchParametersIVF* params); + #ifdef __cplusplus }