All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/tmp/faiss/IVFlib.cpp
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
// -*- c++ -*-
#include "IVFlib.h"
#include <memory>
#include "VectorTransform.h"
#include "FaissAssert.h"
namespace faiss { namespace ivflib {
void check_compatible_for_merge (const Index * index0,
const Index * index1)
{
const faiss::IndexPreTransform *pt0 =
dynamic_cast<const faiss::IndexPreTransform *>(index0);
if (pt0) {
const faiss::IndexPreTransform *pt1 =
dynamic_cast<const faiss::IndexPreTransform *>(index1);
FAISS_THROW_IF_NOT_MSG (pt1, "both indexes should be pretransforms");
FAISS_THROW_IF_NOT (pt0->chain.size() == pt1->chain.size());
for (int i = 0; i < pt0->chain.size(); i++) {
FAISS_THROW_IF_NOT (typeid(pt0->chain[i]) == typeid(pt1->chain[i]));
}
index0 = pt0->index;
index1 = pt1->index;
}
FAISS_THROW_IF_NOT (typeid(index0) == typeid(index1));
FAISS_THROW_IF_NOT (index0->d == index1->d &&
index0->metric_type == index1->metric_type);
const faiss::IndexIVF *ivf0 = dynamic_cast<const faiss::IndexIVF *>(index0);
if (ivf0) {
const faiss::IndexIVF *ivf1 =
dynamic_cast<const faiss::IndexIVF *>(index1);
FAISS_THROW_IF_NOT (ivf1);
ivf0->check_compatible_for_merge (*ivf1);
}
// TODO: check as thoroughfully for other index types
}
const IndexIVF * extract_index_ivf (const Index * index)
{
if (auto *pt =
dynamic_cast<const IndexPreTransform *>(index)) {
index = pt->index;
}
auto *ivf = dynamic_cast<const IndexIVF *>(index);
FAISS_THROW_IF_NOT (ivf);
return ivf;
}
IndexIVF * extract_index_ivf (Index * index) {
return const_cast<IndexIVF*> (extract_index_ivf ((const Index*)(index)));
}
void merge_into(faiss::Index *index0, faiss::Index *index1, bool shift_ids) {
check_compatible_for_merge (index0, index1);
IndexIVF * ivf0 = extract_index_ivf (index0);
IndexIVF * ivf1 = extract_index_ivf (index1);
ivf0->merge_from (*ivf1, shift_ids ? ivf0->ntotal : 0);
// useful for IndexPreTransform
index0->ntotal = ivf0->ntotal;
index1->ntotal = ivf1->ntotal;
}
void search_centroid(faiss::Index *index,
const float* x, int n,
idx_t* centroid_ids)
{
std::unique_ptr<float[]> del;
if (auto index_pre = dynamic_cast<faiss::IndexPreTransform*>(index)) {
x = index_pre->apply_chain(n, x);
del.reset((float*)x);
index = index_pre->index;
}
faiss::IndexIVF* index_ivf = dynamic_cast<faiss::IndexIVF*>(index);
assert(index_ivf);
index_ivf->quantizer->assign(n, x, centroid_ids);
}
void search_and_return_centroids(faiss::Index *index,
size_t n,
const float* xin,
long k,
float *distances,
idx_t* labels,
idx_t* query_centroid_ids,
idx_t* result_centroid_ids)
{
const float *x = xin;
std::unique_ptr<float []> del;
if (auto index_pre = dynamic_cast<faiss::IndexPreTransform*>(index)) {
x = index_pre->apply_chain(n, x);
del.reset((float*)x);
index = index_pre->index;
}
faiss::IndexIVF* index_ivf = dynamic_cast<faiss::IndexIVF*>(index);
assert(index_ivf);
size_t nprobe = index_ivf->nprobe;
std::vector<idx_t> cent_nos (n * nprobe);
std::vector<float> cent_dis (n * nprobe);
index_ivf->quantizer->search(
n, x, nprobe, cent_dis.data(), cent_nos.data());
if (query_centroid_ids) {
for (size_t i = 0; i < n; i++)
query_centroid_ids[i] = cent_nos[i * nprobe];
}
index_ivf->search_preassigned (n, x, k,
cent_nos.data(), cent_dis.data(),
distances, labels, true);
for (size_t i = 0; i < n * k; i++) {
idx_t label = labels[i];
if (label < 0) {
if (result_centroid_ids)
result_centroid_ids[i] = -1;
} else {
long list_no = label >> 32;
long list_index = label & 0xffffffff;
if (result_centroid_ids)
result_centroid_ids[i] = list_no;
labels[i] = index_ivf->invlists->get_single_id(list_no, list_index);
}
}
}
SlidingIndexWindow::SlidingIndexWindow (Index *index): index (index) {
n_slice = 0;
IndexIVF* index_ivf = const_cast<IndexIVF*>(extract_index_ivf (index));
ils = dynamic_cast<ArrayInvertedLists *> (index_ivf->invlists);
nlist = ils->nlist;
FAISS_THROW_IF_NOT_MSG (ils,
"only supports indexes with ArrayInvertedLists");
sizes.resize(nlist);
}
template<class T>
static void shift_and_add (std::vector<T> & dst,
size_t remove,
const std::vector<T> & src)
{
if (remove > 0)
memmove (dst.data(), dst.data() + remove,
(dst.size() - remove) * sizeof (T));
size_t insert_point = dst.size() - remove;
dst.resize (insert_point + src.size());
memcpy (dst.data() + insert_point, src.data (), src.size() * sizeof(T));
}
template<class T>
static void remove_from_begin (std::vector<T> & v,
size_t remove)
{
if (remove > 0)
v.erase (v.begin(), v.begin() + remove);
}
void SlidingIndexWindow::step(const Index *sub_index, bool remove_oldest) {
FAISS_THROW_IF_NOT_MSG (!remove_oldest || n_slice > 0,
"cannot remove slice: there is none");
const ArrayInvertedLists *ils2 = nullptr;
if(sub_index) {
check_compatible_for_merge (index, sub_index);
ils2 = dynamic_cast<const ArrayInvertedLists*>(
extract_index_ivf (sub_index)->invlists);
FAISS_THROW_IF_NOT_MSG (ils2, "supports only ArrayInvertedLists");
}
IndexIVF *index_ivf = extract_index_ivf (index);
if (remove_oldest && ils2) {
for (int i = 0; i < nlist; i++) {
std::vector<size_t> & sizesi = sizes[i];
size_t amount_to_remove = sizesi[0];
index_ivf->ntotal += ils2->ids[i].size() - amount_to_remove;
shift_and_add (ils->ids[i], amount_to_remove, ils2->ids[i]);
shift_and_add (ils->codes[i], amount_to_remove * ils->code_size,
ils2->codes[i]);
for (int j = 0; j + 1 < n_slice; j++) {
sizesi[j] = sizesi[j + 1] - amount_to_remove;
}
sizesi[n_slice - 1] = ils->ids[i].size();
}
} else if (ils2) {
for (int i = 0; i < nlist; i++) {
index_ivf->ntotal += ils2->ids[i].size();
shift_and_add (ils->ids[i], 0, ils2->ids[i]);
shift_and_add (ils->codes[i], 0, ils2->codes[i]);
sizes[i].push_back(ils->ids[i].size());
}
n_slice++;
} else if (remove_oldest) {
for (int i = 0; i < nlist; i++) {
size_t amount_to_remove = sizes[i][0];
index_ivf->ntotal -= amount_to_remove;
remove_from_begin (ils->ids[i], amount_to_remove);
remove_from_begin (ils->codes[i],
amount_to_remove * ils->code_size);
for (int j = 0; j + 1 < n_slice; j++) {
sizes[i][j] = sizes[i][j + 1] - amount_to_remove;
}
sizes[i].resize(sizes[i].size() - 1);
}
n_slice--;
} else {
FAISS_THROW_MSG ("nothing to do???");
}
index->ntotal = index_ivf->ntotal;
}
// Get a subset of inverted lists [i0, i1). Works on IndexIVF's and
// IndexIVF's embedded in a IndexPreTransform
ArrayInvertedLists *
get_invlist_range (const Index *index, long i0, long i1)
{
const IndexIVF *ivf = extract_index_ivf (index);
FAISS_THROW_IF_NOT (0 <= i0 && i0 <= i1 && i1 <= ivf->nlist);
const InvertedLists *src = ivf->invlists;
ArrayInvertedLists * il = new ArrayInvertedLists(i1 - i0, src->code_size);
for (long i = i0; i < i1; i++) {
il->add_entries(i - i0, src->list_size(i),
InvertedLists::ScopedIds (src, i).get(),
InvertedLists::ScopedCodes (src, i).get());
}
return il;
}
void set_invlist_range (Index *index, long i0, long i1,
ArrayInvertedLists * src)
{
IndexIVF *ivf = extract_index_ivf (index);
FAISS_THROW_IF_NOT (0 <= i0 && i0 <= i1 && i1 <= ivf->nlist);
ArrayInvertedLists *dst = dynamic_cast<ArrayInvertedLists *>(ivf->invlists);
FAISS_THROW_IF_NOT_MSG (dst, "only ArrayInvertedLists supported");
FAISS_THROW_IF_NOT (src->nlist == i1 - i0 &&
dst->code_size == src->code_size);
size_t ntotal = index->ntotal;
for (long i = i0 ; i < i1; i++) {
ntotal -= dst->list_size (i);
ntotal += src->list_size (i - i0);
std::swap (src->codes[i - i0], dst->codes[i]);
std::swap (src->ids[i - i0], dst->ids[i]);
}
ivf->ntotal = index->ntotal = ntotal;
}
void search_with_parameters (const Index *index,
idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
IVFSearchParameters *params)
{
FAISS_THROW_IF_NOT (params);
const float *prev_x = x;
ScopeDeleter<float> del;
if (auto ip = dynamic_cast<const IndexPreTransform *> (index)) {
x = ip->apply_chain (n, x);
if (x != prev_x) {
del.set(x);
}
index = ip->index;
}
std::vector<idx_t> Iq(params->nprobe * n);
std::vector<float> Dq(params->nprobe * n);
const IndexIVF *index_ivf = dynamic_cast<const IndexIVF *>(index);
FAISS_THROW_IF_NOT (index_ivf);
index_ivf->quantizer->search(n, x, params->nprobe,
Dq.data(), Iq.data());
index_ivf->search_preassigned(n, x, k, Iq.data(), Dq.data(),
distances, labels,
false, params);
}
} } // namespace faiss::ivflib
faiss::IndexIVF
Definition: IndexIVF.h:91
faiss::IndexPreTransform::index
Index * index
! chain of tranforms
Definition: VectorTransform.h:256
faiss::IndexIVF::search_preassigned
virtual void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const IVFSearchParameters *params=nullptr) const
Definition: IndexIVF.cpp:189
faiss::ArrayInvertedLists
simple (default) implementation as an array of inverted lists
Definition: InvertedLists.h:168
faiss::IndexIVF::check_compatible_for_merge
void check_compatible_for_merge(const IndexIVF &other) const
Definition: IndexIVF.cpp:461
faiss::IndexIVF::nprobe
size_t nprobe
number of probes at query time
Definition: IndexIVF.h:98
faiss::Index::assign
void assign(idx_t n, const float *x, idx_t *labels, idx_t k=1)
Definition: Index.cpp:35
faiss::InvertedLists::list_size
virtual size_t list_size(size_t list_no) const =0
get the size of a list
faiss::ivflib::SlidingIndexWindow::nlist
size_t nlist
same as index-&gt;nlist
Definition: IVFlib.h:95
faiss::IndexPreTransform
Definition: VectorTransform.h:253
faiss::InvertedLists
Definition: InvertedLists.h:32
faiss::InvertedLists::get_single_id
virtual idx_t get_single_id(size_t list_no, size_t offset) const
Definition: InvertedLists.cpp:36
faiss::InvertedLists::code_size
size_t code_size
code size per vector in bytes
Definition: InvertedLists.h:36
faiss::ivflib::SlidingIndexWindow::ils
ArrayInvertedLists * ils
InvertedLists of index.
Definition: IVFlib.h:89
faiss::ivflib::SlidingIndexWindow::n_slice
int n_slice
number of slices currently in index
Definition: IVFlib.h:92
faiss::ivflib::SlidingIndexWindow::sizes
std::vector< std::vector< size_t > > sizes
cumulative list sizes at each slice
Definition: IVFlib.h:98
faiss::ivflib::SlidingIndexWindow::SlidingIndexWindow
SlidingIndexWindow(Index *index)
index should be initially empty and trained
Definition: IVFlib.cpp:160
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::Index::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
faiss::ivflib::SlidingIndexWindow::step
void step(const Index *sub_index, bool remove_oldest)
Definition: IVFlib.cpp:191
faiss::InvertedLists::nlist
size_t nlist
number of possible key values
Definition: InvertedLists.h:35
faiss::IndexIVF::invlists
InvertedLists * invlists
Acess to the actual data.
Definition: IndexIVF.h:93
faiss::ArrayInvertedLists::ids
std::vector< std::vector< idx_t > > ids
Inverted lists for indexes.
Definition: InvertedLists.h:170
faiss::Index
Definition: Index.h:62
faiss::Level1Quantizer::quantizer
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:33
faiss::InvertedLists::ScopedIds
Definition: InvertedLists.h:125
faiss::InvertedLists::ScopedCodes
Definition: InvertedLists.h:144