Skip to content

Commit feb467a

Browse files
Michelle Mafacebook-github-bot
Michelle Ma
authored and
facebook-github-bot
committed
Pass in "own_invlists" to ivf index constructor (#4353)
Summary: Pull Request resolved: #4353 Unicorn index maintains its own posting lists and only uses faiss index as quantizer. Therefore invlists allocated within faiss index are unused. Pass in an option "own_invlists" (already a field in IndexIVF) to optionally initializes invlists Reviewed By: mnorris11 Differential Revision: D75012827 fbshipit-source-id: 6dfe0f624747b1abc1b715271f3d2a4dba1b6cc8
1 parent 3d7659e commit feb467a

26 files changed

Lines changed: 283 additions & 123 deletions

‎faiss/IndexBinaryIVF.cpp‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,16 @@
2626

2727
namespace faiss {
2828

29-
IndexBinaryIVF::IndexBinaryIVF(IndexBinary* quantizer, size_t d, size_t nlist)
29+
IndexBinaryIVF::IndexBinaryIVF(
30+
IndexBinary* quantizer,
31+
size_t d,
32+
size_t nlist,
33+
bool own_invlists)
3034
: IndexBinary(d),
31-
invlists(new ArrayInvertedLists(nlist, code_size)),
35+
invlists(
36+
own_invlists ? new ArrayInvertedLists(nlist, code_size)
37+
: nullptr),
38+
own_invlists(own_invlists),
3239
quantizer(quantizer),
3340
nlist(nlist) {
3441
FAISS_THROW_IF_NOT(d == quantizer->d);

‎faiss/IndexBinaryIVF.h‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,11 @@ struct IndexBinaryIVF : IndexBinary {
6868
* identifier. The pointer is borrowed: the quantizer should not
6969
* be deleted while the IndexBinaryIVF is in use.
7070
*/
71-
IndexBinaryIVF(IndexBinary* quantizer, size_t d, size_t nlist);
71+
IndexBinaryIVF(
72+
IndexBinary* quantizer,
73+
size_t d,
74+
size_t nlist,
75+
bool own_invlists = true);
7276

7377
IndexBinaryIVF();
7478

‎faiss/IndexIVF.cpp‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -158,11 +158,14 @@ IndexIVF::IndexIVF(
158158
size_t d,
159159
size_t nlist,
160160
size_t code_size,
161-
MetricType metric)
161+
MetricType metric,
162+
bool own_invlists)
162163
: Index(d, metric),
163164
IndexIVFInterface(quantizer, nlist),
164-
invlists(new ArrayInvertedLists(nlist, code_size)),
165-
own_invlists(true),
165+
invlists(
166+
own_invlists ? new ArrayInvertedLists(nlist, code_size)
167+
: nullptr),
168+
own_invlists(own_invlists),
166169
code_size(code_size) {
167170
FAISS_THROW_IF_NOT(d == quantizer->d);
168171
is_trained = quantizer->is_trained && (quantizer->ntotal == nlist);

‎faiss/IndexIVF.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,8 @@ struct IndexIVF : Index, IndexIVFInterface {
210210
size_t d,
211211
size_t nlist,
212212
size_t code_size,
213-
MetricType metric = METRIC_L2);
213+
MetricType metric = METRIC_L2,
214+
bool own_invlists = true);
214215

215216
void reset() override;
216217

‎faiss/IndexIVFAdditiveQuantizer.cpp‎

Lines changed: 59 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,9 @@ IndexIVFAdditiveQuantizer::IndexIVFAdditiveQuantizer(
2828
Index* quantizer,
2929
size_t d,
3030
size_t nlist,
31-
MetricType metric)
32-
: IndexIVF(quantizer, d, nlist, 0, metric), aq(aq) {
31+
MetricType metric,
32+
bool own_invlists)
33+
: IndexIVF(quantizer, d, nlist, 0, metric, own_invlists), aq(aq) {
3334
by_residual = true;
3435
}
3536

@@ -301,10 +302,20 @@ IndexIVFResidualQuantizer::IndexIVFResidualQuantizer(
301302
size_t nlist,
302303
const std::vector<size_t>& nbits,
303304
MetricType metric,
304-
Search_type_t search_type)
305-
: IndexIVFAdditiveQuantizer(&rq, quantizer, d, nlist, metric),
305+
Search_type_t search_type,
306+
bool own_invlists)
307+
: IndexIVFAdditiveQuantizer(
308+
&rq,
309+
quantizer,
310+
d,
311+
nlist,
312+
metric,
313+
own_invlists),
306314
rq(d, nbits, search_type) {
307-
code_size = invlists->code_size = rq.code_size;
315+
code_size = rq.code_size;
316+
if (invlists) {
317+
invlists->code_size = code_size;
318+
}
308319
}
309320

310321
IndexIVFResidualQuantizer::IndexIVFResidualQuantizer()
@@ -317,14 +328,16 @@ IndexIVFResidualQuantizer::IndexIVFResidualQuantizer(
317328
size_t M, /* number of subquantizers */
318329
size_t nbits, /* number of bit per subvector index */
319330
MetricType metric,
320-
Search_type_t search_type)
331+
Search_type_t search_type,
332+
bool own_invlists)
321333
: IndexIVFResidualQuantizer(
322334
quantizer,
323335
d,
324336
nlist,
325337
std::vector<size_t>(M, nbits),
326338
metric,
327-
search_type) {}
339+
search_type,
340+
own_invlists) {}
328341

