101 lines
3.3 KiB
C++
101 lines
3.3 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 {
|
|
|
|
/**
|
|
* Index that concatenates the results from several sub-indexes
|
|
*/
|
|
template <typename IndexT>
|
|
struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
|
|
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?
|
|
* @param successive_ids should we shift the returned ids by
|
|
* the size of each sub-index or return them
|
|
* as they are?
|
|
*/
|
|
explicit IndexShardsTemplate(bool threaded = false,
|
|
bool successive_ids = true);
|
|
|
|
/**
|
|
* @param threaded do we use one thread per sub_index or do
|
|
* queries sequentially?
|
|
* @param successive_ids should we shift the returned ids by
|
|
* the size of each sub-index or return them
|
|
* as they are?
|
|
*/
|
|
explicit IndexShardsTemplate(idx_t d,
|
|
bool threaded = false,
|
|
bool successive_ids = true);
|
|
|
|
/// int version due to the implicit bool conversion ambiguity of int as
|
|
/// dimension
|
|
explicit IndexShardsTemplate(int d,
|
|
bool threaded = false,
|
|
bool successive_ids = true);
|
|
|
|
/// Alias for addIndex()
|
|
void add_shard(IndexT* index) { this->addIndex(index); }
|
|
|
|
/// Alias for removeIndex()
|
|
void remove_shard(IndexT* index) { this->removeIndex(index); }
|
|
|
|
/// supported only for sub-indices that implement add_with_ids
|
|
void add(idx_t n, const component_t* x) override;
|
|
|
|
/**
|
|
* Cases (successive_ids, xids):
|
|
* - true, non-NULL ERROR: it makes no sense to pass in ids and
|
|
* request them to be shifted
|
|
* - true, NULL OK, but should be called only once (calls add()
|
|
* on sub-indexes).
|
|
* - false, non-NULL OK: will call add_with_ids with passed in xids
|
|
* distributed evenly over shards
|
|
* - false, NULL OK: will call add_with_ids on each sub-index,
|
|
* starting at ntotal
|
|
*/
|
|
void add_with_ids(idx_t n, const component_t* x, const idx_t* xids) override;
|
|
|
|
void search(idx_t n, const component_t* x, idx_t k,
|
|
distance_t* distances, idx_t* labels) const override;
|
|
|
|
void train(idx_t n, const component_t* x) override;
|
|
|
|
// update metric_type and ntotal. Call if you changes something in
|
|
// the shard indexes.
|
|
void sync_with_shard_indexes();
|
|
|
|
bool successive_ids;
|
|
|
|
protected:
|
|
/// Called just after an index is added
|
|
void onAfterAddIndex(IndexT* index) override;
|
|
|
|
/// Called just after an index is removed
|
|
void onAfterRemoveIndex(IndexT* index) override;
|
|
};
|
|
|
|
using IndexShards = IndexShardsTemplate<Index>;
|
|
using IndexBinaryShards = IndexShardsTemplate<IndexBinary>;
|
|
|
|
|
|
} // namespace faiss
|