Add NearestNeighborQueryNode.

8 files changed, 86 insertions, 3 deletions

diff --git a/searchlib/src/vespa/searchlib/query/query_term_simple.h b/searchlib/src/vespa/searchlib/query/query_term_simple.h
index 74728ab1f2e..a79e33dba32 100644
--- a/searchlib/src/vespa/searchlib/query/query_term_simple.h
+++ b/searchlib/src/vespa/searchlib/query/query_term_simple.h
@@ -23,7 +23,8 @@ public:
         SUFFIXTERM = 4,
         REGEXP = 5,
         GEO_LOCATION = 6,
-        FUZZYTERM = 7
+        FUZZYTERM = 7,
+        NEAREST_NEIGHBOR = 8
     };
 
     template <typename N>
@@ -65,6 +66,7 @@ public:
     bool isRegex()         const { return (_type == Type::REGEXP); }
     bool isGeoLoc()        const { return (_type == Type::GEO_LOCATION); }
     bool isFuzzy()         const { return (_type == Type::FUZZYTERM); }
+    bool is_nearest_neighbor() const noexcept { return (_type == Type::NEAREST_NEIGHBOR); }
     bool empty()           const { return _term.empty(); }
     virtual void visitMembers(vespalib::ObjectVisitor &visitor) const;
     vespalib::string getClassName() const;
diff --git a/searchlib/src/vespa/searchlib/query/streaming/CMakeLists.txt b/searchlib/src/vespa/searchlib/query/streaming/CMakeLists.txt
index 27f9870dc18..c71b838fb37 100644
--- a/searchlib/src/vespa/searchlib/query/streaming/CMakeLists.txt
+++ b/searchlib/src/vespa/searchlib/query/streaming/CMakeLists.txt
@@ -1,6 +1,7 @@
 # Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 vespa_add_library(searchlib_query_streaming OBJECT
     SOURCES
+    nearest_neighbor_query_node.cpp
     query.cpp
     querynode.cpp
     querynoderesultbase.cpp
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
new file mode 100644
index 00000000000..fdc513f9617
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.cpp
@@ -0,0 +1,23 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "nearest_neighbor_query_node.h"
+
+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)
+{
+    setUniqueId(id);
+    setWeight(weight);
+}
+
+NearestNeighborQueryNode::~NearestNeighborQueryNode() = default;
+
+NearestNeighborQueryNode*
+NearestNeighborQueryNode::as_nearest_neighbor_query_node() noexcept
+{
+    return this;
+}
+
+}
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
new file mode 100644
index 00000000000..ddc84a4b6d3
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/query/streaming/nearest_neighbor_query_node.h
@@ -0,0 +1,27 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "queryterm.h"
+
+namespace search::streaming {
+
+/*
+ * Nearest neighbor query node.
+ */
+class NearestNeighborQueryNode: public QueryTerm
+{
+    double                _distance_threshold;
+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;
+    NearestNeighborQueryNode & operator = (const NearestNeighborQueryNode &) = delete;
+    NearestNeighborQueryNode(NearestNeighborQueryNode &&) = delete;
+    NearestNeighborQueryNode & operator = (NearestNeighborQueryNode &&) = delete;
+    ~NearestNeighborQueryNode() 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; }
+};
+
+}
diff --git a/searchlib/src/vespa/searchlib/query/streaming/querynode.cpp b/searchlib/src/vespa/searchlib/query/streaming/querynode.cpp
index 6d59886a4f5..226cb92c894 100644
--- a/searchlib/src/vespa/searchlib/query/streaming/querynode.cpp
+++ b/searchlib/src/vespa/searchlib/query/streaming/querynode.cpp
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "query.h"
+#include "nearest_neighbor_query_node.h"
 #include <vespa/searchlib/parsequery/stackdumpiterator.h>
 #include <charconv>
 #include <vespa/log/log.h>
@@ -77,6 +78,9 @@ QueryNode::Build(const QueryNode * parent, const QueryNodeResultFactory & factor
                                          queryRep.getIndexName(),
                                          QueryTerm::Type::GEO_LOCATION);
         break;
+    case ParseItem::ITEM_NEAREST_NEIGHBOR:
+        qn = build_nearest_neighbor_query_node(factory, queryRep);
+        break;
     case ParseItem::ITEM_NUMTERM:
     case ParseItem::ITEM_TERM:
     case ParseItem::ITEM_PREFIXTERM:
@@ -191,4 +195,20 @@ const HitList & QueryNode::evaluateHits(HitList & hl) const
     return hl;
 }
 
