From 4a1a03428d7884c185b05c364c6f220237a9567e Mon Sep 17 00:00:00 2001 From: Areek Zillur Date: Thu, 7 Apr 2016 15:34:39 -0400 Subject: Add bwc support for pre-5.0 completion index This commit adds support for reading and querying completion fields that were indexed in 2.x --- .../suggest/completion/CompletionSuggester.java | 158 +++++++++++++++------ .../completion/CompletionSuggestionBuilder.java | 136 +++++++++++++++++- .../completion/CompletionSuggestionContext.java | 20 +++ 3 files changed, 265 insertions(+), 49 deletions(-) (limited to 'core/src/main/java/org/elasticsearch/search/suggest/completion') 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 @Override protected Suggest.Suggestion> 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 leaves = searcher.getIndexReader().leaves(); - for (TopSuggestDocs.SuggestScoreDoc suggestScoreDoc : collector.get().scoreLookupDocs()) { - TopDocumentsCollector.SuggestDoc suggestDoc = (TopDocumentsCollector.SuggestDoc) suggestScoreDoc; - // collect contexts - Map> contexts = Collections.emptyMap(); - if (fieldType.hasContextMappings() && suggestDoc.getContexts().isEmpty() == false) { - contexts = fieldType.getContextMappings().getNamedContexts(suggestDoc.getContexts()); - } - // collect payloads - final Map> payload = new HashMap<>(0); - List 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 leaves = searcher.getIndexReader().leaves(); + for (TopSuggestDocs.SuggestScoreDoc suggestScoreDoc : collector.get().scoreLookupDocs()) { + TopDocumentsCollector.SuggestDoc suggestDoc = (TopDocumentsCollector.SuggestDoc) suggestScoreDoc; + // collect contexts + Map> contexts = Collections.emptyMap(); + if (fieldType.hasContextMappings() && suggestDoc.getContexts().isEmpty() == false) { + contexts = fieldType.getContextMappings().getNamedContexts(suggestDoc.getContexts()); + } + // collect payloads + final Map> payload = new HashMap<>(0); + List 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 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 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 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 { + @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 { diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java index 2d9307a882..7810d03004 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java @@ -34,6 +34,7 @@ import org.elasticsearch.common.xcontent.XContentType; import org.elasticsearch.index.mapper.MappedFieldType; import org.elasticsearch.index.mapper.MapperService; import org.elasticsearch.index.mapper.core.CompletionFieldMapper; +import org.elasticsearch.index.mapper.core.CompletionFieldMapper2x; import org.elasticsearch.index.query.QueryParseContext; import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.search.suggest.SuggestUtils; @@ -41,6 +42,9 @@ import org.elasticsearch.search.suggest.SuggestionBuilder; import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext; import org.elasticsearch.search.suggest.completion.context.ContextMapping; import org.elasticsearch.search.suggest.completion.context.ContextMappings; +import org.elasticsearch.search.suggest.completion2x.context.CategoryContextMapping; +import org.elasticsearch.search.suggest.completion2x.context.ContextMapping.ContextQuery; +import org.elasticsearch.search.suggest.completion2x.context.GeolocationContextMapping; import java.io.IOException; import java.util.ArrayList; @@ -219,13 +223,116 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder contextQueries = new ArrayList<>(); + + @SuppressWarnings("unchecked") + private Contexts2x addContextQuery(ContextQuery ctx) { + this.contextQueries.add(ctx); + return this; + } + + /** + * Setup a Geolocation for suggestions. See {@link GeolocationContextMapping}. + * @param lat Latitude of the location + * @param lon Longitude of the Location + * @return this + */ + @Deprecated + public Contexts2x addGeoLocation(String name, double lat, double lon, int ... precisions) { + return addContextQuery(GeolocationContextMapping.query(name, lat, lon, precisions)); + } + + /** + * Setup a Geolocation for suggestions. See {@link GeolocationContextMapping}. + * @param lat Latitude of the location + * @param lon Longitude of the Location + * @param precisions precisions as string var-args + * @return this + */ + @Deprecated + public Contexts2x addGeoLocationWithPrecision(String name, double lat, double lon, String ... precisions) { + return addContextQuery(GeolocationContextMapping.query(name, lat, lon, precisions)); + } + + /** + * Setup a Geolocation for suggestions. See {@link GeolocationContextMapping}. + * @param geohash Geohash of the location + * @return this + */ + @Deprecated + public Contexts2x addGeoLocation(String name, String geohash) { + return addContextQuery(GeolocationContextMapping.query(name, geohash)); + } + + /** + * Setup a Category for suggestions. See {@link CategoryContextMapping}. + * @param categories name of the category + * @return this + */ + @Deprecated + public Contexts2x addCategory(String name, CharSequence...categories) { + return addContextQuery(CategoryContextMapping.query(name, categories)); + } + + /** + * Setup a Category for suggestions. See {@link CategoryContextMapping}. + * @param categories name of the category + * @return this + */ + @Deprecated + public Contexts2x addCategory(String name, Iterable categories) { + return addContextQuery(CategoryContextMapping.query(name, categories)); + } + + /** + * Setup a Context Field for suggestions. See {@link CategoryContextMapping}. + * @param fieldvalues name of the category + * @return this + */ + @Deprecated + public Contexts2x addContextField(String name, CharSequence...fieldvalues) { + return addContextQuery(CategoryContextMapping.query(name, fieldvalues)); + } + + /** + * Setup a Context Field for suggestions. See {@link CategoryContextMapping}. + * @param fieldvalues name of the category + * @return this + */ + @Deprecated + public Contexts2x addContextField(String name, Iterable fieldvalues) { + return addContextQuery(CategoryContextMapping.query(name, fieldvalues)); + } + } + private static class InnerBuilder extends CompletionSuggestionBuilder { private String field; @@ -285,7 +392,12 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder> queryContexts = Collections.emptyMap(); private List payloadFields = Collections.emptyList(); + private CompletionFieldMapper2x.CompletionFieldType fieldType2x; + private List contextQueries; CompletionFieldMapper.CompletionFieldType getFieldType() { return this.fieldType; } + CompletionFieldMapper2x.CompletionFieldType getFieldType2x() { + return this.fieldType2x; + } + void setFieldType(CompletionFieldMapper.CompletionFieldType fieldType) { this.fieldType = fieldType; } @@ -116,4 +124,16 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest } return query; } + + public void setFieldType2x(CompletionFieldMapper2x.CompletionFieldType type) { + this.fieldType2x = type; + } + + public void setContextQueries(List contextQueries) { + this.contextQueries = contextQueries; + } + + public List getContextQueries() { + return contextQueries; + } } -- cgit v1.2.3