add NearestNeighborItem with validation

2 files changed, 48 insertions, 1 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/yql/VespaSerializer.java b/container-search/src/main/java/com/yahoo/search/yql/VespaSerializer.java
index 5565c680efb..6ab07d95c76 100644
--- a/container-search/src/main/java/com/yahoo/search/yql/VespaSerializer.java
+++ b/container-search/src/main/java/com/yahoo/search/yql/VespaSerializer.java
@@ -20,6 +20,7 @@ import static com.yahoo.search.yql.YqlParser.HIT_LIMIT;
 import static com.yahoo.search.yql.YqlParser.IMPLICIT_TRANSFORMS;
 import static com.yahoo.search.yql.YqlParser.LABEL;
 import static com.yahoo.search.yql.YqlParser.NEAR;
+import static com.yahoo.search.yql.YqlParser.NEAREST_NEIGHBOR;
 import static com.yahoo.search.yql.YqlParser.NORMALIZE_CASE;
 import static com.yahoo.search.yql.YqlParser.ONEAR;
 import static com.yahoo.search.yql.YqlParser.ORIGIN;
@@ -73,6 +74,7 @@ import com.yahoo.prelude.query.IntItem;
 import com.yahoo.prelude.query.Item;
 import com.yahoo.prelude.query.MarkerWordItem;
 import com.yahoo.prelude.query.NearItem;
+import com.yahoo.prelude.query.NearestNeighborItem;
 import com.yahoo.prelude.query.NotItem;
 import com.yahoo.prelude.query.NullItem;
 import com.yahoo.prelude.query.ONearItem;
@@ -687,6 +689,28 @@ public class VespaSerializer {
 
     }
 
+    private static class NearestNeighborSerializer extends Serializer<NearestNeighborItem> {
+
+        @Override
+        void onExit(StringBuilder destination, NearestNeighborItem item) { }
+
+        @Override
+        boolean serialize(StringBuilder destination, NearestNeighborItem item) {
+            destination.append("[{");
+            int initLen = destination.length();
+            destination.append(leafAnnotations(item));
+            comma(destination, initLen);
+            int targetNumHits = item.getTargetNumHits();
+            destination.append("\"targetNumHits\": ").append(targetNumHits);
+            destination.append("}]");
+            destination.append(NEAREST_NEIGHBOR).append('(');
+            destination.append(item.getIndexName()).append(", ");
+            destination.append(item.getQueryRankFeatureName()).append(')');
+            return false;
+        }
+
+    }
+
     private static class PredicateQuerySerializer extends Serializer<PredicateQueryItem> {
 
         @Override
@@ -1131,6 +1155,7 @@ public class VespaSerializer {
         dispatchBuilder.put(BoolItem.class, new BoolSerializer());
         dispatchBuilder.put(MarkerWordItem.class, new WordSerializer()); // gotcha
         dispatchBuilder.put(NearItem.class, new NearSerializer());
+        dispatchBuilder.put(NearestNeighborItem.class, new NearestNeighborSerializer());
         dispatchBuilder.put(NotItem.class, new NotSerializer());
         dispatchBuilder.put(NullItem.class, new NullSerializer());
         dispatchBuilder.put(ONearItem.class, new ONearSerializer());
diff --git a/container-search/src/main/java/com/yahoo/search/yql/YqlParser.java b/container-search/src/main/java/com/yahoo/search/yql/YqlParser.java
index 306ac4e42f2..8d013e501e8 100644
--- a/container-search/src/main/java/com/yahoo/search/yql/YqlParser.java
+++ b/container-search/src/main/java/com/yahoo/search/yql/YqlParser.java
@@ -35,6 +35,7 @@ import com.yahoo.prelude.query.IntItem;
 import com.yahoo.prelude.query.Item;
 import com.yahoo.prelude.query.Limit;
 import com.yahoo.prelude.query.NearItem;
+import com.yahoo.prelude.query.NearestNeighborItem;
 import com.yahoo.prelude.query.NotItem;
 import com.yahoo.prelude.query.NullItem;
 import com.yahoo.prelude.query.ONearItem;
@@ -151,6 +152,7 @@ public class YqlParser implements Parser {
     static final String IMPLICIT_TRANSFORMS = "implicitTransforms";
     static final String LABEL = "label";
     static final String NEAR = "near";
+    static final String NEAREST_NEIGHBOR = "nearestNeighbor";
     static final String NORMALIZE_CASE = "normalizeCase";
     static final String ONEAR = "onear";
     static final String ORIGIN_LENGTH = "length";
@@ -367,6 +369,8 @@ public class YqlParser implements Parser {
                 return buildWeightedSet(ast);
             case DOT_PRODUCT:
                 return buildDotProduct(ast);
+            case NEAREST_NEIGHBOR:
+                return buildNearestNeighbor(ast);
             case PREDICATE:
                 return buildPredicate(ast);
             case RANK:
@@ -378,7 +382,7 @@ public class YqlParser implements Parser {
             case NON_EMPTY:
                 return ensureNonEmpty(ast);
             default:
-                throw newUnexpectedArgumentException(names.get(0), DOT_PRODUCT,
+                throw newUnexpectedArgumentException(names.get(0), DOT_PRODUCT, NEAREST_NEIGHBOR,
                                                      RANGE, RANK, USER_QUERY, WAND, WEAK_AND, WEIGHTED_SET,
                                                      PREDICATE, USER_INPUT, NON_EMPTY);
         }
@@ -406,6 +410,24 @@ public class YqlParser implements Parser {
         return fillWeightedSet(ast, args.get(1), new DotProductItem(getIndex(args.get(0))));
     }
 
+    private Item buildNearestNeighbor(OperatorNode<ExpressionOperator> ast) {
+        List<OperatorNode<ExpressionOperator>> args = ast.getArgument(1);
+        Preconditions.checkArgument(args.size() == 2, "Expected 2 arguments, got %s.", args.size());
+        String field = fetchFieldRead(args.get(0));
+        String property = fetchFieldRead(args.get(1));
+        NearestNeighborItem item = new NearestNeighborItem(field, property);
+        Integer targetNumHits = getAnnotation(ast, TARGET_NUM_HITS,
+                Integer.class, null, "desired minimum hits to produce");
+        if (targetNumHits != null) {
+            item.setTargetNumHits(targetNumHits);
+        }
+        String label = getAnnotation(ast, LABEL, String.class, null, "item label");
+        if (label != null) {
+                item.setLabel(label);
+        }
+        return item;
+    }
+
     private Item buildPredicate(OperatorNode<ExpressionOperator> ast) {
         List<OperatorNode<ExpressionOperator>> args = ast.getArgument(1);
         Preconditions.checkArgument(args.size() == 3, "Expected 3 arguments, got %s.", args.size());