329342
IndexIVFResidualQuantizer::~IndexIVFResidualQuantizer() = default;
330343

@@ -339,10 +352,20 @@ IndexIVFLocalSearchQuantizer::IndexIVFLocalSearchQuantizer(
339352
size_t M, /* number of subquantizers */
340353
size_t nbits, /* number of bit per subvector index */
341354
MetricType metric,
342-
Search_type_t search_type)
343-
: IndexIVFAdditiveQuantizer(&lsq, quantizer, d, nlist, metric),
355+
Search_type_t search_type,
356+
bool own_invlists)
357+
: IndexIVFAdditiveQuantizer(
358+
&lsq,
359+
quantizer,
360+
d,
361+
nlist,
362+
metric,
363+
own_invlists),
344364
lsq(d, M, nbits, search_type) {
345-
code_size = invlists->code_size = lsq.code_size;
365+
code_size = lsq.code_size;
366+
if (invlists) {
367+
invlists->code_size = code_size;
368+
}
346369
}
347370

348371
IndexIVFLocalSearchQuantizer::IndexIVFLocalSearchQuantizer()
@@ -362,10 +385,20 @@ IndexIVFProductResidualQuantizer::IndexIVFProductResidualQuantizer(
362385
size_t Msub,
363386
size_t nbits,
364387
MetricType metric,
365-
Search_type_t search_type)
366-
: IndexIVFAdditiveQuantizer(&prq, quantizer, d, nlist, metric),
388+
Search_type_t search_type,
389+
bool own_invlists)
390+
: IndexIVFAdditiveQuantizer(
391+
&prq,
392+
quantizer,
393+
d,
394+
nlist,
395+
metric,
396+
own_invlists),
367397
prq(d, nsplits, Msub, nbits, search_type) {
368-
code_size = invlists->code_size = prq.code_size;
398+
code_size = prq.code_size;
399+
if (invlists) {
400+
invlists->code_size = code_size;
401+
}
369402
}
370403

371404
IndexIVFProductResidualQuantizer::IndexIVFProductResidualQuantizer()
@@ -385,10 +418,20 @@ IndexIVFProductLocalSearchQuantizer::IndexIVFProductLocalSearchQuantizer(
385418
size_t Msub,
386419
size_t nbits,
387420
MetricType metric,
388-
Search_type_t search_type)
389-
: IndexIVFAdditiveQuantizer(&plsq, quantizer, d, nlist, metric),
421+
Search_type_t search_type,
422+
bool own_invlists)
423+
: IndexIVFAdditiveQuantizer(
424+
&plsq,
425+
quantizer,
426+
d,
427+
nlist,
428+
metric,
429+
own_invlists),
390430
plsq(d, nsplits, Msub, nbits, search_type) {
391-
code_size = invlists->code_size = plsq.code_size;
431+
code_size = plsq.code_size;
432+
if (invlists) {
433+
invlists->code_size = code_size;
434+
}
392435
}
393436

394437
IndexIVFProductLocalSearchQuantizer::IndexIVFProductLocalSearchQuantizer()

‎faiss/IndexIVFAdditiveQuantizer.h‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
3535
Index* quantizer,
3636
size_t d,
3737
size_t nlist,
38-
MetricType metric = METRIC_L2);
38+
MetricType metric = METRIC_L2,
39+
bool own_invlists = true);
3940

4041
explicit IndexIVFAdditiveQuantizer(AdditiveQuantizer* aq);
4142

