15 #include <gtest/gtest.h>
17 #include <faiss/IndexIVF.h>
18 #include <faiss/AutoTune.h>
19 #include <faiss/index_io.h>
20 #include <faiss/IVFlib.h>
22 using namespace faiss;
49 std::vector<float> make_data(
size_t n)
51 std::vector <float> database (n * d);
52 for (
size_t i = 0; i < n * d; i++) {
53 database[i] = drand48();
58 std::unique_ptr<Index> make_trained_index(
const char *index_type)
60 auto index = std::unique_ptr<Index>(
index_factory(d, index_type));
61 auto xt = make_data(nt * d);
62 index->
train(nt, xt.data());
67 std::vector<idx_t> search_index(
Index *index,
const float *xq) {
69 std::vector<idx_t> I(k * nq);
70 std::vector<float> D(k * nq);
71 index->
search (nq, xq, k, D.data(), I.data());
92 static void * memdup (
const void *m,
size_t size) {
93 if (size == 0)
return nullptr;
94 return memcpy (malloc(size), m, size);
97 size_t list_size(
size_t list_no)
const override {
98 return il->list_size (list_no);
101 const uint8_t * get_codes (
size_t list_no)
const override {
102 return (uint8_t*)memdup (il->get_codes(list_no),
103 list_size(list_no) * code_size);
106 const idx_t * get_ids (
size_t list_no)
const override {
107 return (idx_t*)memdup (il->get_ids(list_no),
108 list_size(list_no) *
sizeof(idx_t));
111 void release_codes (
const uint8_t *codes)
const override {
115 void release_ids (
const idx_t *ids)
const override {
119 const uint8_t * get_single_code (
size_t list_no,
size_t offset)
121 return (uint8_t*)memdup (il->get_single_code(list_no, offset),
125 virtual size_t add_entries (
127 const idx_t*,
const uint8_t *)
override
129 assert(!
"not implemented");
134 virtual void update_entries (
size_t ,
size_t ,
size_t ,
135 const idx_t *,
const uint8_t *)
137 assert(!
"not implemented");
140 virtual void resize (
size_t ,
size_t )
override {
141 assert(!
"not implemented");
144 ~EncapsulateInvertedLists () {}
149 int test_dealloc_invlists (
const char *index_key) {
151 std::unique_ptr<Index> index = make_trained_index(index_key);
152 IndexIVF * index_ivf = ivflib::extract_index_ivf (index.get());
154 auto xb = make_data (nb * d);
155 index->add(nb, xb.data());
157 auto xq = make_data (nq * d);
159 auto ref_res = search_index (index.get(), xq.data());
161 EncapsulateInvertedLists eil(index_ivf->
invlists);
163 index_ivf->own_invlists =
false;
167 auto new_res = search_index (index.get(), xq.data());
173 EXPECT_EQ (ref_res, new_res);
185 TEST(TestIvlistDealloc, IVFFlat) {
186 test_dealloc_invlists (
"IVF32,Flat");
189 TEST(TestIvlistDealloc, IVFSQ) {
190 test_dealloc_invlists (
"IVF32,SQ8");
193 TEST(TestIvlistDealloc, IVFPQ) {
194 test_dealloc_invlists (
"IVF32,PQ4np");
void train(idx_t n, const float *x) override
long idx_t
all indices are this type
void replace_invlists(InvertedLists *il, bool own=false)
replace the inverted lists, old one is deallocated if own_invlists
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
virtual void set_index_parameter(Index *index, const std::string &name, double val) const
set one of the parameters
InvertedLists * invlists
Acess to the actual data.
Index * index_factory(int d, const char *description_in, MetricType metric)