Make suggester implementation pluggable

This patch tries to make the suggester implementation as pluggable as
facets or highlight implementations. The goal is to be able to create
own suggest implementations in a suggest query.

Closes #3089

src/main/java/org/elasticsearch/action/suggest/TransportSuggestAction.java

@@ -65,10 +65,10 @@ public class TransportSuggestAction extends TransportBroadcastOperationAction<Su
 
     @Inject
     public TransportSuggestAction(Settings settings, ThreadPool threadPool, ClusterService clusterService, TransportService transportService,
-                                  IndicesService indicesService) {
+                                  IndicesService indicesService, SuggestPhase suggestPhase) {
         super(settings, threadPool, clusterService, transportService);
         this.indicesService = indicesService;
-        this.suggestPhase = new SuggestPhase(settings);
+        this.suggestPhase = suggestPhase;
     }
 
     @Override

src/main/java/org/elasticsearch/search/SearchModule.java

@@ -36,6 +36,7 @@ import org.elasticsearch.search.fetch.version.VersionFetchSubPhase;
 import org.elasticsearch.search.highlight.HighlightModule;
 import org.elasticsearch.search.highlight.HighlightPhase;
 import org.elasticsearch.search.query.QueryPhase;
+import org.elasticsearch.search.suggest.SuggestModule;
 
 /**
  *
@@ -44,7 +45,7 @@ public class SearchModule extends AbstractModule implements SpawnModules {
 
     @Override
     public Iterable<? extends Module> spawnModules() {
-        return ImmutableList.of(new TransportSearchModule(), new FacetModule(), new HighlightModule());
+        return ImmutableList.of(new TransportSearchModule(), new FacetModule(), new HighlightModule(), new SuggestModule());
     }
 
     @Override

src/main/java/org/elasticsearch/search/suggest/SuggestModule.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to ElasticSearch and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. ElasticSearch licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.search.suggest;
+
+import com.google.common.collect.Lists;
+import org.elasticsearch.common.inject.AbstractModule;
+import org.elasticsearch.common.inject.multibindings.Multibinder;
+import org.elasticsearch.search.suggest.phrase.PhraseSuggester;
+import org.elasticsearch.search.suggest.term.TermSuggester;
+
+import java.util.List;
+
+/**
+ *
+ */
+public class SuggestModule extends AbstractModule {
+
+    private List<Class<? extends Suggester>> suggesters = Lists.newArrayList();
+
+    public SuggestModule() {
+        registerSuggester(PhraseSuggester.class);
+        registerSuggester(TermSuggester.class);
+    }
+
+    public void registerSuggester(Class<? extends Suggester> suggester) {
+        suggesters.add(suggester);
+    }
+
+    @Override
+    protected void configure() {
+        Multibinder<Suggester> suggesterMultibinder = Multibinder.newSetBinder(binder(), Suggester.class);
+        for (Class<? extends Suggester> clazz : suggesters) {
+            suggesterMultibinder.addBinding().to(clazz);
+        }
+
+        bind(SuggestParseElement.class).asEagerSingleton();
+        bind(SuggestPhase.class).asEagerSingleton();
+        bind(Suggesters.class).asEagerSingleton();
+    }
+}

src/main/java/org/elasticsearch/search/suggest/SuggestParseElement.java

