All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/data/users/hoss/faiss/ThreadedIndex-inl.h
/**
* 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 "FaissAssert.h"
#include <exception>
#include <iostream>
namespace faiss {
template <typename IndexT>
ThreadedIndex<IndexT>::ThreadedIndex(bool threaded)
// 0 is default dimension
: ThreadedIndex(0, threaded) {
}
template <typename IndexT>
ThreadedIndex<IndexT>::ThreadedIndex(int d, bool threaded)
: IndexT(d),
own_fields(false),
isThreaded_(threaded) {
}
template <typename IndexT>
ThreadedIndex<IndexT>::~ThreadedIndex() {
for (auto& p : indices_) {
if (isThreaded_) {
// should have worker thread
FAISS_ASSERT((bool) p.second);
// This will also flush all pending work
p.second->stop();
p.second->waitForThreadExit();
} else {
// should not have worker thread
FAISS_ASSERT(!(bool) p.second);
}
if (own_fields) {
delete p.first;
}
}
}
template <typename IndexT>
void ThreadedIndex<IndexT>::addIndex(IndexT* index) {
// We inherit the dimension from the first index added to us if we don't have
// a set dimension
if (indices_.empty() && this->d == 0) {
this->d = index->d;
}
// The new index must match our set dimension
FAISS_THROW_IF_NOT_FMT(this->d == index->d,
"addIndex: dimension mismatch for "
"newly added index; expecting dim %d, "
"new index has dim %d",
this->d, index->d);
if (!indices_.empty()) {
auto& existing = indices_.front().first;
FAISS_THROW_IF_NOT_MSG(index->metric_type == existing->metric_type,
"addIndex: newly added index is "
"of different metric type than old index");
// Make sure this index is not duplicated
for (auto& p : indices_) {
FAISS_THROW_IF_NOT_MSG(p.first != index,
"addIndex: attempting to add index "
"that is already in the collection");
}
}
indices_.emplace_back(
std::make_pair(
index,
std::unique_ptr<WorkerThread>(isThreaded_ ?
new WorkerThread : nullptr)));
onAfterAddIndex(index);
}
template <typename IndexT>
void ThreadedIndex<IndexT>::removeIndex(IndexT* index) {
for (auto it = indices_.begin(); it != indices_.end(); ++it) {
if (it->first == index) {
// This is our index; stop the worker thread before removing it,
// to ensure that it has finished before function exit
if (isThreaded_) {
// should have worker thread
FAISS_ASSERT((bool) it->second);
it->second->stop();
it->second->waitForThreadExit();
} else {
// should not have worker thread
FAISS_ASSERT(!(bool) it->second);
}
indices_.erase(it);
onAfterRemoveIndex(index);
if (own_fields) {
delete index;
}
return;
}
}
// could not find our index
FAISS_THROW_MSG("IndexReplicas::removeIndex: index not found");
}
template <typename IndexT>
void ThreadedIndex<IndexT>::runOnIndex(std::function<void(int, IndexT*)> f) {
if (isThreaded_) {
std::vector<std::future<bool>> v;
for (int i = 0; i < this->indices_.size(); ++i) {
auto& p = this->indices_[i];
auto indexPtr = p.first;
v.emplace_back(p.second->add([f, i, indexPtr](){ f(i, indexPtr); }));
}
waitAndHandleFutures(v);
} else {
// Multiple exceptions may be thrown; gather them as we encounter them,
// while letting everything else run to completion
std::vector<std::pair<int, std::exception_ptr>> exceptions;
for (int i = 0; i < this->indices_.size(); ++i) {
auto& p = this->indices_[i];
try {
f(i, p.first);
} catch (...) {
exceptions.emplace_back(std::make_pair(i, std::current_exception()));
}
}
handleExceptions(exceptions);
}
}
template <typename IndexT>
void ThreadedIndex<IndexT>::runOnIndex(
std::function<void(int, const IndexT*)> f) const {
const_cast<ThreadedIndex<IndexT>*>(this)->runOnIndex(
[f](int i, IndexT* idx){ f(i, idx); });
}
template <typename IndexT>
void ThreadedIndex<IndexT>::reset() {
runOnIndex([](int, IndexT* index){ index->reset(); });
this->ntotal = 0;
this->is_trained = false;
}
template <typename IndexT>
void
ThreadedIndex<IndexT>::onAfterAddIndex(IndexT* index) {
}
template <typename IndexT>
void
ThreadedIndex<IndexT>::onAfterRemoveIndex(IndexT* index) {
}
template <typename IndexT>
void
ThreadedIndex<IndexT>::waitAndHandleFutures(std::vector<std::future<bool>>& v) {
// Blocking wait for completion for all of the indices, capturing any
// exceptions that are generated
std::vector<std::pair<int, std::exception_ptr>> exceptions;
for (int i = 0; i < v.size(); ++i) {
auto& fut = v[i];
try {
fut.get();
} catch (...) {
exceptions.emplace_back(std::make_pair(i, std::current_exception()));
}
}
handleExceptions(exceptions);
}
} // namespace
faiss::ThreadedIndex::removeIndex
void removeIndex(IndexT *index)
Definition: ThreadedIndex-inl.h:88
faiss::handleExceptions
void handleExceptions(std::vector< std::pair< int, std::exception_ptr >> &exceptions)
Definition: FaissException.cpp:35
faiss::ThreadedIndex
Definition: ThreadedIndex.h:21
faiss::ThreadedIndex::addIndex
void addIndex(IndexT *index)
Definition: ThreadedIndex-inl.h:49
faiss::ThreadedIndex::runOnIndex
void runOnIndex(std::function< void(int, IndexT *)> f)
Definition: ThreadedIndex-inl.h:119
faiss::ThreadedIndex::onAfterRemoveIndex
virtual void onAfterRemoveIndex(IndexT *index)
Called just after an index is removed.
Definition: ThreadedIndex-inl.h:169
faiss::WorkerThread
Definition: WorkerThread.h:18
faiss::ThreadedIndex::onAfterAddIndex
virtual void onAfterAddIndex(IndexT *index)
Called just after an index is added.
Definition: ThreadedIndex-inl.h:164
faiss::ProductQuantizer::d
size_t d
size of the input vectors
Definition: ProductQuantizer.h:27
faiss::ThreadedIndex::reset
void reset() override
Definition: ThreadedIndex-inl.h:156