Reject non-finite brute-force search distances

CMakeLists.txt

@@ -28,6 +28,7 @@ set(EXAMPLE_NAMES
     )
 
 set(TEST_NAMES
+    bruteforce_numeric_range_test
     epsilon_search_test
     multiThread_replace_test
     multiThreadLoad_test

README.md

@@ -71,6 +71,8 @@ Note that inner product is not an actual metric. An element can be closer to som
 
 For other spaces use the nmslib library https://github.com/nmslib/nmslib. 
 
+For brute-force search, `BFIndex.knn_query` raises `RuntimeError` if a distance to an allowed label is NaN or infinite, including when finite coordinates overflow during distance calculation. Use smaller magnitudes to keep squared L2 distances within float32 range. Labels excluded by `filter` do not cause this error, and `k=0` returns empty results without evaluating distances. The C++ `BruteforceSearch<float>` and `BruteforceSearch<double>` APIs return an error `Status` through `searchKnnNoExceptions`; their throwing wrappers raise `std::runtime_error` when exceptions are enabled.
+
 #### API description
 * `hnswlib.get_simd()` returns the ISA selected at import (`sse`, `avx`, `avx512`, or `aarch64`). Set `HNSWLIB_SIMD` to force a lower-or-equal level (raises if the CPU cannot run it).
 * `hnswlib.Index(space, dim)` creates a non-initialized index an HNSW in space `space` with integer dimension `dim`.

hnswlib/bruteforce.h

@@ -8,6 +8,9 @@
 #include <algorithm>
 #include <assert.h>
 #include <iostream>
+#include <cstdint>
+#include <cstring>
+#include <limits>
 
 namespace hnswlib {
 template<typename dist_t>
@@ -120,9 +123,18 @@ class BruteforceSearch : public AlgorithmInterface<dist_t> {
     searchKnnNoExceptions(const void *query_data, size_t k, BaseFilterFunctor* isIdAllowed = nullptr) const override {
         assert(k <= cur_element_count);
         std::priority_queue<std::pair<dist_t, labeltype >> topResults;
+        if (k == 0)
+            return topResults;
         dist_t lastdist = std::numeric_limits<dist_t>::max();
         for (int i = 0; i < cur_element_count; i++) {
             dist_t dist = fstdistfunc_(query_data, data_ + size_per_element_ * i, dist_func_param_);
+            if (!distanceIsFinite(dist)) {
+                labeltype label = getExternalLabel(i);
+                if (isIdAllowed && !(*isIdAllowed)(label))
+                    continue;
+                return Status("BruteforceSearch encountered a non-finite distance; "
+                              "check input values or rescale to avoid overflow");
+            }
             if (dist <= lastdist || topResults.size() < k) {
                 labeltype label = getExternalLabel(i);
                 if ((!isIdAllowed) || (*isIdAllowed)(label)) {
@@ -196,5 +208,30 @@ class BruteforceSearch : public AlgorithmInterface<dist_t> {
 
         input.close();
     }
+
+ private:
+    template<typename T>
+    static bool distanceIsFinite(const T&) {
+        return true;
+    }
+
+    // Inspect the exponent bits: std::isfinite may be optimized away by fast-math.
+    static bool distanceIsFinite(float distance) {
+        static_assert(sizeof(float) == sizeof(std::uint32_t) &&
+                      std::numeric_limits<float>::is_iec559,
+                      "BruteforceSearch requires IEEE 754 binary32 floats");
+        std::uint32_t bits;
+        std::memcpy(&bits, &distance, sizeof(bits));
+        return (bits & 0x7f800000u) != 0x7f800000u;
+    }
+
+    static bool distanceIsFinite(double distance) {
+        static_assert(sizeof(double) == sizeof(std::uint64_t) &&
+                      std::numeric_limits<double>::is_iec559,
+                      "BruteforceSearch requires IEEE 754 binary64 doubles");
+        std::uint64_t bits;
+        std::memcpy(&bits, &distance, sizeof(bits));
+        return (bits & 0x7ff0000000000000ULL) != 0x7ff0000000000000ULL;
+    }
 };
 }  // namespace hnswlib

python_bindings/bindings.cpp

@@ -6,6 +6,7 @@
 #include "hnswlib.h"
 #include <thread>
 #include <atomic>
+#include <memory>
 #include <stdlib.h>
 #include <assert.h>
 
@@ -856,6 +857,8 @@ class BFIndex {
         const std::function<bool(hnswlib::labeltype)>& filter = nullptr) {
         py::array_t < dist_t, py::array::c_style | py::array::forcecast > items(input);
         auto buffer = items.request();
+        std::unique_ptr<hnswlib::labeltype[]> labels_owner;
+        std::unique_ptr<dist_t[]> distances_owner;
         hnswlib::labeltype *data_numpy_l;
         dist_t *data_numpy_d;
         size_t rows, features;
@@ -867,8 +870,10 @@ class BFIndex {
             py::gil_scoped_release l;
             get_input_array_shapes(buffer, &rows, &features);
 
-            data_numpy_l = new hnswlib::labeltype[rows * k];
-            data_numpy_d = new dist_t[rows * k];
+            labels_owner.reset(new hnswlib::labeltype[rows * k]);
+            data_numpy_l = labels_owner.get();
+            distances_owner.reset(new dist_t[rows * k]);
+            data_numpy_d = distances_owner.get();
 
             CustomFilterFunctor idFilter(filter);
             CustomFilterFunctor* p_idFilter = filter ? &idFilter : nullptr;
@@ -909,11 +914,13 @@ class BFIndex {
         }
 
         py::capsule free_when_done_l(data_numpy_l, [](void *f) {
-            delete[] f;
+            delete[] static_cast<hnswlib::labeltype*>(f);
         });
+        labels_owner.release();
         py::capsule free_when_done_d(data_numpy_d, [](void *f) {
-            delete[] f;
+            delete[] static_cast<dist_t*>(f);
         });
+        distances_owner.release();
 
 
         return py::make_tuple(