feat: support nearestNeighbor operator

3 files changed, 68 insertions, 0 deletions

diff --git a/client/src/main/java/ai/vespa/client/dsl/Annotation.java b/client/src/main/java/ai/vespa/client/dsl/Annotation.java
index 906f2abcca0..1949bc7d3f9 100644
--- a/client/src/main/java/ai/vespa/client/dsl/Annotation.java
+++ b/client/src/main/java/ai/vespa/client/dsl/Annotation.java
@@ -20,6 +20,10 @@ public class Annotation {
         return this;
     }
 
+    public boolean contains(String key) {
+        return annotations.containsKey(key);
+    }
+
     @Override
     public String toString() {
         return annotations == null || annotations.isEmpty()
diff --git a/client/src/main/java/ai/vespa/client/dsl/NearestNeighbor.java b/client/src/main/java/ai/vespa/client/dsl/NearestNeighbor.java
new file mode 100644
index 00000000000..6c95d2b6fd7
--- /dev/null
+++ b/client/src/main/java/ai/vespa/client/dsl/NearestNeighbor.java
@@ -0,0 +1,52 @@
+package ai.vespa.client.dsl;
+
+import java.util.stream.Collectors;
+
+public class NearestNeighbor extends QueryChain {
+
+    private Annotation annotation;
+    private String docVectorName;
+    private String queryVectorName;
+
+
+    public NearestNeighbor(String docVectorName, String queryVectorName) {
+        this.docVectorName = docVectorName;
+        this.queryVectorName = queryVectorName;
+        this.nonEmpty = true;
+    }
+
+    NearestNeighbor annotate(Annotation annotation) {
+        this.annotation = annotation;
+        return this;
+    }
+
+    @Override
+    boolean hasPositiveSearchField(String fieldName) {
+        return this.docVectorName.equals(fieldName);
+    }
+
+    @Override
+    boolean hasPositiveSearchField(String fieldName, Object value) {
+        return this.docVectorName.equals(fieldName) && queryVectorName.equals(value);
+    }
+
+    @Override
+    boolean hasNegativeSearchField(String fieldName) {
+        return false;
+    }
+
+    @Override
+    boolean hasNegativeSearchField(String fieldName, Object value) {
+        return false;
+    }
+
+    @Override
+    public String toString() {
+        boolean hasAnnotation = A.hasAnnotation(annotation);
+        if (!hasAnnotation || !annotation.contains("targetHits")) {
+            throw new IllegalArgumentException("must specify target hits in nearest neighbor query");
+        }
+        String s = Text.format("nearestNeighbor(%s, %s)", docVectorName, queryVectorName);
+        return Text.format("([%s]%s)", annotation, s);
+    }
+}
diff --git a/client/src/main/java/ai/vespa/client/dsl/Q.java b/client/src/main/java/ai/vespa/client/dsl/Q.java
index 2d957dcfb92..f15ffed1ea9 100644
--- a/client/src/main/java/ai/vespa/client/dsl/Q.java
+++ b/client/src/main/java/ai/vespa/client/dsl/Q.java
@@ -183,4 +183,16 @@ public final class Q {
     public static GeoLocation geoLocation(String field, Double longitude, Double latitude, String radius) {
         return new GeoLocation(field, longitude, latitude, radius);
     }
+
+    /**
+     * NearestNeighbor nearest neighbor
+     * https://docs.vespa.ai/en/reference/query-language-reference.html#nearestneighbor
+     *
+     * @param docVectorName the vector name defined in the vespa schema
+     * @param queryVectorName the vector name in this query
+     * @return the nearest neighbor query
+     */
+    public static NearestNeighbor nearestNeighbor(String docVectorName, String queryVectorName) {
+        return new NearestNeighbor(docVectorName, queryVectorName);
+    }
 }