Skip to content

Commit 8898eab

Browse files
yuhaijun999facebook-github-bot
authored andcommitted
Add feature in IndexIDMap.cpp range_search with Parameters. (facebookresearch#3213)
Summary: for example: #include <cstdio> #include <cstdlib> #include <iostream> #include <random> #include <faiss/IndexFlat.h> #include <faiss/IndexIDMap.h> #include <faiss/impl/AuxIndexStructures.h> #include <faiss/impl/IDSelector.h> // 64-bit int using idx_t = faiss::idx_t; int main() { int d = 64; // dimension int nb = 10000; // database size int nq = 5; // nb of queries std::mt19937 rng; std::uniform_real_distribution<> distrib; float* xb = new float[d * nb]; float* xq = new float[d * nq]; for (int i = 0; i < nb; i++) { for (int j = 0; j < d; j++) xb[d * i + j] = distrib(rng); xb[d * i] += i / 1000.; } for (int i = 0; i < nq; i++) { for (int j = 0; j < d; j++) xq[d * i + j] = distrib(rng); xq[d * i] += i / 1000.; } faiss::IndexFlatL2 index(d); faiss::IndexIDMap2 index_id_map2(&index); idx_t* xids = new idx_t[nb](); for (int i = 0; i < nb; i++) { xids[i] = i + nb; // add ids } index_id_map2.add_with_ids(nb, xb, xids); faiss::SearchParameters params; std::vector<faiss::idx_t> ids; ids.reserve(nb / 2); for (faiss::idx_t i = 0; i < nb / 2; i++) { ids.push_back(i + nb); // search ids } faiss::IDSelectorArray id_selector_array(ids.size(), ids.data()); params.sel = &id_selector_array; // range search with param { float radius = 7.0f; faiss::RangeSearchResult* result = new faiss::RangeSearchResult(nq); index_id_map2.range_search(nq, xb, radius, result, &params); size_t off = 0; for (size_t i = 0; i < result->nq; i++) { size_t n = (result->lims[i + 1] - result->lims[i]); std::cout << "i : " << i << std::endl; for (size_t j = 0; j < n; j++) { std::cout << "\t label : " << result->labels[off + j] << " distance : " << result->distances[off + j] << std::endl; } off += n; } delete result; } delete[] xb; delete[] xq; delete[] xids; return 0; } outputs: server@dingo11 cpp [main] $ ./6-Range-Search i : 0 label : 10000 distance : 0 i : 1 label : 10001 distance : 0 label : 10136 distance : 6.72638 label : 10183 distance : 6.73293 label : 10223 distance : 6.76569 label : 10555 distance : 6.93339 label : 10995 distance : 5.78548 i : 2 label : 10002 distance : 0 label : 10253 distance : 6.84876 label : 10312 distance : 5.07469 i : 3 label : 10003 distance : 0 label : 10983 distance : 6.77275 i : 4 label : 10004 distance : 0 label : 10112 distance : 6.89793 label : 10403 distance : 6.84196 Pull Request resolved: facebookresearch#3213 Reviewed By: mdouze Differential Revision: D53704072 Pulled By: algoriddle fbshipit-source-id: ca7f03f5a474a59089ebdf9685fb83e54ae198b0
1 parent ebb5f84 commit 8898eab

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

faiss/IndexIDMap.cpp

+10-3
Original file line numberDiff line numberDiff line change
@@ -146,9 +146,16 @@ void IndexIDMapTemplate<IndexT>::range_search(
146146
typename IndexT::distance_t radius,
147147
RangeSearchResult* result,
148148
const SearchParameters* params) const {
149-
FAISS_THROW_IF_NOT_MSG(
150-
!params, "search params not supported for this index");
151-
index->range_search(n, x, radius, result);
149+
if (params) {
150+
SearchParameters internal_search_parameters;
151+
IDSelectorTranslated id_selector_translated(id_map, params->sel);
152+
internal_search_parameters.sel = &id_selector_translated;
153+
154+
index->range_search(n, x, radius, result, &internal_search_parameters);
155+
} else {
156+
index->range_search(n, x, radius, result);
157+
}
158+
152159
#pragma omp parallel for
153160
for (idx_t i = 0; i < result->lims[result->nq]; i++) {
154161
result->labels[i] = result->labels[i] < 0 ? result->labels[i]

0 commit comments

Comments
 (0)