facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

generic result handlers for most indexes (#4762)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4762

Some search functions require special handling on the search resutls. For example collecting resutls from diverse groups.

In order to support this, we re-purpose the ResultHandler object to offer an arbitrary callback whenever a new vector should be taken in consideration.

Many indexes did already have such a callback, but in non-uniform formats. This diff consolidates these diverse implementations. This is especially the case for the InvertedListScanner that is now implemented on top of that.

Most indexes now get a `search1` that searches one query vector and takes a ResultCollector as input. This is the object that can be used a a callback to do specific handling on the results.

The current impleentatoin of search1 may not be perfectly optimized, especially for IndexFlatCodes indexes. Another todo is to support BlockResultHandlers to allow multiple queries at a time.

Differential Revision: D91003349

fbshipit-source-id: 3ce98145f9633174bd3d84bb1ae85f36084e41ce
This commit is contained in:
Matthijs Douze
2026-01-21 13:58:27 -08:00
committed by meta-codesync[bot]
parent 49c64f9704
commit 8eeaa92393
22 changed files with 640 additions and 408 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/Index.cpp
+4
View File
@@ -115,6 +115,10 @@ void Index::search_subset(
FAISS_THROW_MSG("search_subset not implemented for this type of index");
}
void Index::search1(const float*, ResultHandler&, SearchParameters*) const {
FAISS_THROW_MSG("search1 not implemented for this type of index");
}
void Index::compute_residual(const float* x, float* residual, idx_t key) const {
reconstruct(key, residual);
for (size_t i = 0; i < d; i++) {
faiss/Index.h
+9 -1
View File
@@ -14,7 +14,6 @@
#include <faiss/impl/FaissAssert.h>
#include <cstdio>
#include <sstream>
#define FAISS_VERSION_MAJOR 1
#define FAISS_VERSION_MINOR 13
@@ -55,6 +54,9 @@ namespace faiss {
struct IDSelector;
struct RangeSearchResult;
struct DistanceComputer;
template <typename T, typename TI>
struct ResultHandlerUnordered;
using ResultHandler = ResultHandlerUnordered<float, idx_t>;
enum NumericType {
Float32,
@@ -216,6 +218,12 @@ struct Index {
}
}
/** search one vector with a custom result handler */
virtual void search1(
const float* x,
ResultHandler& handler,
SearchParameters* params = nullptr) const;
/** query n vectors of dimension d to the index.
*
* return all vectors with distance < radius. Note that many
faiss/IndexBinary.cpp
+1
View File
@@ -12,6 +12,7 @@
#include <cinttypes>
#include <cstring>
#include <typeinfo>
namespace faiss {
faiss/IndexFlatCodes.cpp
+34 -3
View File
@@ -231,9 +231,11 @@ struct Run_search_with_decompress {
struct Run_search_with_decompress_res {
using T = void;
template <class ResultHandler>
void f(ResultHandler& res, const IndexFlatCodes* index, const float* xq) {
Run_search_with_decompress<ResultHandler> r;
template <class BlockResultHandler>
void f(BlockResultHandler& res,
const IndexFlatCodes* index,
const float* xq) {
Run_search_with_decompress<BlockResultHandler> r;
dispatch_VectorDistance(
index->d,
index->metric_type,
@@ -277,4 +279,33 @@ void IndexFlatCodes::range_search(
dispatch_range_ResultHandler(result, radius, metric_type, sel, r, this, x);
}
void IndexFlatCodes::search1(
const float* x,
ResultHandler& handler,
SearchParameters* params) const {
const IDSelector* sel = params ? params->sel : nullptr;
Run_search_with_decompress_res r;
if (sel) {
if (is_similarity_metric(metric_type)) {
SingleQueryBlockResultHandler<CMin<float, idx_t>, true> res(
handler, sel);
r.f(res, this, x);
} else {
SingleQueryBlockResultHandler<CMax<float, idx_t>, true> res(
handler, sel);
r.f(res, this, x);
}
} else {
if (is_similarity_metric(metric_type)) {
SingleQueryBlockResultHandler<CMin<float, idx_t>, false> res(
handler);
r.f(res, this, x);
} else {
SingleQueryBlockResultHandler<CMax<float, idx_t>, false> res(
handler);
r.f(res, this, x);
}
}
}
} // namespace faiss
faiss/IndexFlatCodes.h
+7 -1
View File
@@ -55,7 +55,8 @@ struct IndexFlatCodes : Index {
return get_FlatCodesDistanceComputer();
}
/** Search implemented by decoding */
/** Search implemented by decoding (most index types will have a faster
* implementation) */
void search(
idx_t n,
const float* x,
@@ -71,6 +72,11 @@ struct IndexFlatCodes : Index {
RangeSearchResult* result,
const SearchParameters* params = nullptr) const override;
virtual void search1(
const float* x,
ResultHandler& handler,
SearchParameters* params = nullptr) const override;
// returns a new instance of a CodePacker
CodePacker* get_CodePacker() const;
faiss/IndexHNSW.cpp
+8
View File
@@ -333,6 +333,14 @@ void IndexHNSW::range_search(
}
}
void IndexHNSW::search1(
const float* x,
ResultHandler& handler,
SearchParameters* params) const {
SingleQueryBlockResultHandler<HNSW::C, false> bres(handler);
hnsw_search(this, 1, x, bres, params);
}
void IndexHNSW::add(idx_t n, const float* x) {
FAISS_THROW_IF_NOT_MSG(
storage,
faiss/IndexHNSW.h
+6
View File
@@ -74,6 +74,12 @@ struct IndexHNSW : Index {
RangeSearchResult* result,
const SearchParameters* params = nullptr) const override;
/** search one vector with a custom result handler */
void search1(
const float* x,
ResultHandler& handler,
SearchParameters* params = nullptr) const override;
void reconstruct(idx_t key, float* recons) const override;
void reset() override;
faiss/IndexIDMap.cpp
+1 -1
View File
@@ -12,7 +12,7 @@
#include <cinttypes>
#include <cstdint>
#include <cstdio>
#include "faiss/Index.h"
#include <stdexcept>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
faiss/IndexIVF.cpp
+142 -47
View File
@@ -27,6 +27,7 @@
#include <faiss/impl/CodePacker.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ResultHandler.h>
namespace faiss {
@@ -920,6 +921,52 @@ void IndexIVF::range_search_preassigned(
stats->ndis += ndis;
}
void IndexIVF::search1(
const float* x,
ResultHandler& handler,
SearchParameters* params_in) const {
const IVFSearchParameters* params = nullptr;
const SearchParameters* quantizer_params = nullptr;
if (params_in) {
params = dynamic_cast<const IVFSearchParameters*>(params_in);
FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
quantizer_params = params->quantizer_params;
}
const size_t nprobe =
std::min(nlist, params ? params->nprobe : this->nprobe);
size_t nx = 1;
std::unique_ptr<idx_t[]> keys(new idx_t[nx * nprobe]);
std::unique_ptr<float[]> coarse_dis(new float[nx * nprobe]);
double t0 = getmillisecs();
quantizer->search(
nx, x, nprobe, coarse_dis.get(), keys.get(), quantizer_params);
indexIVF_stats.quantization_time += getmillisecs() - t0;
t0 = getmillisecs();
invlists->prefetch_lists(keys.get(), nx * nprobe);
std::unique_ptr<InvertedListScanner> scanner(
get_InvertedListScanner(false, nullptr, params));
scanner->set_query(x);
for (idx_t i = 0; i < nprobe; i++) {
idx_t key = keys[i];
if (key < 0 || invlists->is_empty(key)) {
continue;
}
scanner->set_list(key, coarse_dis[i]);
InvertedLists::ScopedCodes scodes(invlists, key);
InvertedLists::ScopedIds ids(invlists, key);
size_t list_size = invlists->list_size(key);
scanner->scan_codes(list_size, scodes.get(), ids.get(), handler);
}
indexIVF_stats.search_time += getmillisecs() - t0;
}
InvertedListScanner* IndexIVF::get_InvertedListScanner(
bool /*store_pairs*/,
const IDSelector* /* sel */,
@@ -1298,6 +1345,87 @@ IndexIVFStats indexIVF_stats;
* InvertedListScanner
*************************************************************************/
namespace {
template <typename C, bool store_pairs, bool use_sel>
size_t run_scan_codes1(
const InvertedListScanner& scanner,
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) {
size_t nup = 0;
size_t list_no = scanner.list_no;
size_t code_size = scanner.code_size;
const IDSelector* sel = scanner.sel;
for (size_t j = 0; j < list_size; j++) {
if (use_sel) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
// skip code without computing distance
if (!sel->is_member(id)) {
codes += code_size;
continue;
}
}
float dis = scanner.distance_to_code(codes);
if (C::cmp(handler.threshold, dis)) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
handler.add_result(dis, id);
nup++;
}
codes += code_size;
}
return nup;
}
template <bool store_pairs, bool use_sel>
size_t run_scan_codes(
const InvertedListScanner& scanner,
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) {
if (scanner.keep_max) {
return run_scan_codes1<CMin<float, idx_t>, store_pairs, use_sel>(
scanner, list_size, codes, ids, handler);
} else {
return run_scan_codes1<CMax<float, idx_t>, store_pairs, use_sel>(
scanner, list_size, codes, ids, handler);
}
}
} // anonymous namespace
size_t InvertedListScanner::scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const {
if (sel == nullptr) {
if (store_pairs) {
return run_scan_codes<true, false>(
*this, list_size, codes, ids, handler);
} else {
return run_scan_codes<false, false>(
*this, list_size, codes, ids, handler);
}
} else {
if (store_pairs) {
return run_scan_codes<true, true>(
*this, list_size, codes, ids, handler);
} else {
return run_scan_codes<false, true>(
*this, list_size, codes, ids, handler);
}
}
}
void InvertedListScanner::set_list(idx_t list_no_in, float /* coarse_dis */) {
this->list_no = list_no_in;
}
size_t InvertedListScanner::scan_codes(
size_t list_size,
const uint8_t* codes,
@@ -1305,46 +1433,15 @@ size_t InvertedListScanner::scan_codes(
float* simi,
idx_t* idxi,
size_t k) const {
size_t nup = 0;
if (!keep_max) {
for (size_t j = 0; j < list_size; j++) {
if (sel != nullptr) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
if (!sel->is_member(id)) {
codes += code_size;
continue;
}
}
float dis = distance_to_code(codes);
if (dis < simi[0]) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
maxheap_replace_top(k, simi, idxi, dis, id);
nup++;
}
codes += code_size;
}
using C = CMax<float, idx_t>;
HeapResultHandler<C, false> handler(k, simi, idxi);
return scan_codes(list_size, codes, ids, handler);
} else {
for (size_t j = 0; j < list_size; j++) {
if (sel != nullptr) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
if (!sel->is_member(id)) {
codes += code_size;
continue;
}
}
float dis = distance_to_code(codes);
if (dis > simi[0]) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
minheap_replace_top(k, simi, idxi, dis, id);
nup++;
}
codes += code_size;
}
using C = CMin<float, idx_t>;
HeapResultHandler<C, false> handler(k, simi, idxi);
return scan_codes(list_size, codes, ids, handler);
}
return nup;
}
size_t InvertedListScanner::iterate_codes(
@@ -1386,16 +1483,14 @@ void InvertedListScanner::scan_codes_range(
const idx_t* ids,
float radius,
RangeQueryResult& res) const {
for (size_t j = 0; j < list_size; j++) {
float dis = distance_to_code(codes);
bool keep = !keep_max
? dis < radius
: dis > radius; // TODO templatize to remove this test
if (keep) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
res.add(dis, id);
}
codes += code_size;
if (!keep_max) {
using C = CMax<float, idx_t>;
RangeResultHandler<C, false> handler(&res, radius);
scan_codes(list_size, codes, ids, handler);
} else {
using C = CMin<float, idx_t>;
RangeResultHandler<C, false> handler(&res, radius);
scan_codes(list_size, codes, ids, handler);
}
}
faiss/IndexIVF.h
+14 -4
View File
@@ -11,9 +11,6 @@
#define FAISS_INDEX_IVF_H
#include <stdint.h>
#include <memory>
#include <unordered_map>
#include <vector>
#include <faiss/Clustering.h>
#include <faiss/Index.h>
@@ -325,6 +322,12 @@ struct IndexIVF : Index, IndexIVFInterface {
RangeSearchResult* result,
const SearchParameters* params = nullptr) const override;
/** search one vector with a custom result handler */
void search1(
const float* x,
ResultHandler& handler,
SearchParameters* params = nullptr) const override;
/** Get a scanner for this index (store_pairs means ignore labels)
*
* The default search implementation uses this to compute the distances.
@@ -492,7 +495,7 @@ struct InvertedListScanner {
virtual void set_query(const float* query_vector) = 0;
/// following codes come from this inverted list
virtual void set_list(idx_t list_no, float coarse_dis) = 0;
virtual void set_list(idx_t list_no, float coarse_dis);
/// compute a single query-to-code distance
virtual float distance_to_code(const uint8_t* code) const = 0;
@@ -543,6 +546,13 @@ struct InvertedListScanner {
RangeQueryResult& result,
size_t& list_size) const;
// accumulate results with a ResultHandler
virtual size_t scan_codes(
size_t n,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const;
virtual ~InvertedListScanner() {}
};
faiss/IndexIVFFlat.cpp
+1 -45
View File
@@ -170,54 +170,10 @@ struct IVFFlatScanner : InvertedListScanner {
this->list_no = list_no;
}
float distance_to_code(const uint8_t* code) const override {
float distance_to_code(const uint8_t* code) const final {
const float* yj = (float*)code;
return vd(xi, yj);
}
size_t scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float* simi,
idx_t* idxi,
size_t k) const override {
const float* list_vecs = (const float*)codes;
size_t nup = 0;
for (size_t j = 0; j < list_size; j++) {
const float* yj = list_vecs + vd.d * j;
if (use_sel && !sel->is_member(ids[j])) {
continue;
}
float dis = vd(xi, yj);
if (C::cmp(simi[0], dis)) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
heap_replace_top<C>(k, simi, idxi, dis, id);
nup++;
}
}
return nup;
}
void scan_codes_range(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float radius,
RangeQueryResult& res) const override {
const float* list_vecs = (const float*)codes;
for (size_t j = 0; j < list_size; j++) {
const float* yj = list_vecs + vd.d * j;
if (use_sel && !sel->is_member(ids[j])) {
continue;
}
float dis = vd(xi, yj);
if (C::cmp(radius, dis)) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
res.add(dis, id);
}
}
}
};
struct Run_get_InvertedListScanner {
faiss/IndexIVFFlatPanorama.cpp
+5 -62
View File
@@ -17,6 +17,7 @@
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/PanoramaStats.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/distances.h>
@@ -89,9 +90,7 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float* simi,
idx_t* idxi,
size_t k) const override {
ResultHandler& handler) const override {
size_t nup = 0;
const size_t n_batches =
@@ -121,7 +120,7 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
use_sel,
active_indices,
exact_distances,
simi[0],
handler.threshold,
local_stats);
});
@@ -131,10 +130,10 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
size_t global_idx = batch_start + idx;
float dis = exact_distances[idx];
if (C::cmp(simi[0], dis)) {
if (C::cmp(handler.threshold, dis)) {
int64_t id = store_pairs ? lo_build(list_no, global_idx)
: ids[global_idx];
heap_replace_top<C>(k, simi, idxi, dis, id);
handler.add_result(dis, id);
nup++;
}
}
@@ -143,62 +142,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
indexPanorama_stats.add(local_stats);
return nup;
}
void scan_codes_range(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float radius,
RangeQueryResult& res) const override {
const size_t n_batches =
(list_size + storage->kBatchSize - 1) / storage->kBatchSize;
const float* cum_sums_data = storage->get_cum_sums(list_no);
std::vector<float> exact_distances(storage->kBatchSize);
std::vector<uint32_t> active_indices(storage->kBatchSize);
PanoramaStats local_stats;
local_stats.reset();
// Same progressive filtering as scan_codes, but with fixed radius
// threshold instead of dynamic heap threshold.
for (size_t batch_no = 0; batch_no < n_batches; batch_no++) {
size_t batch_start = batch_no * storage->kBatchSize;
size_t num_active = with_metric_type(metric, [&]<MetricType M>() {
return storage->pano.progressive_filter_batch<C, M>(
codes,
cum_sums_data,
xi,
cum_sums.data(),
batch_no,
list_size,
sel,
ids,
use_sel,
active_indices,
exact_distances,
radius,
local_stats);
});
// Add batch survivors to range result.
for (size_t i = 0; i < num_active; i++) {
uint32_t idx = active_indices[i];
size_t global_idx = batch_start + idx;
float dis = exact_distances[idx];
if (C::cmp(radius, dis)) {
int64_t id = store_pairs ? lo_build(list_no, global_idx)
: ids[global_idx];
res.add(dis, id);
}
}
}
indexPanorama_stats.add(local_stats);
}
};
struct Run_get_InvertedListScanner {
faiss/IndexIVFPQ.cpp
+22 -77
View File
@@ -25,13 +25,11 @@
#include <faiss/utils/hamming.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ProductQuantizer.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/impl/code_distance/code_distance.h>
namespace faiss {
@@ -762,56 +760,35 @@ struct QueryTables {
}
};
// This way of handling the selector is not optimal since all distances
// are computed even if the id would filter it out.
template <class C, bool use_sel>
struct KnnSearchResults {
idx_t key;
struct WrappedSearchResult {
ResultHandler& res;
size_t nup = 0;
idx_t list_no;
const idx_t* ids;
const IDSelector* sel;
// heap params
size_t k;
float* heap_sim;
idx_t* heap_ids;
size_t nup;
WrappedSearchResult(
idx_t list_no,
const idx_t* ids,
const IDSelector* sel,
ResultHandler& res)
: res(res), list_no(list_no), ids(ids), sel(sel) {}
inline bool skip_entry(idx_t j) {
return use_sel && !sel->is_member(ids[j]);
}
inline void add(idx_t j, float dis) {
if (C::cmp(heap_sim[0], dis)) {
idx_t id = ids ? ids[j] : lo_build(key, j);
heap_replace_top<C>(k, heap_sim, heap_ids, dis, id);
if (C::cmp(res.threshold, dis)) {
idx_t id = ids ? ids[j] : lo_build(this->list_no, j);
res.add_result(dis, id);
nup++;
}
}
};
template <class C, bool use_sel>
struct RangeSearchResults {
idx_t key;
const idx_t* ids;
const IDSelector* sel;
// wrapped result structure
float radius;
RangeQueryResult& rres;
inline bool skip_entry(idx_t j) {
return use_sel && !sel->is_member(ids[j]);
}
inline void add(idx_t j, float dis) {
if (C::cmp(radius, dis)) {
idx_t id = ids ? ids[j] : lo_build(key, j);
rres.add(dis, id);
}
}
};
/*****************************************************
* Scaning the codes.
* The scanning functions call their favorite precompute_*
@@ -1239,17 +1216,12 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
size_t ncode,
const uint8_t* codes,
const idx_t* ids,
float* heap_sim,
idx_t* heap_ids,
size_t k) const override {
KnnSearchResults<C, use_sel> res = {
/* key */ this->key,
/* ids */ this->store_pairs ? nullptr : ids,
/* sel */ this->sel,
/* k */ k,
/* heap_sim */ heap_sim,
/* heap_ids */ heap_ids,
/* nup */ 0};
ResultHandler& handler) const override {
WrappedSearchResult<C, use_sel> res(
this->key,
this->store_pairs ? nullptr : ids,
this->sel,
handler);
if (this->polysemous_ht > 0) {
assert(precompute_mode == 2);
@@ -1265,33 +1237,6 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
}
return res.nup;
}
void scan_codes_range(
size_t ncode,
const uint8_t* codes,
const idx_t* ids,
float radius,
RangeQueryResult& rres) const override {
RangeSearchResults<C, use_sel> res = {
/* key */ this->key,
/* ids */ this->store_pairs ? nullptr : ids,
/* sel */ this->sel,
/* radius */ radius,
/* rres */ rres};
if (this->polysemous_ht > 0) {
assert(precompute_mode == 2);
this->scan_list_polysemous(ncode, codes, res);
} else if (precompute_mode == 2) {
this->scan_list_with_table(ncode, codes, res);
} else if (precompute_mode == 1) {
this->scan_list_with_pointer(ncode, codes, res);
} else if (precompute_mode == 0) {
this->scan_on_the_fly_dist(ncode, codes, res);
} else {
FAISS_THROW_MSG("bad precomp mode");
}
}
};
template <class PQDecoder, bool use_sel>
faiss/IndexIVFRaBitQ.cpp
+6 -16
View File
@@ -17,6 +17,7 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/RaBitQuantizer.h>
#include <faiss/impl/ResultHandler.h>
namespace faiss {
@@ -205,15 +206,13 @@ struct RaBitInvertedListScanner : InvertedListScanner {
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float* simi,
idx_t* idxi,
size_t k) const override {
ResultHandler& handler) const override {
size_t ex_bits = ivf_rabitq.rabitq.nb_bits - 1;
// For 1-bit codes, use default implementation
if (ex_bits == 0 || rabitq_dc == nullptr) {
return InvertedListScanner::scan_codes(
list_size, codes, ids, simi, idxi, k);
list_size, codes, ids, handler);
}
// Multi-bit: Two-stage search with adaptive filtering
@@ -254,24 +253,15 @@ struct RaBitInvertedListScanner : InvertedListScanner {
est_distance,
base_fac->f_error,
rabitq_dc->g_error,
simi[0],
handler.threshold,
keep_max);
if (should_refine) {
local_multibit_evaluations++;
// Lower bound is promising, compute full distance
float dis = distance_to_code(codes);
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
// Check if distance improves heap
bool improves_heap =
keep_max ? (dis > simi[0]) : (dis < simi[0]);
if (improves_heap) {
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
if (keep_max) {
minheap_replace_top(k, simi, idxi, dis, id);
} else {
maxheap_replace_top(k, simi, idxi, dis, id);
}
if (handler.add_result(dis, id)) {
nup++;
}
}
faiss/impl/HNSW.cpp
+5 -5
View File
@@ -616,7 +616,7 @@ static inline void extract_search_params(
int search_from_candidates(
const HNSW& hnsw,
DistanceComputer& qdis,
ResultHandler<C>& res,
ResultHandler& res,
MinimaxHeap& candidates,
VisitedTable& vt,
HNSWStats& stats,
@@ -755,7 +755,7 @@ int search_from_candidates_panorama(
const HNSW& hnsw,
const IndexHNSW* index,
DistanceComputer& qdis,
ResultHandler<C>& res,
ResultHandler& res,
MinimaxHeap& candidates,
VisitedTable& vt,
HNSWStats& stats,
@@ -1182,7 +1182,7 @@ using Node = HNSW::Node;
using C = HNSW::C;
// just used as a lower bound for the minmaxheap, but it is set for heap search
int extract_k_from_ResultHandler(ResultHandler<C>& res) {
int extract_k_from_ResultHandler(ResultHandler& res) {
using RH = HeapBlockResultHandler<C>;
if (auto hres = dynamic_cast<RH::SingleResultHandler*>(&res)) {
return hres->k;
@@ -1195,7 +1195,7 @@ int extract_k_from_ResultHandler(ResultHandler<C>& res) {
HNSWStats HNSW::search(
DistanceComputer& qdis,
const IndexHNSW* index,
ResultHandler<C>& res,
ResultHandler& res,
VisitedTable& vt,
const SearchParameters* params) const {
HNSWStats stats;
@@ -1272,7 +1272,7 @@ HNSWStats HNSW::search(
void HNSW::search_level_0(
DistanceComputer& qdis,
ResultHandler<C>& res,
ResultHandler& res,
idx_t nprobe,
const storage_idx_t* nearest_i,
const float* nearest_d,
faiss/impl/HNSW.h
+4 -6
View File
@@ -45,8 +45,6 @@ struct IndexHNSWFlatPanorama;
struct VisitedTable;
struct DistanceComputer; // from AuxIndexStructures
struct HNSWStats;
template <class C>
struct ResultHandler;
struct SearchParametersHNSW : SearchParameters {
int efSearch = 16;
@@ -212,14 +210,14 @@ struct HNSW {
HNSWStats search(
DistanceComputer& qdis,
const IndexHNSW* index,
ResultHandler<C>& res,
ResultHandler& res,
VisitedTable& vt,
const SearchParameters* params = nullptr) const;
/// search only in level 0 from a given vertex
void search_level_0(
DistanceComputer& qdis,
ResultHandler<C>& res,
ResultHandler& res,
idx_t nprobe,
const storage_idx_t* nearest_i,
const float* nearest_d,
@@ -272,7 +270,7 @@ FAISS_API extern HNSWStats hnsw_stats;
int search_from_candidates(
const HNSW& hnsw,
DistanceComputer& qdis,
ResultHandler<HNSW::C>& res,
ResultHandler& res,
HNSW::MinimaxHeap& candidates,
VisitedTable& vt,
HNSWStats& stats,
@@ -288,7 +286,7 @@ int search_from_candidates_panorama(
const HNSW& hnsw,
const IndexHNSW* index,
DistanceComputer& qdis,
ResultHandler<HNSW::C>& res,
ResultHandler& res,
HNSW::MinimaxHeap& candidates,
VisitedTable& vt,
HNSWStats& stats,
faiss/impl/ResultHandler.h
+131 -50
View File
@@ -12,16 +12,42 @@
#pragma once
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/FaissException.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/partitioning.h>
#include <algorithm>
#include <iostream>
namespace faiss {
/* result handlers for a single query */
template <typename T, typename TI>
struct ResultHandlerUnordered {
// if not better than threshold, then not necessary to call add_result
T threshold{};
// return whether threshold was updated
virtual bool add_result(T dis, TI idx) = 0;
virtual ~ResultHandlerUnordered() {}
};
// handler for a single query
template <class C>
struct ResultHandlerT : ResultHandlerUnordered<typename C::T, typename C::TI> {
ResultHandlerT() {
this->threshold = C::neutral();
}
virtual ~ResultHandlerT() {}
};
// generic handler for searching in float indexes
using ResultHandler = ResultHandlerUnordered<float, idx_t>;
/*****************************************************************
* The classes below are intended to be used as template arguments
* they handle results for batches of queries (size nq).
@@ -67,16 +93,46 @@ struct BlockResultHandler {
}
};
// handler for a single query
template <class C>
struct ResultHandler {
// if not better than threshold, then not necessary to call add_result
typename C::T threshold = C::neutral();
/*****************************************************************
* Convenience classes for float Indexes
*****************************************************************/
// return whether threshold was updated
virtual bool add_result(typename C::T dis, typename C::TI idx) = 0;
// fake BlockResultHandler for a single query
template <class C, bool use_sel>
struct SingleQueryBlockResultHandler : BlockResultHandler<C, use_sel> {
using T = typename C::T;
using TI = typename C::TI;
using BlockResultHandler<C, use_sel>::i0;
using BlockResultHandler<C, use_sel>::i1;
virtual ~ResultHandler() {}
ResultHandler& the_handler;
explicit SingleQueryBlockResultHandler(
ResultHandler& the_handler,
const IDSelector* sel = nullptr)
: BlockResultHandler<C, use_sel>(1, sel),
the_handler(the_handler) {}
struct SingleResultHandler : ResultHandlerT<C> {
ResultHandler& the_handler;
using ResultHandlerT<C>::threshold;
explicit SingleResultHandler(SingleQueryBlockResultHandler& hr)
: the_handler(hr.the_handler) {}
/// begin results for query # i
void begin(const size_t qid) {
assert(qid == 0);
}
/// add one result for query i
bool add_result(T dis, TI idx) final {
return the_handler.add_result(dis, idx);
}
/// series of results for query i is done
void end() {}
};
};
/*****************************************************************
@@ -127,9 +183,9 @@ struct Top1BlockResultHandler : TopkBlockResultHandler<C, use_sel> {
: TopkBlockResultHandler<C, use_sel>(nq, dis_tab, ids_tab, 1, sel) {
}
struct SingleResultHandler : ResultHandler<C> {
struct SingleResultHandler : ResultHandlerT<C> {
Top1BlockResultHandler& hr;
using ResultHandler<C>::threshold;
using ResultHandlerT<C>::threshold;
TI min_idx;
size_t current_idx = 0;
@@ -204,6 +260,34 @@ struct Top1BlockResultHandler : TopkBlockResultHandler<C, use_sel> {
* Heap based result handler
*****************************************************************/
template <class C, bool use_sel = false>
struct HeapResultHandler : ResultHandlerT<C> {
using T = typename C::T;
using TI = typename C::TI;
using ResultHandlerT<C>::threshold;
size_t k;
T* heap_dis;
TI* heap_ids;
HeapResultHandler(size_t k, T* heap_dis, TI* heap_ids)
: k(k), heap_dis(heap_dis), heap_ids(heap_ids) {
if (heap_dis) {
this->threshold = heap_dis[0];
}
}
/// add one result for query i
bool add_result(T dis, TI idx) final {
if (C::cmp(threshold, dis)) {
heap_replace_top<C>(k, heap_dis, heap_ids, dis, idx);
threshold = heap_dis[0];
return true;
}
return false;
}
};
template <class C, bool use_sel = false>
struct HeapBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
using T = typename C::T;
@@ -220,43 +304,30 @@ struct HeapBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
const IDSelector* sel = nullptr)
: TopkBlockResultHandler<C, use_sel>(nq, dis_tab, ids_tab, k, sel) {
}
/******************************************************
* API for 1 result at a time (each SingleResultHandler is
* called from 1 thread)
*/
struct SingleResultHandler : ResultHandler<C> {
struct SingleResultHandler : HeapResultHandler<C, use_sel> {
HeapBlockResultHandler& hr;
using ResultHandler<C>::threshold;
size_t k;
T* heap_dis;
TI* heap_ids;
explicit SingleResultHandler(HeapBlockResultHandler& hr)
: hr(hr), k(hr.k) {}
: HeapResultHandler<C, use_sel>(hr.k, nullptr, nullptr),
hr(hr) {}
/// begin results for query # i
void begin(size_t i) {
heap_dis = hr.dis_tab + i * k;
heap_ids = hr.ids_tab + i * k;
heap_heapify<C>(k, heap_dis, heap_ids);
threshold = heap_dis[0];
}
/// add one result for query i
bool add_result(T dis, TI idx) final {
if (C::cmp(threshold, dis)) {
heap_replace_top<C>(k, heap_dis, heap_ids, dis, idx);
threshold = heap_dis[0];
return true;
}
return false;
this->heap_dis = hr.dis_tab + i * this->k;
this->heap_ids = hr.ids_tab + i * this->k;
heap_heapify<C>(this->k, this->heap_dis, this->heap_ids);
this->threshold = this->heap_dis[0];
}
/// series of results for query i is done
void end() {
heap_reorder<C>(k, heap_dis, heap_ids);
heap_reorder<C>(this->k, this->heap_dis, this->heap_ids);
}
};
@@ -312,10 +383,10 @@ struct HeapBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
/// Reservoir for a single query
template <class C>
struct ReservoirTopN : ResultHandler<C> {
struct ReservoirTopN : ResultHandlerT<C> {
using T = typename C::T;
using TI = typename C::TI;
using ResultHandler<C>::threshold;
using ResultHandlerT<C>::threshold;
T* vals;
TI* ids;
@@ -488,6 +559,27 @@ struct ReservoirBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
* Result handler for range searches
*****************************************************************/
template <class C, bool use_sel = false>
struct RangeResultHandler : ResultHandlerT<C> {
using T = typename C::T;
using TI = typename C::TI;
using ResultHandlerT<C>::threshold;
RangeQueryResult* qr = nullptr;
RangeResultHandler(RangeQueryResult* qr, T threshold) : qr(qr) {
this->threshold = threshold;
}
/// add one result for query i
bool add_result(T dis, TI idx) final {
if (C::cmp(threshold, dis)) {
qr->add(dis, idx);
}
return false;
}
};
template <class C, bool use_sel = false>
struct RangeSearchBlockResultHandler : BlockResultHandler<C, use_sel> {
using T = typename C::T;
@@ -511,28 +603,17 @@ struct RangeSearchBlockResultHandler : BlockResultHandler<C, use_sel> {
* called from 1 thread)
******************************************************/
struct SingleResultHandler : ResultHandler<C> {
struct SingleResultHandler : RangeResultHandler<C> {
// almost the same interface as RangeSearchResultHandler
using ResultHandler<C>::threshold;
using ResultHandlerT<C>::threshold;
RangeSearchPartialResult pres;
RangeQueryResult* qr = nullptr;
explicit SingleResultHandler(RangeSearchBlockResultHandler& rh)
: pres(rh.res) {
threshold = rh.radius;
}
: RangeResultHandler<C>(nullptr, rh.radius), pres(rh.res) {}
/// begin results for query # i
void begin(size_t i) {
qr = &pres.new_result(i);
}
/// add one result for query i
bool add_result(T dis, TI idx) final {
if (C::cmp(threshold, dis)) {
qr->add(dis, idx);
}
return false;
this->qr = &pres.new_result(i);
}
/// series of results for query i is done
faiss/impl/ScalarQuantizer.cpp
-87
View File
@@ -2161,50 +2161,6 @@ struct IVFSQScannerIP : InvertedListScanner {
float distance_to_code(const uint8_t* code) const final {
return accu0 + dc.query_to_code(code);
}
size_t scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float* simi,
idx_t* idxi,
size_t k) const override {
size_t nup = 0;
for (size_t j = 0; j < list_size; j++, codes += code_size) {
if (use_sel && !sel->is_member(use_sel == 1 ? ids[j] : j)) {
continue;
}
float accu = accu0 + dc.query_to_code(codes);
if (accu > simi[0]) {
int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
minheap_replace_top(k, simi, idxi, accu, id);
nup++;
}
}
return nup;
}
void scan_codes_range(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float radius,
RangeQueryResult& res) const override {
for (size_t j = 0; j < list_size; j++, codes += code_size) {
if (use_sel && !sel->is_member(use_sel == 1 ? ids[j] : j)) {
continue;
}
float accu = accu0 + dc.query_to_code(codes);
if (accu > radius) {
int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
res.add(accu, id);
}
}
}
};
/* use_sel = 0: don't check selector
@@ -2260,49 +2216,6 @@ struct IVFSQScannerL2 : InvertedListScanner {
float distance_to_code(const uint8_t* code) const final {
return dc.query_to_code(code);
}
size_t scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float* simi,
idx_t* idxi,
size_t k) const override {
size_t nup = 0;
for (size_t j = 0; j < list_size; j++, codes += code_size) {
if (use_sel && !sel->is_member(use_sel == 1 ? ids[j] : j)) {
continue;
}
float dis = dc.query_to_code(codes);
if (dis < simi[0]) {
int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
maxheap_replace_top(k, simi, idxi, dis, id);
nup++;
}
}
return nup;
}
void scan_codes_range(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
float radius,
RangeQueryResult& res) const override {
for (size_t j = 0; j < list_size; j++, codes += code_size) {
if (use_sel && !sel->is_member(use_sel == 1 ? ids[j] : j)) {
continue;
}
float dis = dc.query_to_code(codes);
if (dis < radius) {
int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
res.add(dis, id);
}
}
}
};
template <class DCClass, int use_sel>
tests/CMakeLists.txt
+1
View File
@@ -38,6 +38,7 @@ set(FAISS_TEST_SRC
test_hamming.cpp
test_mmap.cpp
test_zerocopy.cpp
test_custom_result_handler.cpp
)
if(FAISS_ENABLE_SVS)
tests/test_custom_result_handler.cpp
+236
View File
@@ -0,0 +1,236 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <limits>
#include <memory>
#include <unordered_set>
#include <vector>
#include <gtest/gtest.h>
#include <faiss/AutoTune.h>
#include <faiss/IndexIVF.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/index_factory.h>
#include <faiss/utils/random.h>
using namespace faiss;
/** A ResultHandler that just collects all results that presented to it. */
struct CollectAllResultHandler : ResultHandler {
std::vector<float> D;
std::vector<idx_t> I;
bool add_result(float distance, idx_t i) override {
// we never change the default threshold so that all results pass
D.push_back(distance);
I.push_back(i);
return true;
}
// Sort results by distance and return the top k
void get_top_k(
size_t top_k,
std::vector<float>& out_D,
std::vector<idx_t>& out_I,
bool is_max) const {
size_t n_results = D.size();
// Create indices for sorting
std::vector<size_t> indices(n_results);
for (size_t i = 0; i < n_results; i++) {
indices[i] = i;
}
// Make local copies for lambda capture
const std::vector<float>& distances = D;
// Sort by distance (ascending for L2, descending for IP)
if (is_max) {
std::sort(
indices.begin(),
indices.end(),
[&distances](size_t a, size_t b) {
return distances[a] > distances[b];
});
} else {
std::sort(
indices.begin(),
indices.end(),
[&distances](size_t a, size_t b) {
return distances[a] < distances[b];
});
}
// Take top k
size_t n_out = std::min(top_k, n_results);
out_D.resize(n_out);
out_I.resize(n_out);
for (size_t i = 0; i < n_out; i++) {
out_D[i] = D[indices[i]];
out_I[i] = I[indices[i]];
}
}
};
namespace {
// dimension of the vectors to index
constexpr int d = 64;
// nb of training vectors
constexpr size_t nt = 5000;
// size of the database
constexpr size_t nb = 2000;
// nb of queries
constexpr size_t nq = 100;
// k for search
constexpr int k = 10;
/**
* Generate smooth random data using Faiss's rand_smooth_vectors.
* This produces data with intrinsic dimensionality ~10 that is harder to
* index than a subspace but easier than uniform random data.
*/
std::vector<float> make_smooth_data(size_t n, int64_t seed) {
std::vector<float> data(n * d);
rand_smooth_vectors(n, d, data.data(), seed);
return data;
}
/**
* Test helper: trains an index, adds data, and performs a search using
* both the standard search method and the custom handler search1 method.
* Compares the results to ensure they match.
*/
void test_index(
const char* index_key,
MetricType metric,
double min_match_ratio = 1.0) {
// Create index using factory
std::unique_ptr<Index> index(index_factory(d, index_key, metric));
ASSERT_NE(index, nullptr) << "Failed to create IVF index for " << index_key;
// Generate smooth random data for training and database
auto xt = make_smooth_data(nt, 1234);
auto xb = make_smooth_data(nb, 4567);
auto xq = make_smooth_data(nq, 7890);
// Train the index
index->train(nt, xt.data());
// Add database vectors
index->add(nb, xb.data());
// Set nprobe for IVF indexes
if (IndexIVF* index_ivf = dynamic_cast<IndexIVF*>(index.get())) {
index_ivf->nprobe = 4;
}
// Perform reference search using standard search method
std::vector<idx_t> Iref(nq * k);
std::vector<float> Dref(nq * k);
index->search(nq, xq.data(), k, Dref.data(), Iref.data());
// For IP metric, we need to sort in descending order
bool is_max = (metric == METRIC_INNER_PRODUCT);
// Now test search1 with custom handler for each query
for (size_t q = 0; q < nq; q++) {
CollectAllResultHandler handler;
// Set threshold to collect all results
// For L2: use max float (condition: threshold > distance)
// For IP: use min float (condition: threshold < distance)
if (is_max) {
handler.threshold = -std::numeric_limits<float>::max();
} else {
handler.threshold = std::numeric_limits<float>::max();
}
index->search1(xq.data() + q * d, handler);
// Sort the handler results and get top k
std::vector<float> D_handler;
std::vector<idx_t> I_handler;
handler.get_top_k(k, D_handler, I_handler, is_max);
// Compare with reference results for this query using set comparison
const idx_t* Iref_q = Iref.data() + q * k;
// Create sets of IDs for comparison
std::unordered_set<idx_t> ref_set(Iref_q, Iref_q + k);
std::unordered_set<idx_t> handler_set(
I_handler.begin(), I_handler.end());
// Count how many results from handler are in the reference set
int matches = 0;
for (idx_t id : handler_set) {
if (ref_set.count(id) > 0) {
matches++;
}
}
// Check that matches meet the minimum threshold
EXPECT_GE(matches, static_cast<int>(k * min_match_ratio))
<< "Query " << q << ": expected at least "
<< static_cast<int>(k * min_match_ratio) << " matches, got "
<< matches;
}
}
/*************************************************************
* Test cases for different IVF index types
*************************************************************/
TEST(TestIndexTypes, IVFFlat_L2) {
test_index("IVF32,Flat", METRIC_L2);
}
TEST(TestIndexTypes, IVFFlat_IP) {
test_index("IVF32,Flat", METRIC_INNER_PRODUCT);
}
TEST(TestIndexTypes, IVFPQ_L2) {
test_index("IVF32,PQ8np", METRIC_L2);
}
TEST(TestIndexTypes, IVFPQ_IP) {
test_index("IVF32,PQ8np", METRIC_INNER_PRODUCT);
}
TEST(TestIndexTypes, IVFSQ_L2) {
test_index("IVF32,SQ8", METRIC_L2);
}
TEST(TestIndexTypes, IVFRaBitQ_IP) {
test_index("IVF32,RaBitQ", METRIC_INNER_PRODUCT);
}
TEST(TestIndexTypes, IVFRaBitQ4_L2) {
// RaBitQ4 does a topk per invlist, so not exactly the same as top over all
test_index("IVF32,RaBitQ4", METRIC_L2, 0.8);
}
TEST(TestIndexTypes, IVFRaBitQ4_IP) {
// RaBitQ4 does a topk per invlist, so not exactly the same as top over all
test_index("IVF32,RaBitQ4", METRIC_INNER_PRODUCT, 0.8);
}
TEST(TestIndexTypes, HNSW) {
test_index("HNSW32,Flat", METRIC_L2);
}
TEST(TestIndexTypes, SQ8) {
test_index("SQ8", METRIC_L2);
}
} // namespace
tests/test_hnsw.cpp
+1 -1
View File
@@ -249,7 +249,7 @@ class HNSWTest : public testing::Test {
int reference_search_from_candidates(
const faiss::HNSW& hnsw,
faiss::DistanceComputer& qdis,
faiss::ResultHandler<faiss::HNSW::C>& res,
faiss::ResultHandler& res,
faiss::HNSW::MinimaxHeap& candidates,
faiss::VisitedTable& vt,
faiss::HNSWStats& stats,
tests/test_index_accuracy.py
+2 -2
View File
@@ -374,8 +374,8 @@ class TestSQByte(unittest.TestCase):
index.add(xb)
D, I = index.search(xq, 10)
assert np.all(I == Iref)
assert np.all(D == Dref)
np.testing.assert_array_equal(D, Dref)
np.testing.assert_array_equal(I, Iref)
def test_8bit_direct(self):
for quantizer in faiss.ScalarQuantizer.QT_8bit_direct, faiss.ScalarQuantizer.QT_8bit_direct_signed: