diff --git a/docs/development.md b/docs/development.md index b9574089..8cc3b04b 100644 --- a/docs/development.md +++ b/docs/development.md @@ -533,6 +533,10 @@ Use the [path conversion helper](#filesystem-paths) for file access. Preserve Frozen release-produced indexes and reference searches live in `tests/fixtures/persistence/`; keep their provenance and hashes, and do not regenerate them with the current writer. +Index files are trusted input (see the usage guide): validate structure and sizes so +that corrupt or mismatched files are reported before allocation, without treating a +crafted file as a supported input. + Do not silently reinterpret an old file. Add a magic/version discriminator, use fixed-width serialized fields for new formats, validate sizes before allocation, check every read, and retain a compatibility test fixture or an explicit rejection path. Review each affected format; diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index aaef0f75..dd1c365f 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -62,7 +62,7 @@ see [C++ integration](../cpp_integration.md). | `nbits` / C++ `total_bits` | 1–9 total bits including the sign bit; Python defaults to 8. Raw-vector storage is not supported. | | `M`, `ef_construction` | Graph connectivity and construction search budget; Python defaults are 16 and 200. | | `metric`, `random_seed` | Python defaults are `"l2"` and 100; `"ip"` is also supported. | -| `k`, `ef` | Positive result count and search window. Python requires `k <= max_elements`; omitted/zero `ef` uses `max(k, 10)`, and positive `ef` is raised to at least `k`. | +| `k`, `ef` | Positive result count and search window. Python requires `k <= max_elements`; omitted/zero `ef` uses `max(k, 10)`, positive `ef` is raised to at least `k`, and `ef` above the stored point count is clamped to it. | `random_seed` controls graph levels and the FHT/Kac rotation (using its low 32 bits). With identical inputs and parameters, single-threaded construction is reproducible diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index 47c9f7b1..cde6f5fb 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -164,3 +164,9 @@ or unsupported versions. See the Compatibility is format-specific; preserving old readers' ability to open new files is not guaranteed. Consult each index's format documentation before downgrading the library. + +Index files are trusted input. Loaders validate structure and reject truncated, +inconsistent, or unsupported files so that mistakes are reported, but they are +not a security boundary: a crafted file can make `load` allocate its stored +capacity or violate invariants that later `add` and `remove` calls rely on. +Load only files written by this library. diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 3f2946de..f3906664 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -73,8 +73,19 @@ class HierarchicalNSW { return num_stored_points_.load() - num_removed_; } [[nodiscard]] size_t num_stored_points() const { return num_stored_points_; } + // Cluster of a stored record by internal ID (0 through num_stored_points() - 1; + // these coincide with labels only for points added in order). Removed points + // keep their cluster. [[nodiscard]] PID cluster_id_of(PID internal_id) const { - return raw_cluster_id(internal_id) & ~kRemovedMask; + if (data_level0_memory_ == nullptr) { + throw std::logic_error( + "HNSW index must be constructed or loaded before cluster lookup" + ); + } + if (internal_id >= num_stored_points_.load()) { + throw std::invalid_argument("HNSW internal ID is out of range"); + } + return stored_cluster_id(internal_id); } void save(const char* filename) const; @@ -157,6 +168,7 @@ class HierarchicalNSW { // Queries use the original dimension. Each result is (distance, external label), // where distance estimates squared L2 or 1 - dot according to the index metric. + // ef_search is raised to k and clamped to the stored point count. std::vector>> search( const float* queries, size_t num_queries, @@ -288,6 +300,9 @@ class HierarchicalNSW { } // The stored cluster ID with the removal mark, as save and load see it. + [[nodiscard]] PID stored_cluster_id(PID internal_id) const { + return raw_cluster_id(internal_id) & ~kRemovedMask; + } [[nodiscard]] PID raw_cluster_id(PID internal_id) const { return *(reinterpret_cast( data_level0_memory_.get() + (internal_id * level0_record_bytes_) + diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 98d15e44..42638d99 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -571,7 +571,8 @@ class HnswIndex: k : int Number of neighbors requested, between 1 and max_elements. ef : int, optional - Search breadth. 0 (default) selects max(k, 10). + Search breadth. 0 (default) selects max(k, 10); values above the stored + point count are clamped. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 668c98ef..d36ff43f 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -376,7 +376,8 @@ queries : array_like of float, shape (n_queries, dim) k : int Number of neighbors requested, between 1 and max_elements. ef : int, optional - Search breadth. 0 (default) selects max(k, 10). + Search breadth. 0 (default) selects max(k, 10); values above the stored + point count are clamped. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 16719952..54cd04f5 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -954,6 +954,9 @@ std::vector> HierarchicalNSW::search_knn_sorted_with( if (ef_search == 0) { ef_search = kDefaultEfSearch; } + // No candidate pool can hold more distinct points than are stored, so a huge + // ef_search must not turn into a huge allocation. + ef_search = std::min(ef_search, num_stored_points_.load()); SplitSingleQuery query_wrapper( rotated_query, padded_dim_, ex_bits_, query_config_, metric_type_ @@ -989,7 +992,9 @@ std::vector> HierarchicalNSW::search_knn_sorted_with( auto visited = visited_set_pool_->acquire(); // Keep the ef_search best estimates in a bounded sorted buffer. - buffer::SearchBuffer search_pool(std::max(ef_search, k)); + buffer::SearchBuffer search_pool( + std::min(std::max(ef_search, k), num_stored_points_.load()) + ); std::vector> nearest(k); const size_t num_results = search_layers( *this, @@ -1068,7 +1073,7 @@ void HierarchicalNSW::reconstruct_rotated(PID internal_id, float* rotated) const quant::rabitq_impl::unpack_sign_bits(bin.bin_code(), padded_dim_, combined.data()); const auto* centroid = reinterpret_cast(centroids_memory_.get()) + - (cluster_id_of(internal_id) * padded_dim_); + (stored_cluster_id(internal_id) * padded_dim_); if (ex_bits_ == 0) { quant::reconstruct_full_vec( @@ -1118,7 +1123,7 @@ float HierarchicalNSW::quantized_distance( } const auto* centroid = reinterpret_cast(centroids_memory_.get()) + - (cluster_id_of(target) * padded_dim_); + (stored_cluster_id(target) * padded_dim_); const float dist_sqr = euclidean_sqr(cached.rotated.data(), centroid, padded_dim_); const float g_error = std::sqrt(dist_sqr); const float g_add = (metric_type_ == METRIC_IP) diff --git a/src/index/hnsw_search_impl.hpp b/src/index/hnsw_search_impl.hpp index 587c51ae..d62e4ed4 100644 --- a/src/index/hnsw_search_impl.hpp +++ b/src/index/hnsw_search_impl.hpp @@ -56,9 +56,9 @@ inline void HierarchicalNSW::estimate_bin_direct( DistanceEstimate& estimate ) const { if (metric_type_ == METRIC_IP) { - const float query_centroid_ip = q_to_centroids[cluster_id_of(candidate_id)]; + const float query_centroid_ip = q_to_centroids[stored_cluster_id(candidate_id)]; const float residual_norm = - q_to_centroids[cluster_id_of(candidate_id) + num_clusters_]; + q_to_centroids[stored_cluster_id(candidate_id) + num_clusters_]; estimate = estimate_split_single_direct( bin_data(candidate_id), query_wrapper, @@ -68,7 +68,7 @@ inline void HierarchicalNSW::estimate_bin_direct( ); } else { // L2 distance - const float residual_norm = q_to_centroids[cluster_id_of(candidate_id)]; + const float residual_norm = q_to_centroids[stored_cluster_id(candidate_id)]; estimate = estimate_split_single_direct( bin_data(candidate_id), query_wrapper, @@ -87,9 +87,9 @@ inline void HierarchicalNSW::estimate_full_direct( DistanceEstimate& estimate ) const { if (metric_type_ == METRIC_IP) { - const float query_centroid_ip = q_to_centroids[cluster_id_of(candidate_id)]; + const float query_centroid_ip = q_to_centroids[stored_cluster_id(candidate_id)]; const float residual_norm = - q_to_centroids[cluster_id_of(candidate_id) + num_clusters_]; + q_to_centroids[stored_cluster_id(candidate_id) + num_clusters_]; estimate = estimate_split_single_full_direct( bin_data(candidate_id), ex_data(candidate_id), @@ -102,7 +102,7 @@ inline void HierarchicalNSW::estimate_full_direct( ); } else { // L2 distance - const float residual_norm = q_to_centroids[cluster_id_of(candidate_id)]; + const float residual_norm = q_to_centroids[stored_cluster_id(candidate_id)]; estimate = estimate_split_single_full_direct( bin_data(candidate_id), ex_data(candidate_id), diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index a5bd35c3..eaf3fad8 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -560,3 +560,13 @@ def test_legacy_padding_round_trip_and_add(tmp_path): np.testing.assert_array_equal(actual[1], distances) loaded.add(np.ones((1, 65), dtype=np.float32)) assert loaded.num_points == 4 + + +def test_huge_ef_is_clamped_to_the_stored_point_count(): + data = np.random.default_rng(5).normal(size=(80, 64)).astype(np.float32) + index = HnswIndex(64, 80, nbits=4) + index.build(data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32)) + exhaustive_ids, exhaustive_distances = index.search(data[:3], k=5, ef=80) + ids, distances = index.search(data[:3], k=5, ef=2**40) + np.testing.assert_array_equal(ids, exhaustive_ids) + np.testing.assert_array_equal(distances, exhaustive_distances) diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 50d9e044..98bcd986 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -1430,6 +1430,47 @@ TEST(HnswInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) { EXPECT_EQ(index.search(edge_queries.data(), 2, 3, 10, 1).size(), 2U); } +TEST(HnswSearchTest, ClampsEfToStoredPointsAndChecksClusterLookups) { + constexpr size_t kDim = 64; + constexpr size_t kCount = 24; + constexpr size_t kClusters = 3; + std::mt19937 gen(906); + std::normal_distribution normal; + std::vector data(kCount * kDim); + std::vector centroids(kClusters * kDim); + for (auto* values : {&data, ¢roids}) { + for (auto& value : *values) { + value = normal(gen); + } + } + std::vector cluster_ids(kCount); + for (size_t i = 0; i < kCount; ++i) { + cluster_ids[i] = static_cast(i % kClusters); + } + + HierarchicalNSW unbuilt; + EXPECT_THROW(unbuilt.cluster_id_of(0), std::logic_error); + + HierarchicalNSW index(kCount, kDim, 4, 8, 32); + index.construct( + kClusters, centroids.data(), kCount, data.data(), cluster_ids.data(), 1 + ); + for (size_t i = 0; i < kCount; ++i) { + EXPECT_EQ(index.cluster_id_of(static_cast(i)), cluster_ids[i]); + } + EXPECT_THROW(index.cluster_id_of(static_cast(kCount)), std::invalid_argument); + EXPECT_THROW(index.cluster_id_of(kPidMax), std::invalid_argument); + + // A search window beyond the stored point count neither allocates for it nor + // changes the results of an exhaustive window. + const auto exhaustive = index.search(data.data(), 2, 5, kCount, 1); + const auto huge = index.search(data.data(), 2, 5, size_t{1} << 40, 1); + ASSERT_EQ(huge.size(), exhaustive.size()); + for (size_t q = 0; q < huge.size(); ++q) { + EXPECT_EQ(huge[q], exhaustive[q]); + } +} + } // namespace } // namespace rabitqlib::hnsw