@@ -82,7 +83,8 @@ struct IndexIVFResidualQuantizer : IndexIVFAdditiveQuantizer {
8283
size_t nlist,
8384
const std::vector<size_t>& nbits,
8485
MetricType metric = METRIC_L2,
85-
Search_type_t search_type = AdditiveQuantizer::ST_decompress);
86+
Search_type_t search_type = AdditiveQuantizer::ST_decompress,
87+
bool own_invlists = true);
8688

8789
IndexIVFResidualQuantizer(
8890
Index* quantizer,
@@ -91,7 +93,8 @@ struct IndexIVFResidualQuantizer : IndexIVFAdditiveQuantizer {
9193
size_t M, /* number of subquantizers */
9294
size_t nbits, /* number of bit per subvector index */
9395
MetricType metric = METRIC_L2,
94-
Search_type_t search_type = AdditiveQuantizer::ST_decompress);
96+
Search_type_t search_type = AdditiveQuantizer::ST_decompress,
97+
bool own_invlists = true);
9598

9699
IndexIVFResidualQuantizer();
97100

@@ -118,7 +121,8 @@ struct IndexIVFLocalSearchQuantizer : IndexIVFAdditiveQuantizer {
118121
size_t M, /* number of subquantizers */
119122
size_t nbits, /* number of bit per subvector index */
120123
MetricType metric = METRIC_L2,
121-
Search_type_t search_type = AdditiveQuantizer::ST_decompress);
124+
Search_type_t search_type = AdditiveQuantizer::ST_decompress,
125+
bool own_invlists = true);
122126

123127
IndexIVFLocalSearchQuantizer();
124128

@@ -147,7 +151,8 @@ struct IndexIVFProductResidualQuantizer : IndexIVFAdditiveQuantizer {
147151
size_t Msub,
148152
size_t nbits,
149153
MetricType metric = METRIC_L2,
150-
Search_type_t search_type = AdditiveQuantizer::ST_decompress);
154+
Search_type_t search_type = AdditiveQuantizer::ST_decompress,
155+
bool own_invlists = true);
151156

152157
IndexIVFProductResidualQuantizer();
153158

@@ -176,7 +181,8 @@ struct IndexIVFProductLocalSearchQuantizer : IndexIVFAdditiveQuantizer {
176181
size_t Msub,
177182
size_t nbits,
178183
MetricType metric = METRIC_L2,
179-
Search_type_t search_type = AdditiveQuantizer::ST_decompress);
184+
Search_type_t search_type = AdditiveQuantizer::ST_decompress,
185+
bool own_invlists = true);
180186

181187
IndexIVFProductLocalSearchQuantizer();
182188

‎faiss/IndexIVFAdditiveQuantizerFastScan.cpp‎

Lines changed: 33 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -34,18 +34,20 @@ IndexIVFAdditiveQuantizerFastScan::IndexIVFAdditiveQuantizerFastScan(
3434
size_t d,
3535
size_t nlist,
3636
MetricType metric,
37-
int bbs)
38-
: IndexIVFFastScan(quantizer, d, nlist, 0, metric) {
37+
int bbs,
38+
bool own_invlists)
39+
: IndexIVFFastScan(quantizer, d, nlist, 0, metric, own_invlists) {
3940
if (aq != nullptr) {
40-
init(aq, nlist, metric, bbs);
41+
init(aq, nlist, metric, bbs, own_invlists);
4142
}
4243
}
4344

