Completion Suggester: Support returning documents with completions

core/src/main/java/org/elasticsearch/action/ActionModule.java

@@ -176,6 +176,8 @@
 import org.elasticsearch.action.search.type.TransportSearchScrollQueryThenFetchAction;
 import org.elasticsearch.action.search.type.TransportSearchScrollScanAction;
 import org.elasticsearch.action.suggest.SuggestAction;
+import org.elasticsearch.action.suggest.TransportSuggestAndFetchAction;
+import org.elasticsearch.action.suggest.TransportSuggestThenFetchAction;
 import org.elasticsearch.action.suggest.TransportSuggestAction;
 import org.elasticsearch.action.support.ActionFilter;
 import org.elasticsearch.action.support.ActionFilters;
@@ -315,7 +317,7 @@ protected void configure() {
                 TransportShardMultiTermsVectorAction.class);
         registerAction(DeleteAction.INSTANCE, TransportDeleteAction.class);
         registerAction(ExistsAction.INSTANCE, TransportExistsAction.class);
-        registerAction(SuggestAction.INSTANCE, TransportSuggestAction.class);
+        registerAction(SuggestAction.INSTANCE, TransportSuggestAction.class, TransportSuggestThenFetchAction.class, TransportSuggestAndFetchAction.class);
         registerAction(UpdateAction.INSTANCE, TransportUpdateAction.class);
         registerAction(MultiGetAction.INSTANCE, TransportMultiGetAction.class,
                 TransportShardMultiGetAction.class);

core/src/main/java/org/elasticsearch/action/search/type/TransportSearchDfsQueryThenFetchAction.java

@@ -23,10 +23,7 @@
 import org.apache.lucene.search.ScoreDoc;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRunnable;
-import org.elasticsearch.action.search.ReduceSearchPhaseException;
-import org.elasticsearch.action.search.SearchPhaseExecutionException;
-import org.elasticsearch.action.search.SearchRequest;
-import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.search.*;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.cluster.ClusterService;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
@@ -167,7 +164,7 @@ void innerExecuteFetchPhase() throws Exception {
             for (final AtomicArray.Entry<IntArrayList> entry : docIdsToLoad.asList()) {
                 QuerySearchResult queryResult = queryResults.get(entry.index);
                 DiscoveryNode node = nodes.get(queryResult.shardTarget().nodeId());
-                ShardFetchSearchRequest fetchSearchRequest = createFetchRequest(queryResult, entry, lastEmittedDocPerShard);
+                ShardFetchSearchRequest fetchSearchRequest = createFetchRequest(queryResult, entry, null, lastEmittedDocPerShard);
                 executeFetch(entry.index, queryResult.shardTarget(), counter, fetchSearchRequest, node);
             }
         }

core/src/main/java/org/elasticsearch/action/search/type/TransportSearchQueryThenFetchAction.java

@@ -102,7 +102,7 @@ protected void moveToSecondPhase() throws Exception {
             for (AtomicArray.Entry<IntArrayList> entry : docIdsToLoad.asList()) {
                 QuerySearchResultProvider queryResult = firstResults.get(entry.index);
                 DiscoveryNode node = nodes.get(queryResult.shardTarget().nodeId());
-                ShardFetchSearchRequest fetchSearchRequest = createFetchRequest(queryResult.queryResult(), entry, lastEmittedDocPerShard);
+                ShardFetchSearchRequest fetchSearchRequest = createFetchRequest(queryResult.queryResult(), entry, null, lastEmittedDocPerShard);
                 executeFetch(entry.index, queryResult.shardTarget(), counter, fetchSearchRequest, node);
             }
         }

core/src/main/java/org/elasticsearch/action/search/type/TransportSearchTypeAction.java

@@ -347,12 +347,12 @@ protected void sendReleaseSearchContext(long contextId, DiscoveryNode node) {
             }
         }
 
