mirror of
https://github.com/facebookresearch/faiss.git
synced 2025-06-03 21:54:02 +08:00
5840c17332
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2787 Approximate top-k facilities allow to speedup IndexBinaryFlat search / hamming_knn_hc() search a bit, if the number k of needed nearest neighbors is small (no more than NB * BD for `ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_NB_BD`, otherwise an exception is thrown) Use case: ``` approx_mode: int = 0, 1, 2, 3 or 4 subq = faiss.downcast_IndexBinary(index.quantizer) subq.approx_topk_mode subq.approx_topk_mode = approx_mode ``` Please reference `faiss/utils/approx_topk/mode.h` for more detailed explanation of modes. Reviewed By: mdouze Differential Revision: D44398992 fbshipit-source-id: aac7c35c88d567ba73c88614a1c86f3e8b254ef3
117 lines
3.0 KiB
C++
117 lines
3.0 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++ -*-
|
|
|
|
#include <faiss/IndexBinaryFlat.h>
|
|
|
|
#include <faiss/impl/AuxIndexStructures.h>
|
|
#include <faiss/impl/FaissAssert.h>
|
|
#include <faiss/impl/IDSelector.h>
|
|
#include <faiss/utils/Heap.h>
|
|
#include <faiss/utils/hamming.h>
|
|
#include <faiss/utils/utils.h>
|
|
#include <cstring>
|
|
|
|
namespace faiss {
|
|
|
|
IndexBinaryFlat::IndexBinaryFlat(idx_t d) : IndexBinary(d) {}
|
|
|
|
void IndexBinaryFlat::add(idx_t n, const uint8_t* x) {
|
|
xb.insert(xb.end(), x, x + n * code_size);
|
|
ntotal += n;
|
|
}
|
|
|
|
void IndexBinaryFlat::reset() {
|
|
xb.clear();
|
|
ntotal = 0;
|
|
}
|
|
|
|
void IndexBinaryFlat::search(
|
|
idx_t n,
|
|
const uint8_t* x,
|
|
idx_t k,
|
|
int32_t* distances,
|
|
idx_t* labels,
|
|
const SearchParameters* params) const {
|
|
FAISS_THROW_IF_NOT_MSG(
|
|
!params, "search params not supported for this index");
|
|
FAISS_THROW_IF_NOT(k > 0);
|
|
|
|
const idx_t block_size = query_batch_size;
|
|
for (idx_t s = 0; s < n; s += block_size) {
|
|
idx_t nn = block_size;
|
|
if (s + block_size > n) {
|
|
nn = n - s;
|
|
}
|
|
|
|
if (use_heap) {
|
|
// We see the distances and labels as heaps.
|
|
int_maxheap_array_t res = {
|
|
size_t(nn), size_t(k), labels + s * k, distances + s * k};
|
|
|
|
hammings_knn_hc(
|
|
&res,
|
|
x + s * code_size,
|
|
xb.data(),
|
|
ntotal,
|
|
code_size,
|
|
/* ordered = */ true,
|
|
approx_topk_mode);
|
|
} else {
|
|
hammings_knn_mc(
|
|
x + s * code_size,
|
|
xb.data(),
|
|
nn,
|
|
ntotal,
|
|
k,
|
|
code_size,
|
|
distances + s * k,
|
|
labels + s * k);
|
|
}
|
|
}
|
|
}
|
|
|
|
size_t IndexBinaryFlat::remove_ids(const IDSelector& sel) {
|
|
idx_t j = 0;
|
|
for (idx_t i = 0; i < ntotal; i++) {
|
|
if (sel.is_member(i)) {
|
|
// should be removed
|
|
} else {
|
|
if (i > j) {
|
|
memmove(&xb[code_size * j],
|
|
&xb[code_size * i],
|
|
sizeof(xb[0]) * code_size);
|
|
}
|
|
j++;
|
|
}
|
|
}
|
|
long nremove = ntotal - j;
|
|
if (nremove > 0) {
|
|
ntotal = j;
|
|
xb.resize(ntotal * code_size);
|
|
}
|
|
return nremove;
|
|
}
|
|
|
|
void IndexBinaryFlat::reconstruct(idx_t key, uint8_t* recons) const {
|
|
memcpy(recons, &(xb[code_size * key]), sizeof(*recons) * code_size);
|
|
}
|
|
|
|
void IndexBinaryFlat::range_search(
|
|
idx_t n,
|
|
const uint8_t* x,
|
|
int radius,
|
|
RangeSearchResult* result,
|
|
const SearchParameters* params) const {
|
|
FAISS_THROW_IF_NOT_MSG(
|
|
!params, "search params not supported for this index");
|
|
hamming_range_search(x, xb.data(), n, ntotal, radius, code_size, result);
|
|
}
|
|
|
|
} // namespace faiss
|