add distanceThreshold option for nearestNeighbor operator

1 files changed, 15 insertions, 1 deletions

diff --git a/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp b/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
index ad450a91f33..09790e7e360 100644
--- a/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
+++ b/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
@@ -121,11 +121,12 @@ struct Fixture
 };
 
 template <bool strict>
-SimpleResult find_matches(Fixture &env, const Value &qtv) {
+SimpleResult find_matches(Fixture &env, const Value &qtv, double threshold = std::numeric_limits<double>::max()) {
     auto md = MatchData::makeTestInstance(2, 2);
     auto &tfmd = *(md->resolveTermField(0));
     auto &attr = *(env._tensorAttr);
     NearestNeighborDistanceHeap dh(2);
+    dh.set_distance_threshold(env.dist_fun()->convert_threshold(threshold));
     const BitVector *filter = env._global_filter.get();
     auto search = NearestNeighborIterator::create(strict, tfmd, qtv, attr, dh, filter, env.dist_fun());
     if (strict) {
@@ -159,6 +160,19 @@ verify_iterator_returns_expected_results(const vespalib::string& attribute_tenso
     EXPECT_EQUAL(result, farExpect);
     result = find_matches<false>(fixture, *farTensor);
     EXPECT_EQUAL(result, farExpect);
+
+    SimpleResult null_thr5_exp({1,4,6});
+    result = find_matches<true>(fixture, *nullTensor, 5.0);
+    EXPECT_EQUAL(result, null_thr5_exp);
+    result = find_matches<false>(fixture, *nullTensor, 5.0);
+    EXPECT_EQUAL(result, null_thr5_exp);
+
+    SimpleResult far_thr4_exp({2,5});
+    result = find_matches<true>(fixture, *farTensor, 4.0);
+    EXPECT_EQUAL(result, far_thr4_exp);
+    result = find_matches<false>(fixture, *farTensor, 4.0);
+    EXPECT_EQUAL(result, far_thr4_exp);
+
 }
 
 TEST("require that NearestNeighborIterator returns expected results") {