diff options
Diffstat (limited to 'core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java')
| -rw-r--r-- | core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java | 158 |
1 files changed, 114 insertions, 44 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 8bf35a34b2..16a2804b37 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 @@ -18,8 +18,11 @@ */ package org.elasticsearch.search.suggest.completion; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.ReaderUtil; +import org.apache.lucene.index.Terms; import org.apache.lucene.search.BulkScorer; import org.apache.lucene.search.CollectionTerminatedException; import org.apache.lucene.search.IndexSearcher; @@ -29,7 +32,9 @@ import org.apache.lucene.search.suggest.document.CompletionQuery; import org.apache.lucene.search.suggest.document.TopSuggestDocs; import org.apache.lucene.search.suggest.document.TopSuggestDocsCollector; import org.apache.lucene.util.CharsRefBuilder; +import org.apache.lucene.util.CollectionUtil; import org.apache.lucene.util.PriorityQueue; +import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.text.Text; import org.elasticsearch.index.fielddata.AtomicFieldData; @@ -42,10 +47,12 @@ import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.search.suggest.Suggest; import org.elasticsearch.search.suggest.Suggester; import org.elasticsearch.search.suggest.SuggestionBuilder; +import org.elasticsearch.search.suggest.completion2x.Completion090PostingsFormat; import java.io.IOException; import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; @@ -61,56 +68,119 @@ public class CompletionSuggester extends Suggester<CompletionSuggestionContext> @Override protected Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> innerExecute(String name, final CompletionSuggestionContext suggestionContext, final IndexSearcher searcher, CharsRefBuilder spare) throws IOException { - final CompletionFieldMapper.CompletionFieldType fieldType = suggestionContext.getFieldType(); - if (fieldType == null) { - throw new IllegalArgumentException("field [" + suggestionContext.getField() + "] is not a completion field"); - } - CompletionSuggestion completionSuggestion = new CompletionSuggestion(name, suggestionContext.getSize()); - spare.copyUTF8Bytes(suggestionContext.getText()); - CompletionSuggestion.Entry completionSuggestEntry = new CompletionSuggestion.Entry(new Text(spare.toString()), 0, spare.length()); - completionSuggestion.addTerm(completionSuggestEntry); - TopSuggestDocsCollector collector = new TopDocumentsCollector(suggestionContext.getSize()); - suggest(searcher, suggestionContext.toQuery(), collector); - int numResult = 0; - List<LeafReaderContext> leaves = searcher.getIndexReader().leaves(); - for (TopSuggestDocs.SuggestScoreDoc suggestScoreDoc : collector.get().scoreLookupDocs()) { - TopDocumentsCollector.SuggestDoc suggestDoc = (TopDocumentsCollector.SuggestDoc) suggestScoreDoc; - // collect contexts - Map<String, Set<CharSequence>> contexts = Collections.emptyMap(); - if (fieldType.hasContextMappings() && suggestDoc.getContexts().isEmpty() == false) { - contexts = fieldType.getContextMappings().getNamedContexts(suggestDoc.getContexts()); - } - // collect payloads - final Map<String, List<Object>> payload = new HashMap<>(0); - List<String> payloadFields = suggestionContext.getPayloadFields(); - if (payloadFields.isEmpty() == false) { - final int readerIndex = ReaderUtil.subIndex(suggestDoc.doc, leaves); - final LeafReaderContext subReaderContext = leaves.get(readerIndex); - final int subDocId = suggestDoc.doc - subReaderContext.docBase; - for (String field : payloadFields) { - MapperService mapperService = suggestionContext.getShardContext().getMapperService(); - MappedFieldType payloadFieldType = mapperService.fullName(field); - if (payloadFieldType != null) { - QueryShardContext shardContext = suggestionContext.getShardContext(); - final AtomicFieldData data = shardContext.getForField(payloadFieldType) + if (suggestionContext.getFieldType() != null) { + final CompletionFieldMapper.CompletionFieldType fieldType = suggestionContext.getFieldType(); + CompletionSuggestion completionSuggestion = new CompletionSuggestion(name, suggestionContext.getSize()); + spare.copyUTF8Bytes(suggestionContext.getText()); + CompletionSuggestion.Entry completionSuggestEntry = new CompletionSuggestion.Entry( + new Text(spare.toString()), 0, spare.length()); + completionSuggestion.addTerm(completionSuggestEntry); + TopSuggestDocsCollector collector = new TopDocumentsCollector(suggestionContext.getSize()); + suggest(searcher, suggestionContext.toQuery(), collector); + int numResult = 0; + List<LeafReaderContext> leaves = searcher.getIndexReader().leaves(); + for (TopSuggestDocs.SuggestScoreDoc suggestScoreDoc : collector.get().scoreLookupDocs()) { + TopDocumentsCollector.SuggestDoc suggestDoc = (TopDocumentsCollector.SuggestDoc) suggestScoreDoc; + // collect contexts + Map<String, Set<CharSequence>> contexts = Collections.emptyMap(); + if (fieldType.hasContextMappings() && suggestDoc.getContexts().isEmpty() == false) { + contexts = fieldType.getContextMappings().getNamedContexts(suggestDoc.getContexts()); + } + // collect payloads + final Map<String, List<Object>> payload = new HashMap<>(0); + List<String> payloadFields = suggestionContext.getPayloadFields(); + if (payloadFields.isEmpty() == false) { + final int readerIndex = ReaderUtil.subIndex(suggestDoc.doc, leaves); + final LeafReaderContext subReaderContext = leaves.get(readerIndex); + final int subDocId = suggestDoc.doc - subReaderContext.docBase; + for (String field : payloadFields) { + MapperService mapperService = suggestionContext.getShardContext().getMapperService(); + MappedFieldType payloadFieldType = mapperService.fullName(field); + if (payloadFieldType != null) { + QueryShardContext shardContext = suggestionContext.getShardContext(); + final AtomicFieldData data = shardContext.getForField(payloadFieldType) .load(subReaderContext); - final ScriptDocValues scriptValues = data.getScriptValues(); - scriptValues.setNextDocId(subDocId); - payload.put(field, new ArrayList<>(scriptValues.getValues())); - } else { - throw new IllegalArgumentException("payload field [" + field + "] does not exist"); + final ScriptDocValues scriptValues = data.getScriptValues(); + scriptValues.setNextDocId(subDocId); + payload.put(field, new ArrayList<>(scriptValues.getValues())); + } else { + throw new IllegalArgumentException("payload field [" + field + "] does not exist"); + } } } - } - if (numResult++ < suggestionContext.getSize()) { - CompletionSuggestion.Entry.Option option = new CompletionSuggestion.Entry.Option( + if (numResult++ < suggestionContext.getSize()) { + CompletionSuggestion.Entry.Option option = new CompletionSuggestion.Entry.Option( new Text(suggestDoc.key.toString()), suggestDoc.score, contexts, payload); - completionSuggestEntry.addOption(option); - } else { - break; + completionSuggestEntry.addOption(option); + } else { + break; + } } + return completionSuggestion; + } else if (suggestionContext.getFieldType2x() != null) { + final IndexReader indexReader = searcher.getIndexReader(); + org.elasticsearch.search.suggest.completion2x.CompletionSuggestion completionSuggestion = + new org.elasticsearch.search.suggest.completion2x.CompletionSuggestion(name, suggestionContext.getSize()); + spare.copyUTF8Bytes(suggestionContext.getText()); + + org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry completionSuggestEntry = + new org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry(new Text(spare.toString()), 0, spare.length()); + completionSuggestion.addTerm(completionSuggestEntry); + + String fieldName = suggestionContext.getField(); + Map<String, org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option> results = + new HashMap<>(indexReader.leaves().size() * suggestionContext.getSize()); + for (LeafReaderContext atomicReaderContext : indexReader.leaves()) { + LeafReader atomicReader = atomicReaderContext.reader(); + Terms terms = atomicReader.fields().terms(fieldName); + if (terms instanceof Completion090PostingsFormat.CompletionTerms) { + final Completion090PostingsFormat.CompletionTerms lookupTerms = (Completion090PostingsFormat.CompletionTerms) terms; + final Lookup lookup = lookupTerms.getLookup(suggestionContext.getFieldType2x(), suggestionContext); + if (lookup == null) { + // we don't have a lookup for this segment.. this might be possible if a merge dropped all + // docs from the segment that had a value in this segment. + continue; + } + List<Lookup.LookupResult> lookupResults = lookup.lookup(spare.get(), false, suggestionContext.getSize()); + for (Lookup.LookupResult res : lookupResults) { + + final String key = res.key.toString(); + final float score = res.value; + final org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option value = results.get(key); + if (value == null) { + final org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option option = + new org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option(new Text(key), score, + res.payload == null ? null : new BytesArray(res.payload)); + results.put(key, option); + } else if (value.getScore() < score) { + value.setScore(score); + value.setPayload(res.payload == null ? null : new BytesArray(res.payload)); + } + } + } + } + final List<org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option> options = + new ArrayList<>(results.values()); + CollectionUtil.introSort(options, scoreComparator); + + int optionCount = Math.min(suggestionContext.getSize(), options.size()); + for (int i = 0; i < optionCount; i++) { + completionSuggestEntry.addOption(options.get(i)); + } + + return completionSuggestion; + } + return null; + } + + private static final ScoreComparator scoreComparator = new ScoreComparator(); + public static class ScoreComparator implements + Comparator<org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option> { + @Override + public int compare(org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option o1, + org.elasticsearch.search.suggest.completion2x.CompletionSuggestion.Entry.Option o2) { + return Float.compare(o2.getScore(), o1.getScore()); } - return completionSuggestion; } private static void suggest(IndexSearcher searcher, CompletionQuery query, TopSuggestDocsCollector collector) throws IOException { |