Skip to content

Commit 68c3c36

Browse files
alexanderguzhvaalexanderguzhva
committed
modify IndexIVF interface
Signed-off-by: Alexandr Guzhva <[email protected]>
1 parent 2bca503 commit 68c3c36

14 files changed

Lines changed: 30 additions & 40 deletions

‎faiss/IndexIVF.cpp‎

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -455,7 +455,7 @@ void IndexIVF::search_preassigned(
455455
#pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis, nheap)
456456
{
457457
std::unique_ptr<InvertedListScanner> scanner(
458-
get_InvertedListScanner_2(store_pairs, sel, params));
458+
get_InvertedListScanner(store_pairs, sel, params));
459459

460460
/*****************************************************
461461
* Depending on parallel_mode, there are two possible ways
@@ -796,7 +796,7 @@ void IndexIVF::range_search_preassigned(
796796
{
797797
RangeSearchPartialResult pres(result);
798798
std::unique_ptr<InvertedListScanner> scanner(
799-
get_InvertedListScanner_2(store_pairs, sel, params));
799+
get_InvertedListScanner(store_pairs, sel, params));
800800
FAISS_THROW_IF_NOT(scanner.get());
801801
all_pres[omp_get_thread_num()] = &pres;
802802

@@ -912,19 +912,11 @@ void IndexIVF::range_search_preassigned(
912912

913913
InvertedListScanner* IndexIVF::get_InvertedListScanner(
914914
bool /*store_pairs*/,
915-
const IDSelector* /* sel */) const {
915+
const IDSelector* /* sel */,
916+
const IVFSearchParameters* /* params */) const {
916917
FAISS_THROW_MSG("get_InvertedListScanner not implemented");
917918
}
918919

919-
InvertedListScanner* IndexIVF::get_InvertedListScanner_2(
920-
bool store_pairs,
921-
const IDSelector* sel,
922-
const IVFSearchParameters* /* search_params */) const {
923-
// By default, search_params are ignored, and the call
924-
// is forwarded to get_InvertedListScanner().
925-
return get_InvertedListScanner(store_pairs, sel);
926-
}
927-
928920
void IndexIVF::reconstruct(idx_t key, float* recons) const {
929921
idx_t lo = direct_map.get(key);
930922
reconstruct_from_offset(lo_listno(lo), lo_offset(lo), recons);

‎faiss/IndexIVF.h‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -311,26 +311,14 @@ struct IndexIVF : Index, IndexIVFInterface {
311311
const SearchParameters* params = nullptr) const override;
312312

313313
/** Get a scanner for this index (store_pairs means ignore labels)
314-
*
315-
* The default search implementation uses this to compute the distances
316-
* through a call to get_InvertedListScanner_2().
317-
*/
318-
virtual InvertedListScanner* get_InvertedListScanner(
319-
bool store_pairs = false,
320-
const IDSelector* sel = nullptr) const;
321-
322-
/** Get a scanner for this index (store_pairs means ignore labels).
323-
* This extended version allows passing custom parameters for the search
324-
* into a generated InvertedListScanner object instance.
325314
*
326315
* The default search implementation uses this to compute the distances.
327-
* The default implementation of this function just forward the call
328-
* to get_InvertedListScanner() function.
316+
* Use sel instead of params->sel, because sel can get overriden.
329317
*/
330-
virtual InvertedListScanner* get_InvertedListScanner_2(
318+
virtual InvertedListScanner* get_InvertedListScanner(
331319
bool store_pairs = false,
332320
const IDSelector* sel = nullptr,
333-
const IVFSearchParameters* search_params = nullptr) const;
321+
const IVFSearchParameters* params = nullptr) const;
334322

335323
/** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2
336324
*/

‎faiss/IndexIVFAdditiveQuantizer.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,8 @@ struct AQInvertedListScannerLUT : AQInvertedListScanner {
253253

254254
InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
255255
bool store_pairs,
256-
const IDSelector* sel) const {
256+
const IDSelector* sel,
257+
const IVFSearchParameters*) const {
257258
FAISS_THROW_IF_NOT(!sel);
258259
if (metric_type == METRIC_INNER_PRODUCT) {
259260
if (aq->search_type == AdditiveQuantizer::ST_decompress) {

‎faiss/IndexIVFAdditiveQuantizer.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,8 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
5252

5353
InvertedListScanner* get_InvertedListScanner(
5454
bool store_pairs,
55-
const IDSelector* sel) const override;
55+
const IDSelector* sel,
56+
const IVFSearchParameters* params) const override;
5657

5758
void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
5859

‎faiss/IndexIVFFlat.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,8 @@ InvertedListScanner* get_InvertedListScanner1(
223223

224224
InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
225225
bool store_pairs,
226-
const IDSelector* sel) const {
226+
const IDSelector* sel,
227+
const IVFSearchParameters*) const {
227228
if (sel) {
228229
return get_InvertedListScanner1<true>(this, store_pairs, sel);
229230
} else {

‎faiss/IndexIVFFlat.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@ struct IndexIVFFlat : IndexIVF {
4444

4545
InvertedListScanner* get_InvertedListScanner(
4646
bool store_pairs,
47-
const IDSelector* sel) const override;
47+
const IDSelector* sel,
48+
const IVFSearchParameters* params) const override;
4849

4950
void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
5051
const override;

‎faiss/IndexIVFPQ.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1321,7 +1321,8 @@ InvertedListScanner* get_InvertedListScanner2(
13211321

13221322
InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
13231323
bool store_pairs,
1324-
const IDSelector* sel) const {
1324+
const IDSelector* sel,
1325+
const IVFSearchParameters*) const {
13251326
if (sel) {
13261327
return get_InvertedListScanner2<true>(*this, store_pairs, sel);
13271328
} else {

‎faiss/IndexIVFPQ.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,8 @@ struct IndexIVFPQ : IndexIVF {
134134

135135
InvertedListScanner* get_InvertedListScanner(
136136
bool store_pairs,
137-
const IDSelector* sel) const override;
137+
const IDSelector* sel,
138+
const IVFSearchParameters* params) const override;
138139

139140
/// build precomputed table
140141
void precompute_table();

‎faiss/IndexIVFRaBitQ.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
171171
}
172172
};
173173

174-
InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner_2(
174+
InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
175175
bool store_pairs,
176176
const IDSelector* sel,
177177
const IVFSearchParameters* search_params_in) const {

‎faiss/IndexIVFRaBitQ.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,10 @@ struct IndexIVFRaBitQ : IndexIVF {
4646
const idx_t* precomputed_idx,
4747
void* inverted_list_context = nullptr) override;
4848

49-
InvertedListScanner* get_InvertedListScanner_2(
49+
InvertedListScanner* get_InvertedListScanner(
5050
bool store_pairs,
5151
const IDSelector* sel,
52-
const IVFSearchParameters* search_params) const override;
52+
const IVFSearchParameters* params) const override;
5353

5454
void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
5555
const override;

0 commit comments

Comments
 (0)