Merge pull request #16309 from cbuescher/phraseSuggestion-smoothingModel

With SmoothingModels and CandidateGenerator now beeing able to be serialized via NamedWritable infrastructure, PhraseSuggestionBuilder can now properly delegate read/write and equals/hashCode to their implementations. Also adding randomization of smoothing model and candidate generators to tests.

core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java

@@ -39,6 +39,7 @@ import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder;
 import org.elasticsearch.search.rescore.RescoreBuilder;
 import org.elasticsearch.search.suggest.SuggestionBuilder;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.SmoothingModel;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 
@@ -706,6 +707,13 @@ public abstract class StreamInput extends InputStream {
         return readNamedWriteable(ScoreFunctionBuilder.class);
     }
 
+    /**
+     * Reads a {@link SmoothingModel} from the current stream
+     */
+    public SmoothingModel readPhraseSuggestionSmoothingModel() throws IOException {
+        return readNamedWriteable(SmoothingModel.class);
+    }
+
     /**
      * Reads a list of objects
      */

core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java

@@ -38,6 +38,7 @@ import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder;
 import org.elasticsearch.search.rescore.RescoreBuilder;
 import org.elasticsearch.search.suggest.SuggestionBuilder;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.SmoothingModel;
 import org.joda.time.ReadableInstant;
 
 import java.io.EOFException;
@@ -670,6 +671,13 @@ public abstract class StreamOutput extends OutputStream {
         writeNamedWriteable(scoreFunctionBuilder);
     }
 
+    /**
+     * Writes the given {@link SmoothingModel} to the stream
+     */
+    public void writePhraseSuggestionSmoothingModel(SmoothingModel smoothinModel) throws IOException {
+        writeNamedWriteable(smoothinModel);
+    }
+
     /**
      * Writes the given {@link GeoPoint} to the stream
      */

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

@@ -25,7 +25,6 @@ import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
@@ -42,7 +41,7 @@ import java.util.Set;
 import java.util.function.Consumer;
 
 public final class DirectCandidateGeneratorBuilder
-        implements Writeable<DirectCandidateGeneratorBuilder>, CandidateGenerator {
+        implements CandidateGenerator {
 
     private static final String TYPE = "direct_generator";
     static final DirectCandidateGeneratorBuilder PROTOTYPE = new DirectCandidateGeneratorBuilder("_na_");

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

@@ -27,6 +27,7 @@ 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;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
@@ -216,6 +217,13 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
         return this;
     }
 
+    /**
+     * Gets the {@link SmoothingModel}
+     */
+    public SmoothingModel smoothingModel() {
+        return this.model;
+    }
+
     public PhraseSuggestionBuilder tokenLimit(int tokenLimit) {
         this.tokenLimit = tokenLimit;
         return this;
@@ -391,8 +399,8 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
          * Default discount parameter for {@link StupidBackoff} smoothing
          */
         public static final double DEFAULT_BACKOFF_DISCOUNT = 0.4;
-        private double discount = DEFAULT_BACKOFF_DISCOUNT;
         static final StupidBackoff PROTOTYPE = new StupidBackoff(DEFAULT_BACKOFF_DISCOUNT);
+        private double discount = DEFAULT_BACKOFF_DISCOUNT;
         private static final String NAME = "stupid_backoff";
         private static final ParseField DISCOUNT_FIELD = new ParseField("discount");
 
@@ -743,7 +751,11 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
         out.writeOptionalFloat(realWordErrorLikelihood);
         out.writeOptionalFloat(confidence);
         out.writeOptionalVInt(gramSize);
-        // NORELEASE model.writeTo();
+        boolean hasModel = model != null;
+        out.writeBoolean(hasModel);
+        if (hasModel) {
+            out.writePhraseSuggestionSmoothingModel(model);
+        }
         out.writeOptionalBoolean(forceUnigrams);
         out.writeOptionalVInt(tokenLimit);
         out.writeOptionalString(preTag);
@@ -757,7 +769,15 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
         }
         out.writeMap(collateParams);
         out.writeOptionalBoolean(collatePrune);
-        // NORELEASE write Map<String, List<CandidateGenerator>> generators = new HashMap<>();
+        out.writeVInt(this.generators.size());
+        for (Entry<String, List<CandidateGenerator>> entry : this.generators.entrySet()) {
+            out.writeString(entry.getKey());
+            List<CandidateGenerator> generatorsList = entry.getValue();
+            out.writeVInt(generatorsList.size());
+            for (CandidateGenerator generator : generatorsList) {
+                generator.writeTo(out);
+            }
+        }
     }
 
     @Override