-        protected ShardFetchSearchRequest createFetchRequest(QuerySearchResult queryResult, AtomicArray.Entry<IntArrayList> entry, ScoreDoc[] lastEmittedDocPerShard) {
+        protected ShardFetchSearchRequest createFetchRequest(QuerySearchResult queryResult, AtomicArray.Entry<IntArrayList> entry, Map<String, IntArrayList> namedDocIds, ScoreDoc[] lastEmittedDocPerShard) {
             if (lastEmittedDocPerShard != null) {
                 ScoreDoc lastEmittedDoc = lastEmittedDocPerShard[entry.index];
-                return new ShardFetchSearchRequest(request, queryResult.id(), entry.value, lastEmittedDoc);
+                return new ShardFetchSearchRequest(request, queryResult.id(), entry.value, namedDocIds, lastEmittedDoc);
             } else {
-                return new ShardFetchSearchRequest(request, queryResult.id(), entry.value);
+                return new ShardFetchSearchRequest(request, queryResult.id(), entry.value, namedDocIds);
             }
         }
 

core/src/main/java/org/elasticsearch/action/suggest/SuggestRequest.java

@@ -19,17 +19,13 @@
 
 package org.elasticsearch.action.suggest;
 
-import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.support.broadcast.BroadcastRequest;
-import org.elasticsearch.client.Requests;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentHelper;
-import org.elasticsearch.search.suggest.SuggestBuilder;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -39,8 +35,7 @@
  * {@link org.elasticsearch.client.Requests#suggestRequest(String...)}.
  * <p/>
  * <p>The request requires the suggest query source to be set either using
- * {@link #suggest(org.elasticsearch.common.bytes.BytesReference)} / {@link #suggest(org.elasticsearch.common.bytes.BytesReference)}
- * or by using {@link #suggest(org.elasticsearch.search.suggest.SuggestBuilder)}
+ * {@link #suggest(org.elasticsearch.common.bytes.BytesReference)}
  * (Best created using the {link @org.elasticsearch.search.suggest.SuggestBuilders)}).
  *
  * @see SuggestResponse
@@ -69,12 +64,6 @@ public SuggestRequest(String... indices) {
         super(indices);
     }
 
-    @Override
-    public ActionRequestValidationException validate() {
-        ActionRequestValidationException validationException = super.validate();
-        return validationException;
-    }
-
     /**
      * The Phrase to get correction suggestions for 
      */
@@ -90,26 +79,6 @@ public SuggestRequest suggest(BytesReference suggestSource) {
         return this;
     }
 
-    /**
-     * set a new source using a {@link org.elasticsearch.search.suggest.SuggestBuilder}
-     * for phrase and term suggestion lookup
-     */
-    public SuggestRequest suggest(SuggestBuilder suggestBuilder) {
-        return suggest(suggestBuilder.buildAsBytes(Requests.CONTENT_TYPE));
-    }
-
-    /**
-     * set a new source using a {@link org.elasticsearch.search.suggest.SuggestBuilder.SuggestionBuilder}
-     * for completion suggestion lookup
-     */
-    public SuggestRequest suggest(SuggestBuilder.SuggestionBuilder suggestionBuilder) {
-        return suggest(suggestionBuilder.buildAsBytes(Requests.CONTENT_TYPE));
-    }
-
-    public SuggestRequest suggest(String source) {
-        return suggest(new BytesArray(source));
-    }
-
     /**
      * A comma separated list of routing values to control the shards the search will be executed on.
      */
@@ -147,7 +116,7 @@ public void readFrom(StreamInput in) throws IOException {
         super.readFrom(in);
         routing = in.readOptionalString();
         preference = in.readOptionalString();
-        suggest(in.readBytesReference());
+        suggestSource = in.readBytesReference();
     }
 
     @Override

core/src/main/java/org/elasticsearch/action/suggest/SuggestRequestBuilder.java

@@ -23,9 +23,12 @@
 import org.elasticsearch.action.support.broadcast.BroadcastOperationRequestBuilder;
 import org.elasticsearch.client.ElasticsearchClient;
 import org.elasticsearch.client.Requests;
+import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.fetch.source.FetchSourceContext;
 import org.elasticsearch.search.suggest.SuggestBuilder;
 import org.elasticsearch.search.suggest.SuggestBuilder.SuggestionBuilder;
 
@@ -36,7 +39,7 @@
  */
 public class SuggestRequestBuilder extends BroadcastOperationRequestBuilder<SuggestRequest, SuggestResponse, SuggestRequestBuilder> {
 
-    final SuggestBuilder suggest = new SuggestBuilder();
+    final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
 
     public SuggestRequestBuilder(ElasticsearchClient client, SuggestAction action) {
         super(client, action, new SuggestRequest());
@@ -46,7 +49,7 @@ public SuggestRequestBuilder(ElasticsearchClient client, SuggestAction action) {
      * Add a definition for suggestions to the request
      */
     public <T> SuggestRequestBuilder addSuggestion(SuggestionBuilder<T> suggestion) {
-        suggest.addSuggestion(suggestion);
+        searchSourceBuilder.suggest().addSuggestion(suggestion);
         return this;
     }
 
@@ -59,7 +62,7 @@ public SuggestRequestBuilder setRouting(String routing) {
     }
 
     public SuggestRequestBuilder setSuggestText(String globalText) {
-        this.suggest.setText(globalText);
+        searchSourceBuilder.suggest().setText(globalText);
         return this;
     }
 
@@ -82,11 +85,62 @@ public SuggestRequestBuilder setRouting(String... routing) {
         return this;
     }
 
+    /**
+     * Indicates whether the response should contain the stored _source for
+     * every hit
+     */
+    public SuggestRequestBuilder fetchSource(boolean fetch) {
+        searchSourceBuilder.fetchSource(fetch);
+        return this;
+    }
+
+    /**
+     * Indicate that _source should be returned with every hit, with an
+     * "include" and/or "exclude" set which can include simple wildcard
+     * elements.
+     *
+     * @param include
+     *            An optional include (optionally wildcarded) pattern to filter
+     *            the returned _source
+     * @param exclude
+     *            An optional exclude (optionally wildcarded) pattern to filter
+     *            the returned _source
+     */
+    public SuggestRequestBuilder fetchSource(@Nullable String include, @Nullable String exclude) {
+        searchSourceBuilder.fetchSource(include, exclude);
+        return this;
+    }
+
+    /**
+     * Indicate that _source should be returned with every hit, with an
+     * "include" and/or "exclude" set which can include simple wildcard
+     * elements.
+     *
+     * @param includes
+     *            An optional list of include (optionally wildcarded) pattern to
+     *            filter the returned _source
+     * @param excludes
+     *            An optional list of exclude (optionally wildcarded) pattern to
+     *            filter the returned _source
+     */
+    public SuggestRequestBuilder fetchSource(@Nullable String[] includes, @Nullable String[] excludes) {
+        searchSourceBuilder.fetchSource(includes, excludes);
+        return this;
+    }
+
+    /**
+     * Indicate how the _source should be fetched.
+     */
+    public SuggestRequestBuilder fetchSource(@Nullable FetchSourceContext fetchSourceContext) {
+        searchSourceBuilder.fetchSource(fetchSourceContext);
+        return this;
+    }
+
     @Override
     protected SuggestRequest beforeExecute(SuggestRequest request) {
         try {
             XContentBuilder builder = XContentFactory.contentBuilder(Requests.CONTENT_TYPE);
-            suggest.toXContent(builder, ToXContent.EMPTY_PARAMS);
+            searchSourceBuilder.toXContent(builder, ToXContent.EMPTY_PARAMS);
             request.suggest(builder.bytes());
         } catch (IOException e) {
             throw new ElasticsearchException("Unable to build suggestion request", e);