77 lines
2.4 KiB
C++
77 lines
2.4 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.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "Index.h"
|
|
#include "IndexBinary.h"
|
|
#include "ThreadedIndex.h"
|
|
|
|
namespace faiss {
|
|
|
|
/// Takes individual faiss::Index instances, and splits queries for
|
|
/// sending to each Index instance, and joins the results together
|
|
/// when done.
|
|
/// Each index is managed by a separate CPU thread.
|
|
template <typename IndexT>
|
|
class IndexReplicasTemplate : public ThreadedIndex<IndexT> {
|
|
public:
|
|
using idx_t = typename IndexT::idx_t;
|
|
using component_t = typename IndexT::component_t;
|
|
using distance_t = typename IndexT::distance_t;
|
|
|
|
/// The dimension that all sub-indices must share will be the dimension of the
|
|
/// first sub-index added
|
|
/// @param threaded do we use one thread per sub-index or do queries
|
|
/// sequentially?
|
|
explicit IndexReplicasTemplate(bool threaded = true);
|
|
|
|
/// @param d the dimension that all sub-indices must share
|
|
/// @param threaded do we use one thread per sub index or do queries
|
|
/// sequentially?
|
|
explicit IndexReplicasTemplate(idx_t d, bool threaded = true);
|
|
|
|
/// int version due to the implicit bool conversion ambiguity of int as
|
|
/// dimension
|
|
explicit IndexReplicasTemplate(int d, bool threaded = true);
|
|
|
|
/// Alias for addIndex()
|
|
void add_replica(IndexT* index) { this->addIndex(index); }
|
|
|
|
/// Alias for removeIndex()
|
|
void remove_replica(IndexT* index) { this->removeIndex(index); }
|
|
|
|
/// faiss::Index API
|
|
/// All indices receive the same call
|
|
void train(idx_t n, const component_t* x) override;
|
|
|
|
/// faiss::Index API
|
|
/// All indices receive the same call
|
|
void add(idx_t n, const component_t* x) override;
|
|
|
|
/// faiss::Index API
|
|
/// Query is partitioned into a slice for each sub-index
|
|
/// split by ceil(n / #indices) for our sub-indices
|
|
void search(idx_t n,
|
|
const component_t* x,
|
|
idx_t k,
|
|
distance_t* distances,
|
|
idx_t* labels) const override;
|
|
|
|
/// reconstructs from the first index
|
|
void reconstruct(idx_t, component_t *v) const override;
|
|
|
|
protected:
|
|
/// Called just after an index is added
|
|
void onAfterAddIndex(IndexT* index) override;
|
|
};
|
|
|
|
using IndexReplicas = IndexReplicasTemplate<Index>;
|
|
using IndexBinaryReplicas = IndexReplicasTemplate<IndexBinary>;
|
|
|
|
} // namespace
|