From 43b76205b8622f07069b96bd4d061a219154c705 Mon Sep 17 00:00:00 2001
From: yehzu <yehzu2@gmail.com>
Date: Wed, 8 Sep 2021 22:29:13 +0800
Subject: feat: support nearestNeighbor operator

---
 .../main/java/ai/vespa/client/dsl/Annotation.java  |  4 ++
 .../java/ai/vespa/client/dsl/NearestNeighbor.java  | 52 ++++++++++++++++++++++
 client/src/main/java/ai/vespa/client/dsl/Q.java    | 12 +++++
 3 files changed, 68 insertions(+)
 create mode 100644 client/src/main/java/ai/vespa/client/dsl/NearestNeighbor.java

(limited to 'client/src/main')

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);
+    }
 }
-- 
cgit v1.2.3