diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index f2bfc11f..ce646a3b 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -442,6 +442,9 @@ inline HierarchicalNSW::~HierarchicalNSW() { free_memory(); } inline void HierarchicalNSW::save(const char* filename) const { std::ofstream output(filename, std::ios::binary); + if (!output.is_open()) { + throw std::runtime_error("HNSW: cannot open index file for writing"); + } output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); output.write(reinterpret_cast(&cur_element_count_), sizeof(size_t)); @@ -491,6 +494,9 @@ inline void HierarchicalNSW::save(const char* filename) const { rotator_->save(output); output.close(); + if (!output) { + throw std::runtime_error("HNSW: failed to write index file"); + } } inline void HierarchicalNSW::load(const char* filename) { @@ -644,6 +650,7 @@ inline void HierarchicalNSW::add_point( std::unique_lock lock_label(get_lable_op_mutex(label)); int level = -1; + int curlevel = 0; PID cur_c = 0; { std::unique_lock lock_table(label_lookup_lock_); @@ -661,10 +668,10 @@ inline void HierarchicalNSW::add_point( cur_c = cur_element_count_; cur_element_count_++; label_lookup_[label] = cur_c; + curlevel = get_random_level(mult_); } std::unique_lock lock_el(link_list_locks_[cur_c]); - int curlevel = get_random_level(mult_); if (level > 0) { curlevel = level; } diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index 2e2c1de1..e6707dbb 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -130,13 +130,12 @@ inline void one_bit_code_with_factor( // We use unnormalized vector to get error factor. To be more specific, // sqrt((1 - ^2) / ^2) / sqrt(dim - 1) = 3rd item in following // expression - T tmp_error = - l2_norm * kConstEpsilon * - std::sqrt( - (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - - 1) / - (dim - 1) - ); + // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero. + const T normalized_error = + (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - 1 + ) / + (dim - 1); + T tmp_error = l2_norm * kConstEpsilon * std::sqrt(std::max(normalized_error, T{0})); // 3 factors used for distance estimation, please refer to document for more info. // For f_rescale and 2nd item of f_add, we use the dot product of raw residual (rather @@ -611,13 +610,12 @@ inline void ex_bits_code_with_factor( // A nonzero residual and its quantized code have a strictly positive inner product. assert(ip_resi_xucb > 0); - T tmp_error = - l2_norm * kConstEpsilon * - std::sqrt( - (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - - 1) / - (dim - 1) - ); + // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero. + const T normalized_error = + (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - 1 + ) / + (dim - 1); + T tmp_error = l2_norm * kConstEpsilon * std::sqrt(std::max(normalized_error, T{0})); if (metric_type == METRIC_L2) { f_add_ex = l2_sqr + (2 * l2_sqr * ip_cent_xucb / ip_resi_xucb); diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 9a54bf0b..48ce4484 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include "bindings_common.hpp" @@ -65,9 +66,9 @@ class SymqgIndex { throw std::invalid_argument("data dimension does not match index dim"); } - num_points_ = static_cast(data_array.shape(0)); - index_ = std::make_unique>( - num_points_, + const size_t num_points = static_cast(data_array.shape(0)); + auto index = std::make_unique>( + num_points, dim_, max_degree_, metric_, @@ -78,13 +79,15 @@ class SymqgIndex { const auto initialization = init == "pipnn" ? symqg::QGInitialization::PiPNN : symqg::QGInitialization::Random; symqg::QGBuilder builder( - *index_, + *index, static_cast(ef_construction), data_array.data(), num_threads, initialization ); builder.build(); + index_ = std::move(index); + num_points_ = num_points; built_ = true; } diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index 6806c7bd..9e8d15bb 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -1,5 +1,7 @@ """Tests for HnswIndex: construction, search, properties, error handling, save/load.""" +from pathlib import Path + import numpy as np import pytest from conftest import DIM, N_CLUSTERS, N_QUERIES, N_VECTORS, brute_force_knn, recall_at_k @@ -42,6 +44,15 @@ def test_fast_quantization_builds(base_data, clusters): assert idx.is_built +def test_parallel_build_preserves_self_retrieval(base_data, clusters): + idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + centroids, cluster_ids = clusters + idx.build(base_data, centroids, cluster_ids, num_threads=8) + ids, distances = idx.search(base_data[:10], k=1, ef=N_VECTORS) + np.testing.assert_array_equal(ids[:, 0], np.arange(10)) + assert np.isfinite(distances).all() + + # ── search output shape and dtype ───────────────────────────────────────────── @@ -124,6 +135,17 @@ def test_search_before_build_raises(query_data): # ── save / load roundtrip ───────────────────────────────────────────────────── +def test_save_rejects_unopenable_destination(built_hnsw, tmp_path): + with pytest.raises(RuntimeError, match="HNSW: cannot open index file for writing"): + built_hnsw.save(str(tmp_path / "missing" / "hnsw.index")) + + +@pytest.mark.skipif(not Path("/dev/full").exists(), reason="/dev/full is unavailable") +def test_save_reports_write_or_close_failure(built_hnsw): + with pytest.raises(RuntimeError, match="HNSW: failed to write index file"): + built_hnsw.save("/dev/full") + + def test_save_load_roundtrip(built_hnsw, query_data, tmp_path): path = str(tmp_path / "hnsw.index") built_hnsw.save(path) diff --git a/tests/python/test_quantization_factors.py b/tests/python/test_quantization_factors.py new file mode 100644 index 00000000..39af396d --- /dev/null +++ b/tests/python/test_quantization_factors.py @@ -0,0 +1,23 @@ +"""Python boundary coverage for finite quantization factors and pruning.""" + +import numpy as np +import pytest +from rabitqlib import IvfIndex + + +@pytest.mark.parametrize("nbits", [4, 9, 32]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_collinear_ivf_query_keeps_finite_distances(nbits, metric): + data = np.full((1, 512), 0.1, dtype=np.float32) + centroid = np.zeros_like(data) + index = IvfIndex(data.shape[1], len(data), 1, nbits=nbits, metric=metric) + index.build(data, centroid, np.zeros(len(data), dtype=np.uint32), num_threads=1) + + # The C++ regression fixes the rotated residual to trigger negative roundoff. + # This exercises the normal Python boundary with its randomly sampled rotation. + # A zero query avoids query-LUT approximation in the distance oracle. + ids, distances = index.search(centroid, k=1, nprobe=1, num_threads=1) + np.testing.assert_array_equal(ids, [[0]]) + assert np.all(np.isfinite(distances)) + expected = np.sum(data.astype(np.float64) ** 2) if metric == "l2" else 1.0 + np.testing.assert_allclose(distances, expected, rtol=2e-5, atol=2e-5) diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index d406202c..117f14bc 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -117,6 +117,70 @@ def test_search_rejects_ef_smaller_than_k(built_symqg, query_data): built_symqg.search(query_data, k=10, ef=9) +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("replacement_points", [0, _MAX_DEGREE]) +def test_failed_rebuild_preserves_index( + base_data, query_data, tmp_path, bits, metric, replacement_points +): + data = base_data[:40] + queries = query_data[:3] + index = SymqgIndex( + DIM, max_degree=_MAX_DEGREE, metric=metric, quantization_bits=bits + ) + index.build(data, ef_construction=_EF_BUILD) + expected_ids, expected_distances = index.search(queries, k=5, ef=_EF) + before_path = tmp_path / "before.index" + index.save(str(before_path)) + + with pytest.raises( + ValueError, + match="QuantizedGraph degree bound must be smaller than the number of points", + ): + index.build(data[:replacement_points], ef_construction=_EF_BUILD) + + assert index.is_built + assert index.num_points == len(data) + assert index.dim == DIM + assert index.max_degree == _MAX_DEGREE + assert index.metric == metric + assert index.quantization_bits == bits + ids, distances = index.search(queries, k=5, ef=_EF) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) + + after_path = tmp_path / "after.index" + index.save(str(after_path)) + assert after_path.read_bytes() == before_path.read_bytes() + restored = SymqgIndex.load(str(after_path)) + assert restored.num_points == len(data) + restored_ids, restored_distances = restored.search(queries, k=5, ef=_EF) + np.testing.assert_array_equal(restored_ids, expected_ids) + np.testing.assert_array_equal(restored_distances, expected_distances) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +def test_successful_rebuild_replaces_index(base_data, tmp_path, bits): + index = SymqgIndex(DIM, max_degree=_MAX_DEGREE, quantization_bits=bits) + index.build(base_data[:40], ef_construction=_EF_BUILD) + replacement = base_data[40:88] + 10 + index.build(replacement, ef_construction=_EF_BUILD) + + assert index.is_built + assert index.num_points == len(replacement) + assert index.quantization_bits == bits + ids, distances = index.search(replacement[:3], k=1, ef=_EF) + np.testing.assert_array_equal(ids[:, 0], np.arange(3)) + + path = str(tmp_path / "rebuilt.index") + index.save(path) + restored = SymqgIndex.load(path) + assert restored.num_points == len(replacement) + restored_ids, restored_distances = restored.search(replacement[:3], k=1, ef=_EF) + np.testing.assert_array_equal(restored_ids, ids) + np.testing.assert_array_equal(restored_distances, distances) + + def test_invalid_degree_raises(): with pytest.raises(Exception): SymqgIndex(DIM, max_degree=16) diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp index 90814951..d2f68f7d 100644 --- a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -17,6 +17,7 @@ #include "rabitqlib/index/lut.hpp" #include "rabitqlib/index/query.hpp" #include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/quantization/rabitq.hpp" #include "rabitqlib/simd/estimator_dispatch.hpp" #include "rabitqlib/simd/fastscan_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" @@ -235,6 +236,57 @@ TEST(FastScanHighAccuracyTest, Avx512TransferAcceptsUnalignedOutput) { EXPECT_TRUE(std::equal(expected.begin(), expected.end(), storage.begin() + 1)); } +TEST(BatchEstimatorTest, CollinearResidualKeepsFiniteLowerBounds) { + if (!cpu::has_avx2() && !cpu::has_avx512_core()) { + GTEST_SKIP() << "FastScan requires AVX2/FMA or AVX512"; + } + using Estimator = decltype(&rabitqlib::simd::split_batch_estdist); + std::vector backends{ + rabitqlib::simd::split_batch_estdist, rabitqlib::simd::split_batch_estdist_generic}; + if (cpu::has_avx2()) { + backends.push_back(rabitqlib::simd::split_batch_estdist_avx2); + } + if (cpu::has_avx512_core()) { + backends.push_back(rabitqlib::simd::split_batch_estdist_avx512); + } + + constexpr size_t kDim = 512; + std::array data; + data.fill(0.1F); + const std::array centroid{}; + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + std::vector storage(BatchDataMap::data_bytes(kDim), 0); + quant::quantize_one_batch( + data.data(), centroid.data(), 1, kDim, storage.data(), metric + ); + for (bool hacc : {false, true}) { + SCOPED_TRACE( + ::testing::Message() + << "metric=" << static_cast(metric) << " hacc=" << hacc + ); + SplitBatchQuery query(data.data(), kDim, 3, metric, hacc); + query.set_g_add(std::sqrt(static_cast(kDim)) * 0.1F); + for (auto backend : backends) { + std::array distance{}, lower{}, ip{}; + backend( + storage.data(), + query, + kDim, + distance.data(), + lower.data(), + ip.data(), + hacc + ); + EXPECT_TRUE(std::isfinite(distance[0])); + EXPECT_TRUE(std::isfinite(lower[0])); + EXPECT_LE(lower[0], distance[0]); + // IVF must be able to admit this candidate into an empty search buffer. + EXPECT_LT(lower[0], std::numeric_limits::max()); + } + } + } +} + TEST(BatchEstimatorTest, BackendsMatchScalarCorrectionAcrossChunksAndTailBatches) { if (!cpu::has_avx2() && !cpu::has_avx512_core()) { GTEST_SKIP() << "FastScan requires AVX2/FMA or AVX512"; diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 0026403f..c2e1f811 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -3,6 +3,11 @@ #include #include +#include +#include +#include +#include +#include #include namespace rabitqlib::hnsw { @@ -47,5 +52,87 @@ TEST(HnswOneBitSearchTest, UsesBinaryEstimateForEntryPoint) { EXPECT_NEAR(results[0][0].first, exact_distance, exact_distance * 0.1F); } +TEST(HnswConstructionTest, ParallelConstructionProducesSearchableIndex) { + constexpr size_t kDim = 64; + constexpr size_t kCount = 512; + constexpr size_t kQueries = 16; + std::mt19937 generator(42); + std::uniform_real_distribution distribution(-1, 1); + std::vector data(kCount * kDim); + for (auto& value : data) { + value = distribution(generator); + } + std::vector centroid(kDim); + std::vector cluster_ids(kCount, 0); + HierarchicalNSW index(kCount, kDim, 4, 8, 50); + index.construct(1, centroid.data(), kCount, data.data(), cluster_ids.data(), 8, false); + + const auto results = index.search(data.data(), kQueries, 1, kCount, 1); + ASSERT_EQ(results.size(), kQueries); + for (size_t i = 0; i < kQueries; ++i) { + ASSERT_EQ(results[i].size(), 1U); + EXPECT_EQ(results[i][0].second, i); + EXPECT_TRUE(std::isfinite(results[i][0].first)); + } +} + +class HnswSaveTest : public ::testing::Test { + protected: + static constexpr size_t kDim = 64; + static constexpr size_t kCount = 8; + std::vector data_{std::vector(kCount * kDim)}; + std::vector centroid_{std::vector(kDim)}; + std::vector cluster_ids_{std::vector(kCount, 0)}; + HierarchicalNSW index_{kCount, kDim, 4, 4, 10}; + std::string path_; + + void SetUp() override { + path_ = ::testing::TempDir() + "rabitq_hnsw_" + + ::testing::UnitTest::GetInstance()->current_test_info()->name() + ".index"; + for (size_t i = 0; i < data_.size(); ++i) { + data_[i] = static_cast((i * 37) % 127) / 32.0F; + } + index_.construct( + 1, centroid_.data(), kCount, data_.data(), cluster_ids_.data(), 1, false + ); + } + + void TearDown() override { std::remove(path_.c_str()); } +}; + +TEST_F(HnswSaveTest, RejectsUnopenableDestination) { + try { + index_.save((path_ + "/index").c_str()); + FAIL() << "Saving to a missing directory must fail"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "HNSW: cannot open index file for writing"); + } +} + +TEST_F(HnswSaveTest, ReportsWriteOrCloseFailure) { + if (!std::filesystem::exists("/dev/full")) { + GTEST_SKIP() << "/dev/full is unavailable"; + } + try { + index_.save("/dev/full"); + FAIL() << "Saving to a full destination must fail"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "HNSW: failed to write index file"); + } +} + +TEST_F(HnswSaveTest, SuccessfulSavePreservesSearchAfterLoading) { + ASSERT_NO_THROW(index_.save(path_.c_str())); + HierarchicalNSW loaded; + ASSERT_NO_THROW(loaded.load(path_.c_str())); + EXPECT_EQ(loaded.dimension(), index_.dimension()); + EXPECT_EQ(loaded.num_clusters(), index_.num_clusters()); + EXPECT_EQ(loaded.nbits(), index_.nbits()); + EXPECT_EQ( + loaded.search(data_.data(), kCount, 2, kCount, 1), + index_.search(data_.data(), kCount, 2, kCount, 1) + ); +} + } // namespace } // namespace rabitqlib::hnsw diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index bef869ab..8e93b5d9 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -356,6 +356,43 @@ TEST(RabitqDegenerateInputTest, ExtraBitFactorsAreFiniteForZeroResidual) { EXPECT_FLOAT_EQ(f_error, 0.0F); } +TEST(RabitqDegenerateInputTest, FullFactorsAreFiniteForCollinearResidual) { + constexpr size_t kDim = 512; + std::array data; + data.fill(0.1F); + const std::array centroid{}; + std::array code{}; + + // Equal coordinates make the residual and quantized direction collinear. + // Their theoretically zero error must remain finite despite norm/dot rounding. + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits = 1; bits <= 9; ++bits) { + SCOPED_TRACE( + ::testing::Message() + << "metric=" << static_cast(metric) << " bits=" << bits + ); + float f_add = 0; + float f_rescale = 0; + float f_error = 0; + quantize_full_single( + data.data(), + centroid.data(), + kDim, + bits, + code.data(), + f_add, + f_rescale, + f_error, + metric + ); + EXPECT_TRUE(std::isfinite(f_add)); + EXPECT_TRUE(std::isfinite(f_rescale)); + EXPECT_TRUE(std::isfinite(f_error)); + EXPECT_GE(f_error, 0.0F); + } + } +} + TEST(RabitqDegenerateInputTest, ScalarQuantizationReconstructsZeroVector) { constexpr size_t kDim = 64; std::array data{};