Skip to content

Commit

Permalink
Add bounds checking to hnsw nb_neighbors (#4185)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #4185

Based on this users comment it seems like we should do bound checking: #4177

Reviewed By: mnorris11

Differential Revision: D69497295

fbshipit-source-id: 97025cf29c464afb0f85aa98f4b303489b7fc989
  • Loading branch information
bshethmeta authored and facebook-github-bot committed Feb 14, 2025
1 parent f0e3832 commit 657c563
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 0 deletions.
1 change: 1 addition & 0 deletions faiss/impl/HNSW.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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];
}
Expand Down
10 changes: 10 additions & 0 deletions tests/test_hnsw.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<faiss::idx_t> I(k * nq);
Expand Down

0 comments on commit 657c563

Please sign in to comment.