+std::unique_ptr<QueryNode>
+QueryNode::build_nearest_neighbor_query_node(const QueryNodeResultFactory& factory, SimpleQueryStackDumpIterator& query_rep)
+{
+    vespalib::stringref query_tensor_name = query_rep.getTerm();
+    vespalib::stringref field_name = query_rep.getIndexName();
+    int32_t id = query_rep.getUniqueId();
+    search::query::Weight weight = query_rep.GetWeight();
+    double distance_threshold = query_rep.getDistanceThreshold();
+    return std::make_unique<NearestNeighborQueryNode>(factory.create(),
+                                                      query_tensor_name,
+                                                      field_name,
+                                                      id,
+                                                      weight,
+                                                      distance_threshold);
+}
+
 }
diff --git a/searchlib/src/vespa/searchlib/query/streaming/querynode.h b/searchlib/src/vespa/searchlib/query/streaming/querynode.h
index 574a3c16ca3..c3fa2b63f69 100644
--- a/searchlib/src/vespa/searchlib/query/streaming/querynode.h
+++ b/searchlib/src/vespa/searchlib/query/streaming/querynode.h
@@ -28,6 +28,7 @@ using ConstQueryTermList = std::vector<const QueryTerm *>;
 */
 class QueryNode
 {
+    static std::unique_ptr<QueryNode> build_nearest_neighbor_query_node(const QueryNodeResultFactory& factory, SimpleQueryStackDumpIterator& queryRep);
  public:
   using UP = std::unique_ptr<QueryNode>;
 
@@ -54,7 +55,7 @@ class QueryNode
   virtual size_t depth() const { return 1; }
   /// Return the width of this tree.
   virtual size_t width() const { return 1; }
-  static UP Build(const QueryNode * parent, const QueryNodeResultFactory & org, SimpleQueryStackDumpIterator & queryRep, bool allowRewrite);
+  static UP Build(const QueryNode * parent, const QueryNodeResultFactory& factory, SimpleQueryStackDumpIterator & queryRep, bool allowRewrite);
 };
 
 /// A list conating the QuerNode objects. With copy/assignment.
diff --git a/searchlib/src/vespa/searchlib/query/streaming/queryterm.cpp b/searchlib/src/vespa/searchlib/query/streaming/queryterm.cpp
index 83f4410a520..11557bf1dcc 100644
--- a/searchlib/src/vespa/searchlib/query/streaming/queryterm.cpp
+++ b/searchlib/src/vespa/searchlib/query/streaming/queryterm.cpp
@@ -92,4 +92,10 @@ void QueryTerm::add(unsigned pos, unsigned context, uint32_t elemId, int32_t wei
     _hitList.emplace_back(pos, context, elemId, weight_);
 }
 
+NearestNeighborQueryNode*
+QueryTerm::as_nearest_neighbor_query_node() noexcept
+{
+    return nullptr;
+}
+
 }
diff --git a/searchlib/src/vespa/searchlib/query/streaming/queryterm.h b/searchlib/src/vespa/searchlib/query/streaming/queryterm.h
index dd9f56b11e1..51987225692 100644
--- a/searchlib/src/vespa/searchlib/query/streaming/queryterm.h
+++ b/searchlib/src/vespa/searchlib/query/streaming/queryterm.h
@@ -12,6 +12,8 @@
 
 namespace search::streaming {
 
+class NearestNeighborQueryNode;
+
 /**
    This is a leaf in the Query tree. All terms are leafs.
    A QueryTerm has the index for where to find the term. The term is a string,
@@ -57,7 +59,7 @@ public:
     QueryTerm & operator = (const QueryTerm &) = delete;
     QueryTerm(QueryTerm &&) = delete;
     QueryTerm & operator = (QueryTerm &&) = delete;
-    ~QueryTerm();
+    ~QueryTerm() override;
     bool evaluate() const override;
     const HitList & evaluateHits(HitList & hl) const override;
     void reset() override;
@@ -87,6 +89,7 @@ public:
     const string & getIndex() const override { return _index; }
     void setFuzzyMaxEditDistance(uint32_t fuzzyMaxEditDistance) { _fuzzyMaxEditDistance = fuzzyMaxEditDistance; }
     void setFuzzyPrefixLength(uint32_t fuzzyPrefixLength) { _fuzzyPrefixLength = fuzzyPrefixLength; }
+    virtual NearestNeighborQueryNode* as_nearest_neighbor_query_node() noexcept;
 protected:
     using QueryNodeResultBaseContainer = std::unique_ptr<QueryNodeResultBase>;
     string                       _index;