From 4d1462f0df37a56e12d9caed3f958d340154e8b7 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 18:15:49 +0800 Subject: [PATCH 1/6] docs(ivf): say the hnsw initializer suits nprobe far below the cluster count Graph construction can leave centroids unreachable, so points in those clusters are never returned and an nprobe close to the cluster count raises instead of probing every cluster. State the intended regime and point at flat and flat_rabitq for exhaustive probing. --- docs/docs/index/ivf.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index f43dea45..f196cf40 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -146,7 +146,12 @@ auto routing = ivf.initializer_type(); using twice the standard RaBitQ error margin, and recomputes exact distances for the retained candidates. Returned distances are exact for those centroids, but routing is approximate: pruning can miss a true nearest centroid. -- `hnsw` searches a graph of float32 centroids approximately. +- `hnsw` searches a graph of float32 centroids approximately. It suits workloads + whose `nprobe` is much smaller than the cluster count. Graph construction can + leave a few centroids unreachable, so points stored in those clusters are not + returned at any `nprobe`, and an `nprobe` close to the cluster count raises an + error instead of probing every cluster. Use `flat` or `flat_rabitq` when a + query should probe most clusters. All choices support L2 and inner product. Flat RaBitQ uses the IVF-RaBitQ rotation and padded domain, retains float32 centroids for refinement, and stores its mean and From ad7fa1003f5729458a636b9ac2a006677832632a Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 18:15:49 +0800 Subject: [PATCH 2/6] fix(fastscan): keep estimator corrections out of shared Eigen helpers The batch correction in estimator_kernels.hpp and qg_estimator.hpp instantiated Eigen assignment loops from the generic, AVX2 and AVX-512 translation units. Below -O3 those helpers survive as identically named weak symbols with different instruction sets, and only the static-archive member order kept the generic copy in front of the ISA copies. Write the 32-lane corrections as plain loops so each backend vectorizes them with its own flags and no ISA source emits a shared Eigen helper. check-isa-symbols.sh now passes at -O1 and -O2 -g as well as -O3; search outputs on identical saved indexes are bit-identical before and after. --- docs/development.md | 13 +++--- .../fastscan/detail/qg_estimator.hpp | 29 ++++++------- src/fastscan/estimator_kernels.hpp | 42 ++++++++----------- 3 files changed, 37 insertions(+), 47 deletions(-) diff --git a/docs/development.md b/docs/development.md index 36b0b170..cd589669 100644 --- a/docs/development.md +++ b/docs/development.md @@ -165,11 +165,14 @@ Clang: It recompiles the library sources from `compile_commands.json` with that build's flags (extra arguments such as `-O1` are appended) and fails when copies of a weak function -differ between ISA flag sets. Only the Release (`-O3`) configuration is clean and checked. -Lower optimization levels emit helpers the compiler otherwise inlines: Debug (`-O0`) -emits many standard-library and Eigen functions per source, and GCC at `-O1`/`-O2` -still shares two Eigen assignment loops between the FastScan estimator sources. The -script only sees `src/`, needs GNU binutils, and does not cover MSVC or ARM-only builds. +differ between ISA flag sets. CI checks the Release (`-O3`) configuration; with GCC 14 +the check also passes at `-O1` and `-O2 -g`, which is what the sanitizer and +RelWithDebInfo configurations compile. Debug (`-O0`) still emits standard-library +helpers (`std::min`, `std::round`, `std::fill`) per source and is not clean. The FastScan +estimator correction is written as plain loops rather than Eigen expressions so that no +ISA source instantiates a shared Eigen helper; keep it that way, since a static library +only hides such sharing while its generic objects precede the ISA objects. The script +only sees `src/`, needs GNU binutils, and does not cover MSVC or ARM-only builds. ### Static analysis diff --git a/include/rabitqlib/fastscan/detail/qg_estimator.hpp b/include/rabitqlib/fastscan/detail/qg_estimator.hpp index f350aaaa..a578e40e 100644 --- a/include/rabitqlib/fastscan/detail/qg_estimator.hpp +++ b/include/rabitqlib/fastscan/detail/qg_estimator.hpp @@ -4,14 +4,14 @@ #include #include -#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/fastscan/data_layout.hpp" #include "rabitqlib/fastscan/pack_codes.hpp" #include "rabitqlib/fastscan/query.hpp" namespace rabitqlib::fastscan::detail { // Share QG accumulation and correction while callers select the accumulator. -// TU-local linkage keeps each ISA's Eigen expression compiled with its own flags. +// The correction is a plain loop so each ISA translation unit vectorizes it with +// its own flags without instantiating shared weak helper functions. template static inline void qg_batch_estdist_impl( const char* batch_data, @@ -25,25 +25,20 @@ static inline void qg_batch_estdist_impl( cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim ); - ConstRowMajorArrayMap ip_arr( - accumulated_values.data(), 1, fastscan::kBatchSize - ); std::array f_add_values; std::array f_rescale_values; cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size()); - ConstRowMajorArrayMap f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize); - ConstRowMajorArrayMap f_rescale_arr( - f_rescale_values.data(), 1, fastscan::kBatchSize - ); - - RowMajorArrayMap estimated_distances_arr( - estimated_distances, 1, fastscan::kBatchSize - ); - estimated_distances_arr = - f_add_arr + query_wrapper.g_add() + - (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + - query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); + const T delta = query_wrapper.delta(); + const T sum_vl_lut = query_wrapper.sum_vl_lut(); + const T g_add = query_wrapper.g_add(); + const T k1xsumq = query_wrapper.k1xsumq(); + for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { + const T ip = + delta * static_cast(accumulated_values[lane]) + sum_vl_lut + k1xsumq; + estimated_distances[lane] = + f_add_values[lane] + g_add + f_rescale_values[lane] * ip; + } } } // namespace rabitqlib::fastscan::detail diff --git a/src/fastscan/estimator_kernels.hpp b/src/fastscan/estimator_kernels.hpp index f046ce7b..d78fe836 100644 --- a/src/fastscan/estimator_kernels.hpp +++ b/src/fastscan/estimator_kernels.hpp @@ -4,7 +4,6 @@ #include #include -#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/fastscan/data_layout.hpp" #include "rabitqlib/fastscan/pack_codes.hpp" #include "rabitqlib/fastscan/query.hpp" @@ -12,7 +11,10 @@ namespace rabitqlib::simd { namespace { // Compile the complete batch estimator at each ISA so accumulation conversion and -// factor correction use the same vector width as the scanning kernels. +// factor correction use the same vector width as the scanning kernels. The +// correction is written as plain loops so each backend vectorizes it with its own +// flags: an Eigen expression here would instantiate identically named weak helper +// functions in every ISA translation unit, and the linker keeps one copy for all. template inline void split_batch_estdist_impl( const char* batch_data, @@ -25,9 +27,6 @@ inline void split_batch_estdist_impl( ) { ConstBatchDataMap cur_batch(batch_data, padded_dim); std::array accumulated_values{}; - RowMajorArrayMap accumulated_arr( - accumulated_values.data(), 1, fastscan::kBatchSize - ); if (use_hacc) { AccumulateHacc( cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim @@ -44,27 +43,20 @@ inline void split_batch_estdist_impl( cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size()); cur_batch.f_error().copy_to(f_error_values.data(), f_error_values.size()); - ConstRowMajorArrayMap f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize); - ConstRowMajorArrayMap f_rescale_arr( - f_rescale_values.data(), 1, fastscan::kBatchSize - ); - ConstRowMajorArrayMap f_error_arr( - f_error_values.data(), 1, fastscan::kBatchSize - ); - RowMajorArrayMap estimated_distances_arr( - estimated_distances, 1, fastscan::kBatchSize - ); - RowMajorArrayMap ip_x0_qr_arr(ip_x0_qr, 1, fastscan::kBatchSize); - RowMajorArrayMap lower_bounds_arr(lower_bounds, 1, fastscan::kBatchSize); - - ip_x0_qr_arr = query_wrapper.delta() * (accumulated_arr.template cast()) + - query_wrapper.sum_vl_lut(); - - estimated_distances_arr = f_add_arr + query_wrapper.g_add() + - f_rescale_arr * (ip_x0_qr_arr + query_wrapper.k1xsumq()); - - lower_bounds_arr = estimated_distances_arr - f_error_arr * query_wrapper.g_error(); + const float delta = query_wrapper.delta(); + const float sum_vl_lut = query_wrapper.sum_vl_lut(); + const float g_add = query_wrapper.g_add(); + const float k1xsumq = query_wrapper.k1xsumq(); + const float g_error = query_wrapper.g_error(); + for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { + const float ip = delta * static_cast(accumulated_values[lane]) + sum_vl_lut; + const float estimated = + f_add_values[lane] + g_add + f_rescale_values[lane] * (ip + k1xsumq); + ip_x0_qr[lane] = ip; + estimated_distances[lane] = estimated; + lower_bounds[lane] = estimated - f_error_values[lane] * g_error; + } } } // namespace From 04d89fdb0dfd45158519eba08dc27b2c9712ae80 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 18:15:49 +0800 Subject: [PATCH 3/6] ci: fail wheel gates when the release eligibility check fails Each wheel gate was skipped whenever its build job was skipped, and the build job is skipped when the prepare job fails. A skipped required check satisfies branch protection, so a failed eligibility check waived the wheel builds, the installed-wheel suite and the typing checks on a pull request. Run the gates with always() and decide from the prepare, changes and build results, matching the other aggregator jobs; the contract test now exercises every input combination of the shared verdict script. --- .github/scripts/test_ci_workflow_contracts.py | 68 ++++++++-- .github/workflows/release.yml | 128 ++++++++++++++++-- 2 files changed, 171 insertions(+), 25 deletions(-) diff --git a/.github/scripts/test_ci_workflow_contracts.py b/.github/scripts/test_ci_workflow_contracts.py index e3905b47..b4b9d2c3 100644 --- a/.github/scripts/test_ci_workflow_contracts.py +++ b/.github/scripts/test_ci_workflow_contracts.py @@ -89,6 +89,7 @@ def selections(build, event): ] artifacts = set() + verdicts = set() for prefix, tag, check_name in ( ("linux", "manylinux_x86_64", "Linux x86-64 wheels"), ( @@ -131,21 +132,64 @@ def selections(build, event): self.assertNotIn("continue-on-error", entry) gate = job(workflow, f"{prefix}-wheels") self.assertEqual(gate["name"], check_name) - self.assertEqual(needs(gate), [f"{prefix}-wheel-build"]) + # The gate always runs so that a build skipped because the + # eligibility check failed cannot satisfy the required check. + self.assertEqual( + needs(gate), ["prepare", "changes", f"{prefix}-wheel-build"] + ) + self.assertEqual(gate["if"], "${{ always() }}") (verdict_step,) = steps(gate) self.assertEqual( - verdict_step["env"]["BUILD_RESULT"], - f"${{{{ needs.{prefix}-wheel-build.result }}}}", + verdict_step["env"], + { + "PREPARE_RESULT": "${{ needs.prepare.result }}", + "BUILD_SELECTED": "${{ needs.prepare.outputs.build }}", + "CHANGES_RESULT": "${{ needs.changes.result }}", + "WHEELS_SELECTED": "${{ needs.changes.outputs.wheels }}", + "BUILD_RESULT": f"${{{{ needs.{prefix}-wheel-build.result }}}}", + }, ) - for result in ("success", "failure", "cancelled", "skipped", ""): - with self.subTest(result=result): - verdict = subprocess.run( - ["bash", "-e", "-c", verdict_step["run"]], - env=dict(os.environ, BUILD_RESULT=result), - capture_output=True, - check=False, - ) - self.assertEqual(verdict.returncode == 0, result == "success") + verdicts.add(verdict_step["run"]) + # One shared verdict script, exercised over every input combination. + (command,) = verdicts + results = ("success", "failure", "cancelled", "skipped", "") + for prepare in results: + for selected in ("true", "false", ""): + for changes in results: + for wheels in ("true", "false", ""): + for build in results: + if build == "success": + expected = True + elif build != "skipped": + expected = False + elif prepare == "skipped": + expected = True # excluded by the workflow_run filter + else: + expected = prepare == "success" and ( + selected == "false" + or (changes == "success" and wheels == "false") + ) + with self.subTest( + prepare=prepare, + selected=selected, + changes=changes, + wheels=wheels, + build=build, + ): + verdict = subprocess.run( + ["bash", "-e", "-c", command], + env=dict( + os.environ, + PREPARE_RESULT=prepare, + BUILD_SELECTED=selected, + CHANGES_RESULT=changes, + WHEELS_SELECTED=wheels, + BUILD_RESULT=build, + ), + capture_output=True, + check=False, + ) + self.assertEqual(verdict.returncode == 0, expected) def test_installed_typing_checks_are_required(self): config = tomllib.loads((ROOT / "pyproject.toml").read_text()) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a9b1fd92..f3f55c19 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -85,16 +85,43 @@ jobs: if-no-files-found: error linux-wheels: - # Preserve the required check name and include every selected ABI. + # Preserve the required check name and include every selected ABI. The gate + # always runs: a skipped build passes only when nothing selected it, never + # because the eligibility check failed. name: Linux x86-64 wheels - needs: linux-wheel-build - if: ${{ always() && !cancelled() && needs.linux-wheel-build.result != 'skipped' }} + needs: [prepare, changes, linux-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.linux-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 linux-arm64-wheel-build: name: Linux ARM64 wheel builds @@ -129,14 +156,39 @@ jobs: linux-arm64-wheels: name: Linux ARM64 wheels (CPython 3.11–3.14) - needs: linux-arm64-wheel-build - if: ${{ always() && !cancelled() && needs.linux-arm64-wheel-build.result != 'skipped' }} + needs: [prepare, changes, linux-arm64-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.linux-arm64-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 windows-wheel-build: # Windows wheels take longest, so each ABI builds on its own runner. @@ -174,14 +226,39 @@ jobs: windows-wheels: name: Windows x86-64 wheels - needs: windows-wheel-build - if: ${{ always() && !cancelled() && needs.windows-wheel-build.result != 'skipped' }} + needs: [prepare, changes, windows-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.windows-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 macos-wheel-build: name: macOS arm64 wheel builds @@ -218,14 +295,39 @@ jobs: macos-wheels: # Preserve the protected check name; PRs test the oldest and newest ABI. name: macOS arm64 wheels (CPython 3.11–3.14) - needs: macos-wheel-build - if: ${{ always() && !cancelled() && needs.macos-wheel-build.result != 'skipped' }} + needs: [prepare, changes, macos-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.macos-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 release-tag: name: Create version tag From ce0cd45644c8e1622b151e91d250cac1e434087d Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 18:15:49 +0800 Subject: [PATCH 4/6] fix: reject non-finite and oversized coordinates at every index entry point IVF, HNSW and SymphonyQG accepted NaN, infinite and arbitrarily large coordinates in data, centroids and queries. The core quantizer then converted NaN to int (undefined behaviour, reported by UBSan's float-cast-overflow check), stored NaN factors, and in one-bit high-accuracy IVF search returned a poisoned point first with a NaN distance even after remove(). The stable core API and the clustering APIs already rejected such input. Add utils/input_checks.hpp with a finiteness scan and the clustering coordinate bound, call it at construct, add and search entry points before any state changes, and reject non-finite queries in Lut. Enable float-cast-overflow and float-divide-by-zero in the sanitizer build, since -fsanitize=undefined does not include them. Document the contract in the usage guide and the Python docstrings, and cover each entry point with C++ and Python regression tests. --- CMakeLists.txt | 6 +- docs/development.md | 10 ++ docs/docs/index/usage.md | 8 ++ include/rabitqlib/core/encoder.hpp | 2 + include/rabitqlib/fastscan/detail/lut.hpp | 5 + include/rabitqlib/utils/input_checks.hpp | 67 +++++++++++ python_bindings/__init__.pyi | 12 ++ python_bindings/hnsw_bindings.cpp | 4 +- python_bindings/ivf_bindings.cpp | 2 + python_bindings/symqg_bindings.cpp | 4 +- src/core/api.cpp | 15 +-- src/index/hnsw.cpp | 10 ++ src/index/ivf.cpp | 6 + src/index/qg.cpp | 4 + src/index/qg_builder.cpp | 2 + tests/python/test_input_boundaries.py | 59 ++++++++++ .../rabitqlib/clustering/qgkmeans_test.cpp | 15 ++- .../rabitqlib/fastscan/lut_prepare_test.cpp | 32 ++++++ tests/unit/rabitqlib/index/hnsw_test.cpp | 106 ++++++++++++++++++ tests/unit/rabitqlib/index/ivf_test.cpp | 104 +++++++++++++++++ tests/unit/rabitqlib/index/qg_test.cpp | 76 +++++++++++++ .../rabitqlib/utils/input_checks_test.cpp | 92 +++++++++++++++ 22 files changed, 620 insertions(+), 21 deletions(-) create mode 100644 include/rabitqlib/utils/input_checks.hpp create mode 100644 tests/unit/rabitqlib/utils/input_checks_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index ad0097c0..deb5132e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -63,11 +63,13 @@ endif() if(RABITQ_ENABLE_SANITIZERS) # Keep sanitizer instrumentation consistent across public inline code and # compiled components, including their runtime linkage. + # float-cast-overflow and float-divide-by-zero are not part of "undefined"; + # quantization converts floats to integers, so enable them explicitly. target_compile_options(rabitq_compile_options INTERFACE - "$<$:-fsanitize=address,undefined;-fno-omit-frame-pointer>" + "$<$:-fsanitize=address,undefined,float-cast-overflow,float-divide-by-zero;-fno-omit-frame-pointer>" ) target_link_options(rabitq_compile_options INTERFACE - $<$:-fsanitize=address,undefined> + $<$:-fsanitize=address,undefined,float-cast-overflow,float-divide-by-zero> ) endif() diff --git a/docs/development.md b/docs/development.md index cd589669..b9574089 100644 --- a/docs/development.md +++ b/docs/development.md @@ -334,6 +334,16 @@ Keep validation at the public operation boundary and reuse only invariants estab for that operation. Buffer contents and per-record factors may change between calls: validate them again even when reusing a query or workspace. +**Index inputs** + +Every index entry point that accepts float data, centroids or queries calls +`check_coordinates` from `utils/input_checks.hpp` after its shape checks and before +mutating state: non-finite values and magnitudes above the clustering bound raise +`std::invalid_argument`. The core `quantize_*` encoders and `Lut` rely on that (the +LUT additionally rejects non-finite rotated queries) because quantization converts +floats to integers, which is undefined for NaN. Keep the check at the boundary rather +than inside per-vector kernels. + **Python conversion** `ensure_*` returns a normalized NumPy owner retained across GIL release; diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index 894d25d8..47c9f7b1 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -14,6 +14,14 @@ as the indexed data. Arrays whose data pointers are not aligned for float32 are copied into aligned storage before native processing; inputs remain unchanged. +Data, centroids and queries must be finite. Coordinates are also limited to the +magnitude the clustering APIs accept, `sqrt(sqrt(FLT_MAX) / (64 * dim))` (about +`5e7` for `dim = 128`), so that squared distances and their products stay inside +float32. `build`, `add` and `search` raise `ValueError` (C++: +`std::invalid_argument`) for NaN, infinite, or larger coordinates before changing +the index. The C++ `quantize_*` encoders in `core/encoder.hpp` assume their +callers have validated the input. + These limits apply to the Python indexes. The C++ [stable core API](../rabitq/core_api.md) accepts working dimensions from 2, with multiples of 32 required for `SplitPacked`. C++ [MatrixRotator](../rabitq/rotator.md) has separate dimension and storage requirements; diff --git a/include/rabitqlib/core/encoder.hpp b/include/rabitqlib/core/encoder.hpp index 25a77f87..32d8bdaf 100644 --- a/include/rabitqlib/core/encoder.hpp +++ b/include/rabitqlib/core/encoder.hpp @@ -122,6 +122,8 @@ inline void quantize_compact_ex_bits( cur_ex_data.f_rescale_ex() = f_rescale_ex; } +// The per-vector encoders assume finite inputs; callers validate once per batch +// (see utils/input_checks.hpp) because quantization converts floats to integers. inline void quantize_split_single( const float* data, const float* centroid, diff --git a/include/rabitqlib/fastscan/detail/lut.hpp b/include/rabitqlib/fastscan/detail/lut.hpp index bd8ecf74..4b13e22a 100644 --- a/include/rabitqlib/fastscan/detail/lut.hpp +++ b/include/rabitqlib/fastscan/detail/lut.hpp @@ -12,6 +12,7 @@ #include "rabitqlib/fastscan/detail/lut_prepare.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib { @@ -68,6 +69,10 @@ class Lut { "Lut::reset requires uint16 scratch for high-accuracy lookup tables" ); } + // Table quantization converts query-derived values to integers. + if (!all_finite(rotated_query, padded_dim)) { + throw std::invalid_argument("Lut::reset requires a finite query"); + } lut_.resize(shape.storage_length); table_length_ = shape.table_length; build_lut(rotated_query, padded_dim, lut_scratch, lut_u16_scratch, use_hacc); diff --git a/include/rabitqlib/utils/input_checks.hpp b/include/rabitqlib/utils/input_checks.hpp new file mode 100644 index 00000000..ba883018 --- /dev/null +++ b/include/rabitqlib/utils/input_checks.hpp @@ -0,0 +1,67 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace rabitqlib { + +// True when every float32 value is finite. The integer reduction lets portable +// builds check several values per instruction without relaxing floating-point +// compiler options, and never allocates. Internal linkage keeps every ISA-flagged +// translation unit that includes this header on its own copy. +static inline bool all_finite(const float* values, size_t count) { + static_assert( + sizeof(float) == sizeof(uint32_t) && std::numeric_limits::is_iec559 + ); + uint32_t nonfinite = 0; + for (size_t i = 0; i < count; ++i) { + uint32_t bits; + std::memcpy(&bits, values + i, sizeof(bits)); + nonfinite |= static_cast((bits & 0x7f800000U) == 0x7f800000U); + } + return nonfinite == 0; +} + +template +static inline bool all_finite(const T* values, size_t count) { + size_t nonfinite = 0; + for (size_t i = 0; i < count; ++i) { + nonfinite += static_cast(!std::isfinite(values[i])); + } + return nonfinite == 0; +} + +// Largest coordinate magnitude whose squared distances, their products in graph +// pruning, and rotated residuals stay inside float32; shared with clustering. +static inline float max_coordinate_magnitude(size_t dim) { + return static_cast(std::sqrt( + std::sqrt(static_cast(std::numeric_limits::max())) / + (64.0 * static_cast(dim)) + )); +} + +// Reject vectors the indexes cannot quantize. `what` names the input in the message. +static inline void check_coordinates( + const float* values, size_t count, size_t dim, const char* what +) { + if (!all_finite(values, count)) { + throw std::invalid_argument(std::string(what) + " must contain only finite values"); + } + const float bound = max_coordinate_magnitude(dim); + size_t out_of_range = 0; + for (size_t i = 0; i < count; ++i) { + out_of_range += static_cast(!(std::fabs(values[i]) <= bound)); + } + if (out_of_range != 0) { + throw std::invalid_argument( + std::string(what) + " exceeds the safe float32 coordinate range" + ); + } +} + +} // namespace rabitqlib diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 4823f37e..98d15e44 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -94,6 +94,8 @@ class IvfIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. Examples -------- @@ -140,6 +142,8 @@ class IvfIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. Examples -------- @@ -388,6 +392,8 @@ class HnswIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __init__( @@ -424,6 +430,8 @@ class HnswIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __repr__(self) -> str: @@ -641,6 +649,8 @@ class SymqgIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __init__( @@ -669,6 +679,8 @@ class SymqgIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __repr__(self) -> str: diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index ae47f2c3..668c98ef 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -311,7 +311,9 @@ random_seed : int, optional Notes ----- Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc"; +supported; an overlapping mutation raises RuntimeError. +Data, centroids and queries must be finite float32 values within the index +coordinate range; other values raise ValueError.)doc"; py::class_(m, "HnswIndex", kIndexDoc) .def( diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index e03f8b52..b2bf292d 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -353,6 +353,8 @@ Notes ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. +Data, centroids and queries must be finite float32 values within the index +coordinate range; other values raise ValueError. Examples -------- diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index ee1d3fd1..e144b5c3 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -232,7 +232,9 @@ quantization_bits : {0, 4, 8}, optional Notes ----- Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc"; +supported; an overlapping mutation raises RuntimeError. +Data, centroids and queries must be finite float32 values within the index +coordinate range; other values raise ValueError.)doc"; static constexpr char kSearchDoc[] = R"doc(Search a batch of query vectors. diff --git a/src/core/api.cpp b/src/core/api.cpp index fc8737d8..88cfdf9e 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -4,7 +4,6 @@ #include #include #include -#include #include #include #include @@ -17,6 +16,7 @@ #include "rabitqlib/core/encoder.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/pack_excode.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib::quant::v1 { @@ -94,18 +94,7 @@ void check_input(const float* input, size_t dim) { check_buffer( {input, checked_multiply(dim, sizeof(float))}, checked_multiply(dim, sizeof(float)) ); - static_assert( - sizeof(float) == sizeof(uint32_t) && std::numeric_limits::is_iec559 - ); - // An integer reduction lets portable builds check several float32 values - // per instruction without relaxing floating-point compiler options. - uint32_t nonfinite = 0; - for (size_t i = 0; i < dim; ++i) { - uint32_t bits; - std::memcpy(&bits, input + i, sizeof(bits)); - nonfinite |= static_cast((bits & 0x7f800000U) == 0x7f800000U); - } - if (nonfinite != 0) { + if (!all_finite(input, dim)) { throw std::invalid_argument("Input vectors must contain only finite values"); } } diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index d70bc774..984638ee 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -35,6 +35,7 @@ #include "rabitqlib/index/detail/visited_pool.hpp" #include "rabitqlib/index/hnsw/hnsw_quant.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/path.hpp" #include "rabitqlib/utils/space.hpp" @@ -518,6 +519,8 @@ void HierarchicalNSW::construct( throw std::invalid_argument("HNSW cluster ID is out of range"); } } + check_coordinates(data, num_points * dim_, dim_, "HNSW construction data"); + check_coordinates(centroids, num_clusters * dim_, dim_, "HNSW construction centroids"); try { num_clusters_ = num_clusters; @@ -913,6 +916,12 @@ std::vector>> HierarchicalNSW::search( if (k == 0) { throw std::invalid_argument("HNSW k must be greater than zero"); } + if (num_queries != 0) { + if (queries == nullptr) { + throw std::invalid_argument("HNSW search queries must not be null"); + } + check_coordinates(queries, num_queries * dim_, dim_, "HNSW search queries"); + } std::vector>> results(num_queries); index_detail::parallel_for( 0, @@ -1290,6 +1299,7 @@ std::vector HierarchicalNSW::add( if (num_points > buffer::kSearchBufferMaxPointCount - num_stored_points_) { throw std::invalid_argument("HNSW add point count exceeds the supported ID range"); } + check_coordinates(data, num_points * dim_, dim_, "HNSW add data"); std::vector labels(num_points); // Either the caller names the cluster of every point, or each one goes to its // nearest centroid, ordered the way a query is routed. diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 4c7c08ae..e906a006 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -38,6 +38,7 @@ #include "rabitqlib/index/ivf/cluster.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/path.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -210,6 +211,8 @@ void IVF::construct( id_lists[cluster_id].push_back(static_cast(i)); cluster_sizes[cluster_id] += 1; } + check_coordinates(data, max_elements_ * dim_, dim_, "IVF construction data"); + check_coordinates(centroids, num_clusters_ * dim_, dim_, "IVF construction centroids"); allocate_storage(cluster_sizes); create_cluster_views(cluster_sizes); @@ -356,6 +359,7 @@ void IVF::add( if (exceeds_storage(new_num_points)) { throw std::invalid_argument("IVF configuration exceeds addressable storage"); } + check_coordinates(data, num_points * dim_, dim_, "IVF add data"); // Assign every new point to a cluster std::vector assigned(num_points); @@ -977,6 +981,7 @@ void IVF::search( bool use_hacc ) const { validate_search(query, k, nprobe, result_ids); + check_coordinates(query, dim_, dim_, "IVF search query"); nprobe = std::min(nprobe, num_clusters_); // Probe at most every cluster std::vector rotated_query(padded_dim_); @@ -1039,6 +1044,7 @@ void IVF::search_batch( if (num_queries > max_values / dim_ || num_queries > max_values / k) { throw std::length_error("IVF query batch is too large"); } + check_coordinates(queries, num_queries * dim_, dim_, "IVF search queries"); nprobe = std::min(nprobe, num_clusters_); const bool hacc = use_hacc.value_or(!raw_reranking_ && ex_bits_ > 2); const size_t num_workers = diff --git a/src/index/qg.cpp b/src/index/qg.cpp index 1aeb5551..e3253e09 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -32,6 +32,7 @@ #include "rabitqlib/index/symqg/detail/search_scratch.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/bitops.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/path.hpp" @@ -452,6 +453,7 @@ void QuantizedGraph::search_impl( const float* query, uint32_t k, uint32_t* result_ids, float* distances, size_t ef_search ) { validate_search(query, k, result_ids, distances, ef_search); + check_coordinates(query, dim_, dim_, "QuantizedGraph search query"); detail::SearchScratch scratch(padded_dim_, max_degree_, ef_search, k); thread_local VisitedSet visited; thread_local size_t visited_size = 0; @@ -498,6 +500,7 @@ void QuantizedGraph::search_batch_with_ef( search_impl(queries, k, result_ids, distances, ef_search); return; } + check_coordinates(queries, num_queries * dim_, dim_, "QuantizedGraph search queries"); const int num_workers = rabitqlib::detail::openmp_team_size(num_threads, num_queries); index_detail::ParallelException failure; @@ -808,6 +811,7 @@ void detail::search_with_workspace( const float* k1xsumq ) { graph.validate_search(query, k, result_ids, distances, graph.default_ef_search_); + check_coordinates(query, graph.dim_, graph.dim_, "QuantizedGraph search query"); graph.search_with_scratch_impl(query, k, result_ids, distances, scratch, hint, k1xsumq); } diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp index 338a5efd..dcf1bbc4 100644 --- a/src/index/qg_builder.cpp +++ b/src/index/qg_builder.cpp @@ -25,6 +25,7 @@ #include "rabitqlib/index/symqg/detail/pipnn.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -100,6 +101,7 @@ void QGBuilder::reset(const float* data) { if (data == nullptr) { throw std::invalid_argument("QuantizedGraph build data must not be null"); } + check_coordinates(data, num_points_ * dim_, dim_, "QuantizedGraph build data"); if (initialization_ == QGInitialization::PiPNN) { auto initial = detail::build_initial_graph( data, num_points_, dim_, max_degree_, qg_.metric_type_, num_workers_ diff --git a/tests/python/test_input_boundaries.py b/tests/python/test_input_boundaries.py index f9c5ff26..1fbb4dc4 100644 --- a/tests/python/test_input_boundaries.py +++ b/tests/python/test_input_boundaries.py @@ -277,3 +277,62 @@ def test_array_conversion_reports_target_dtype(built_index, value, description): == "queries must be convertible to a NumPy array of dtype float32; got " + description ) + + +@pytest.mark.parametrize( + "value,message", + [ + (np.nan, "must contain only finite values"), + (np.inf, "must contain only finite values"), + (-np.inf, "must contain only finite values"), + (1e30, "exceeds the safe float32 coordinate range"), + ], +) +def test_nonfinite_or_oversized_coordinates_are_rejected(built_index, value, message): + index, queries = built_index + before = index.num_points + bad_queries = np.array(queries, copy=True) + bad_queries[1, 3] = value + with pytest.raises(ValueError, match=message): + index.search(bad_queries, 3, 32) + extra = np.random.default_rng(7).normal(size=(2, 64)).astype(np.float32) + extra[0, 0] = value + if isinstance(index, IvfIndex): + with pytest.raises(ValueError, match=message): + index.add(extra) + elif isinstance(index, HnswIndex): + # The shared fixture is at capacity, which is checked before coordinates. + data = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + spare = HnswIndex(64, 82, nbits=4) + spare.build( + data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32) + ) + with pytest.raises(ValueError, match=message): + spare.add(extra) + assert spare.num_points == 80 + assert index.num_points == before + ids, distances = index.search(queries, 3, 32) + assert ids.shape == (2, 3) and np.all(np.isfinite(distances)) + + +@pytest.mark.parametrize("index_kind", ["ivf", "hnsw", "symqg"]) +@pytest.mark.parametrize("value", [np.nan, np.inf, 1e30]) +def test_build_rejects_bad_coordinates_without_building(index_kind, value): + data = np.random.default_rng(3).normal(size=(80, 64)).astype(np.float32) + bad_data = data.copy() + bad_data[5, 3] = value + if index_kind == "symqg": + index = SymqgIndex(64, 32) + with pytest.raises(ValueError): + index.build(bad_data, 40) + else: + centroid = data.mean(axis=0, keepdims=True) + ids = np.zeros(80, dtype=np.uint32) + index = IvfIndex(64, 80, 1, 4) if index_kind == "ivf" else HnswIndex(64, 80) + with pytest.raises(ValueError): + index.build(bad_data, centroid, ids) + bad_centroid = centroid.copy() + bad_centroid[0, 0] = value + with pytest.raises(ValueError): + index.build(data, bad_centroid, ids) + assert not index.is_built diff --git a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp index 8d2401c5..a529b9c1 100644 --- a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp @@ -1289,17 +1289,24 @@ TEST(QGKMeansTest, RejectsInvalidNativeParametersBeforeAllocation) { namespace rabitqlib::qgkmeans { TEST(QGAssignerTest, PropagatesSearchFailureFromWorkers) { + // Out-of-range queries are rejected by the per-query search entry that each + // worker calls, so the failure originates inside the parallel region. const std::vector centroids(33 * 64, 0.0F); const std::vector data(128 * 64, 1e19F); std::vector assignments(128, kPidMax); std::vector distances(128, -1.0F); QGAssigner assigner(64, 33, 32, 33, 33, 1, 2); - EXPECT_THROW( + try { assigner.assign( centroids.data(), data.data(), 128, assignments.data(), distances.data() - ), - std::runtime_error - ); + ); + ADD_FAILURE() << "expected the worker failure to propagate"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ( + error.what(), + "QuantizedGraph search query exceeds the safe float32 coordinate range" + ); + } EXPECT_EQ(assignments.front(), kPidMax); EXPECT_EQ(distances.front(), -1.0F); } diff --git a/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp index 11b664be..bf1cf3ad 100644 --- a/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp +++ b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include "rabitqlib/core/detail/simd/space_dispatch.hpp" @@ -262,5 +263,36 @@ TEST(FastScanLutPreparationTest, FloatLutPreservesStandardAndHaccCodes) { TEST(FastScanLutPreparationTest, DoubleLutRetainsOriginalPreparation) { check_prepared_lut(); } + +TEST(FastScanLutPreparationTest, RejectsNonFiniteQueriesForFloatAndDouble) { + constexpr size_t kPaddedDim = 64; + const std::vector query(kPaddedDim, 0.5F); + const std::vector double_query(kPaddedDim, 0.5); + const float poison[] = { + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + }; + for (float value : poison) { + auto bad = query; + bad[7] = value; + EXPECT_THROW((Lut(bad.data(), kPaddedDim)), std::invalid_argument); + EXPECT_THROW((Lut(bad.data(), kPaddedDim, true)), std::invalid_argument); + auto bad_double = double_query; + bad_double[7] = static_cast(value); + EXPECT_THROW((Lut(bad_double.data(), kPaddedDim)), std::invalid_argument); + EXPECT_THROW( + (Lut(bad_double.data(), kPaddedDim, true)), std::invalid_argument + ); + } + // A rejected reset keeps the previously prepared table. + Lut lut(query.data(), kPaddedDim); + std::vector scratch(kPaddedDim * 4); + auto bad = query; + bad[0] = std::numeric_limits::quiet_NaN(); + EXPECT_THROW(lut.reset(bad.data(), kPaddedDim, scratch.data()), std::invalid_argument); + EXPECT_EQ(lut.padded_dim(), kPaddedDim); +} + } // namespace } // namespace rabitqlib::fastscan diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 17ec4226..50d9e044 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -11,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -25,6 +27,7 @@ #include "rabitqlib/index/hnsw/hnsw_quant.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" #include "temp_file.hpp" @@ -1324,6 +1327,109 @@ TEST(HnswRemoveTest, LoadRejectsAnOutOfRangeClusterEvenWithTheMark) { } } +TEST(HnswInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) { + constexpr size_t kDim = 64; + constexpr size_t kCount = 16; + constexpr size_t kClusters = 2; + std::mt19937 gen(905); + std::normal_distribution normal; + std::vector data(kCount * kDim); + std::vector centroids(kClusters * kDim); + std::vector extra(2 * kDim); + std::vector queries(2 * kDim); + for (auto* values : {&data, ¢roids, &extra, &queries}) { + 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); + } + const float bound = max_coordinate_magnitude(kDim); + const std::array rejected{ + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 2.0F * bound, + }; + const auto message_for = [](float value, const char* what) { + return std::string(what) + (std::isfinite(value) + ? " exceeds the safe float32 coordinate range" + : " must contain only finite values"); + }; + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + + HierarchicalNSW index(kCount + 2, kDim, 4, 8, 32); + for (float value : rejected) { + auto bad_data = data; + bad_data[kDim + 3] = value; + expect_message( + [&] { + index.construct( + kClusters, + centroids.data(), + kCount, + bad_data.data(), + cluster_ids.data(), + 1 + ); + }, + message_for(value, "HNSW construction data") + ); + auto bad_centroids = centroids; + bad_centroids.back() = value; + expect_message( + [&] { + index.construct( + kClusters, + bad_centroids.data(), + kCount, + data.data(), + cluster_ids.data(), + 1 + ); + }, + message_for(value, "HNSW construction centroids") + ); + EXPECT_EQ(index.num_stored_points(), 0U); + } + ASSERT_NO_THROW(index.construct( + kClusters, centroids.data(), kCount, data.data(), cluster_ids.data(), 1 + )); + ASSERT_EQ(index.num_stored_points(), kCount); + + for (float value : rejected) { + auto bad_extra = extra; + bad_extra[kDim] = value; + expect_message( + [&] { index.add(bad_extra.data(), 2, nullptr, false, 1); }, + message_for(value, "HNSW add data") + ); + EXPECT_EQ(index.num_stored_points(), kCount); + auto bad_queries = queries; + bad_queries[kDim + 1] = value; + expect_message( + [&] { index.search(bad_queries.data(), 2, 3, 10, 1); }, + message_for(value, "HNSW search queries") + ); + } + // The bound itself is accepted everywhere. + auto edge = extra; + edge[0] = bound; + edge[kDim] = -bound; + ASSERT_NO_THROW(index.add(edge.data(), 2, nullptr, false, 1)); + EXPECT_EQ(index.num_stored_points(), kCount + 2); + auto edge_queries = queries; + edge_queries[0] = bound; + EXPECT_EQ(index.search(edge_queries.data(), 2, 3, 10, 1).size(), 2U); +} + } // namespace } // namespace rabitqlib::hnsw diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 6bcc640a..621a9d66 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -27,6 +27,7 @@ #include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "temp_file.hpp" namespace rabitqlib::ivf { @@ -1444,6 +1445,109 @@ TEST(IvfRemoveTest, RejectsBadInputAndRemovesNothingOnFailure) { } } +TEST(IvfInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) { + constexpr size_t kCount = 40; + const auto points = random_rows(kCount, 901); + const auto centroids = random_rows(kDynNumClusters, 902); + std::vector assigned(kCount); + for (size_t i = 0; i < kCount; ++i) { + assigned[i] = static_cast(i % kDynNumClusters); + } + const float bound = max_coordinate_magnitude(kDynDim); + const std::array rejected{ + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 2.0F * bound, + }; + const auto message_for = [](float value, const char* what) { + return std::string(what) + (std::isfinite(value) + ? " exceeds the safe float32 coordinate range" + : " must contain only finite values"); + }; + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + + IVF index(kCount, kDynDim, kDynNumClusters, 4); + for (float value : rejected) { + auto bad_points = points; + bad_points[kDynDim + 3] = value; + expect_message( + [&] { + index.construct( + bad_points.data(), centroids.data(), assigned.data(), false, 1 + ); + }, + message_for(value, "IVF construction data") + ); + auto bad_centroids = centroids; + bad_centroids.back() = value; + expect_message( + [&] { + index.construct( + points.data(), bad_centroids.data(), assigned.data(), false, 1 + ); + }, + message_for(value, "IVF construction centroids") + ); + EXPECT_EQ(index.num_points(), 0U); + } + ASSERT_NO_THROW( + index.construct(points.data(), centroids.data(), assigned.data(), false, 1) + ); + ASSERT_EQ(index.num_points(), kCount); + + const auto extra = random_rows(2, 903); + const auto queries = random_rows(2, 904); + std::vector result_ids(2 * 3); + std::vector distances(2 * 3); + for (float value : rejected) { + auto bad_extra = extra; + bad_extra[kDynDim] = value; + expect_message( + [&] { index.add(bad_extra.data(), 2, nullptr, false, 1); }, + message_for(value, "IVF add data") + ); + EXPECT_EQ(index.num_points(), kCount); + auto bad_query = queries; + bad_query[0] = value; + expect_message( + [&] { + index.search(bad_query.data(), 3, 1, result_ids.data(), distances.data()); + }, + message_for(value, "IVF search query") + ); + auto bad_batch = queries; + bad_batch[kDynDim + 1] = value; + expect_message( + [&] { + index.search_batch( + bad_batch.data(), 2, 3, 1, result_ids.data(), distances.data() + ); + }, + message_for(value, "IVF search queries") + ); + } + // The bound itself is accepted everywhere. + auto edge = extra; + edge[0] = bound; + edge[kDynDim] = -bound; + ASSERT_NO_THROW(index.add(edge.data(), 2, nullptr, false, 1)); + EXPECT_EQ(index.num_points(), kCount + 2); + auto edge_query = queries; + edge_query[0] = bound; + ASSERT_NO_THROW( + index.search(edge_query.data(), 3, 1, result_ids.data(), distances.data()) + ); + ASSERT_NO_THROW( + index.search_batch(edge_query.data(), 2, 3, 1, result_ids.data(), distances.data()) + ); +} + } // namespace } // namespace rabitqlib::ivf diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index 2359a5ed..32dae1a3 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -40,6 +40,7 @@ #include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" #include "temp_file.hpp" @@ -2116,3 +2117,78 @@ TEST(QgSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { } } } // namespace rabitqlib::symqg + +namespace rabitqlib::symqg { +TEST(QgInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) { + constexpr size_t kCount = 33, kDim = 64; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.17F); + } + const float bound = max_coordinate_magnitude(kDim); + const std::array rejected{ + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 2.0F * bound, + }; + const auto message_for = [](float value, const char* what) { + return std::string(what) + (std::isfinite(value) + ? " exceeds the safe float32 coordinate range" + : " must contain only finite values"); + }; + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + + for (float value : rejected) { + auto bad_data = data; + bad_data[kDim + 5] = value; + for (auto init : {QGInitialization::Random, QGInitialization::PiPNN}) { + QuantizedGraph graph(kCount, kDim, 32); + expect_message( + [&] { QGBuilder(graph, 32, bad_data.data(), 1, init); }, + message_for(value, "QuantizedGraph build data") + ); + EXPECT_EQ(graph.num_points(), 0U); + } + } + + QuantizedGraph graph(kCount, kDim, 32); + QGBuilder builder(graph, 32, data.data(), 1, QGInitialization::Random); + builder.build(2); + graph.set_ef(10); + const std::vector queries(data.begin(), data.begin() + 2 * kDim); + std::vector result_ids(2 * 3); + std::vector distances(2 * 3); + for (float value : rejected) { + auto bad_query = queries; + bad_query[0] = value; + expect_message( + [&] { graph.search(bad_query.data(), 3, result_ids.data(), distances.data()); }, + message_for(value, "QuantizedGraph search query") + ); + auto bad_batch = queries; + bad_batch[kDim + 1] = value; + expect_message( + [&] { + graph.search_batch( + bad_batch.data(), 2, 3, result_ids.data(), distances.data(), 1 + ); + }, + message_for(value, "QuantizedGraph search queries") + ); + } + auto edge_query = queries; + edge_query[0] = bound; + ASSERT_NO_THROW( + graph.search(edge_query.data(), 3, result_ids.data(), distances.data()) + ); + ASSERT_NO_THROW( + graph.search_batch(edge_query.data(), 2, 3, result_ids.data(), distances.data(), 1) + ); +} +} // namespace rabitqlib::symqg diff --git a/tests/unit/rabitqlib/utils/input_checks_test.cpp b/tests/unit/rabitqlib/utils/input_checks_test.cpp new file mode 100644 index 00000000..f0abd877 --- /dev/null +++ b/tests/unit/rabitqlib/utils/input_checks_test.cpp @@ -0,0 +1,92 @@ +#include "rabitqlib/utils/input_checks.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace rabitqlib { +namespace { + +TEST(InputChecksTest, DetectsEveryNonFiniteValueForFloatAndDouble) { + const std::vector finite{ + 0.0F, + -0.0F, + std::numeric_limits::denorm_min(), + std::numeric_limits::max(), + -std::numeric_limits::max(), + 1.5F, + }; + EXPECT_TRUE(all_finite(finite.data(), finite.size())); + EXPECT_TRUE(all_finite(finite.data(), 0)); + const std::vector finite_double(finite.begin(), finite.end()); + EXPECT_TRUE(all_finite(finite_double.data(), finite_double.size())); + + const float poison[] = { + std::numeric_limits::quiet_NaN(), + std::numeric_limits::signaling_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + }; + for (float value : poison) { + for (size_t position : {size_t{0}, finite.size() / 2, finite.size() - 1}) { + auto values = finite; + values[position] = value; + EXPECT_FALSE(all_finite(values.data(), values.size())); + const std::vector doubles(values.begin(), values.end()); + EXPECT_FALSE(all_finite(doubles.data(), doubles.size())); + } + } +} + +TEST(InputChecksTest, CoordinateBoundMatchesClusteringAndAcceptsItself) { + constexpr size_t kDim = 128; + const double expected = std::sqrt( + std::sqrt(static_cast(std::numeric_limits::max())) / + (64.0 * static_cast(kDim)) + ); + const float bound = max_coordinate_magnitude(kDim); + EXPECT_NEAR(static_cast(bound), expected, expected * 1e-6); + EXPECT_LT(max_coordinate_magnitude(2 * kDim), bound); + + std::vector values(kDim, 0.25F); + values[0] = bound; + values[1] = -bound; + EXPECT_NO_THROW(check_coordinates(values.data(), values.size(), kDim, "data")); + EXPECT_NO_THROW(check_coordinates(values.data(), 0, kDim, "data")); + + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + values[kDim - 1] = std::nextafter(bound, std::numeric_limits::infinity()); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "data"); }, + "data exceeds the safe float32 coordinate range" + ); + values[kDim - 1] = -std::nextafter(bound, std::numeric_limits::infinity()); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "queries"); }, + "queries exceeds the safe float32 coordinate range" + ); + // Non-finite values are reported as such even when the range is also exceeded. + values[kDim - 1] = std::numeric_limits::infinity(); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "data"); }, + "data must contain only finite values" + ); + values[kDim - 1] = std::numeric_limits::quiet_NaN(); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "data"); }, + "data must contain only finite values" + ); +} + +} // namespace +} // namespace rabitqlib From a7688acbaed2288802058e52463c47d3353f09a1 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 18:38:56 +0800 Subject: [PATCH 5/6] build: install the input validation header lut.hpp includes utils/input_checks.hpp, so installed consumers of the FastScan headers need it too. --- CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index deb5132e..579dd400 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -283,6 +283,7 @@ install(DIRECTORY include/rabitqlib/core DESTINATION "${rabitq_include_dir}") install(FILES include/rabitqlib/utils/bitops.hpp include/rabitqlib/utils/cpu_features.hpp + include/rabitqlib/utils/input_checks.hpp include/rabitqlib/utils/memory.hpp include/rabitqlib/utils/openmp.hpp include/rabitqlib/utils/packed.hpp From b6ff0916ca60039cf02850a3b2e54d023d5e650d Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 19:23:55 +0800 Subject: [PATCH 6/6] test(consumer): expect the extended sanitizer flag from installed exports The installed-consumer options check hard-coded -fsanitize=address,undefined; the sanitizer build now also enables float-cast-overflow and float-divide-by-zero. --- tests/consumer/check_options.cmake | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/consumer/check_options.cmake b/tests/consumer/check_options.cmake index 5df01dc5..d01252d2 100644 --- a/tests/consumer/check_options.cmake +++ b/tests/consumer/check_options.cmake @@ -7,8 +7,10 @@ include("${OPTIONS_FILE}") # Catch optimization, architecture and warnings leaking through any level of # the target graph, including installed exports. set(required_options ${openmp_options}) +# Keep in sync with RABITQ_ENABLE_SANITIZERS in the top-level CMakeLists.txt. +set(sanitizer_flag -fsanitize=address,undefined,float-cast-overflow,float-divide-by-zero) if(expect_sanitizers) - list(APPEND required_options -fsanitize=address,undefined -fno-omit-frame-pointer) + list(APPEND required_options ${sanitizer_flag} -fno-omit-frame-pointer) endif() foreach(flag IN LISTS consumer_options) if(NOT flag STREQUAL "" AND NOT flag IN_LIST required_options) @@ -35,6 +37,6 @@ if(expect_msvc) endif() endforeach() endif() -if(expect_sanitizers AND NOT "-fsanitize=address,undefined" IN_LIST consumer_link_options) +if(expect_sanitizers AND NOT sanitizer_flag IN_LIST consumer_link_options) message(FATAL_ERROR "The instrumented library must propagate sanitizer runtime linkage") endif()