@@ -767,7 +787,9 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
         builder.realWordErrorLikelihood = in.readOptionalFloat();
         builder.confidence = in.readOptionalFloat();
         builder.gramSize = in.readOptionalVInt();
-        // NORELEASE read model
+        if (in.readBoolean()) {
+            builder.model = in.readPhraseSuggestionSmoothingModel();
+        }
         builder.forceUnigrams = in.readOptionalBoolean();
         builder.tokenLimit = in.readOptionalVInt();
         builder.preTag = in.readOptionalString();
@@ -778,7 +800,17 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
         }
         builder.collateParams = in.readMap();
         builder.collatePrune = in.readOptionalBoolean();
-        // NORELEASE read Map<String, List<CandidateGenerator>> generators;
+        int generatorsEntries = in.readVInt();
+        for (int i = 0; i < generatorsEntries; i++) {
+            String type = in.readString();
+            int numberOfGenerators = in.readVInt();
+            List<CandidateGenerator> generatorsList = new ArrayList<>(numberOfGenerators);
+            for (int g = 0; g < numberOfGenerators; g++) {
+                DirectCandidateGeneratorBuilder generator = DirectCandidateGeneratorBuilder.PROTOTYPE.readFrom(in);
+                generatorsList.add(generator);
+            }
+            builder.generators.put(type, generatorsList);
+        }
         return builder;
     }
 
@@ -788,9 +820,9 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
                 Objects.equals(separator, other.separator) &&
                 Objects.equals(realWordErrorLikelihood, other.realWordErrorLikelihood) &&
                 Objects.equals(confidence, other.confidence) &&
-                // NORELEASE Objects.equals(generator, other.generator) &&
+                Objects.equals(generators, other.generators) &&
                 Objects.equals(gramSize, other.gramSize) &&
-                // NORELEASE Objects.equals(model, other.model) &&
+                Objects.equals(model, other.model) &&
                 Objects.equals(forceUnigrams, other.forceUnigrams) &&
                 Objects.equals(tokenLimit, other.tokenLimit) &&
                 Objects.equals(preTag, other.preTag) &&
@@ -803,17 +835,14 @@ public final class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSugge
     @Override
     protected int doHashCode() {
         return Objects.hash(maxErrors, separator, realWordErrorLikelihood, confidence,
-                /** NORELEASE generators, */
+                generators, gramSize, model, forceUnigrams, tokenLimit, preTag, postTag,
-                gramSize,
-                /** NORELEASE model, */
-                forceUnigrams, tokenLimit, preTag, postTag,
                 collateQuery, collateParams, collatePrune);
     }
 
     /**
      * {@link CandidateGenerator} interface.
      */