4445
void IndexIVFAdditiveQuantizerFastScan::init(
4546
AdditiveQuantizer* aq,
4647
size_t nlist,
4748
MetricType metric,
48-
int bbs) {
49+
int bbs,
50+
bool own_invlists) {
4951
FAISS_THROW_IF_NOT(aq != nullptr);
5052
FAISS_THROW_IF_NOT(!aq->nbits.empty());
5153
FAISS_THROW_IF_NOT(aq->nbits[0] == 4);
@@ -66,7 +68,7 @@ void IndexIVFAdditiveQuantizerFastScan::init(
6668
} else {
6769
M = aq->M;
6870
}
69-
init_fastscan(aq, M, 4, nlist, metric, bbs);
71+
init_fastscan(aq, M, 4, nlist, metric, bbs, own_invlists);
7072

7173
max_train_points = 1024 * ksub * M;
7274
by_residual = true;
@@ -80,17 +82,20 @@ IndexIVFAdditiveQuantizerFastScan::IndexIVFAdditiveQuantizerFastScan(
8082
orig.d,
8183
orig.nlist,
8284
0,
83-
orig.metric_type),
85+
orig.metric_type,
86+
orig.own_invlists),
8487
aq(orig.aq) {
8588
FAISS_THROW_IF_NOT(
8689
metric_type == METRIC_INNER_PRODUCT || !orig.by_residual);
8790

88-
init(aq, nlist, metric_type, bbs);
91+
init(aq, nlist, metric_type, bbs, own_invlists);
8992

9093
is_trained = orig.is_trained;
9194
ntotal = orig.ntotal;
9295
nprobe = orig.nprobe;
93-
96+
if (!orig.own_invlists) {
97+
return; // skip packing codes below
98+
}
9499
for (size_t i = 0; i < nlist; i++) {
95100
size_t nb = orig.invlists->list_size(i);
96101
size_t nb2 = roundup(nb, bbs);
@@ -448,17 +453,19 @@ IndexIVFLocalSearchQuantizerFastScan::IndexIVFLocalSearchQuantizerFastScan(
448453
size_t nbits,
449454
MetricType metric,
450455
Search_type_t search_type,
451-
int bbs)
456+
int bbs,
457+
bool own_invlists)
452458
: IndexIVFAdditiveQuantizerFastScan(
453459
quantizer,
454460
nullptr,
455461
d,
456462
nlist,
457463
metric,
458-
bbs),
464+
bbs,
465+
own_invlists),
459466
lsq(d, M, nbits, search_type) {
460467
FAISS_THROW_IF_NOT(nbits == 4);
461-
init(&lsq, nlist, metric, bbs);
468+
init(&lsq, nlist, metric, bbs, own_invlists);
462469
}
463470

464471
IndexIVFLocalSearchQuantizerFastScan::IndexIVFLocalSearchQuantizerFastScan() {
@@ -474,17 +481,19 @@ IndexIVFResidualQuantizerFastScan::IndexIVFResidualQuantizerFastScan(
474481
size_t nbits,
475482
MetricType metric,
476483
Search_type_t search_type,
477-
int bbs)
484+
int bbs,
485+
bool own_invlists)
478486
: IndexIVFAdditiveQuantizerFastScan(
479487
quantizer,
480488
nullptr,
481489
d,
482490
nlist,
483491
metric,
484-
bbs),
492+
bbs,
493+
own_invlists),
485494
rq(d, M, nbits, search_type) {
486495
FAISS_THROW_IF_NOT(nbits == 4);
487-
init(&rq, nlist, metric, bbs);
496+
init(&rq, nlist, metric, bbs, own_invlists);
488497
}
489498

490499
IndexIVFResidualQuantizerFastScan::IndexIVFResidualQuantizerFastScan() {
@@ -502,17 +511,19 @@ IndexIVFProductLocalSearchQuantizerFastScan::
502511
size_t nbits,
503512
MetricType metric,
504513
Search_type_t search_type,
505-
int bbs)
514+
int bbs,
515+
bool own_invlists)
506516
: IndexIVFAdditiveQuantizerFastScan(
507517
quantizer,
508518
nullptr,
509519
d,
510520
nlist,
511521
metric,
512-
bbs),
522+
bbs,
523+
own_invlists),
513524
plsq(d, nsplits, Msub, nbits, search_type) {
514525
FAISS_THROW_IF_NOT(nbits == 4);
515-
init(&plsq, nlist, metric, bbs);
526+
init(&plsq, nlist, metric, bbs, own_invlists);
516527
}
517528

518529
IndexIVFProductLocalSearchQuantizerFastScan::
@@ -531,17 +542,19 @@ IndexIVFProductResidualQuantizerFastScan::
531542
size_t nbits,
532543
MetricType metric,
533544
Search_type_t search_type,
534-
int bbs)
545+
int bbs,
546+
bool own_invlists)
535547
: IndexIVFAdditiveQuantizerFastScan(
536548
quantizer,
537549
nullptr,
538550
d,
539551
nlist,
540552
metric,
541-
bbs),
553+
bbs,
554+
own_invlists),
542555
prq(d, nsplits, Msub, nbits, search_type) {
543556
FAISS_THROW_IF_NOT(nbits == 4);
544-
init(&prq, nlist, metric, bbs);
557+
init(&prq, nlist, metric, bbs, own_invlists);
545558
}
546559

547560
IndexIVFProductResidualQuantizerFastScan::

0 commit comments

Comments
 (0)