mirror of
https://github.com/facebookresearch/faiss.git
synced 2025-06-03 21:54:02 +08:00
ebb5f8401b
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/3247 add a context parameter to be passed to InvertedLists and InvertedListsIterator. - add a context field in `SearchParametersIVF`, the context will be passed to `InvertedLists::get_iterator`. The user can create `InvertedListsIterator` with the context object - add a context parameter in `IndexIVF::add_core` method. the context will be passed to `InvertedLists::add_entry`. The user can use the context object to pass storage handlers, store error codes from storage layer, logging information, etc. Reviewed By: mdouze Differential Revision: D53113911 fbshipit-source-id: ff31d247d3dc949d0bb50bcaffc3142efd027089
111 lines
2.9 KiB
C++
111 lines
2.9 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.
|
|
*/
|
|
|
|
// -*- c++ -*-
|
|
|
|
#ifndef FAISS_INDEX_SCALAR_QUANTIZER_H
|
|
#define FAISS_INDEX_SCALAR_QUANTIZER_H
|
|
|
|
#include <stdint.h>
|
|
#include <vector>
|
|
|
|
#include <faiss/IndexFlatCodes.h>
|
|
#include <faiss/IndexIVF.h>
|
|
#include <faiss/impl/ScalarQuantizer.h>
|
|
|
|
namespace faiss {
|
|
|
|
/**
|
|
* Flat index built on a scalar quantizer.
|
|
*/
|
|
struct IndexScalarQuantizer : IndexFlatCodes {
|
|
/// Used to encode the vectors
|
|
ScalarQuantizer sq;
|
|
|
|
/** Constructor.
|
|
*
|
|
* @param d dimensionality of the input vectors
|
|
* @param M number of subquantizers
|
|
* @param nbits number of bit per subvector index
|
|
*/
|
|
IndexScalarQuantizer(
|
|
int d,
|
|
ScalarQuantizer::QuantizerType qtype,
|
|
MetricType metric = METRIC_L2);
|
|
|
|
IndexScalarQuantizer();
|
|
|
|
void train(idx_t n, const float* x) override;
|
|
|
|
void search(
|
|
idx_t n,
|
|
const float* x,
|
|
idx_t k,
|
|
float* distances,
|
|
idx_t* labels,
|
|
const SearchParameters* params = nullptr) const override;
|
|
|
|
FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
|
|
|
|
/* standalone codec interface */
|
|
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
|
|
|
|
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
|
|
};
|
|
|
|
/** An IVF implementation where the components of the residuals are
|
|
* encoded with a scalar quantizer. All distance computations
|
|
* are asymmetric, so the encoded vectors are decoded and approximate
|
|
* distances are computed.
|
|
*/
|
|
|
|
struct IndexIVFScalarQuantizer : IndexIVF {
|
|
ScalarQuantizer sq;
|
|
|
|
IndexIVFScalarQuantizer(
|
|
Index* quantizer,
|
|
size_t d,
|
|
size_t nlist,
|
|
ScalarQuantizer::QuantizerType qtype,
|
|
MetricType metric = METRIC_L2,
|
|
bool by_residual = true);
|
|
|
|
IndexIVFScalarQuantizer();
|
|
|
|
void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
|
|
|
|
idx_t train_encoder_num_vectors() const override;
|
|
|
|
void encode_vectors(
|
|
idx_t n,
|
|
const float* x,
|
|
const idx_t* list_nos,
|
|
uint8_t* codes,
|
|
bool include_listnos = false) const override;
|
|
|
|
void add_core(
|
|
idx_t n,
|
|
const float* x,
|
|
const idx_t* xids,
|
|
const idx_t* precomputed_idx,
|
|
void* inverted_list_context = nullptr) override;
|
|
|
|
InvertedListScanner* get_InvertedListScanner(
|
|
bool store_pairs,
|
|
const IDSelector* sel) const override;
|
|
|
|
void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
|
|
const override;
|
|
|
|
/* standalone codec interface */
|
|
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
|
|
};
|
|
|
|
} // namespace faiss
|
|
|
|
#endif
|