mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Alibek Zhakubayev
d492af446e
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5249 Before this change, Top1BlockResultHandler::begin_multiple initialized dis_tab[i] to C::neutral() but left ids_tab[i] uninitialized for every query slot. This was inconsistent with two related code paths in the same class: 1. SingleResultHandler::begin() (the per-query API) explicitly sets min_idx = -1 before accumulating results. 2. HeapBlockResultHandler::begin_multiple() calls heap_heapify<C>() which sets ids_tab[i*k + j] = -1 for all k result slots. The inconsistency means that any caller invoking begin_multiple + end_multiple without an intervening add_results — for example, when the database is empty and there are no results to report — reads uninitialized memory from ids_tab instead of the expected -1 sentinel. Fix: add ids_tab[i] = -1 inside the begin_multiple loop, making the block API consistent with the single-query API and with the Heap handler. This is a strictly additive, zero-risk change: in the normal path where add_results is called and covers all queries, the -1 is overwritten with the actual winner index before end_multiple returns. Test: TestTop1BlockResultHandler/BeginMultipleInitializesIds — calls begin_multiple(0, nq) followed by end_multiple() with no add_results, then verifies every ids_tab[i] == -1 and dis_tab[i] == C::neutral(). Reviewed By: junjieqi Differential Revision: D106359898 fbshipit-source-id: f864016703aa98b0c3ff6226968bb61f361a45f7
80 lines
2.6 KiB
C++
80 lines
2.6 KiB
C++
/*
|
|
* 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/impl/ResultHandler.h>
|
|
#include <faiss/utils/Heap.h>
|
|
#include <gtest/gtest.h>
|
|
#include <algorithm>
|
|
#include <numeric>
|
|
|
|
using namespace faiss;
|
|
|
|
// Verify that Top1BlockResultHandler::begin_multiple initializes ids_tab to -1.
|
|
// The single-query path (SingleResultHandler::begin) explicitly sets
|
|
// min_idx=-1; begin_multiple must be consistent so that
|
|
// begin_multiple/end_multiple without an intervening add_results produces the
|
|
// "no result found" sentinel -1 in ids_tab rather than leaving it
|
|
// uninitialized.
|
|
TEST(Top1BlockResultHandler, BeginMultipleInitializesIds) {
|
|
constexpr size_t nq = 16;
|
|
std::vector<float> dis_tab(nq, 12345.0f);
|
|
std::vector<int64_t> ids_tab(nq, 42);
|
|
|
|
Top1BlockResultHandler<CMax<float, int64_t>> handler(
|
|
nq, dis_tab.data(), ids_tab.data());
|
|
|
|
handler.begin_multiple(0, nq);
|
|
handler.end_multiple();
|
|
|
|
const std::vector<int64_t> expected_ids(nq, -1);
|
|
EXPECT_EQ(ids_tab, expected_ids);
|
|
|
|
const std::vector<float> expected_dis(nq, CMax<float, int64_t>::neutral());
|
|
EXPECT_EQ(dis_tab, expected_dis);
|
|
}
|
|
|
|
TEST(Heap, addn_with_ids) {
|
|
size_t n = 1000;
|
|
size_t k = 1;
|
|
std::vector<int64_t> heap_labels(n, -1);
|
|
std::vector<float> heap_distances(n, 0);
|
|
float_minheap_array_t heaps = {
|
|
n, k, heap_labels.data(), heap_distances.data()};
|
|
heaps.heapify();
|
|
std::vector<int64_t> labels(n, 1);
|
|
std::vector<float> distances(n, 0.0f);
|
|
std::vector<int64_t> subset(n);
|
|
std::iota(subset.begin(), subset.end(), 0);
|
|
heaps.addn_with_ids(1, distances.data(), labels.data(), 1);
|
|
heaps.reorder();
|
|
EXPECT_TRUE(
|
|
std::all_of(heap_labels.begin(), heap_labels.end(), [](int64_t i) {
|
|
return i == 1;
|
|
}));
|
|
}
|
|
|
|
TEST(Heap, addn_query_subset_with_ids) {
|
|
size_t n = 20000000; // more than 2^24
|
|
size_t k = 1;
|
|
std::vector<int64_t> heap_labels(n, -1);
|
|
std::vector<float> heap_distances(n, 0);
|
|
float_minheap_array_t heaps = {
|
|
n, k, heap_labels.data(), heap_distances.data()};
|
|
heaps.heapify();
|
|
std::vector<int64_t> labels(n, 1);
|
|
std::vector<float> distances(n, 0.0f);
|
|
std::vector<int64_t> subset(n);
|
|
std::iota(subset.begin(), subset.end(), 0);
|
|
heaps.addn_query_subset_with_ids(
|
|
n, subset.data(), 1, distances.data(), labels.data(), 1);
|
|
heaps.reorder();
|
|
EXPECT_TRUE(
|
|
std::all_of(heap_labels.begin(), heap_labels.end(), [](int64_t i) {
|
|
return i == 1;
|
|
}));
|
|
}
|