Add unpack_match_data member function to

search::streaming::NearestNeighborQueryNode.

2 files changed, 17 insertions, 0 deletions

diff --git a/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.cpp b/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.cpp
index 1317d1c0651..07d16c7bddc 100644
--- a/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.cpp
+++ b/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.cpp
@@ -1,6 +1,8 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "nearest_neighbor_query_node.h"
+#include <vespa/searchlib/fef/itermdata.h>
+#include <vespa/searchlib/fef/matchdata.h>
 #include <cassert>
 
 namespace search::streaming {
@@ -49,4 +51,18 @@ NearestNeighborQueryNode::get_raw_score() const
     return std::nullopt;
 }
 
+void
+NearestNeighborQueryNode::unpack_match_data(uint32_t docid, const fef::ITermData& td, fef::MatchData& match_data)
+{
+    auto raw_score = get_raw_score();
+    if (raw_score.has_value()) {
+        if (td.numFields() == 1u) {
+            auto& tfd = td.field(0u);
+            auto tmd = match_data.resolveTermField(tfd.getHandle());
+            assert(tmd != nullptr);
+            tmd->setRawScore(docid, raw_score.value());
+        }
+    }
+}
+
 }
diff --git a/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.h b/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.h
index c3eaad45031..277956e27a4 100644
--- a/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.h
+++ b/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.h
@@ -52,6 +52,7 @@ public:
     const std::optional<double>& get_distance() const { return _distance; }
     // This is used during unpacking, and also signals to the RawScoreCalculator that the entire document was a match.
     std::optional<double> get_raw_score() const;
+    void unpack_match_data(uint32_t docid, const fef::ITermData& td, fef::MatchData& match_data) override;
 };
 
 }