@@ -22,21 +22,24 @@ import java.io.IOException;
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ElasticSearchIllegalArgumentException;
+import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
-import org.elasticsearch.search.suggest.phrase.PhraseSuggestParser;
-import org.elasticsearch.search.suggest.term.TermSuggestParser;
 
 /**
  *
  */
 public class SuggestParseElement implements SearchParseElement {
-    private final SuggestContextParser termSuggestParser = new TermSuggestParser();
-    private final SuggestContextParser phraseSuggestParser = new PhraseSuggestParser();
-    
+    private Suggesters suggesters;
+
+    @Inject
+    public SuggestParseElement(Suggesters suggesters) {
+        this.suggesters = suggesters;
+    }
+
     @Override
     public void parse(XContentParser parser, SearchContext context) throws Exception {
         SuggestionSearchContext suggestionSearchContext = parseInternal(parser, context.mapperService());
@@ -74,16 +77,11 @@ public class SuggestParseElement implements SearchParseElement {
                         if (suggestionName == null) {
                             throw new ElasticSearchIllegalArgumentException("Suggestion must have name");
                         }
-                        final SuggestContextParser contextParser;
-                        if ("term".equals(fieldName)) {
-                            contextParser = termSuggestParser;
-                        } else if ("phrase".equals(fieldName)) {
-                            contextParser = phraseSuggestParser;
-                        } else {
+                        if (suggesters.get(fieldName) == null) {
                             throw new ElasticSearchIllegalArgumentException("Suggester[" + fieldName + "] not supported");
                         }
+                        final SuggestContextParser contextParser = suggesters.get(fieldName).getContextParser();
                         parseAndVerify(parser, mapperService, suggestionSearchContext, globalText, suggestionName, suggestText, contextParser);
-                        
                     }
                 }
             }

src/main/java/org/elasticsearch/search/suggest/SuggestPhase.java

@@ -45,9 +45,9 @@ public class SuggestPhase extends AbstractComponent implements SearchPhase {
     private final SuggestParseElement parseElement;
 
     @Inject
-    public SuggestPhase(Settings settings) {
+    public SuggestPhase(Settings settings, SuggestParseElement suggestParseElement) {
         super(settings);
-        this.parseElement = new SuggestParseElement();
+        this.parseElement = suggestParseElement;
     }
 
     @Override

src/main/java/org/elasticsearch/search/suggest/Suggester.java

@@ -26,6 +26,11 @@ import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry.Option;
 
 public interface Suggester<T extends SuggestionSearchContext.SuggestionContext> {
-    public abstract Suggestion<? extends Entry<? extends Option>> execute(String name, T suggestion, IndexReader indexReader, CharsRef spare)
+
+    public Suggestion<? extends Entry<? extends Option>> execute(String name, T suggestion, IndexReader indexReader, CharsRef spare)
             throws IOException;
+
+    public String[] names();
+
+    public SuggestContextParser getContextParser();
 }

src/main/java/org/elasticsearch/search/suggest/Suggesters.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed to ElasticSearch and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. ElasticSearch licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.search.suggest;
+
+import com.google.common.collect.ImmutableMap;
+import org.elasticsearch.common.collect.MapBuilder;
+import org.elasticsearch.common.inject.Inject;
+
+import java.util.Set;
+
+/**
+ *
+ */
+public class Suggesters {
+    private final ImmutableMap<String, Suggester> parsers;
+
+    @Inject
+    public Suggesters(Set<Suggester> suggesters) {
+        MapBuilder<String, Suggester> builder = MapBuilder.newMapBuilder();
+        for (Suggester suggester : suggesters) {
+            for (String type : suggester.names()) {
+                builder.put(type, suggester);
+            }
+        }
+        this.parsers = builder.immutableMap();
+    }
+
+    public Suggester get(String type) {
+        return parsers.get(type);
+    }
+
+}

src/main/java/org/elasticsearch/search/suggest/SuggestionSearchContext.java

@@ -55,7 +55,7 @@ public class SuggestionSearchContext {
             this.text = text;
         }
         
-        protected SuggestionContext(Suggester suggester) {
+        public SuggestionContext(Suggester suggester) {
             this.suggester = suggester;
         }
         

src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggester.java

@@ -33,10 +33,11 @@ import org.elasticsearch.common.text.Text;
 import org.elasticsearch.search.suggest.Suggest.Suggestion;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry.Option;
+import org.elasticsearch.search.suggest.SuggestContextParser;
 import org.elasticsearch.search.suggest.SuggestUtils;
 import org.elasticsearch.search.suggest.Suggester;
 
-final class PhraseSuggester implements Suggester<PhraseSuggestionContext> {
+public final class PhraseSuggester implements Suggester<PhraseSuggestionContext> {
     private final BytesRef SEPARATOR = new BytesRef(" ");
     
     /*
@@ -81,4 +82,14 @@ final class PhraseSuggester implements Suggester<PhraseSuggestionContext> {
         return response;
     }
 
+    @Override
+    public String[] names() {
+        return new String[] {"phrase"};
+    }
+
+    @Override
+    public SuggestContextParser getContextParser() {
+        return new PhraseSuggestParser();
+    }
+
 }

src/main/java/org/elasticsearch/search/suggest/term/TermSuggester.java

@@ -32,11 +32,12 @@ import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.text.BytesText;
 import org.elasticsearch.common.text.StringText;
 import org.elasticsearch.common.text.Text;
+import org.elasticsearch.search.suggest.SuggestContextParser;
 import org.elasticsearch.search.suggest.SuggestUtils;
 import org.elasticsearch.search.suggest.Suggester;
 import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
 
-final class TermSuggester implements Suggester<TermSuggestionContext> {
+public final class TermSuggester implements Suggester<TermSuggestionContext> {
 
     @Override
     public TermSuggestion execute(String name, TermSuggestionContext suggestion, IndexReader indexReader, CharsRef spare) throws IOException {
@@ -62,7 +63,17 @@ final class TermSuggester implements Suggester<TermSuggestionContext> {
         return response;
     }
 
-   
+    @Override
+    public String[] names() {
+        return new String[] {"term"};
+    }
+
+    @Override
+    public SuggestContextParser getContextParser() {
+        return new TermSuggestParser();
+    }
+
+
     private List<Token> queryTerms(SuggestionContext suggestion, CharsRef spare) throws IOException {
         final List<Token> result = new ArrayList<TermSuggester.Token>();
         final String field = suggestion.getField();

src/test/java/org/elasticsearch/test/integration/search/suggest/CustomSuggester.java

@@ -0,0 +1,88 @@
+/*
+ * Licensed to ElasticSearch and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. ElasticSearch licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.test.integration.search.suggest;
+
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.util.CharsRef;
+import org.elasticsearch.common.text.StringText;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.search.suggest.Suggest;
+import org.elasticsearch.search.suggest.SuggestContextParser;
+import org.elasticsearch.search.suggest.Suggester;
+import org.elasticsearch.search.suggest.SuggestionSearchContext;
+
+import java.io.IOException;
+import java.util.Locale;
+import java.util.Map;
+
+/**
+ *
+ */
+public class CustomSuggester implements Suggester<CustomSuggester.CustomSuggestionsContext> {
+
+
+    // This is a pretty dumb implementation which returns the original text + fieldName + custom config option + 12 or 123
+    @Override
+    public Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> execute(String name, CustomSuggestionsContext suggestion, IndexReader indexReader, CharsRef spare) throws IOException {
+        // Get the suggestion context
+        String text = suggestion.getText().utf8ToString();
+
+        // create two suggestions with 12 and 123 appended
+        Suggest.Suggestion<Suggest.Suggestion.Entry<Suggest.Suggestion.Entry.Option>> response = new Suggest.Suggestion<Suggest.Suggestion.Entry<Suggest.Suggestion.Entry.Option>>(name, suggestion.getSize());
+
+        String firstSuggestion = String.format(Locale.ROOT, "%s-%s-%s-%s", text, suggestion.getField(), suggestion.options.get("suffix"), "12");
+        Suggest.Suggestion.Entry<Suggest.Suggestion.Entry.Option> resultEntry12 = new Suggest.Suggestion.Entry<Suggest.Suggestion.Entry.Option>(new StringText(firstSuggestion), 0, text.length() + 2);
+        response.addTerm(resultEntry12);
+
+        String secondSuggestion = String.format(Locale.ROOT, "%s-%s-%s-%s", text, suggestion.getField(), suggestion.options.get("suffix"), "123");
+        Suggest.Suggestion.Entry<Suggest.Suggestion.Entry.Option> resultEntry123 = new Suggest.Suggestion.Entry<Suggest.Suggestion.Entry.Option>(new StringText(secondSuggestion), 0, text.length() + 3);
+        response.addTerm(resultEntry123);
+
+        return response;
+    }
+
+    @Override
+    public String[] names() {
+        return new String[] {"custom"};
+    }
+
+    @Override
+    public SuggestContextParser getContextParser() {
+        return new SuggestContextParser() {
+            @Override
+            public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, MapperService mapperService) throws IOException {
+                Map<String, Object> options = parser.map();
+                CustomSuggestionsContext suggestionContext = new CustomSuggestionsContext(CustomSuggester.this, options);
+                suggestionContext.setField((String) options.get("field"));
+                return suggestionContext;
+            }
+        };
+    }
+
+    public static class CustomSuggestionsContext extends SuggestionSearchContext.SuggestionContext {
+
+        public Map<String, Object> options;
+
+        public CustomSuggestionsContext(Suggester suggester, Map<String, Object> options) {
+            super(suggester);
+            this.options = options;
+        }
+    }
+}

src/test/java/org/elasticsearch/test/integration/search/suggest/CustomSuggesterPlugin.java

@@ -0,0 +1,43 @@
+/*
+ * Licensed to ElasticSearch and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. ElasticSearch licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.test.integration.search.suggest;
+
+import org.elasticsearch.plugins.AbstractPlugin;
+import org.elasticsearch.search.suggest.SuggestModule;
+
+/**
+ *
+ */
+public class CustomSuggesterPlugin extends AbstractPlugin {
+
+    @Override
+    public String name() {
+        return "test-plugin-custom-suggester";
+    }
+
+    @Override
+    public String description() {
+        return "Custom suggester to test pluggable implementation";
+    }
+
+    public void onModule(SuggestModule suggestModule) {
+        suggestModule.registerSuggester(CustomSuggester.class);
+    }
+
+}

src/test/java/org/elasticsearch/test/integration/search/suggest/CustomSuggesterSearchTests.java

@@ -0,0 +1,98 @@
+/*
+ * Licensed to ElasticSearch and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. ElasticSearch licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.test.integration.search.suggest;
+
+import com.google.common.collect.Lists;
+import org.elasticsearch.action.search.SearchRequestBuilder;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.common.Priority;
+import org.elasticsearch.common.RandomStringGenerator;
+import org.elasticsearch.common.settings.ImmutableSettings;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.search.suggest.Suggest;
+import org.elasticsearch.test.integration.AbstractNodesTests;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.List;
+import java.util.Locale;
+
+import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder;
+import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.is;
+
+/**
+ *
+ */
+public class CustomSuggesterSearchTests extends AbstractNodesTests {
+
+    private Client client;
+
+    @BeforeClass
+    public void createNodes() throws Exception {
+        ImmutableSettings.Builder settings = settingsBuilder().put("plugin.types", CustomSuggesterPlugin.class.getName());
+        startNode("server1", settings);
+        client = client("server1");
+
+        client.prepareIndex("test", "test", "1").setSource(jsonBuilder()
+                .startObject()
+                .field("name", "arbitrary content")
+                .endObject())
+                .setRefresh(true).execute().actionGet();
+        client.admin().cluster().prepareHealth().setWaitForEvents(Priority.LANGUID).setWaitForYellowStatus().execute().actionGet();
+    }
+
+    @AfterClass
+    public void closeNodes() {
+        client.close();
+        closeAllNodes();
+    }
+
+    @Test
+    public void testThatCustomSuggestersCanBeRegisteredAndWork() throws Exception {
+        String randomText = RandomStringGenerator.randomAlphanumeric(10);
+        String randomField = RandomStringGenerator.randomAlphanumeric(10);
+        String randomSuffix = RandomStringGenerator.randomAlphanumeric(10);
+        SearchRequestBuilder searchRequestBuilder = client.prepareSearch("test").setTypes("test").setFrom(0).setSize(1);
+        XContentBuilder query = jsonBuilder().startObject()
+                .startObject("suggest")
+                    .startObject("someName")
+                        .field("text", randomText)
+                        .startObject("custom")
+                            .field("field", randomField)
+                            .field("suffix", randomSuffix)
+                        .endObject()
+                    .endObject()
+                .endObject()
+            .endObject();
+        searchRequestBuilder.setExtraSource(query.bytes());
+
+        SearchResponse searchResponse = searchRequestBuilder.execute().actionGet();
+
+        List<Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> suggestions = Lists.newArrayList(searchResponse.getSuggest().getSuggestion("someName").iterator());
+        assertThat(suggestions, hasSize(2));
+        assertThat(suggestions.get(0).getText().string(), is(String.format(Locale.ROOT, "%s-%s-%s-12", randomText, randomField, randomSuffix)));
+        assertThat(suggestions.get(1).getText().string(), is(String.format(Locale.ROOT, "%s-%s-%s-123", randomText, randomField, randomSuffix)));
+    }
+
+}