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

3 files changed, 265 insertions, 49 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 {
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<CompletionSug
                 contentBuilder.endArray();
             }
             contentBuilder.endObject();
-            contextBytes = contentBuilder.bytes();
-            return this;
+            return contexts(contentBuilder);
         } catch (IOException e) {
             throw new IllegalArgumentException(e);
         }
     }
 
+    private CompletionSuggestionBuilder contexts(XContentBuilder contextBuilder) {
+        contextBytes = contextBuilder.bytes();
+        return this;
+    }
+
+    public CompletionSuggestionBuilder contexts(Contexts2x contexts2x) {
+        Objects.requireNonNull(contexts2x, "contexts must not be null");
+        try {
+            XContentBuilder contentBuilder = XContentFactory.jsonBuilder();
+            contentBuilder.startObject();
+            for (ContextQuery contextQuery : contexts2x.contextQueries) {
+                contextQuery.toXContent(contentBuilder, EMPTY_PARAMS);
+            }
+            contentBuilder.endObject();
+            return contexts(contentBuilder);
+        } catch (IOException e) {
+            throw new IllegalArgumentException(e);
+        }
+    }
+
+    // for 2.x context support
+    public static class Contexts2x {
+        private List<ContextQuery> 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<? extends CharSequence> 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<? extends CharSequence> 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<CompletionSug
         suggestionContext.setFuzzyOptions(fuzzyOptions);
         suggestionContext.setRegexOptions(regexOptions);
         MappedFieldType mappedFieldType = mapperService.fullName(suggestionContext.getField());
-        if (mappedFieldType != null && mappedFieldType instanceof CompletionFieldMapper.CompletionFieldType) {
+        if (mappedFieldType == null ||
+            (mappedFieldType instanceof CompletionFieldMapper.CompletionFieldType == false
+                && mappedFieldType instanceof CompletionFieldMapper2x.CompletionFieldType == false)) {
+            throw new IllegalArgumentException("Field [" + suggestionContext.getField() + "] is not a completion suggest field");
+        }
+        if (mappedFieldType instanceof CompletionFieldMapper.CompletionFieldType) {
             CompletionFieldMapper.CompletionFieldType type = (CompletionFieldMapper.CompletionFieldType) mappedFieldType;
             suggestionContext.setFieldType(type);
             if (type.hasContextMappings() && contextBytes != null) {
@@ -310,9 +422,23 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
             } else if (contextBytes != null) {
                 throw new IllegalArgumentException("suggester [" + type.name() + "] doesn't expect any context");
             }
-        } else {
-            throw new IllegalArgumentException("Field [" + suggestionContext.getField() + "] is not a completion suggest field");
+        } else if (mappedFieldType instanceof CompletionFieldMapper2x.CompletionFieldType) {
+            CompletionFieldMapper2x.CompletionFieldType type = ((CompletionFieldMapper2x.CompletionFieldType) mappedFieldType);
+            suggestionContext.setFieldType2x(type);
+            if (type.requiresContext()) {
+                if (contextBytes != null) {
+                    try (XContentParser contextParser = XContentFactory.xContent(contextBytes).createParser(contextBytes)) {
+                        contextParser.nextToken();
+                        suggestionContext.setContextQueries(ContextQuery.parseQueries(type.getContextMapping(), contextParser));
+                    }
+                } else {
+                    throw new IllegalArgumentException("suggester [completion] requires context to be setup");
+                }
+            } else if (contextBytes != null) {
+                throw new IllegalArgumentException("suggester [completion] doesn't expect any context");
+            }
         }
+        assert suggestionContext.getFieldType() != null || suggestionContext.getFieldType2x() != null : "no completion field type set";
         return suggestionContext;
     }
 
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java
index 1941bc9fb8..268e0553ff 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java
@@ -21,10 +21,12 @@ package org.elasticsearch.search.suggest.completion;
 import org.apache.lucene.search.suggest.document.CompletionQuery;
 import org.elasticsearch.common.unit.Fuzziness;
 import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
+import org.elasticsearch.index.mapper.core.CompletionFieldMapper2x;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.suggest.SuggestionSearchContext;
 import org.elasticsearch.search.suggest.completion.context.ContextMapping;
 import org.elasticsearch.search.suggest.completion.context.ContextMappings;
+import org.elasticsearch.search.suggest.completion2x.context.ContextMapping.ContextQuery;
 
 import java.util.Collections;
 import java.util.List;
@@ -44,11 +46,17 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
     private RegexOptions regexOptions;
     private Map<String, List<ContextMapping.InternalQueryContext>> queryContexts = Collections.emptyMap();
     private List<String> payloadFields = Collections.emptyList();
+    private CompletionFieldMapper2x.CompletionFieldType fieldType2x;
+    private List<ContextQuery> 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<ContextQuery> contextQueries) {
+        this.contextQueries = contextQueries;
+    }
+
+    public List<ContextQuery> getContextQueries() {
+        return contextQueries;
+    }
 }