From 6ee3eb6cf5e365c3562a4c1ffc22b401706d2c86 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 27 Oct 2026 22:48:57 +0800 Subject: [PATCH] test: strengthen regression oracles and consolidate duplicate coverage --- .github/scripts/test_prepare_release.py | 15 ++- tests/python/test_array_alignment.py | 64 +++++++--- tests/python/test_historical_persistence.py | 20 ++- tests/python/test_ivf.py | 13 +- tests/python/test_qgkmeans.py | 6 - .../rabitqlib/quantization/core_api_test.cpp | 115 ++++++++++++++---- .../quantization/xy_quantization_test.cpp | 50 +++----- tests/unit/rabitqlib/utils/bitops_test.cpp | 3 - tests/unit/rabitqlib/utils/space_test.cpp | 1 + 9 files changed, 189 insertions(+), 98 deletions(-) diff --git a/.github/scripts/test_prepare_release.py b/.github/scripts/test_prepare_release.py index f3b54311..94576a4b 100644 --- a/.github/scripts/test_prepare_release.py +++ b/.github/scripts/test_prepare_release.py @@ -41,11 +41,18 @@ def fake_command(*args): if args[:2] == ("git", "rev-list"): return tag_sha if "/actions/workflows/" in args[-1]: - conclusion = ( - ci.get(args[-1].split("/workflows/")[1].split("/")[0]) - if isinstance(ci, dict) - else ci + workflow = args[-1].split("/workflows/")[1].split("/")[0] + self.assertIn(workflow, ("test.yaml", "python.yml")) + self.assertEqual( + args, + ( + "gh", + "api", + f"repos/{REPO}/actions/workflows/{workflow}/runs" + f"?event=push&head_sha={SHA}&per_page=100", + ), ) + conclusion = ci.get(workflow) if isinstance(ci, dict) else ci runs = ( [] if conclusion == "missing" diff --git a/tests/python/test_array_alignment.py b/tests/python/test_array_alignment.py index 7b31caf5..bb083f5d 100644 --- a/tests/python/test_array_alignment.py +++ b/tests/python/test_array_alignment.py @@ -34,7 +34,7 @@ def test_unaligned_training_matches_aligned_input(model_type, offset): @pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) @pytest.mark.parametrize("offset", [1, 2, 3]) -def test_indexes_accept_unaligned_build_and_query_arrays(kind, offset): +def test_indexes_accept_unaligned_build_and_query_arrays(kind, offset, tmp_path): rng = np.random.default_rng(43) data = rng.standard_normal((65, 65), dtype=np.float32) centroids = data[:2].copy() @@ -43,23 +43,59 @@ def test_indexes_accept_unaligned_build_and_query_arrays(kind, offset): unaligned_data = unaligned_copy(data, offset) unaligned_centroids = unaligned_copy(centroids, offset) unaligned_queries = unaligned_copy(queries, offset) + + def make_index(): + if kind == "ivf": + return IvfIndex(65, len(data), 2, 4) + if kind == "hnsw": + return HnswIndex( + 65, len(data), M=8, ef_construction=32, nbits=4, random_seed=42 + ) + return SymqgIndex(65, 32) + if kind == "ivf": - index = IvfIndex(65, len(data), 2, 4) parameters = {"nprobe": 2} - elif kind == "hnsw": - index = HnswIndex(65, len(data), M=8, ef_construction=32, nbits=4) - parameters = {"ef": 65} else: - index = SymqgIndex(65, 32) parameters = {"ef": 65} - if kind == "qg": - index.build(unaligned_data, ef_construction=32, num_threads=1) - else: - index.build(unaligned_data, unaligned_centroids, labels, num_threads=1) - expected = index.search(queries, k=5, num_threads=1, **parameters) - actual = index.search(unaligned_queries, k=5, num_threads=1, **parameters) - for result, reference in zip(actual, expected, strict=True): - np.testing.assert_array_equal(result, reference) + + if kind != "qg": + reference = make_index() + reference.build(data, centroids, labels, num_threads=1) + reference_path = tmp_path / "aligned.index" + reference.save(str(reference_path)) + expected_bytes = reference_path.read_bytes() + expected = reference.search(queries, k=5, num_threads=1, **parameters) + + # Isolate each conversion so a damaged centroid cannot hide behind data damage. + inputs = [(unaligned_data, centroids)] + if kind != "qg": + inputs += [(data, unaligned_centroids), (unaligned_data, unaligned_centroids)] + for vectors, centers in inputs: + # IVF samples a rotation at construction; rebuild a loaded reference to + # reuse it. HNSW's explicit seed reproduces its rotation and serial graph. + index = IvfIndex.load(str(reference_path)) if kind == "ivf" else make_index() + if kind == "qg": + index.build(vectors, ef_construction=32, num_threads=1) + ids, distances = index.search(queries, k=len(data), **parameters) + for row in ids: + np.testing.assert_array_equal(np.sort(row), np.arange(len(data))) + exact = np.sum( + (queries.astype(np.float64)[:, None] - data.astype(np.float64)) ** 2, + axis=2, + ) + np.testing.assert_allclose( + distances, np.take_along_axis(exact, ids, axis=1), rtol=2e-6, atol=1e-5 + ) + expected = index.search(queries, k=5, num_threads=1, **parameters) + else: + index.build(vectors, centers, labels, num_threads=1) + actual_path = tmp_path / "unaligned.index" + index.save(str(actual_path)) + assert actual_path.read_bytes() == expected_bytes + for values in (queries, unaligned_queries): + actual = index.search(values, k=5, num_threads=1, **parameters) + for result, reference_result in zip(actual, expected, strict=True): + np.testing.assert_array_equal(result, reference_result) np.testing.assert_array_equal(unaligned_data, data) np.testing.assert_array_equal(unaligned_centroids, centroids) np.testing.assert_array_equal(unaligned_queries, queries) diff --git a/tests/python/test_historical_persistence.py b/tests/python/test_historical_persistence.py index c822e218..7cb9c006 100644 --- a/tests/python/test_historical_persistence.py +++ b/tests/python/test_historical_persistence.py @@ -35,9 +35,23 @@ def reference(name): def test_fixture_hashes_match_pinned_manifest(): - manifest = json.loads((FIXTURES / "manifest.json").read_text()) - for filename, digest in manifest["sha256"].items(): - assert hashlib.sha256((FIXTURES / filename).read_bytes()).hexdigest() == digest + registered = set() + for path in sorted(FIXTURES.rglob("manifest.json")): + manifest = json.loads(path.read_text()) + assert manifest["sha256"], path + for filename, digest in manifest["sha256"].items(): + assert filename not in registered, filename + registered.add(filename) + assert ( + hashlib.sha256((FIXTURES / filename).read_bytes()).hexdigest() == digest + ), filename + expected = { + path.relative_to(FIXTURES).as_posix() + for suffix in ("*.index", "*.expected") + for path in FIXTURES.rglob(suffix) + } + assert expected + assert registered == expected @pytest.mark.parametrize("name", NAMES) diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 91c643ea..a83f0c75 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -70,7 +70,7 @@ def test_batch_search_matches_individual_queries(bits, metric): individual = [index.search(q[None, :], **kwargs) for q in queries] expected_ids = np.concatenate([result[0] for result in individual]) expected_distances = np.concatenate([result[1] for result in individual]) - for threads in [1, 4]: + for threads in [1, 2, 4]: ids, distances = index.search(queries, num_threads=threads, **kwargs) np.testing.assert_array_equal(ids, expected_ids) np.testing.assert_array_equal(distances, expected_distances) @@ -186,17 +186,6 @@ def test_full_quantization_centroid_distances_roundtrip(nbits, metric, tmp_path) np.testing.assert_array_equal(loaded_distances, distances) -def test_multithreaded_search_matches_single(built_ivf, query_data): - ids1, dists1 = built_ivf.search( - query_data, k=_TOPK, nprobe=_NPROBE_ALL, num_threads=1 - ) - ids2, dists2 = built_ivf.search( - query_data, k=_TOPK, nprobe=_NPROBE_ALL, num_threads=2 - ) - np.testing.assert_array_equal(ids1, ids2) - np.testing.assert_allclose(dists1, dists2, rtol=1e-5) - - # ── error handling ──────────────────────────────────────────────────────────── diff --git a/tests/python/test_qgkmeans.py b/tests/python/test_qgkmeans.py index a0f10699..d50f001b 100644 --- a/tests/python/test_qgkmeans.py +++ b/tests/python/test_qgkmeans.py @@ -54,12 +54,6 @@ def test_constructor_accepts_faiss_style_keyword_parameters(): assert clustering.obj is None -@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) -def test_unknown_parameter_raises(model_type): - with pytest.raises(AttributeError): - model_type(DIM, K, unknown_parameter=True) - - def test_train_populates_faiss_style_outputs(base_data): clustering = qgkmeans() final_obj = clustering.train(base_data) diff --git a/tests/unit/rabitqlib/quantization/core_api_test.cpp b/tests/unit/rabitqlib/quantization/core_api_test.cpp index 783ea5ed..e7a9ae89 100644 --- a/tests/unit/rabitqlib/quantization/core_api_test.cpp +++ b/tests/unit/rabitqlib/quantization/core_api_test.cpp @@ -2,9 +2,11 @@ #include #include +#include #include #include #include +#include #include #include #include @@ -540,34 +542,99 @@ TEST(CoreApi, ScoringRetainsBufferAndNumericalChecks) { } TEST(CoreApi, ConcurrentCallsUseOnlyTheirOwnWorkspace) { - constexpr size_t dim = 96; - std::array data{}, query{}; - for (size_t i = 0; i < dim; ++i) { - data[i] = std::sin(static_cast(i)); - query[i] = std::cos(static_cast(i)); + struct Case { + api::Parameters parameters; + std::vector data, query; + std::vector signs, extra, unpacked; + api::Factors factors; + float score = 0; + }; + std::array cases; + for (size_t i = 0; i < cases.size(); ++i) { + auto& c = cases[i]; + c.parameters = { + 32 + i * 64, + i == 0 ? 1U + : i == 1 ? 4U + : 9U, + i % 2 == 0 ? METRIC_L2 : METRIC_IP}; + const size_t dim = c.parameters.dim; + c.data.resize(dim); + c.query.resize(dim); + for (size_t j = 0; j < dim; ++j) { + c.data[j] = std::sin(static_cast(j + i * 7)) * static_cast(i + 1); + c.query[j] = std::cos(static_cast(j + i * 11)); + } + const auto r = api::requirements(c.parameters, api::CodeLayout::SplitPacked); + Workspace scratch(r.workspace_bytes); + c.signs.resize(r.code_bytes); + c.extra.resize(r.extra_bytes); + c.unpacked.resize(dim * sizeof(uint16_t)); + api::CodeOutput out{ + api::CodeLayout::SplitPacked, + {c.signs.data(), c.signs.size()}, + {c.extra.data(), c.extra.size()}}; + c.factors = api::encode(c.parameters, c.data.data(), nullptr, out, scratch.view()); + c.score = api::compute( + api::prepare_query(c.parameters, c.query.data(), nullptr), view(out), c.factors + ); + api::convert_codes( + c.parameters, + view(out), + {api::CodeLayout::UnpackedU16, {c.unpacked.data(), c.unpacked.size()}, {}}, + scratch.view() + ); } - api::Parameters p{dim, 9}; - auto r = api::requirements(p, api::CodeLayout::SplitPacked); - auto q = api::prepare_query(p, query.data(), nullptr); - std::array results{}; - std::vector workers; - for (size_t i = 0; i < results.size(); ++i) { - workers.emplace_back([&, i] { - Workspace scratch(r.workspace_bytes); - std::vector sign(r.code_bytes), extra(r.extra_bytes); - api::CodeOutput out{ - api::CodeLayout::SplitPacked, - {sign.data(), sign.size()}, - {extra.data(), extra.size()}}; + + std::promise start; + const auto gate = start.get_future().share(); + std::atomic ready{0}; + std::vector> workers; + workers.reserve(cases.size()); + for (size_t worker = 0; worker < cases.size(); ++worker) { + workers.push_back(std::async(std::launch::async, [&, worker] { + ++ready; + gate.wait(); for (size_t run = 0; run < 20; ++run) { - auto f = api::encode(p, data.data(), nullptr, out, scratch.view()); - results[i] = api::compute(q, view(out), f); + // Each round uses distinct inputs across workers, and changes shape + // within each worker. Compare every result to a serial oracle. + const auto& c = cases[(worker + run) % cases.size()]; + const auto r = + api::requirements(c.parameters, api::CodeLayout::SplitPacked); + Workspace scratch(r.workspace_bytes); + std::vector signs(r.code_bytes), extra(r.extra_bytes); + std::vector unpacked(c.unpacked.size()); + api::CodeOutput out{ + api::CodeLayout::SplitPacked, + {signs.data(), signs.size()}, + {extra.data(), extra.size()}}; + const auto f = + api::encode(c.parameters, c.data.data(), nullptr, out, scratch.view()); + const auto query = + api::prepare_query(c.parameters, c.query.data(), nullptr); + if (signs != c.signs || extra != c.extra || f.add != c.factors.add || + f.rescale != c.factors.rescale || f.error != c.factors.error || + api::compute(query, view(out), f) != c.score) + return false; + api::CodeOutput raw{ + api::CodeLayout::UnpackedU16, {unpacked.data(), unpacked.size()}, {}}; + api::convert_codes(c.parameters, view(out), raw, scratch.view()); + if (unpacked != c.unpacked) + return false; + std::fill(signs.begin(), signs.end(), 0xA5); + std::fill(extra.begin(), extra.end(), 0xA5); + api::convert_codes(c.parameters, view(raw), out, scratch.view()); + if (signs != c.signs || extra != c.extra || + api::compute(query, view(out), f) != c.score) + return false; } - }); + return true; + })); } + while (ready.load() != cases.size()) + std::this_thread::yield(); + start.set_value(); for (auto& worker : workers) - worker.join(); - for (float result : results) - EXPECT_EQ(result, results[0]); + EXPECT_TRUE(worker.get()); } } // namespace diff --git a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp index 2697e637..e68c20a7 100644 --- a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp +++ b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp @@ -368,34 +368,6 @@ TEST(XyQuantization, CollinearResidualGivesFiniteErrorFactor) { EXPECT_FLOAT_EQ(base_map.f_error(), 0.0F); } -TEST(XyQuantization, SplitInnerProductRecombinesExactly) { - constexpr size_t kDim = 64; - constexpr size_t kBaseBits = 3; - constexpr size_t kExtraBits = 4; - - std::mt19937 gen(7); - std::vector data = RandomVec(kDim, gen); - std::vector centroid(kDim, 0.0F); - std::vector query = RandomVec(kDim, gen); - - XySplitResult res = SplitCode(data, centroid, kDim, kBaseBits, kExtraBits); - - double split_ip = 0; - double direct_ip = 0; - for (size_t i = 0; i < kDim; ++i) { - uint32_t total = (static_cast(res.base_code[i]) << kExtraBits) | - static_cast(res.extra_code[i]); - direct_ip += static_cast(query[i]) * static_cast(total); - } - for (size_t i = 0; i < kDim; ++i) { - split_ip += static_cast(query[i]) * static_cast(res.base_code[i]) * - static_cast(1U << kExtraBits); - split_ip += static_cast(query[i]) * static_cast(res.extra_code[i]); - } - - EXPECT_NEAR(split_ip, direct_ip, 1e-6 * std::abs(direct_ip) + 1e-6); -} - // The base code must be stored exactly once: total storage is // (base_bits + extra_bits) bits per dimension plus 5 floats (3 base factors + // 2 full factors), never (2*base_bits + extra_bits) as an independently @@ -593,18 +565,32 @@ TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { XySplitResult res = SplitCode(data, centroid, kDim, kBaseBits, kExtraBits); + // Encode the unsplit reference independently of XY splitting and packed reads. + std::vector full_code(kDim); + float full_add = 0, full_rescale = 0, full_error = 0; + quant::quantize_full_single( + data.data(), + centroid.data(), + kDim, + kBaseBits + kExtraBits, + full_code.data(), + full_add, + full_rescale, + full_error + ); + double total_ip = 0; double sumq = 0; for (size_t i = 0; i < kDim; ++i) { uint32_t total = (static_cast(res.base_code[i]) << kExtraBits) | static_cast(res.extra_code[i]); - total_ip += static_cast(query[i]) * static_cast(total); + EXPECT_EQ(total, full_code[i]) << "dim " << i; + total_ip += static_cast(query[i]) * full_code[i]; sumq += query[i]; } double c_b = -(static_cast(1 << (kBaseBits + kExtraBits)) - 1) / 2.0; - float ref_est = static_cast( - res.f_add_full + kGAdd + (res.f_rescale_full * (total_ip + (sumq * c_b))) - ); + float ref_est = + static_cast(full_add + kGAdd + (full_rescale * (total_ip + (sumq * c_b)))); std::vector base_data(BaseDataMap::data_bytes(kDim, kBaseBits)); std::vector extra_data(ExDataMap::data_bytes(kDim, kExtraBits)); diff --git a/tests/unit/rabitqlib/utils/bitops_test.cpp b/tests/unit/rabitqlib/utils/bitops_test.cpp index b76f9a21..c02d2427 100644 --- a/tests/unit/rabitqlib/utils/bitops_test.cpp +++ b/tests/unit/rabitqlib/utils/bitops_test.cpp @@ -11,9 +11,6 @@ TEST(BitopsTest, CountsBitsAndTrailingZeros) { EXPECT_EQ(popcount32(0), 0U); EXPECT_EQ(popcount32(0xFFFFFFFFU), 32U); EXPECT_EQ(popcount32(0x80000001U), 2U); - EXPECT_EQ(popcount64(0), 0U); - EXPECT_EQ(popcount64(0xFFFFFFFFFFFFFFFFULL), 64U); - EXPECT_EQ(popcount64(0x8000000000000001ULL), 2U); for (unsigned bit = 0; bit < 32; ++bit) { EXPECT_EQ(countr_zero32(uint32_t{1} << bit), bit); } diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index 95e5c77c..b423a401 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -40,6 +40,7 @@ TEST(BitOps, Popcount64MatchesBitReference) { }; check(0); check(std::numeric_limits::max()); + check(0x8000000000000001ULL); for (size_t bit = 0; bit < 64; ++bit) { check(uint64_t{1} << bit); check(~(uint64_t{1} << bit));