diff --git a/faiss/impl/NNDescent.cpp b/faiss/impl/NNDescent.cpp index d0c41d694..e27955754 100644 --- a/faiss/impl/NNDescent.cpp +++ b/faiss/impl/NNDescent.cpp @@ -31,25 +31,46 @@ Nhood::Nhood(int /* l */, int s, std::mt19937& rng, int N) { /// Copy operator Nhood& Nhood::operator=(const Nhood& other) { - M = other.M; - std::copy( - other.nn_new.begin(), - other.nn_new.end(), - std::back_inserter(nn_new)); - nn_new.reserve(other.nn_new.capacity()); - pool.reserve(other.pool.capacity()); + if (this != &other) { + M = other.M; + nn_new = other.nn_new; + nn_old = other.nn_old; + rnn_new = other.rnn_new; + rnn_old = other.rnn_old; + pool = other.pool; + } return *this; } /// Copy constructor -Nhood::Nhood(const Nhood& other) { - M = other.M; - std::copy( - other.nn_new.begin(), - other.nn_new.end(), - std::back_inserter(nn_new)); - nn_new.reserve(other.nn_new.capacity()); - pool.reserve(other.pool.capacity()); +Nhood::Nhood(const Nhood& other) + : pool(other.pool), + M(other.M), + nn_old(other.nn_old), + nn_new(other.nn_new), + rnn_old(other.rnn_old), + rnn_new(other.rnn_new) {} + +/// Move constructor +Nhood::Nhood(Nhood&& other) noexcept + : pool(std::move(other.pool)), + M(other.M), + nn_old(std::move(other.nn_old)), + nn_new(std::move(other.nn_new)), + rnn_old(std::move(other.rnn_old)), + rnn_new(std::move(other.rnn_new)) {} + +/// Move assignment operator +Nhood& Nhood::operator=(Nhood&& other) noexcept { + if (this != &other) { + M = other.M; + nn_new = std::move(other.nn_new); + nn_old = std::move(other.nn_old); + rnn_new = std::move(other.rnn_new); + rnn_old = std::move(other.rnn_old); + pool = std::move(other.pool); + } + return *this; } /// Insert a point into the candidate pool @@ -90,6 +111,22 @@ void Nhood::join(C callback) const { } void gen_random(std::mt19937& rng, int* addr, const int size, const int N) { + FAISS_THROW_IF_NOT_FMT( + size > 0 && size <= N, + "gen_random: size (%d) must be > 0 and <= N (%d)", + size, + N); + if (size == N) { + // Special case: return all indices in random order + for (int i = 0; i < size; ++i) { + addr[i] = i; + } + for (int i = size - 1; i > 0; --i) { + int j = rng() % (i + 1); + std::swap(addr[i], addr[j]); + } + return; + } for (int i = 0; i < size; ++i) { addr[i] = rng() % (N - size); } @@ -294,7 +331,7 @@ void NNDescent::nndescent(DistanceComputer& qdis, bool verbose) { int num_eval_points = std::min(NUM_EVAL_POINTS, ntotal); std::vector eval_points(num_eval_points); std::vector> acc_eval_set(num_eval_points); - std::mt19937 rng(random_seed * 6577 + omp_get_thread_num()); + std::mt19937 rng(random_seed * 6577); gen_random(rng, eval_points.data(), eval_points.size(), ntotal); generate_eval_set(qdis, eval_points, acc_eval_set, ntotal); for (int it = 0; it < iter; it++) { diff --git a/faiss/impl/NNDescent.h b/faiss/impl/NNDescent.h index 60c273df7..d6fafd048 100644 --- a/faiss/impl/NNDescent.h +++ b/faiss/impl/NNDescent.h @@ -80,6 +80,10 @@ struct Nhood { Nhood(const Nhood& other); + Nhood(Nhood&& other) noexcept; + + Nhood& operator=(Nhood&& other) noexcept; + void insert(int id, float dist); template diff --git a/tests/test_graph_based.py b/tests/test_graph_based.py index c7ce64ccf..15a61f76a 100644 --- a/tests/test_graph_based.py +++ b/tests/test_graph_based.py @@ -608,6 +608,29 @@ class TestNNDescent(unittest.TestCase): np.testing.assert_array_equal(indices, gt) +class TestNNDescentGenRandom(unittest.TestCase): + """Regression tests for gen_random edge cases in NNDescent.""" + + def test_search_L_equals_ntotal(self): + """gen_random(size, N) crashed with division by zero when size == N. + + In search(), L_2 = max(search_L, topk). When search_L >= ntotal, + gen_random is called with size == N, causing rng() % 0. + """ + d = 32 + nb = 200 # just above NUM_EVAL_POINTS=100 + xb = np.random.default_rng(42).random((nb, d)).astype('float32') + xq = np.random.default_rng(43).random((10, d)).astype('float32') + + index = faiss.IndexNNDescentFlat(d, 32) + index.nndescent.search_L = nb # triggers gen_random(size=nb, N=nb) + index.train(xb) + index.add(xb) + + # This crashed with division by zero before the fix + D, I = index.search(xq, k=1) + + class TestNNDescentKNNG(unittest.TestCase): def test_knng_L2(self): diff --git a/tests/test_nndescent.cpp b/tests/test_nndescent.cpp new file mode 100644 index 000000000..4c78978ac --- /dev/null +++ b/tests/test_nndescent.cpp @@ -0,0 +1,87 @@ +/* + * 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 +#include +#include + +using namespace faiss::nndescent; + +/// Helper: build a Nhood with known data in all fields. +static Nhood make_populated_nhood() { + std::mt19937 rng(42); + int N = 200; + Nhood nh(/*l=*/50, /*s=*/10, rng, N); + + nh.pool.clear(); + nh.pool.push_back(Neighbor(1, 0.1f, true)); + nh.pool.push_back(Neighbor(2, 0.2f, false)); + nh.pool.push_back(Neighbor(3, 0.3f, true)); + + nh.nn_old = {10, 20, 30}; + nh.rnn_new = {40, 50}; + nh.rnn_old = {60, 70, 80}; + + return nh; +} + +TEST(NhoodCopy, CopyConstructorPreservesAllFields) { + Nhood original = make_populated_nhood(); + Nhood copy(original); + + EXPECT_EQ(copy.M, original.M); + EXPECT_EQ(copy.pool.size(), original.pool.size()); + EXPECT_EQ(copy.nn_new, original.nn_new); + EXPECT_EQ(copy.nn_old, original.nn_old); + EXPECT_EQ(copy.rnn_new, original.rnn_new); + EXPECT_EQ(copy.rnn_old, original.rnn_old); +} + +TEST(NhoodCopy, CopyAssignmentPreservesAllFields) { + Nhood original = make_populated_nhood(); + Nhood assigned; + assigned = original; + + EXPECT_EQ(assigned.M, original.M); + EXPECT_EQ(assigned.pool.size(), original.pool.size()); + EXPECT_EQ(assigned.nn_new, original.nn_new); + EXPECT_EQ(assigned.nn_old, original.nn_old); + EXPECT_EQ(assigned.rnn_new, original.rnn_new); + EXPECT_EQ(assigned.rnn_old, original.rnn_old); +} + +TEST(NhoodCopy, CopyAssignmentSelfAssign) { + Nhood nh = make_populated_nhood(); + auto expected_pool_size = nh.pool.size(); + auto expected_nn_new = nh.nn_new; + + // Use a reference to avoid -Wself-assign-overloaded. + Nhood& ref = nh; + nh = ref; + + EXPECT_EQ(nh.pool.size(), expected_pool_size); + EXPECT_EQ(nh.nn_new, expected_nn_new); +} + +/// Simulates std::vector reallocation during push_back. +TEST(NhoodCopy, VectorReallocationPreservesData) { + std::vector vec; + // Do NOT reserve — force reallocation during push_back + for (int i = 0; i < 20; i++) { + Nhood nh = make_populated_nhood(); + nh.pool[0].id = i; + vec.push_back(std::move(nh)); + } + + for (int i = 0; i < 20; i++) { + EXPECT_EQ(vec[i].pool[0].id, i) << "pool lost at index " << i; + EXPECT_EQ(vec[i].pool.size(), 3) << "pool truncated at index " << i; + EXPECT_EQ(vec[i].nn_old.size(), 3) << "nn_old lost at index " << i; + EXPECT_EQ(vec[i].rnn_new.size(), 2) << "rnn_new lost at index " << i; + EXPECT_EQ(vec[i].rnn_old.size(), 3) << "rnn_old lost at index " << i; + } +}