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
Files
94a4331fee6524272a108d8eb5c176133a29c6b2
faiss/faiss/IndexIVFSQFastScan.cpp
T
Michael Norris 94a4331fee Add reranked types to IndexSQFastScan and fix flat RaBitQ filtering (#5661) 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5661

Two gaps in the flat fast-scan indexes, both found while benchmarking
IDSelector filtering.

`IndexSQFastScan` accepted only the native 4-bit quantizers. It now accepts
the reranked types as well, the same set `IndexIVFSQFastScan` handles:
QT_6bit, QT_8bit, QT_8bit_uniform, QT_8bit_direct and QT_8bit_direct_signed.
Each code is SPLIT: the SIMD scan reads the top 4 bits, and `lo_codes` keeps
the rest for rerank. The total is the quantizer's own width, so `SQ8fs` holds
8 bits per dimension, the same as plain `SQ8`. A flat index needs no direct
map for this, unlike the IVF form, because the position is the id. `SQ8fs` and
`SQ6fs` now resolve in `index_factory`.

`IndexRaBitQFastScan` forwarded an `IDSelector` through every layer but its
kernel never applied it, so a filtered search returned ids outside the
admitted set. `RaBitQHeapHandler::handle` now tests the selector per code,
after the SIMD block, which is where the PQ and SQ handlers test theirs. The
entry point no longer refuses a parameter object.

The type predicates and the nibble packing move to `impl/sq_fastscan_utils.h`.
The flat and IVF classes then cannot disagree about which quantizers rerank,
or about how a nibble is laid out.

The `ISfs` payload gains the rerank factor and the original codes. Both are
written for every quantizer, and the codes are empty for a native 4-bit index.
`ISfs` has not been released, so no stored index is affected.

Reviewed By: junjieqi

Differential Revision: D121036639

fbshipit-source-id: 40e51fc348a483cf5da947fa7db56ff10f1af773
2026-09-25 13:47:19 -07:00

802 lines
26 KiB
C++
Raw Blame History

/*
* 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 <faiss/IndexIVFSQFastScan.h>
#include <faiss/impl/sq_fastscan_utils.h>
#include <omp.h>
#include <algorithm>
#include <cstring>
#include <memory>
#include <vector>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/impl/fast_scan/fast_scan.h>
#include <faiss/invlists/BlockInvertedLists.h>
#include <faiss/invlists/DirectMap.h>
#include <faiss/invlists/InvertedLists.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/hamming.h>
#include <faiss/utils/utils.h>
namespace faiss {
using namespace faiss::sq_fastscan;
// -----------------------------------------------------------------------
// Constructors
// -----------------------------------------------------------------------
IndexIVFSQFastScan::IndexIVFSQFastScan(
Index* quantizer_in,
size_t d_in,
size_t nlist_in,
ScalarQuantizer::QuantizerType qtype,
MetricType metric,
int bbs_in,
bool by_residual_in)
: IndexIVFFastScan(quantizer_in, d_in, nlist_in, 0, metric, false),
sq(d_in, qtype) {
if (is_fallback(qtype)) {
by_residual = by_residual_in;
if (ScalarQuantizer::TurboQuantRefine::is_turboq_full(qtype)) {
by_residual = false;
}
} else {
// FastScan path: by_residual=false because our 2D LUT cannot
// handle per-probe residuals (same rationale as IndexIVFPQFastScan).
by_residual = false;
}
if (is_fallback(qtype)) {
// Fallback: use standard ArrayInvertedLists with full SQ codes.
// Search uses IndexIVF's InvertedListScanner path.
code_size = sq.code_size;
replace_invlists(new ArrayInvertedLists(nlist_in, code_size), true);
} else {
size_t M_in = d_in;
M = M_in;
nbits = 4;
bbs = bbs_in;
ksub = 16;
M2 = roundup(M_in, 2);
code_size = M2 / 2;
replace_invlists(
new BlockInvertedLists(nlist_in, get_CodePacker()), true);
if (needs_rerank(qtype)) {
orig_codes_invlists =
new ArrayInvertedLists(nlist_in, sq.code_size);
direct_map.set_type(DirectMap::Hashtable, invlists, 0);
}
}
}
IndexIVFSQFastScan::IndexIVFSQFastScan() = default;
IndexIVFSQFastScan::IndexIVFSQFastScan(
const IndexIVFScalarQuantizer& orig,
int bbs_in)
: IndexIVFFastScan(
orig.quantizer,
orig.d,
orig.nlist,
0,
orig.metric_type,
false),
sq(orig.sq) {
ntotal = orig.ntotal;
is_trained = orig.is_trained;
nprobe = orig.nprobe;
if (is_fallback(sq.qtype)) {
by_residual = orig.by_residual;
code_size = sq.code_size;
replace_invlists(new ArrayInvertedLists(nlist, code_size), true);
for (size_t list_no = 0; list_no < nlist; list_no++) {
size_t list_size = orig.invlists->list_size(list_no);
if (list_size == 0) {
continue;
}
InvertedLists::ScopedCodes codes(orig.invlists, list_no);
InvertedLists::ScopedIds ids(orig.invlists, list_no);
invlists->add_entries(list_no, list_size, ids.get(), codes.get());
}
return;
}
FAISS_THROW_IF_NOT_MSG(
!orig.by_residual,
"IndexIVFSQFastScan: conversion from IndexIVFScalarQuantizer "
"with by_residual=true is not supported for native/rerank "
"types. Set orig.by_residual=false and retrain, or construct "
"IndexIVFSQFastScan directly.");
by_residual = false;
size_t M_in = d;
M = M_in;
nbits = 4;
bbs = bbs_in;
ksub = 16;
M2 = roundup(M_in, 2);
code_size = M2 / 2;
replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true);
if (needs_rerank(sq.qtype)) {
orig_codes_invlists = new ArrayInvertedLists(nlist, sq.code_size);
direct_map.set_type(DirectMap::Hashtable, invlists, 0);
}
for (size_t list_no = 0; list_no < nlist; list_no++) {
size_t list_size = orig.invlists->list_size(list_no);
if (list_size == 0) {
continue;
}
InvertedLists::ScopedCodes orig_codes(orig.invlists, list_no);
InvertedLists::ScopedIds orig_ids(orig.invlists, list_no);
if (needs_rerank(sq.qtype)) {
orig_codes_invlists->add_entries(
list_no, list_size, orig_ids.get(), orig_codes.get());
for (size_t j = 0; j < list_size; j++) {
direct_map.add_single_id(orig_ids.get()[j], list_no, j);
}
}
// Decode to float and re-quantize to 4-bit.
std::vector<float> recon(list_size * d);
sq.decode(orig_codes.get(), recon.data(), list_size);
std::vector<uint8_t> flat4(list_size * M2 / 2, 0);
if (is_native_4bit(sq.qtype)) {
sq.compute_codes(recon.data(), flat4.data(), list_size);
} else {
float_to_4bit_nibbles(
recon.data(), flat4.data(), list_size, d, M2, sq);
}
BlockInvertedLists* bil = dynamic_cast<BlockInvertedLists*>(invlists);
FAISS_THROW_IF_NOT(bil);
size_t nb2 = roundup(list_size, bbs);
AlignedTable<uint8_t> packed(nb2 * M2 / 2);
pq4_pack_codes(flat4.data(), list_size, M, nb2, bbs, M2, packed.get());
bil->add_entries(list_no, list_size, orig_ids.get(), packed.get());
}
}
IndexIVFSQFastScan::~IndexIVFSQFastScan() {
delete orig_codes_invlists;
}
// -----------------------------------------------------------------------
// Virtual method implementations
// -----------------------------------------------------------------------
size_t IndexIVFSQFastScan::fast_scan_code_size() const {
return M2 / 2;
}
void IndexIVFSQFastScan::train_encoder(
idx_t n,
const float* x,
const idx_t* /*assign*/) {
sq.train(n, x);
}
idx_t IndexIVFSQFastScan::train_encoder_num_vectors() const {
return 100000;
}
void IndexIVFSQFastScan::encode_vectors(
idx_t n,
const float* x,
const idx_t* list_nos,
uint8_t* codes_out,
bool include_listnos) const {
std::vector<float> residuals;
const float* to_encode = x;
if (by_residual) {
residuals.resize(n * d);
for (idx_t i = 0; i < n; i++) {
if (list_nos[i] >= 0) {
quantizer->compute_residual(
x + i * d, residuals.data() + i * d, list_nos[i]);
} else {
memset(residuals.data() + i * d, 0, sizeof(float) * d);
}
}
to_encode = residuals.data();
}
if (is_fallback(sq.qtype)) {
sq.compute_codes(to_encode, codes_out, n);
} else if (is_native_4bit(sq.qtype)) {
sq.compute_codes(to_encode, codes_out, n);
} else {
float_to_4bit_nibbles(to_encode, codes_out, n, d, M2, sq);
}
if (include_listnos) {
size_t coarse_size = coarse_code_size();
for (idx_t i = n - 1; i >= 0; i--) {
uint8_t* code = codes_out + i * (coarse_size + code_size);
memmove(code + coarse_size, codes_out + i * code_size, code_size);
encode_listno(list_nos[i], code);
}
}
}
bool IndexIVFSQFastScan::lookup_table_is_3d() const {
// For SQ, the LUT doesn't depend on which list we're scanning
// (unlike PQ with residuals and precomputed tables).
return false;
}
void IndexIVFSQFastScan::compute_LUT(
size_t n,
const float* x,
const CoarseQuantized& /*cq*/,
AlignedTable<float>& dis_tables,
AlignedTable<float>& /*biases*/,
const FastScanDistancePostProcessing&) const {
size_t dim = d;
size_t ksub_val = 16;
size_t dim12 = dim * ksub_val;
// Build reconstruction table: for each dimension, the 16 centroid values.
bool is_uniform = is_uniform_range(sq.qtype) ||
sq.qtype == ScalarQuantizer::QT_4bit_uniform;
std::vector<float> recon_table(dim * ksub_val);
if (is_uniform) {
float vmin, vdiff;
get_uniform_range(sq, vmin, vdiff);
for (size_t c = 0; c < ksub_val; c++) {
float recon = vmin + ((c + 0.5f) / 15.0f) * vdiff;
for (size_t m = 0; m < dim; m++) {
recon_table[m * ksub_val + c] = recon;
}
}
} else {
const float* vmin = sq.trained.data();
const float* vdiff = sq.trained.data() + dim;
for (size_t m = 0; m < dim; m++) {
for (size_t c = 0; c < ksub_val; c++) {
recon_table[m * ksub_val + c] =
vmin[m] + ((c + 0.5f) / 15.0f) * vdiff[m];
}
}
}
// by_residual is always false for the fast-scan path (non-fallback
// types), so codes store raw quantized values and the LUT computes
// distances directly from the raw query. No biases needed.
dis_tables.resize(n * dim12);
if (metric_type == METRIC_L2) {
for (size_t i = 0; i < n; i++) {
const float* xi = x + i * dim;
float* lut_i = dis_tables.get() + i * dim12;
for (size_t m = 0; m < dim; m++) {
float qi = xi[m];
const float* recon_m = recon_table.data() + m * ksub_val;
float* lut_m = lut_i + m * ksub_val;
for (size_t c = 0; c < ksub_val; c++) {
float diff = qi - recon_m[c];
lut_m[c] = diff * diff;
}
}
}
} else {
for (size_t i = 0; i < n; i++) {
const float* xi = x + i * dim;
float* lut_i = dis_tables.get() + i * dim12;
for (size_t m = 0; m < dim; m++) {
float qi = xi[m];
const float* recon_m = recon_table.data() + m * ksub_val;
float* lut_m = lut_i + m * ksub_val;
for (size_t c = 0; c < ksub_val; c++) {
lut_m[c] = qi * recon_m[c];
}
}
}
}
}
void IndexIVFSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
const {
size_t coarse_size = coarse_code_size();
#pragma omp parallel if (n > 1)
{
std::vector<float> residual(d);
#pragma omp for
for (idx_t i = 0; i < n; i++) {
const uint8_t* code = bytes + i * (code_size + coarse_size);
int64_t list_no = decode_listno(code);
float* xi = x + i * d;
sq.decode(code + coarse_size, xi, 1);
if (by_residual) {
quantizer->reconstruct(list_no, residual.data());
for (int j = 0; j < d; j++) {
xi[j] += residual[j];
}
}
}
}
}
void IndexIVFSQFastScan::reconstruct_from_offset(
int64_t list_no,
int64_t offset,
float* recons) const {
if (is_fallback(sq.qtype)) {
// Fallback: codes stored directly in invlists.
InvertedLists::ScopedCodes codes(invlists, list_no);
sq.decode(codes.get() + offset * sq.code_size, recons, 1);
} else if (orig_codes_invlists) {
// Rerank: use original full-precision codes for reconstruction.
InvertedLists::ScopedCodes codes(orig_codes_invlists, list_no);
sq.decode(codes.get() + offset * sq.code_size, recons, 1);
} else {
// Native 4-bit: unpack from block inverted lists.
std::vector<uint8_t> code(M2 / 2, 0);
InvertedLists::ScopedCodes list_codes(invlists, list_no);
BitstringWriter bsw(code.data(), M2 / 2);
for (size_t m = 0; m < M; m++) {
uint8_t c = pq4_get_packed_element(
list_codes.get(), bbs, M2, offset, m);
bsw.write(c, 4);
}
sq.decode(code.data(), recons, 1);
}
if (by_residual) {
std::vector<float> centroid(d);
quantizer->reconstruct(list_no, centroid.data());
for (int j = 0; j < d; j++) {
recons[j] += centroid[j];
}
}
}
InvertedListScanner* IndexIVFSQFastScan::get_InvertedListScanner(
bool store_pairs,
const IDSelector* sel,
const IVFSearchParameters*) const {
return sq.select_InvertedListScanner(
metric_type, quantizer, store_pairs, sel, by_residual);
}
// -----------------------------------------------------------------------
// search — fallback types use IndexIVF's scanner-based path
// -----------------------------------------------------------------------
void IndexIVFSQFastScan::search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels,
const SearchParameters* params) const {
if (is_fallback(sq.qtype)) {
IndexIVF::search(n, x, k, distances, labels, params);
} else if (needs_rerank(sq.qtype)) {
idx_t k_coarse = idx_t(double(k) * rerank_factor);
FAISS_THROW_IF_NOT(k_coarse >= k);
FAISS_THROW_IF_NOT_MSG(
n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64");
std::vector<float> coarse_dis(n * k_coarse);
std::vector<idx_t> coarse_ids(n * k_coarse);
IndexIVFFastScan::search(
n, x, k_coarse, coarse_dis.data(), coarse_ids.data(), params);
// Rerank using original full-precision codes
#pragma omp parallel
{
std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
sq.get_distance_computer(metric_type));
#pragma omp for
for (idx_t i = 0; i < n; i++) {
const float* xi = x + i * d;
dc->set_query(xi);
float* heap_dis = distances + i * k;
idx_t* heap_ids = labels + i * k;
if (metric_type == METRIC_L2) {
maxheap_heapify(k, heap_dis, heap_ids);
} else {
minheap_heapify(k, heap_dis, heap_ids);
}
for (idx_t j = 0; j < k_coarse; j++) {
idx_t id = coarse_ids[i * k_coarse + j];
if (id < 0) {
continue;
}
idx_t lo = direct_map.get(id);
idx_t list_no = lo_listno(lo);
idx_t offset = lo_offset(lo);
InvertedLists::ScopedCodes codes(
orig_codes_invlists, list_no);
float dis = dc->distance_to_code(
codes.get() + offset * sq.code_size);
if (metric_type == METRIC_L2) {
if (dis < heap_dis[0]) {
maxheap_replace_top(k, heap_dis, heap_ids, dis, id);
}
} else {
if (dis > heap_dis[0]) {
minheap_replace_top(k, heap_dis, heap_ids, dis, id);
}
}
}
if (metric_type == METRIC_L2) {
maxheap_reorder(k, heap_dis, heap_ids);
} else {
minheap_reorder(k, heap_dis, heap_ids);
}
}
}
} else {
IndexIVFFastScan::search(n, x, k, distances, labels, params);
}
}
void IndexIVFSQFastScan::search_preassigned(
idx_t n,
const float* x,
idx_t k,
const idx_t* assign,
const float* centroid_dis,
float* distances,
idx_t* labels,
bool store_pairs,
const IVFSearchParameters* params,
IndexIVFStats* stats) const {
if (is_fallback(sq.qtype)) {
IndexIVF::search_preassigned(
n,
x,
k,
assign,
centroid_dis,
distances,
labels,
store_pairs,
params,
stats);
} else if (needs_rerank(sq.qtype)) {
idx_t k_coarse = idx_t(double(k) * rerank_factor);
FAISS_THROW_IF_NOT(k_coarse >= k);
FAISS_THROW_IF_NOT_MSG(
n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64");
std::vector<float> coarse_dis(n * k_coarse);
std::vector<idx_t> coarse_ids(n * k_coarse);
IndexIVFFastScan::search_preassigned(
n,
x,
k_coarse,
assign,
centroid_dis,
coarse_dis.data(),
coarse_ids.data(),
false,
params,
stats);
#pragma omp parallel
{
std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
sq.get_distance_computer(metric_type));
#pragma omp for
for (idx_t i = 0; i < n; i++) {
const float* xi = x + i * d;
dc->set_query(xi);
float* heap_dis = distances + i * k;
idx_t* heap_ids = labels + i * k;
if (metric_type == METRIC_L2) {
maxheap_heapify(k, heap_dis, heap_ids);
} else {
minheap_heapify(k, heap_dis, heap_ids);
}
for (idx_t j = 0; j < k_coarse; j++) {
idx_t id = coarse_ids[i * k_coarse + j];
if (id < 0) {
continue;
}
idx_t lo = direct_map.get(id);
idx_t list_no = lo_listno(lo);
idx_t offset = lo_offset(lo);
InvertedLists::ScopedCodes codes(
orig_codes_invlists, list_no);
float dis = dc->distance_to_code(
codes.get() + offset * sq.code_size);
if (metric_type == METRIC_L2) {
if (dis < heap_dis[0]) {
maxheap_replace_top(k, heap_dis, heap_ids, dis, id);
}
} else {
if (dis > heap_dis[0]) {
minheap_replace_top(k, heap_dis, heap_ids, dis, id);
}
}
}
if (metric_type == METRIC_L2) {
maxheap_reorder(k, heap_dis, heap_ids);
} else {
minheap_reorder(k, heap_dis, heap_ids);
}
}
}
} else {
IndexIVFFastScan::search_preassigned(
n,
x,
k,
assign,
centroid_dis,
distances,
labels,
store_pairs,
params,
stats);
}
}
void IndexIVFSQFastScan::range_search(
idx_t n,
const float* x,
float radius,
RangeSearchResult* result,
const SearchParameters* params) const {
if (is_fallback(sq.qtype)) {
IndexIVF::range_search(n, x, radius, result, params);
} else if (needs_rerank(sq.qtype)) {
const IVFSearchParameters* ivf_params = nullptr;
if (params) {
ivf_params = dynamic_cast<const IVFSearchParameters*>(params);
}
size_t cur_nprobe = ivf_params ? ivf_params->nprobe : nprobe;
const IDSelector* sel = params ? params->sel : nullptr;
std::vector<idx_t> assign(n * cur_nprobe);
std::vector<float> centroid_dis(n * cur_nprobe);
quantizer->search(n, x, cur_nprobe, centroid_dis.data(), assign.data());
std::vector<RangeSearchPartialResult*> partial_results(
omp_get_max_threads());
#pragma omp parallel
{
int rank = omp_get_thread_num();
RangeSearchPartialResult* pres =
new RangeSearchPartialResult(result);
partial_results[rank] = pres;
std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
sq.get_distance_computer(metric_type));
#pragma omp for
for (idx_t i = 0; i < n; i++) {
dc->set_query(x + i * d);
RangeQueryResult& qres = pres->new_result(i);
for (size_t p = 0; p < cur_nprobe; p++) {
idx_t list_no = assign[i * cur_nprobe + p];
if (list_no < 0) {
continue;
}
size_t list_size = orig_codes_invlists->list_size(list_no);
if (list_size == 0) {
continue;
}
InvertedLists::ScopedCodes codes(
orig_codes_invlists, list_no);
InvertedLists::ScopedIds ids(orig_codes_invlists, list_no);
for (size_t j = 0; j < list_size; j++) {
if (sel && !sel->is_member(ids[j])) {
continue;
}
float dis = dc->distance_to_code(
codes.get() + j * sq.code_size);
if (metric_type == METRIC_L2) {
if (dis < radius) {
qres.add(dis, ids[j]);
}
} else {
if (dis > radius) {
qres.add(dis, ids[j]);
}
}
}
}
}
}
RangeSearchPartialResult::merge(partial_results);
} else {
IndexIVFFastScan::range_search(n, x, radius, result, params);
}
}
// -----------------------------------------------------------------------
// add_with_ids — custom implementation to store both packed and original codes
// -----------------------------------------------------------------------
void IndexIVFSQFastScan::add_with_ids(
idx_t n,
const float* x,
const idx_t* xids) {
FAISS_THROW_IF_NOT(is_trained);
constexpr idx_t bs = 65536;
if (n > bs) {
for (idx_t i0 = 0; i0 < n; i0 += bs) {
idx_t i1 = std::min(n, i0 + bs);
add_with_ids(i1 - i0, x + i0 * d, xids ? xids + i0 : nullptr);
}
return;
}
// Assign to coarse clusters
std::unique_ptr<idx_t[]> idx(new idx_t[n]);
quantizer->assign(n, x, idx.get());
if (is_fallback(sq.qtype)) {
// Fallback: store SQ codes directly in invlists (same as
// IndexIVFScalarQuantizer).
std::vector<uint8_t> flat_codes(n * sq.code_size);
encode_vectors(n, x, idx.get(), flat_codes.data());
for (idx_t i = 0; i < n; i++) {
if (idx[i] < 0) {
continue;
}
idx_t id = xids ? xids[i] : ntotal + i;
invlists->add_entry(
idx[i], id, flat_codes.data() + i * sq.code_size);
}
ntotal += n;
return;
}
// Produce 4-bit packed codes for the SIMD scan,
// and (for rerank types) also store original codes.
AlignedTable<uint8_t> flat_codes(n * code_size);
encode_vectors(n, x, idx.get(), flat_codes.get());
// Store original SQ codes for rerank types
std::vector<uint8_t> orig_flat_codes;
if (needs_rerank(sq.qtype)) {
orig_flat_codes.resize(n * sq.code_size);
for (idx_t i = 0; i < n; i++) {
if (idx[i] < 0) {
continue;
}
if (by_residual) {
std::vector<float> residual(d);
quantizer->compute_residual(x + i * d, residual.data(), idx[i]);
sq.compute_codes(
residual.data(),
orig_flat_codes.data() + i * sq.code_size,
1);
} else {
sq.compute_codes(
x + i * d,
orig_flat_codes.data() + i * sq.code_size,
1);
}
}
}
// Sort by list assignment
std::vector<idx_t> order(n);
for (idx_t i = 0; i < n; i++) {
order[i] = i;
}
std::stable_sort(order.begin(), order.end(), [&idx](idx_t a, idx_t b) {
return idx[a] < idx[b];
});
BlockInvertedLists* bil = dynamic_cast<BlockInvertedLists*>(invlists);
FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported");
idx_t i0 = 0;
while (i0 < n) {
idx_t list_no = idx[order[i0]];
idx_t i1 = i0 + 1;
while (i1 < n && idx[order[i1]] == list_no) {
i1++;
}
if (list_no == -1) {
i0 = i1;
continue;
}
size_t list_size = bil->list_size(list_no);
bil->resize(list_no, list_size + i1 - i0);
AlignedTable<uint8_t> list_codes((i1 - i0) * code_size);
for (idx_t i = i0; i < i1; i++) {
size_t ofs = list_size + i - i0;
idx_t id = xids ? xids[order[i]] : ntotal + order[i];
bil->ids[list_no][ofs] = id;
memcpy(list_codes.data() + (i - i0) * code_size,
flat_codes.data() + order[i] * code_size,
code_size);
if (needs_rerank(sq.qtype)) {
orig_codes_invlists->add_entry(
list_no,
id,
orig_flat_codes.data() + order[i] * sq.code_size);
direct_map.add_single_id(id, list_no, ofs);
}
}
pq4_pack_codes_range(
list_codes.data(),
M,
list_size,
list_size + i1 - i0,
bbs,
M2,
bil->codes[list_no].data(),
0,
get_block_stride());
i0 = i1;
}
ntotal += n;
}
void IndexIVFSQFastScan::reset() {
IndexIVFFastScan::reset();
if (orig_codes_invlists) {
for (size_t i = 0; i < nlist; i++) {
orig_codes_invlists->resize(i, 0);
}
direct_map.clear();
}
}
} // namespace faiss
Reference in New Issue View Git Blame Copy Permalink