Merge pull request #26773 from vespa-engine/geirst/raw-score-nearest-neighbor-node

Add raw score that also signals a match.

2 files changed, 25 insertions, 4 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 fdc513f9617..d1c37cd6dcd 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
@@ -6,7 +6,8 @@ namespace search::streaming {
 
 NearestNeighborQueryNode::NearestNeighborQueryNode(std::unique_ptr<QueryNodeResultBase> resultBase, const string& term, const string& index, int32_t id, search::query::Weight weight, double distance_threshold)
     : QueryTerm(std::move(resultBase), term, index, Type::NEAREST_NEIGHBOR),
-      _distance_threshold(distance_threshold)
+      _distance_threshold(distance_threshold),
+      _raw_score()
 {
     setUniqueId(id);
     setWeight(weight);
@@ -14,6 +15,18 @@ NearestNeighborQueryNode::NearestNeighborQueryNode(std::unique_ptr<QueryNodeResu
 
 NearestNeighborQueryNode::~NearestNeighborQueryNode() = default;
 
+bool
+NearestNeighborQueryNode::evaluate() const
+{
+    return _raw_score.has_value();
+}
+
+void
+NearestNeighborQueryNode::reset()
+{
+    _raw_score.reset();
+}
+
 NearestNeighborQueryNode*
 NearestNeighborQueryNode::as_nearest_neighbor_query_node() noexcept
 {
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 ddc84a4b6d3..0beb130c53d 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
@@ -3,15 +3,19 @@
 #pragma once
 
 #include "queryterm.h"
+#include <optional>
 
 namespace search::streaming {
 
 /*
  * Nearest neighbor query node.
  */
-class NearestNeighborQueryNode: public QueryTerm
-{
-    double                _distance_threshold;
+class NearestNeighborQueryNode: public QueryTerm {
+private:
+    double _distance_threshold;
+    // When this value is set it also indicates a match
+    std::optional<double> _raw_score;
+
 public:
     NearestNeighborQueryNode(std::unique_ptr<QueryNodeResultBase> resultBase, const string& term, const string& index, int32_t id, search::query::Weight weight, double distance_threshold);
     NearestNeighborQueryNode(const NearestNeighborQueryNode &) = delete;
@@ -19,9 +23,13 @@ public:
     NearestNeighborQueryNode(NearestNeighborQueryNode &&) = delete;
     NearestNeighborQueryNode & operator = (NearestNeighborQueryNode &&) = delete;
     ~NearestNeighborQueryNode() override;
+    bool evaluate() const override;
+    void reset() override;
     NearestNeighborQueryNode* as_nearest_neighbor_query_node() noexcept override;
     const vespalib::string& get_query_tensor_name() const { return getTermString(); }
     double get_distance_threshold() const { return _distance_threshold; }
+    void set_raw_score(double value) { _raw_score = value; }
+    const std::optional<double>& get_raw_score() const noexcept { return _raw_score; }
 };
 
 }