Propagating annotations for fuzzy query

4 files changed, 106 insertions, 19 deletions

diff --git a/container-search/src/main/java/com/yahoo/prelude/query/FuzzyItem.java b/container-search/src/main/java/com/yahoo/prelude/query/FuzzyItem.java
index 74c31a5d1f0..fda96aa6ecc 100644
--- a/container-search/src/main/java/com/yahoo/prelude/query/FuzzyItem.java
+++ b/container-search/src/main/java/com/yahoo/prelude/query/FuzzyItem.java
@@ -1,7 +1,10 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.prelude.query;
 
+import com.yahoo.compress.IntegerCompressor;
+
 import java.nio.ByteBuffer;
+import java.util.Objects;
 
 /**
  * Fuzzy search term
@@ -11,9 +14,37 @@ import java.nio.ByteBuffer;
 public class FuzzyItem extends TermItem {
     private String term;
 
-    public FuzzyItem(String indexName, boolean isFromQuery, String term) {
+    private int maxEditDistance;
+    private int prefixLength;
+
+    public static int DefaultMaxEditDistance = 2;
+    public static int DefaultPrefixLength = 0;
+
+    public FuzzyItem(String indexName, boolean isFromQuery, String term, int maxEditDistance, int prefixLength) {
         super(indexName, isFromQuery, null);
         setValue(term);
+        setMaxEditDistance(maxEditDistance);
+        setPrefixLength(prefixLength);
+    }
+
+    public void setMaxEditDistance(int maxEditDistance) {
+        if (maxEditDistance < 0)
+            throw new IllegalArgumentException("Can not use negative maxEditDistance " + maxEditDistance);
+        this.maxEditDistance = maxEditDistance;
+    }
+
+    public void setPrefixLength(int prefixLength) {
+        if (prefixLength < 0)
+            throw new IllegalArgumentException("Can not use negative prefixLength " + prefixLength);
+        this.prefixLength = prefixLength;
+    }
+
+    public int getPrefixLength() {
+        return this.prefixLength;
+    }
+
+    public int getMaxEditDistance() {
+        return this.maxEditDistance;
     }
 
     @Override
@@ -73,35 +104,36 @@ public class FuzzyItem extends TermItem {
             return false;
         }
         FuzzyItem other = (FuzzyItem) obj;
-        if (term == null) {
-            if (other.term != null) {
-                return false;
-            }
-        } else if (!term.equals(other.term)) {
-            return false;
-        }
+        if (!this.term.equals(other.term)) return false;
+        if (this.maxEditDistance != other.maxEditDistance) return false;
+        if (this.prefixLength != other.prefixLength) return false;
         return true;
     }
 
     @Override
     public int hashCode() {
-        final int prime = 31;
-        int result = super.hashCode();
-        result = prime * result + ((term == null) ? 0 : term.hashCode());
-        return result;
+        return Objects.hash(super.hashCode(), term, maxEditDistance, prefixLength);
     }
 
     @Override
-    public String toString() {
-        StringBuilder builder = new StringBuilder();
-        builder.append("FuzzyItem [term=").append(term).append("]");
-        return builder.toString();
+    protected void appendHeadingString(StringBuilder buffer) {
+        buffer.append(getName());
+        buffer.append("(");
+        buffer.append(this.term);
+        buffer.append(",");
+        buffer.append(this.maxEditDistance);
+        buffer.append(",");
+        buffer.append(this.prefixLength);
+        buffer.append(")");
+        buffer.append(" ");
     }
 
     @Override
     protected void encodeThis(ByteBuffer buffer) {
         super.encodeThis(buffer);
         putString(getIndexedString(), buffer);
+        IntegerCompressor.putCompressedPositiveNumber(this.maxEditDistance, buffer);
+        IntegerCompressor.putCompressedPositiveNumber(this.prefixLength, buffer);
     }
 }
 
diff --git a/container-search/src/main/java/com/yahoo/search/query/SelectParser.java b/container-search/src/main/java/com/yahoo/search/query/SelectParser.java
index 320148ec01d..c1d1cdd566b 100644
--- a/container-search/src/main/java/com/yahoo/search/query/SelectParser.java
+++ b/container-search/src/main/java/com/yahoo/search/query/SelectParser.java
@@ -61,6 +61,8 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
+import static com.yahoo.search.yql.YqlParser.MAX_EDIT_DISTANCE;
+import static com.yahoo.search.yql.YqlParser.PREFIX_LENGTH;
 import static com.yahoo.slime.Type.ARRAY;
 import static com.yahoo.slime.Type.DOUBLE;
 import static com.yahoo.slime.Type.LONG;
@@ -1161,9 +1163,16 @@ public class SelectParser implements Parser {
 
     private Item instantiateFuzzyItem(String field, String key, Inspector value) {
         HashMap<Integer, Inspector> children = childMap(value);
+        HashMap<String, Inspector> annotations = getAnnotationMap(value);
+
         Preconditions.checkArgument(children.size() == 1, "Expected 1 argument, got %s.", children.size());
+
         String wordData = children.get(0).asString();
-        FuzzyItem fuzzy = new FuzzyItem(field, true, wordData);
+
+        Integer maxEditDistance = getIntegerAnnotation(MAX_EDIT_DISTANCE, annotations, FuzzyItem.DefaultMaxEditDistance);
+        Integer prefixLength = getIntegerAnnotation(PREFIX_LENGTH, annotations, FuzzyItem.DefaultPrefixLength);
+
+        FuzzyItem fuzzy = new FuzzyItem(field, true, wordData, maxEditDistance, prefixLength);
 
         return leafStyleSettings(getAnnotations(value), fuzzy);
     }
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 e778798b0e5..4b511df5e5f 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
@@ -21,6 +21,7 @@ import static com.yahoo.search.yql.YqlParser.GEO_LOCATION;
 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.MAX_EDIT_DISTANCE;
 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;
@@ -31,6 +32,7 @@ import static com.yahoo.search.yql.YqlParser.ORIGIN_OFFSET;
 import static com.yahoo.search.yql.YqlParser.ORIGIN_ORIGINAL;
 import static com.yahoo.search.yql.YqlParser.PHRASE;
 import static com.yahoo.search.yql.YqlParser.PREFIX;
+import static com.yahoo.search.yql.YqlParser.PREFIX_LENGTH;
 import static com.yahoo.search.yql.YqlParser.RANGE;
 import static com.yahoo.search.yql.YqlParser.RANK;
 import static com.yahoo.search.yql.YqlParser.RANKED;
@@ -526,7 +528,8 @@ public class VespaSerializer {
 
         @Override
         boolean serialize(StringBuilder destination, FuzzyItem fuzzy) {
-            String annotations = leafAnnotations(fuzzy);
+            String annotations = fuzzyAnnotations(fuzzy);
+
             destination.append(normalizeIndexName(fuzzy.getIndexName())).append(" contains ");
 
             if (annotations.length() > 0) {
@@ -543,6 +546,30 @@ public class VespaSerializer {
             }
             return false;
         }
+
+        static String fuzzyAnnotations(FuzzyItem fuzzyItem) {
+            boolean isMaxEditDistanceSet = fuzzyItem.getMaxEditDistance() != FuzzyItem.DefaultMaxEditDistance;
+            boolean isPrefixLengthSet = fuzzyItem.getPrefixLength() != FuzzyItem.DefaultPrefixLength;
+            boolean anyAnnotationSet = isMaxEditDistanceSet || isPrefixLengthSet;
+
+            StringBuilder builder = new StringBuilder();
+            if (anyAnnotationSet) {
+                builder.append("[{");
+            }
+            if (isMaxEditDistanceSet) {
+                builder.append(MAX_EDIT_DISTANCE + ":").append(fuzzyItem.getMaxEditDistance());
+            }
+            if (isMaxEditDistanceSet && isPrefixLengthSet) {
+                builder.append(",");
+            }
+            if (isPrefixLengthSet) {
+                builder.append(PREFIX_LENGTH + ":").append(fuzzyItem.getPrefixLength());
+            }
+            if (anyAnnotationSet) {
+                builder.append("}]");
+            }
+            return builder.toString();
+        }
     }
 
     private static class ONearSerializer extends Serializer<ONearItem> {
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 06ee0e706f3..fcb19dde10d 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
@@ -126,6 +126,8 @@ public class YqlParser implements Parser {
     private static final String RANKED_DESCRIPTION = "setting for whether to use term for ranking";
     private static final String STEM_DESCRIPTION = "setting for whether to use stem if field implies it";
     private static final String USE_POSITION_DATA_DESCRIPTION = "setting for whether to use position data for ranking this item";
+    private static final String MAX_EDIT_DISTANCE_DESCRIPTION = "setting for an inclusive upper bound for a fuzzy edit-distance search";
+    private static final String PREFIX_LENGTH_DESCRIPTION = "setting for a prefix length that is considered frozen for a fuzzy search";
     private static final String USER_INPUT_ALLOW_EMPTY = "allowEmpty";
     private static final String USER_INPUT_DEFAULT_INDEX = "defaultIndex";
     private static final String USER_INPUT_GRAMMAR = "grammar";
@@ -194,6 +196,9 @@ public class YqlParser implements Parser {
     public static final String WEIGHT = "weight";
     public static final String WEIGHTED_SET = "weightedSet";
     public static final String FUZZY = "fuzzy";
+    public static final String MAX_EDIT_DISTANCE = "maxEditDistance";
+    public static final String PREFIX_LENGTH = "prefixLength";
+
 
     private final IndexFacts indexFacts;
     private final List<ConnectedItem> connectedItems = new ArrayList<>();
@@ -1313,7 +1318,21 @@ public class YqlParser implements Parser {
 
         String wordData = getStringContents(args.get(0));
 
-        FuzzyItem fuzzy = new FuzzyItem(field, true, wordData);
+        Integer maxEditDistance = getAnnotation(
+                ast,
+                MAX_EDIT_DISTANCE,
+                Integer.class,
+                FuzzyItem.DefaultMaxEditDistance,
+                MAX_EDIT_DISTANCE_DESCRIPTION);
+
+        Integer prefixLength = getAnnotation(
+                ast,
+                PREFIX_LENGTH,
+                Integer.class,
+                FuzzyItem.DefaultPrefixLength,
+                PREFIX_LENGTH_DESCRIPTION);
+
+        FuzzyItem fuzzy = new FuzzyItem(field, true, wordData, maxEditDistance, prefixLength);
         return leafStyleSettings(ast, fuzzy);
     }