From fc0daf9ce48b18e3bd0bdee1e034244f31d33bde Mon Sep 17 00:00:00 2001 From: Bhavik Sheth Date: Wed, 12 Feb 2025 17:05:53 -0800 Subject: [PATCH] Add bounds checking to hnsw nb_neighbors Summary: Based on this users comment it seems like we should do bound checking: https://github.com/facebookresearch/faiss/issues/4177 Differential Revision: D69497295 --- faiss/impl/HNSW.cpp | 1 + tests/test_hnsw.cpp | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp index ece0281221..61afa00231 100644 --- a/faiss/impl/HNSW.cpp +++ b/faiss/impl/HNSW.cpp @@ -32,6 +32,7 @@ namespace faiss { **************************************************************/ int HNSW::nb_neighbors(int layer_no) const { + FAISS_THROW_IF_NOT(layer_no + 1 < cum_nneighbor_per_level.size()); return cum_nneighbor_per_level[layer_no + 1] - cum_nneighbor_per_level[layer_no]; } diff --git a/tests/test_hnsw.cpp b/tests/test_hnsw.cpp index b3c93a861e..9424bd3499 100644 --- a/tests/test_hnsw.cpp +++ b/tests/test_hnsw.cpp @@ -582,6 +582,16 @@ TEST_F(HNSWTest, TEST_search_neighbors_to_add) { } } +TEST_F(HNSWTest, TEST_nb_neighbors_bound) { + omp_set_num_threads(1); + EXPECT_EQ(index->hnsw.nb_neighbors(0), 8); + EXPECT_EQ(index->hnsw.nb_neighbors(1), 4); + EXPECT_EQ(index->hnsw.nb_neighbors(2), 4); + EXPECT_EQ(index->hnsw.nb_neighbors(3), 4); + // picking a large number to trigger an exception based on checking bounds + EXPECT_THROW(index->hnsw.nb_neighbors(100), faiss::FaissException); +} + TEST_F(HNSWTest, TEST_search_level_0) { omp_set_num_threads(1); std::vector I(k * nq);