diff options
Diffstat (limited to 'core/src/main/java/org/elasticsearch/search/suggest/completion')
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()); |