Reject non-finite cosine inputs

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 cosine indexes (`Index` and `BFIndex`), `add_items` and `knn_query` reject NaN and infinity with a `ValueError` after conversion to float32. The whole batch is checked before insertion or search; an insertion batch rejected for non-finite values leaves the index unchanged.
+
 #### 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`.

python_bindings/bindings.cpp

@@ -1,3 +1,4 @@
+#include <cmath>
 #include <iostream>
 #include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
@@ -114,6 +115,17 @@ inline void get_input_array_shapes(const py::buffer_info& buffer, size_t* rows,
 }
 
 
+inline void check_cosine_input(const py::buffer_info& buffer) {
+    const float* data = (const float*)buffer.ptr;
+    // An integer reduction allows vectorization of the finite-input path.
+    int all_finite = 1;
+    for (py::ssize_t i = 0; i < buffer.size; i++)
+        all_finite &= std::isfinite(data[i]);
+    if (!all_finite)
+        throw std::invalid_argument("Cosine vectors must contain only finite values");
+}
+
+
 inline std::vector<size_t> get_input_ids_and_check_shapes(const py::object& ids_, size_t feature_rows) {
     std::vector<size_t> ids;
     if (!ids_.is_none()) {
@@ -269,6 +281,9 @@ class Index {
 
         std::vector<size_t> ids = get_input_ids_and_check_shapes(ids_, rows);
 
+        if (normalize)
+            check_cosine_input(buffer);
+
         {
             int start = 0;
             if (!ep_added) {
@@ -629,6 +644,9 @@ class Index {
             py::gil_scoped_release l;
             get_input_array_shapes(buffer, &rows, &features);
 
+            if (normalize)
+                check_cosine_input(buffer);
+
             // avoid using threads when the number of searches is small:
             if (rows <= num_threads * 4) {
                 num_threads = 1;
@@ -812,6 +830,9 @@ class BFIndex {
 
         std::vector<size_t> ids = get_input_ids_and_check_shapes(ids_, rows);
 
+        if (normalize)
+            check_cosine_input(buffer);
+
         {
             for (size_t row = 0; row < rows; row++) {
                 size_t id = ids.size() ? ids.at(row) : cur_l + row;
@@ -867,6 +888,9 @@ class BFIndex {
             py::gil_scoped_release l;
             get_input_array_shapes(buffer, &rows, &features);
 
+            if (normalize)
+                check_cosine_input(buffer);
+
             data_numpy_l = new hnswlib::labeltype[rows * k];
             data_numpy_d = new dist_t[rows * k];
 

tests/python/bindings_test_spaces.py

@@ -1,4 +1,7 @@
+import itertools
+import tempfile
 import unittest
+import warnings
 
 import numpy as np
 
@@ -37,3 +40,156 @@ def testRandomSelf(self):
 
                     diff=np.mean(np.abs(distances-expected_distances))
                     self.assertAlmostEqual(diff, 0, delta=1e-3)
+
+
+class CosineFiniteInputTestCase(unittest.TestCase):
+    index_types = (hnswlib.Index, hnswlib.BFIndex)
+    non_finite_values = (np.nan, np.inf, -np.inf)
+
+    def make_index(self, index_type, num_threads=1):
+        index = index_type(space='cosine', dim=4)
+        index.init_index(max_elements=64)
+        index.set_num_threads(num_threads)
+        if index_type is hnswlib.Index:
+            index.set_ef(64)
+        return index
+
+    def testNonFiniteComponentsRejected(self):
+        cases = itertools.product(self.index_types, ('add_items', 'knn_query'),
+                                  (1, 2), (0, 1, 3), self.non_finite_values)
+        for index_type, operation, ndim, component, value in cases:
+            with self.subTest(index=index_type.__name__, operation=operation,
+                              ndim=ndim, component=component, value=value):
+                index = self.make_index(index_type)
+                if operation == 'knn_query':
+                    index.add_items([[1, 0, 0, 0]])
+                data = np.array([1, 0, 0, 0], dtype=np.float32)
+                data[component] = value
+                if ndim == 2:
+                    data = data.reshape(1, 4)
+                with self.assertRaises(ValueError):
+                    getattr(index, operation)(data)
+
+    def testRejectedBatchLeavesExistingItemsUnchanged(self):
+        cases = itertools.product(self.index_types, (1, 4), (0, 16, 31),
+                                  self.non_finite_values)
+        for index_type, num_threads, invalid_row, value in cases:
+            with self.subTest(index=index_type.__name__, num_threads=num_threads,
+                              invalid_row=invalid_row, value=value):
+                index = self.make_index(index_type, num_threads)
+                index.add_items([[1, 0, 0, 0], [-1, 0, 0, 0]], [10, 11])
+                data = np.tile(np.array([0, 1, 0, 0], dtype=np.float32), (32, 1))
+                # A valid first row would overwrite label 10 if added prematurely.
+                data[0] = [-1, 0, 0, 0]
+                data[invalid_row, 3] = value
+                ids = np.arange(100, 132)
+                ids[0] = 10
+                with self.assertRaises(ValueError):
+                    index.add_items(data, ids)
+                self.assertEqual(index.get_current_count(), 2)
+                labels, distances = index.knn_query([[1, 0, 0, 0]], k=2)
+                np.testing.assert_array_equal(labels, [[10, 11]])
+                np.testing.assert_allclose(distances, [[0, 2]], atol=2e-6, rtol=0)
+
+    def testRejectedInitialBatchPreservesAutomaticLabels(self):
+        cases = itertools.product(self.index_types, (1, 4), self.non_finite_values)
+        for index_type, num_threads, value in cases:
+            with self.subTest(index=index_type.__name__, num_threads=num_threads,
+                              value=value):
+                index = self.make_index(index_type, num_threads)
+                data = np.tile(np.array([1, 0, 0, 0], dtype=np.float32), (32, 1))
+                data[-1, 3] = value
+                with self.assertRaises(ValueError):
+                    index.add_items(data)
+                self.assertEqual(index.get_current_count(), 0)
+                index.add_items([[1, 0, 0, 0]])
+                labels, distances = index.knn_query([[1, 0, 0, 0]])
+                np.testing.assert_array_equal(labels, [[0]])
+                np.testing.assert_allclose(distances, [[0]], atol=2e-6, rtol=0)
+
+    def testRejectedQueryDoesNotCallFilter(self):
+        cases = itertools.product(self.index_types, (1, 4), (0, 16, 31),
+                                  self.non_finite_values)
+        for index_type, num_threads, invalid_row, value in cases:
+            with self.subTest(index=index_type.__name__, num_threads=num_threads,
+                              invalid_row=invalid_row, value=value):
+                index = self.make_index(index_type, num_threads)
+                index.add_items([[1, 0, 0, 0], [-1, 0, 0, 0]], [10, 11])
+                queries = np.tile(np.array([1, 0, 0, 0], dtype=np.float32), (32, 1))
+                queries[invalid_row, 3] = value
+                visited_labels = []
+
+                def allow_label(label):
+                    visited_labels.append(label)
+                    return True
+
+                with self.assertRaises(ValueError):
+                    index.knn_query(queries, num_threads=num_threads,
+                                    filter=allow_label)
+                self.assertEqual(visited_labels, [])
+                labels, distances = index.knn_query([[1, 0, 0, 0]], k=2,
+                                                   num_threads=1)
+                np.testing.assert_array_equal(labels, [[10, 11]])
+                np.testing.assert_allclose(distances, [[0, 2]], atol=2e-6, rtol=0)
+
+    def testFloat64OverflowRejected(self):
+        cases = itertools.product(self.index_types, ('add_items', 'knn_query'),
+                                  (-1, 1))
+        for index_type, operation, sign in cases:
+            with self.subTest(index=index_type.__name__, operation=operation,
+                              sign=sign):
+                index = self.make_index(index_type)
+                if operation == 'knn_query':
+                    index.add_items([[1, 0, 0, 0]])
+                too_large = sign * float(np.finfo(np.float32).max) * 2
+                data = np.array([[1, too_large, 0, 0]], dtype=np.float64)
+                self.assertTrue(np.isfinite(data).all())
+                with warnings.catch_warnings():
+                    warnings.filterwarnings('ignore', message='overflow encountered in cast',
+                                            category=RuntimeWarning)
+                    with self.assertRaises(ValueError):
+                        getattr(index, operation)(data)
+
+    def testFiniteStridedAndZeroInputs(self):
+        cases = itertools.product(self.index_types, (1, 2))
+        for index_type, ndim in cases:
+            with self.subTest(index=index_type.__name__, ndim=ndim):
+                index = self.make_index(index_type)
+                backing = np.zeros((3, 8), dtype=np.float32)
+                backing[0, 0] = 1
+                backing[1, 0] = -1
+                data = backing[:, ::2]
+                index.add_items(data, [0, 1, 2])
+                queries = data[0] if ndim == 1 else data
+                labels, distances = index.knn_query(queries, k=3)
+                # Literal distances: same/opposite axes are 0/2; zero is 1.
+                expected = np.array([[0, 2, 1], [2, 0, 1], [1, 1, 1]])
+                for row in range(labels.shape[0]):
+                    np.testing.assert_array_equal(np.sort(labels[row]), [0, 1, 2])
+                    np.testing.assert_allclose(distances[row], expected[row, labels[row]],
+                                               atol=2e-6, rtol=0)
+
+    def testSaveLoadAfterRejectedInsertion(self):
+        cases = itertools.product(self.index_types, self.non_finite_values)
+        for index_type, value in cases:
+            with self.subTest(index=index_type.__name__, value=value):
+                index = self.make_index(index_type)
+                index.add_items([[1, 0, 0, 0], [-1, 0, 0, 0]], [10, 11])
+                invalid = np.array([[0, 1, value, 0]], dtype=np.float32)
+                with self.assertRaises(ValueError):
+                    index.add_items(invalid, [12])
+                with tempfile.TemporaryDirectory() as directory:
+                    filename = directory + '/cosine.bin'
+                    index.save_index(filename)
+                    restored = index_type(space='cosine', dim=4)
+                    restored.load_index(filename)
+                    restored.set_num_threads(1)
+                    if index_type is hnswlib.Index:
+                        restored.set_ef(64)
+                    self.assertEqual(restored.get_current_count(), 2)
+                    with self.assertRaises(ValueError):
+                        restored.add_items(invalid, [12])
+                    restored.add_items([[0, 1, 0, 0]], [12])
+                    labels, distances = restored.knn_query([[1, 0, 0, 0]], k=3)
+                    np.testing.assert_array_equal(labels, [[10, 12, 11]])
+                    np.testing.assert_allclose(distances, [[0, 1, 2]], atol=2e-6, rtol=0)