mirror of
https://github.com/facebookresearch/faiss.git
synced 2025-06-03 21:54:02 +08:00
a8118acbc5
Changelog: - changed license: BSD+Patents -> MIT - propagates exceptions raised in sub-indexes of IndexShards and IndexReplicas - support for searching several inverted lists in parallel (parallel_mode != 0) - better support for PQ codes where nbit != 8 or 16 - IVFSpectralHash implementation: spectral hash codes inside an IVF - 6-bit per component scalar quantizer (4 and 8 bit were already supported) - combinations of inverted lists: HStackInvertedLists and VStackInvertedLists - configurable number of threads for OnDiskInvertedLists prefetching (including 0=no prefetch) - more test and demo code compatible with Python 3 (print with parentheses) - refactored benchmark code: data loading is now in a single file
124 lines
3.8 KiB
C++
124 lines
3.8 KiB
C++
/**
|
|
* Copyright (c) Facebook, Inc. and its affiliates.
|
|
*
|
|
* This source code is licensed under the MIT license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
#include "IndexReplicas.h"
|
|
#include "FaissAssert.h"
|
|
|
|
namespace faiss {
|
|
|
|
template <typename IndexT>
|
|
IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(bool threaded)
|
|
: ThreadedIndex<IndexT>(threaded) {
|
|
}
|
|
|
|
template <typename IndexT>
|
|
IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(idx_t d, bool threaded)
|
|
: ThreadedIndex<IndexT>(d, threaded) {
|
|
}
|
|
|
|
template <typename IndexT>
|
|
IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(int d, bool threaded)
|
|
: ThreadedIndex<IndexT>(d, threaded) {
|
|
}
|
|
|
|
template <typename IndexT>
|
|
void
|
|
IndexReplicasTemplate<IndexT>::onAfterAddIndex(IndexT* index) {
|
|
// Make sure that the parameters are the same for all prior indices, unless
|
|
// we're the first index to be added
|
|
if (this->count() > 0 && this->at(0) != index) {
|
|
auto existing = this->at(0);
|
|
|
|
FAISS_THROW_IF_NOT_FMT(index->ntotal == existing->ntotal,
|
|
"IndexReplicas: newly added index does "
|
|
"not have same number of vectors as prior index; "
|
|
"prior index has %ld vectors, new index has %ld",
|
|
existing->ntotal, index->ntotal);
|
|
|
|
FAISS_THROW_IF_NOT_MSG(index->is_trained == existing->is_trained,
|
|
"IndexReplicas: newly added index does "
|
|
"not have same train status as prior index");
|
|
} else {
|
|
// Set our parameters based on the first index we're adding
|
|
// (dimension is handled in ThreadedIndex)
|
|
this->ntotal = index->ntotal;
|
|
this->verbose = index->verbose;
|
|
this->is_trained = index->is_trained;
|
|
this->metric_type = index->metric_type;
|
|
}
|
|
}
|
|
|
|
template <typename IndexT>
|
|
void
|
|
IndexReplicasTemplate<IndexT>::train(idx_t n, const component_t* x) {
|
|
this->runOnIndex([n, x](int, IndexT* index){ index->train(n, x); });
|
|
}
|
|
|
|
template <typename IndexT>
|
|
void
|
|
IndexReplicasTemplate<IndexT>::add(idx_t n, const component_t* x) {
|
|
this->runOnIndex([n, x](int, IndexT* index){ index->add(n, x); });
|
|
this->ntotal += n;
|
|
}
|
|
|
|
template <typename IndexT>
|
|
void
|
|
IndexReplicasTemplate<IndexT>::reconstruct(idx_t n, component_t* x) const {
|
|
FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
|
|
|
|
// Just pass to the first replica
|
|
this->at(0)->reconstruct(n, x);
|
|
}
|
|
|
|
template <typename IndexT>
|
|
void
|
|
IndexReplicasTemplate<IndexT>::search(idx_t n,
|
|
const component_t* x,
|
|
idx_t k,
|
|
distance_t* distances,
|
|
idx_t* labels) const {
|
|
FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
|
|
|
|
if (n == 0) {
|
|
return;
|
|
}
|
|
|
|
auto dim = this->d;
|
|
size_t componentsPerVec =
|
|
sizeof(component_t) == 1 ? (dim + 7) / 8 : dim;
|
|
|
|
// Partition the query by the number of indices we have
|
|
faiss::Index::idx_t queriesPerIndex =
|
|
(faiss::Index::idx_t) (n + this->count() - 1) /
|
|
(faiss::Index::idx_t) this->count();
|
|
FAISS_ASSERT(n / queriesPerIndex <= this->count());
|
|
|
|
auto fn =
|
|
[queriesPerIndex, componentsPerVec,
|
|
n, x, k, distances, labels](int i, const IndexT* index) {
|
|
faiss::Index::idx_t base = (faiss::Index::idx_t) i * queriesPerIndex;
|
|
|
|
if (base < n) {
|
|
auto numForIndex = std::min(queriesPerIndex, n - base);
|
|
|
|
index->search(numForIndex,
|
|
x + base * componentsPerVec,
|
|
k,
|
|
distances + base * k,
|
|
labels + base * k);
|
|
}
|
|
};
|
|
|
|
this->runOnIndex(fn);
|
|
}
|
|
|
|
// explicit instantiations
|
|
template struct IndexReplicasTemplate<Index>;
|
|
template struct IndexReplicasTemplate<IndexBinary>;
|
|
|
|
} // namespace
|