diff --git a/README.md b/README.md index f8e04571..dc557a17 100644 --- a/README.md +++ b/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`. diff --git a/python_bindings/bindings.cpp b/python_bindings/bindings.cpp index 9d2c1082..3e5d576a 100644 --- a/python_bindings/bindings.cpp +++ b/python_bindings/bindings.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -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 get_input_ids_and_check_shapes(const py::object& ids_, size_t feature_rows) { std::vector ids; if (!ids_.is_none()) { @@ -269,6 +281,9 @@ class Index { std::vector 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 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]; diff --git a/tests/python/bindings_test_spaces.py b/tests/python/bindings_test_spaces.py index 901cadcc..49b7a8cd 100644 --- a/tests/python/bindings_test_spaces.py +++ b/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)