diff --git a/faiss/impl/AuxIndexStructures.cpp b/faiss/impl/AuxIndexStructures.cpp index c95f34ba5..94dd215b4 100644 --- a/faiss/impl/AuxIndexStructures.cpp +++ b/faiss/impl/AuxIndexStructures.cpp @@ -35,6 +35,9 @@ RangeSearchResult::RangeSearchResult(idx_t nq, bool alloc_lims) : nq(nq) { /// called when lims contains the nb of elements result entries /// for each query void RangeSearchResult::do_allocation() { + // works only if all the partial results are aggregated + // simulatenously + FAISS_THROW_IF_NOT(labels == nullptr && distances == nullptr); size_t ofs = 0; for (int i = 0; i < nq; i++) { size_t n = lims[i]; diff --git a/faiss/utils/distances.cpp b/faiss/utils/distances.cpp index e39b55e80..6c1b34dcc 100644 --- a/faiss/utils/distances.cpp +++ b/faiss/utils/distances.cpp @@ -104,34 +104,25 @@ void exhaustive_inner_product_seq( size_t nx, size_t ny, ResultHandler& res) { - size_t check_period = InterruptCallback::get_period_hint(ny * d); - - check_period *= omp_get_max_threads(); - using SingleResultHandler = typename ResultHandler::SingleResultHandler; - for (size_t i0 = 0; i0 < nx; i0 += check_period) { - size_t i1 = std::min(i0 + check_period, nx); - #pragma omp parallel - { - SingleResultHandler resi(res); + { + SingleResultHandler resi(res); #pragma omp for - for (int64_t i = i0; i < i1; i++) { - const float* x_i = x + i * d; - const float* y_j = y; + for (int64_t i = 0; i < nx; i++) { + const float* x_i = x + i * d; + const float* y_j = y; - resi.begin(i); + resi.begin(i); - for (size_t j = 0; j < ny; j++) { - float ip = fvec_inner_product(x_i, y_j, d); - resi.add_result(ip, j); - y_j += d; - } - resi.end(); + for (size_t j = 0; j < ny; j++) { + float ip = fvec_inner_product(x_i, y_j, d); + resi.add_result(ip, j); + y_j += d; } + resi.end(); } - InterruptCallback::check(); } } @@ -143,32 +134,25 @@ void exhaustive_L2sqr_seq( size_t nx, size_t ny, ResultHandler& res) { - size_t check_period = InterruptCallback::get_period_hint(ny * d); - check_period *= omp_get_max_threads(); using SingleResultHandler = typename ResultHandler::SingleResultHandler; - for (size_t i0 = 0; i0 < nx; i0 += check_period) { - size_t i1 = std::min(i0 + check_period, nx); - #pragma omp parallel - { - SingleResultHandler resi(res); + { + SingleResultHandler resi(res); #pragma omp for - for (int64_t i = i0; i < i1; i++) { - const float* x_i = x + i * d; - const float* y_j = y; - resi.begin(i); - for (size_t j = 0; j < ny; j++) { - float disij = fvec_L2sqr(x_i, y_j, d); - resi.add_result(disij, j); - y_j += d; - } - resi.end(); + for (int64_t i = 0; i < nx; i++) { + const float* x_i = x + i * d; + const float* y_j = y; + resi.begin(i); + for (size_t j = 0; j < ny; j++) { + float disij = fvec_L2sqr(x_i, y_j, d); + resi.add_result(disij, j); + y_j += d; } + resi.end(); } - InterruptCallback::check(); } -}; +} /** Find the nearest neighbors for nx queries in a set of ny vectors */ template diff --git a/tests/test_index.py b/tests/test_index.py index 2ec92439a..3231c6968 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -1002,3 +1002,24 @@ class TestValidIndexParams(unittest.TestCase): self.assertEqual(D.shape[0], nq) self.assertEqual(D.shape[1], k) + + +class TestLargeRangeSearch(unittest.TestCase): + + def test_range_search(self): + # test for https://github.com/facebookresearch/faiss/issues/1889 + d = 256 + nq = 16 + nb = 1000000 + + # faiss.cvar.distance_compute_blas_threshold = 10 + faiss.omp_set_num_threads(1) + + index = faiss.IndexFlatL2(d) + xb = np.zeros((nb, d), dtype="float32") + index.add(xb) + + xq = np.zeros((nq, d), dtype="float32") + lims, D, I = index.range_search(xq, 1.0) + + assert len(D) == len(xb) * len(xq)