Merge pull request #12281 from vespa-engine/arnej/add-nns-iterator

Arnej/add nns iterator

2 files changed, 71 insertions, 2 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 7bc582ab442..691e80aeb9f 100644
--- a/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
+++ b/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
@@ -12,6 +12,7 @@
 #include <vespa/searchlib/queryeval/simpleresult.h>
 #include <vespa/searchlib/tensor/dense_tensor_attribute.h>
 #include <vespa/vespalib/test/insertion_operators.h>
+#include <vespa/searchlib/queryeval/nns_index_iterator.h>
 
 #include <vespa/log/log.h>
 LOG_SETUP("nearest_neighbor_test");
@@ -190,4 +191,70 @@ TEST("require that NearestNeighborIterator sets expected rawscore") {
     TEST_DO(verify_iterator_sets_expected_rawscore(denseSpecFloat, denseSpecDouble));
 }
 
+TEST("require that NnsIndexIterator works as expected") {
+    std::vector<NnsIndexIterator::Hit> hits{{2,4.0}, {3,9.0}, {5,1.0}, {8,16.0}, {9,36.0}};
+    auto md = MatchData::makeTestInstance(2, 2);
+    auto &tfmd = *(md->resolveTermField(0));
+    auto search = NnsIndexIterator::create(true, tfmd, hits);
+    uint32_t docid = 1;
+    search->initFullRange();
+    bool match = search->seek(docid);
+    EXPECT_FALSE(match);
+    EXPECT_FALSE(search->isAtEnd());
+    EXPECT_EQUAL(2u, search->getDocId());
+    docid = 2;
+    match = search->seek(docid);
+    EXPECT_TRUE(match);
+    EXPECT_FALSE(search->isAtEnd());
+    EXPECT_EQUAL(docid, search->getDocId());
+    search->unpack(docid);
+    EXPECT_EQUAL(2.0, tfmd.getRawScore());
+
+    docid = 3;
+    match = search->seek(docid);
+    EXPECT_TRUE(match);
+    EXPECT_FALSE(search->isAtEnd());
+    EXPECT_EQUAL(docid, search->getDocId());
+    search->unpack(docid);
+    EXPECT_EQUAL(3.0, tfmd.getRawScore());
+
+    docid = 4;
+    match = search->seek(docid);
+    EXPECT_FALSE(match);
+    EXPECT_FALSE(search->isAtEnd());
+    EXPECT_EQUAL(5u, search->getDocId());
+
+    docid = 6;
+    match = search->seek(docid);
+    EXPECT_FALSE(match);
+    EXPECT_FALSE(search->isAtEnd());
+    EXPECT_EQUAL(8u, search->getDocId());
+    docid = 8;
+    search->unpack(docid);
+    EXPECT_EQUAL(4.0, tfmd.getRawScore());
+    docid = 9;
+    match = search->seek(docid);
+    EXPECT_TRUE(match);
+    EXPECT_FALSE(search->isAtEnd());
+    docid = 10;
+    match = search->seek(docid);
+    EXPECT_FALSE(match);
+    EXPECT_TRUE(search->isAtEnd());
+
+    docid = 4;
+    search->initRange(docid, 7);
+    match = search->seek(docid);
+    EXPECT_FALSE(match);
+    EXPECT_FALSE(search->isAtEnd());
+    EXPECT_EQUAL(5u, search->getDocId());
+    docid = 5;
+    search->unpack(docid);
+    EXPECT_EQUAL(1.0, tfmd.getRawScore());
+    EXPECT_FALSE(search->isAtEnd());
+    docid = 6;
+    match = search->seek(docid);
+    EXPECT_FALSE(match);
+    EXPECT_TRUE(search->isAtEnd());
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
index cd0d4bcaad0..1204ae1e9bc 100644
--- a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
+++ b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
@@ -102,8 +102,10 @@ public:
         if (exp_hits.size() == k) {
             std::vector<uint32_t> expected_by_docid = exp_hits;
             std::sort(expected_by_docid.begin(), expected_by_docid.end());
-            std::vector<uint32_t> got_by_docid = index->find_top_k(k, qv, k);
-            EXPECT_EQ(expected_by_docid, got_by_docid);
+            auto got_by_docid = index->find_top_k(k, qv, k);
+            for (idx = 0; idx < k; ++idx) {
+                EXPECT_EQ(expected_by_docid[idx], got_by_docid[idx].docid);
+            }
         }
     }
 };