Commit 8898eab
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, ¶ms);
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: ca7f03f5a474a59089ebdf9685fb83e54ae198b01 parent ebb5f84 commit 8898eab
1 file changed
+10
-3
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
146 | 146 | | |
147 | 147 | | |
148 | 148 | | |
149 | | - | |
150 | | - | |
151 | | - | |
| 149 | + | |
| 150 | + | |
| 151 | + | |
| 152 | + | |
| 153 | + | |
| 154 | + | |
| 155 | + | |
| 156 | + | |
| 157 | + | |
| 158 | + | |
152 | 159 | | |
153 | 160 | | |
154 | 161 | | |
| |||
0 commit comments