-    public interface CandidateGenerator extends ToXContent {
+    public interface CandidateGenerator extends Writeable<CandidateGenerator>, ToXContent {
         String getType();
 
         CandidateGenerator fromXContent(QueryParseContext parseContext) throws IOException;

core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java

@@ -38,7 +38,7 @@ import static org.hamcrest.Matchers.not;
 public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBuilder<SB>> extends ESTestCase {
 
     private static final int NUMBER_OF_TESTBUILDERS = 20;
-    private static NamedWriteableRegistry namedWriteableRegistry;
+    protected static NamedWriteableRegistry namedWriteableRegistry;
 
     /**
      * setup for the whole base test class

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

@@ -28,6 +28,11 @@ public class LaplaceModelTests extends SmoothingModelTestCase {
 
     @Override
     protected SmoothingModel createTestModel() {
+        return createRandomModel();
+    }
+
+
+    static SmoothingModel createRandomModel() {
         return new Laplace(randomDoubleBetween(0.0, 10.0, false));
     }
 

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

@@ -28,6 +28,10 @@ public class LinearInterpolationModelTests extends SmoothingModelTestCase {
 
     @Override
     protected SmoothingModel createTestModel() {
+        return createRandomModel();
+    }
+
+    static LinearInterpolation createRandomModel() {
         double trigramLambda = randomDoubleBetween(0.0, 10.0, false);
         double bigramLambda = randomDoubleBetween(0.0, 10.0, false);
         double unigramLambda = randomDoubleBetween(0.0, 10.0, false);

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

@@ -21,6 +21,11 @@ package org.elasticsearch.search.suggest.phrase;
 
 import org.elasticsearch.script.Template;
 import org.elasticsearch.search.suggest.AbstractSuggestionBuilderTestCase;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.Laplace;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.LinearInterpolation;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.SmoothingModel;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder.StupidBackoff;
+import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -28,6 +33,13 @@ import java.util.Map;
 
 public class PhraseSuggestionBuilderTests extends AbstractSuggestionBuilderTestCase<PhraseSuggestionBuilder> {
 
+    @BeforeClass
+    public static void initSmoothingModels() {
+        namedWriteableRegistry.registerPrototype(SmoothingModel.class, Laplace.PROTOTYPE);
+        namedWriteableRegistry.registerPrototype(SmoothingModel.class, LinearInterpolation.PROTOTYPE);
+        namedWriteableRegistry.registerPrototype(SmoothingModel.class, StupidBackoff.PROTOTYPE);
+    }
+
     @Override
     protected PhraseSuggestionBuilder randomSuggestionBuilder() {
         PhraseSuggestionBuilder testBuilder = new PhraseSuggestionBuilder(randomAsciiOfLength(10));
@@ -50,18 +62,36 @@ public class PhraseSuggestionBuilderTests extends AbstractSuggestionBuilderTestC
             testBuilder.collateParams(collateParams );
         }
         if (randomBoolean()) {
-            // NORELEASE add random model
+            testBuilder.smoothingModel(randomSmoothingModel());
         }
-
         if (randomBoolean()) {
-            // NORELEASE add random generator
+            int numGenerators = randomIntBetween(1, 5);
+            for (int i = 0; i < numGenerators; i++) {
+                testBuilder.addCandidateGenerator(DirectCandidateGeneratorTests.randomCandidateGenerator());
+            }
         }
         return testBuilder;
     }
 
+    private static SmoothingModel randomSmoothingModel() {
+        SmoothingModel model = null;
+        switch (randomIntBetween(0,2)) {
+        case 0:
+            model = LaplaceModelTests.createRandomModel();
+            break;
+        case 1:
+            model = StupidBackoffModelTests.createRandomModel();
+            break;
+        case 2:
+            model = LinearInterpolationModelTests.createRandomModel();
+            break;
+        }
+        return model;
+    }
+
     @Override
     protected void mutateSpecificParameters(PhraseSuggestionBuilder builder) throws IOException {
-        switch (randomIntBetween(0, 7)) {
+        switch (randomIntBetween(0, 12)) {
         case 0:
             builder.maxErrors(randomValueOtherThan(builder.maxErrors(), () -> randomFloat()));
             break;
@@ -105,9 +135,16 @@ public class PhraseSuggestionBuilderTests extends AbstractSuggestionBuilderTestC
             builder.forceUnigrams(builder.forceUnigrams() == null ? randomBoolean() : ! builder.forceUnigrams());
             break;
         case 10:
-            builder.collateParams().put(randomAsciiOfLength(5), randomAsciiOfLength(5));
+            Map<String, Object> collateParams = builder.collateParams() == null ? new HashMap<>(1) : builder.collateParams();
+            collateParams.put(randomAsciiOfLength(5), randomAsciiOfLength(5));
+            builder.collateParams(collateParams);
+            break;
+        case 11:
+            builder.smoothingModel(randomValueOtherThan(builder.smoothingModel(), PhraseSuggestionBuilderTests::randomSmoothingModel));
+            break;
+        case 12:
+            builder.addCandidateGenerator(DirectCandidateGeneratorTests.randomCandidateGenerator());
             break;
-        // TODO mutate random Model && generator
         }
     }
 

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

@@ -28,6 +28,10 @@ public class StupidBackoffModelTests extends SmoothingModelTestCase {
 
     @Override
     protected SmoothingModel createTestModel() {
+        return createRandomModel();
+    }
+
+    static SmoothingModel createRandomModel() {
         return new StupidBackoff(randomDoubleBetween(0.0, 10.0, false));
     }