Add support for returning documents with completion suggester

This commit enables completion suggester to return documents
associated with suggestions. Now the document source is returned
with every suggestion, which respects source filtering options.

In case of suggest queries spanning more than one shard, the
suggest is executed in two phases, where the last phase fetches
the relevant documents from shards, implying executing suggest
requests against a single shard is more performant due to the
document fetch overhead when the suggest spans multiple shards.

2 files changed, 107 insertions, 26 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java
index 759ab8d206..c27f378915 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java
@@ -109,7 +109,7 @@ public class CompletionSuggester extends Suggester<CompletionSuggestionContext>
                     }
                 }
                 if (numResult++ < suggestionContext.getSize()) {
-                    CompletionSuggestion.Entry.Option option = new CompletionSuggestion.Entry.Option(
+                    CompletionSuggestion.Entry.Option option = new CompletionSuggestion.Entry.Option(suggestDoc.doc,
                         new Text(suggestDoc.key.toString()), suggestDoc.score, contexts, payload);
                     completionSuggestEntry.addOption(option);
                 } else {
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java
index 0c209e00a7..a92cbfe1e2 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java
@@ -18,11 +18,16 @@
  */
 package org.elasticsearch.search.suggest.completion;
 
+import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.suggest.Lookup;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.text.Text;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.search.internal.InternalSearchHit;
+import org.elasticsearch.search.internal.InternalSearchHits;
+import org.elasticsearch.search.internal.InternalSearchHits.StreamContext.ShardTargetType;
 import org.elasticsearch.search.suggest.Suggest;
 
 import java.io.IOException;
@@ -35,6 +40,8 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import static org.elasticsearch.search.suggest.Suggest.COMPARATOR;
+
 /**
  * Suggestion response for {@link CompletionSuggester} results
  *
@@ -62,6 +69,25 @@ public final class CompletionSuggestion extends Suggest.Suggestion<CompletionSug
         super(name, size);
     }
 
+    /**
+     * @return the result options for the suggestion
+     */
+    public List<Entry.Option> getOptions() {
+        if (entries.isEmpty() == false) {
+            assert entries.size() == 1 : "CompletionSuggestion must have only one entry";
+            return entries.get(0).getOptions();
+        } else {
+            return Collections.emptyList();
+        }
+    }
+
+    /**
+     * @return whether there is any hits for the suggestion
+     */
+    public boolean hasScoreDocs() {
+        return getOptions().size() > 0;
+    }
+
     private static final class OptionPriorityQueue extends org.apache.lucene.util.PriorityQueue<Entry.Option> {
 
         private final Comparator<Suggest.Suggestion.Entry.Option> comparator;
@@ -90,30 +116,54 @@ public final class CompletionSuggestion extends Suggest.Suggestion<CompletionSug
         }
     }
 
-    @Override
-    public Suggest.Suggestion<Entry> reduce(List<Suggest.Suggestion<Entry>> toReduce) {
-        if (toReduce.size() == 1) {
-            return toReduce.get(0);
+    /**
+     * Reduces suggestions to a single suggestion containing at most
+     * top {@link CompletionSuggestion#getSize()} options across <code>toReduce</code>
+     */
+    public static CompletionSuggestion reduceTo(List<Suggest.Suggestion<Entry>> toReduce) {
+        if (toReduce.isEmpty()) {
+            return null;
         } else {
-            // combine suggestion entries from participating shards on the coordinating node
-            // the global top <code>size</code> entries are collected from the shard results
-            // using a priority queue
-            OptionPriorityQueue priorityQueue = new OptionPriorityQueue(size, sortComparator());
-            for (Suggest.Suggestion<Entry> entries : toReduce) {
-                assert entries.getEntries().size() == 1 : "CompletionSuggestion must have only one entry";
-                for (Entry.Option option : entries.getEntries().get(0)) {
-                    if (option == priorityQueue.insertWithOverflow(option)) {
-                        // if the current option has overflown from pq,
-                        // we can assume all of the successive options
-                        // from this shard result will be overflown as well
-                        break;
+            final CompletionSuggestion leader = (CompletionSuggestion) toReduce.get(0);
+            final Entry leaderEntry = leader.getEntries().get(0);
+            final String name = leader.getName();
+            if (toReduce.size() == 1) {
+                return leader;
+            } else {
+                // combine suggestion entries from participating shards on the coordinating node
+                // the global top <code>size</code> entries are collected from the shard results
+                // using a priority queue
+                OptionPriorityQueue priorityQueue = new OptionPriorityQueue(leader.getSize(), COMPARATOR);
+                for (Suggest.Suggestion<Entry> suggestion : toReduce) {
+                    assert suggestion.getName().equals(name) : "name should be identical across all suggestions";
+                    for (Entry.Option option : ((CompletionSuggestion) suggestion).getOptions()) {
+                        if (option == priorityQueue.insertWithOverflow(option)) {
+                            // if the current option has overflown from pq,
+                            // we can assume all of the successive options
+                            // from this shard result will be overflown as well
+                            break;
+                        }
                     }
                 }
+                final CompletionSuggestion suggestion = new CompletionSuggestion(leader.getName(), leader.getSize());
+                final Entry entry = new Entry(leaderEntry.getText(), leaderEntry.getOffset(), leaderEntry.getLength());
+                Collections.addAll(entry.getOptions(), priorityQueue.get());
+                suggestion.addTerm(entry);
+                return suggestion;
+            }
+        }
+    }
+
+    @Override
+    public Suggest.Suggestion<Entry> reduce(List<Suggest.Suggestion<Entry>> toReduce) {
+        return reduceTo(toReduce);
+    }
+
+    public void setShardIndex(int shardIndex) {
+        if (entries.isEmpty() == false) {
+            for (Entry.Option option : getOptions()) {
+                option.setShardIndex(shardIndex);
             }
-            Entry options = this.entries.get(0);
-            options.getOptions().clear();
-            Collections.addAll(options.getOptions(), priorityQueue.get());
-            return this;
         }
     }
 
@@ -145,9 +195,12 @@ public final class CompletionSuggestion extends Suggest.Suggestion<CompletionSug
         public static class Option extends Suggest.Suggestion.Entry.Option {
             private Map<String, Set<CharSequence>> contexts;
             private Map<String, List<Object>> payload;
+            private ScoreDoc doc;
+            private InternalSearchHit hit;
 
-            public Option(Text text, float score, Map<String, Set<CharSequence>> contexts, Map<String, List<Object>> payload) {
+            public Option(int docID, Text text, float score, Map<String, Set<CharSequence>> contexts, Map<String, List<Object>> payload) {
                 super(text, score);
+                this.doc = new ScoreDoc(docID, score);
                 this.payload = payload;
                 this.contexts = contexts;
             }
@@ -171,14 +224,30 @@ public final class CompletionSuggestion extends Suggest.Suggestion<CompletionSug
                 return contexts;
             }
 
-            @Override
-            public void setScore(float score) {
-                super.setScore(score);
+            public ScoreDoc getDoc() {
+                return doc;
+            }
+
+            public InternalSearchHit getHit() {
+                return hit;
+            }
+
+            public void setShardIndex(int shardIndex) {
+                this.doc.shardIndex = shardIndex;
+            }
+
+            public void setHit(InternalSearchHit hit) {
+                this.hit = hit;
             }
 
             @Override
             protected XContentBuilder innerToXContent(XContentBuilder builder, Params params) throws IOException {
-                super.innerToXContent(builder, params);
+                builder.field("text", getText());
+                if (hit != null) {
+                    hit.toInnerXContent(builder, params);
+                } else {
+                    builder.field("score", getScore());
+                }
                 if (payload.size() > 0) {
                     builder.startObject("payload");
                     for (Map.Entry<String, List<Object>> entry : payload.entrySet()) {
@@ -207,6 +276,11 @@ public final class CompletionSuggestion extends Suggest.Suggestion<CompletionSug
             @Override
             public void readFrom(StreamInput in) throws IOException {
                 super.readFrom(in);
+                this.doc = Lucene.readScoreDoc(in);
+                if (in.readBoolean()) {
+                    this.hit = InternalSearchHit.readSearchHit(in,
+                        InternalSearchHits.streamContext().streamShardTarget(ShardTargetType.STREAM));
+                }
                 int payloadSize = in.readInt();
                 this.payload = new LinkedHashMap<>(payloadSize);
                 for (int i = 0; i < payloadSize; i++) {
@@ -234,6 +308,13 @@ public final class CompletionSuggestion extends Suggest.Suggestion<CompletionSug
             @Override
             public void writeTo(StreamOutput out) throws IOException {
                 super.writeTo(out);
+                Lucene.writeScoreDoc(out, doc);
+                if (hit != null) {
+                    out.writeBoolean(true);
+                    hit.writeTo(out, InternalSearchHits.streamContext().streamShardTarget(ShardTargetType.STREAM));
+                } else {
+                    out.writeBoolean(false);
+                }
                 out.writeInt(payload.size());
                 for (Map.Entry<String, List<Object>> entry : payload.entrySet()) {
                     out.writeString(entry.getKey());