Adding builder method to SmoothingModel implementations

Adds a method that emits a WordScorerFactory to all of the
three SmoothingModel implementatins that will be needed when
we switch to parsing the PhraseSuggestion on the coordinating
node and need to delay creating the WordScorer on the shards.

Author: cbuescher

core/src/main/java/org/elasticsearch/search/suggest/phrase/LaplaceScorer.java

@@ -27,22 +27,26 @@
 import java.io.IOException;
 //TODO public for tests
 public final class LaplaceScorer extends WordScorer {
-    
+
     public static final WordScorerFactory FACTORY = new WordScorer.WordScorerFactory() {
         @Override
         public WordScorer newScorer(IndexReader reader, Terms terms, String field, double realWordLikelyhood, BytesRef separator) throws IOException {
             return new LaplaceScorer(reader, terms, field, realWordLikelyhood, separator, 0.5);
         }
     };
-    
+
     private double alpha;
 
     public LaplaceScorer(IndexReader reader, Terms terms, String field,
             double realWordLikelyhood, BytesRef separator, double alpha) throws IOException {
         super(reader, terms, field, realWordLikelyhood, separator);
         this.alpha = alpha;
     }
-    
+
+    double alpha() {
+        return this.alpha;
+    }
+
     @Override
     protected double scoreBigram(Candidate word, Candidate w_1) throws IOException {
         SuggestUtils.join(separator, spare, w_1.term, word.term);

core/src/main/java/org/elasticsearch/search/suggest/phrase/LinearInterpoatingScorer.java

@@ -41,7 +41,19 @@ public LinearInterpoatingScorer(IndexReader reader, Terms terms, String field,
         this.bigramLambda = bigramLambda / sum;
         this.trigramLambda = trigramLambda / sum;
     }
-    
+
+    double trigramLambda() {
+        return this.trigramLambda;
+    }
+
+    double bigramLambda() {
+        return this.bigramLambda;
+    }
+
+    double unigramLambda() {
+        return this.unigramLambda;
+    }
+
     @Override
     protected double scoreBigram(Candidate word, Candidate w_1) throws IOException {
         SuggestUtils.join(separator, spare, w_1.term, word.term);

core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionBuilder.java

@@ -18,8 +18,12 @@
  */
 package org.elasticsearch.search.suggest.phrase;
 
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.Terms;
+import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.ParseFieldMatcher;
+import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -30,6 +34,7 @@
 import org.elasticsearch.index.query.QueryParseContext;
 import org.elasticsearch.script.Template;
 import org.elasticsearch.search.suggest.SuggestBuilder.SuggestionBuilder;
+import org.elasticsearch.search.suggest.phrase.WordScorer.WordScorerFactory;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -50,7 +55,7 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
     private Float confidence;
     private final Map<String, List<CandidateGenerator>> generators = new HashMap<>();
     private Integer gramSize;
-    private SmoothingModel<?> model;
+    private SmoothingModel model;
     private Boolean forceUnigrams;
     private Integer tokenLimit;
     private String preTag;
@@ -159,7 +164,7 @@ public PhraseSuggestionBuilder forceUnigrams(boolean forceUnigrams) {
      * Sets an explicit smoothing model used for this suggester. The default is
      * {@link PhraseSuggestionBuilder.StupidBackoff}.
      */
-    public PhraseSuggestionBuilder smoothingModel(SmoothingModel<?> model) {
+    public PhraseSuggestionBuilder smoothingModel(SmoothingModel model) {
         this.model = model;
         return this;
     }
@@ -292,7 +297,7 @@ public static DirectCandidateGenerator candidateGenerator(String field) {
      * Smoothing</a> for details.
      * </p>
      */
-    public static final class StupidBackoff extends SmoothingModel<StupidBackoff> {
+    public static final class StupidBackoff extends SmoothingModel {
         /**
          * Default discount parameter for {@link StupidBackoff} smoothing
          */
@@ -341,8 +346,9 @@ public StupidBackoff readFrom(StreamInput in) throws IOException {
         }
 
         @Override
-        protected boolean doEquals(StupidBackoff other) {
-            return Objects.equals(discount, other.discount);
+        protected boolean doEquals(SmoothingModel other) {
+            StupidBackoff otherModel = (StupidBackoff) other;
+            return Objects.equals(discount, otherModel.discount);
         }
 
         @Override
@@ -351,7 +357,7 @@ public final int hashCode() {
         }
 
         @Override
-        public StupidBackoff fromXContent(QueryParseContext parseContext) throws IOException {
+        public SmoothingModel fromXContent(QueryParseContext parseContext) throws IOException {
             XContentParser parser = parseContext.parser();
             XContentParser.Token token;
             String fieldName = null;
@@ -366,6 +372,12 @@ public StupidBackoff fromXContent(QueryParseContext parseContext) throws IOExcep
             }
             return new StupidBackoff(discount);
         }
+
+        @Override
+        public WordScorerFactory buildWordScorerFactory() {
+            return (IndexReader reader, Terms terms, String field, double realWordLikelyhood, BytesRef separator)
+                    -> new StupidBackoffScorer(reader, terms, field, realWordLikelyhood, separator, discount);
+        }
     }
 
     /**
@@ -377,7 +389,7 @@ public StupidBackoff fromXContent(QueryParseContext parseContext) throws IOExcep
      * Smoothing</a> for details.
      * </p>
      */
-    public static final class Laplace extends SmoothingModel<Laplace> {
+    public static final class Laplace extends SmoothingModel {
         private double alpha = DEFAULT_LAPLACE_ALPHA;
         private static final String NAME = "laplace";
         private static final ParseField ALPHA_FIELD = new ParseField("alpha");
@@ -419,13 +431,14 @@ public void writeTo(StreamOutput out) throws IOException {
         }
 
         @Override
-        public Laplace readFrom(StreamInput in) throws IOException {
+        public SmoothingModel readFrom(StreamInput in) throws IOException {
             return new Laplace(in.readDouble());
         }
 
         @Override
-        protected boolean doEquals(Laplace other) {
-            return Objects.equals(alpha, other.alpha);
+        protected boolean doEquals(SmoothingModel other) {
+            Laplace otherModel = (Laplace) other;
+            return Objects.equals(alpha, otherModel.alpha);
         }
 
         @Override
@@ -434,7 +447,7 @@ public final int hashCode() {
         }
 
         @Override
-        public Laplace fromXContent(QueryParseContext parseContext) throws IOException {
+        public SmoothingModel fromXContent(QueryParseContext parseContext) throws IOException {
             XContentParser parser = parseContext.parser();
             XContentParser.Token token;
             String fieldName = null;
@@ -449,10 +462,16 @@ public Laplace fromXContent(QueryParseContext parseContext) throws IOException {
             }
             return new Laplace(alpha);
         }
+
+        @Override
+        public WordScorerFactory buildWordScorerFactory() {
+            return (IndexReader reader, Terms terms, String field, double realWordLikelyhood, BytesRef separator)
+                    -> new LaplaceScorer(reader, terms,  field, realWordLikelyhood, separator, alpha);
+        }
     }
 
 
-    public static abstract class SmoothingModel<SM extends SmoothingModel<?>> implements NamedWriteable<SM>, ToXContent {
+    public static abstract class SmoothingModel implements NamedWriteable<SmoothingModel>, ToXContent {
 
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
@@ -471,16 +490,18 @@ public final boolean equals(Object obj) {
                 return false;
             }
             @SuppressWarnings("unchecked")
-            SM other = (SM) obj;
+            SmoothingModel other = (SmoothingModel) obj;
             return doEquals(other);
         }
 
-        public abstract SM fromXContent(QueryParseContext parseContext) throws IOException;
+        public abstract SmoothingModel fromXContent(QueryParseContext parseContext) throws IOException;
+
+        public abstract WordScorerFactory buildWordScorerFactory();
 
         /**
          * subtype specific implementation of "equals".
          */
-        protected abstract boolean doEquals(SM other);
+        protected abstract boolean doEquals(SmoothingModel other);
 
         protected abstract XContentBuilder innerToXContent(XContentBuilder builder, Params params) throws IOException;
     }
@@ -493,7 +514,7 @@ public final boolean equals(Object obj) {
      * Smoothing</a> for details.
      * </p>
      */
-    public static final class LinearInterpolation extends SmoothingModel<LinearInterpolation> {
+    public static final class LinearInterpolation extends SmoothingModel {
         private static final String NAME = "linear";
         static final LinearInterpolation PROTOTYPE = new LinearInterpolation(0.8, 0.1, 0.1);
         private final double trigramLambda;
@@ -563,10 +584,11 @@ public LinearInterpolation readFrom(StreamInput in) throws IOException {
         }
 
         @Override
-        protected boolean doEquals(LinearInterpolation other) {
-            return Objects.equals(trigramLambda, other.trigramLambda) &&
-                    Objects.equals(bigramLambda, other.bigramLambda) &&
-                    Objects.equals(unigramLambda, other.unigramLambda);
+        protected boolean doEquals(SmoothingModel other) {
+            final LinearInterpolation otherModel = (LinearInterpolation) other;
+            return Objects.equals(trigramLambda, otherModel.trigramLambda) &&
+                    Objects.equals(bigramLambda, otherModel.bigramLambda) &&
+                    Objects.equals(unigramLambda, otherModel.unigramLambda);
         }
 
         @Override
@@ -579,35 +601,45 @@ public LinearInterpolation fromXContent(QueryParseContext parseContext) throws I
             XContentParser parser = parseContext.parser();
             XContentParser.Token token;
             String fieldName = null;
-            final double[] lambdas = new double[3];
+            double trigramLambda = 0.0;
+            double bigramLambda = 0.0;
+            double unigramLambda = 0.0;
             ParseFieldMatcher matcher = parseContext.parseFieldMatcher();
             while ((token = parser.nextToken()) != Token.END_OBJECT) {
                 if (token == XContentParser.Token.FIELD_NAME) {
                     fieldName = parser.currentName();
-                }
-                if (token.isValue()) {
+                } else if (token.isValue()) {
                     if (matcher.match(fieldName, TRIGRAM_FIELD)) {
-                        lambdas[0] = parser.doubleValue();
-                        if (lambdas[0] < 0) {
+                        trigramLambda = parser.doubleValue();
+                        if (trigramLambda < 0) {
                             throw new IllegalArgumentException("trigram_lambda must be positive");
                         }
                     } else if (matcher.match(fieldName, BIGRAM_FIELD)) {
-                        lambdas[1] = parser.doubleValue();
-                        if (lambdas[1] < 0) {
+                        bigramLambda = parser.doubleValue();
+                        if (bigramLambda < 0) {
                             throw new IllegalArgumentException("bigram_lambda must be positive");
                         }
                     } else if (matcher.match(fieldName, UNIGRAM_FIELD)) {
-                        lambdas[2] = parser.doubleValue();
-                        if (lambdas[2] < 0) {
+                        unigramLambda = parser.doubleValue();
+                        if (unigramLambda < 0) {
                             throw new IllegalArgumentException("unigram_lambda must be positive");
                         }
                     } else {
                         throw new IllegalArgumentException(
                                 "suggester[phrase][smoothing][linear] doesn't support field [" + fieldName + "]");
                     }
+                } else {
+                    throw new ParsingException(parser.getTokenLocation(), "[" + NAME + "] unknown token [" + token + "] after [" + fieldName + "]");
                 }
             }
-            return new LinearInterpolation(lambdas[0], lambdas[1], lambdas[2]);
+            return new LinearInterpolation(trigramLambda, bigramLambda, unigramLambda);
+        }
+
+        @Override
+        public WordScorerFactory buildWordScorerFactory() {
+            return (IndexReader reader, Terms terms, String field, double realWordLikelyhood, BytesRef separator) ->
+                        new LinearInterpoatingScorer(reader, terms, field, realWordLikelyhood, separator, trigramLambda, bigramLambda,
+                            unigramLambda);
         }
     }
 

core/src/main/java/org/elasticsearch/search/suggest/phrase/StupidBackoffScorer.java

@@ -42,6 +42,10 @@ public StupidBackoffScorer(IndexReader reader, Terms terms,String field, double
         this.discount = discount;
     }
 
+    double discount() {
+        return this.discount;
+    }
+
     @Override
     protected double scoreBigram(Candidate word, Candidate w_1) throws IOException {
         SuggestUtils.join(separator, spare, w_1.term, word.term);

core/src/test/java/org/elasticsearch/search/suggest/phrase/LaplaceModelTests.java

@@ -20,19 +20,30 @@
 package org.elasticsearch.search.suggest.phrase;
 
 import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.Laplace;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.SmoothingModel;
 
-public class LaplaceModelTests extends SmoothingModelTest<Laplace> {
+import static org.hamcrest.Matchers.instanceOf;
+
+public class LaplaceModelTests extends SmoothingModelTestCase {
 
     @Override
-    protected Laplace createTestModel() {
+    protected SmoothingModel createTestModel() {
         return new Laplace(randomDoubleBetween(0.0, 10.0, false));
     }
 
     /**
      * mutate the given model so the returned smoothing model is different
      */
     @Override
-    protected Laplace createMutation(Laplace original) {
+    protected Laplace createMutation(SmoothingModel input) {
+        Laplace original = (Laplace) input;
         return new Laplace(original.getAlpha() + 0.1);
     }
+
+    @Override
+    void assertWordScorer(WordScorer wordScorer, SmoothingModel input) {
+        Laplace model = (Laplace) input;
+        assertThat(wordScorer, instanceOf(LaplaceScorer.class));
+        assertEquals(model.getAlpha(), ((LaplaceScorer) wordScorer).alpha(), Double.MIN_VALUE);
+    }
 }

core/src/test/java/org/elasticsearch/search/suggest/phrase/LinearInterpolationModelTests.java

@@ -20,15 +20,18 @@
 package org.elasticsearch.search.suggest.phrase;
 
 import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.LinearInterpolation;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.SmoothingModel;
 
-public class LinearInterpolationModelTests extends SmoothingModelTest<LinearInterpolation> {
+import static org.hamcrest.Matchers.instanceOf;
+
+public class LinearInterpolationModelTests extends SmoothingModelTestCase {
 
     @Override
-    protected LinearInterpolation createTestModel() {
+    protected SmoothingModel createTestModel() {
         double trigramLambda = randomDoubleBetween(0.0, 10.0, false);
         double bigramLambda = randomDoubleBetween(0.0, 10.0, false);
         double unigramLambda = randomDoubleBetween(0.0, 10.0, false);
-        // normalize
+        // normalize so parameters sum to 1
         double sum = trigramLambda + bigramLambda + unigramLambda;
         return new LinearInterpolation(trigramLambda / sum, bigramLambda / sum, unigramLambda / sum);
     }
@@ -37,7 +40,8 @@ protected LinearInterpolation createTestModel() {
      * mutate the given model so the returned smoothing model is different
      */
     @Override
-    protected LinearInterpolation createMutation(LinearInterpolation original) {
+    protected LinearInterpolation createMutation(SmoothingModel input) {
+        LinearInterpolation original = (LinearInterpolation) input;
         // swap two values permute original lambda values
         switch (randomIntBetween(0, 2)) {
         case 0:
@@ -52,4 +56,14 @@ protected LinearInterpolation createMutation(LinearInterpolation original) {
             return new LinearInterpolation(original.getUnigramLambda(), original.getBigramLambda(), original.getTrigramLambda());
         }
     }
+
+    @Override
+    void assertWordScorer(WordScorer wordScorer, SmoothingModel in) {
+        LinearInterpolation testModel = (LinearInterpolation) in;
+        LinearInterpoatingScorer testScorer = (LinearInterpoatingScorer) wordScorer;
+        assertThat(wordScorer, instanceOf(LinearInterpoatingScorer.class));
+        assertEquals(testModel.getTrigramLambda(), (testScorer).trigramLambda(), 1e-15);
+        assertEquals(testModel.getBigramLambda(), (testScorer).bigramLambda(), 1e-15);
+        assertEquals(testModel.getUnigramLambda(), (testScorer).unigramLambda(), 1e-15);
+    }
 }