1 files changed, 276 insertions, 0 deletions
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMappings.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMappings.java
new file mode 100644
index 0000000000..76eb5a0ffd
--- /dev/null
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMappings.java
@@ -0,0 +1,276 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.suggest.completion.context;
+
+import org.apache.lucene.search.suggest.document.CompletionQuery;
+import org.apache.lucene.search.suggest.document.ContextQuery;
+import org.apache.lucene.search.suggest.document.ContextSuggestField;
+import org.apache.lucene.util.CharsRefBuilder;
+import org.elasticsearch.ElasticsearchParseException;
+import org.elasticsearch.Version;
+import org.elasticsearch.common.xcontent.ToXContent;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.mapper.DocumentMapperParser;
+import org.elasticsearch.index.mapper.ParseContext;
+import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
+
+import java.io.IOException;
+import java.util.*;
+
+import static org.elasticsearch.search.suggest.completion.context.ContextMapping.*;
+
+/**
+ * ContextMappings indexes context-enabled suggestion fields
+ * and creates context queries for defined {@link ContextMapping}s
+ * for a {@link CompletionFieldMapper}
+ */
+public class ContextMappings implements ToXContent {
+    private final List<ContextMapping> contextMappings;
+    private final Map<String, ContextMapping> contextNameMap;
+
+    public ContextMappings(List<ContextMapping> contextMappings) {
+        if (contextMappings.size() > 255) {
+            // we can support more, but max of 255 (1 byte) unique context types per suggest field
+            // seems reasonable?
+            throw new UnsupportedOperationException("Maximum of 10 context types are supported was: " + contextMappings.size());
+        }
+        this.contextMappings = contextMappings;
+        contextNameMap = new HashMap<>(contextMappings.size());
+        for (ContextMapping mapping : contextMappings) {
+            contextNameMap.put(mapping.name(), mapping);
+        }
+    }
+
+    /**
+     * @return number of context mappings
+     * held by this instance
+     */
+    public int size() {
+        return contextMappings.size();
+    }
+
+    /**
+     * Returns a context mapping by its name
+     */
+    public ContextMapping get(String name) {
+        ContextMapping contextMapping = contextNameMap.get(name);
+        if (contextMapping == null) {
+            throw new IllegalArgumentException("Unknown context name[" + name + "], must be one of " + contextNameMap.size());
+        }
+        return contextMapping;
+    }
+
+    /**
+     * Adds a context-enabled field for all the defined mappings to <code>document</code>
+     * see {@link org.elasticsearch.search.suggest.completion.context.ContextMappings.TypedContextField}
+     */
+    public void addField(ParseContext.Document document, String name, String input, int weight, Map<String, Set<CharSequence>> contexts) {
+        document.add(new TypedContextField(name, input, weight, contexts, document));
+    }
+
+    /**
+     * Field prepends context values with a suggestion
+     * Context values are associated with a type, denoted by
+     * a type id, which is prepended to the context value.
+     *
+     * Every defined context mapping yields a unique type id (index of the
+     * corresponding context mapping in the context mappings list)
+     * for all its context values
+     *
+     * The type, context and suggestion values are encoded as follows:
+     * <p>
+     *     TYPE_ID | CONTEXT_VALUE | CONTEXT_SEP | SUGGESTION_VALUE
+     * </p>
+     *
+     * Field can also use values of other indexed fields as contexts
+     * at index time
+     */
+    private class TypedContextField extends ContextSuggestField {
+        private final Map<String, Set<CharSequence>> contexts;
+        private final ParseContext.Document document;
+
+        public TypedContextField(String name, String value, int weight, Map<String, Set<CharSequence>> contexts,
+                                 ParseContext.Document document) {
+            super(name, value, weight);
+            this.contexts = contexts;
+            this.document = document;
+        }
+
+        @Override
+        protected Iterable<CharSequence> contexts() {
+            Set<CharSequence> typedContexts = new HashSet<>();
+            final CharsRefBuilder scratch = new CharsRefBuilder();
+            scratch.grow(1);
+            for (int typeId = 0; typeId < contextMappings.size(); typeId++) {
+                scratch.setCharAt(0, (char) typeId);
+                scratch.setLength(1);
+                ContextMapping mapping = contextMappings.get(typeId);
+                Set<CharSequence> contexts = new HashSet<>(mapping.parseContext(document));
+                if (this.contexts.get(mapping.name()) != null) {
+                    contexts.addAll(this.contexts.get(mapping.name()));
+                }
+                for (CharSequence context : contexts) {
+                    scratch.append(context);
+                    typedContexts.add(scratch.toCharsRef());
+                    scratch.setLength(1);
+                }
+            }
+            return typedContexts;
+        }
+    }
+
+    /**
+     * Wraps a {@link CompletionQuery} with context queries,
+     * individual context mappings adds query contexts using
+     * {@link ContextMapping#getQueryContexts(List)}s
+     *
+     * @param query base completion query to wrap
+     * @param queryContexts a map of context mapping name and collected query contexts
+     * @return a context-enabled query
+     */
+    public ContextQuery toContextQuery(CompletionQuery query, Map<String, List<CategoryQueryContext>> queryContexts) {
+        ContextQuery typedContextQuery = new ContextQuery(query);
+        if (queryContexts.isEmpty() == false) {
+            CharsRefBuilder scratch = new CharsRefBuilder();
+            scratch.grow(1);
+            for (int typeId = 0; typeId < contextMappings.size(); typeId++) {
+                scratch.setCharAt(0, (char) typeId);
+                scratch.setLength(1);
+                ContextMapping mapping = contextMappings.get(typeId);
+                List<CategoryQueryContext> queryContext = queryContexts.get(mapping.name());
+                if (queryContext != null) {
+                    for (CategoryQueryContext context : mapping.getQueryContexts(queryContext)) {
+                        scratch.append(context.context);
+                        typedContextQuery.addContext(scratch.toCharsRef(), context.boost, !context.isPrefix);
+                        scratch.setLength(1);
+                    }
+                }
+            }
+        }
+        return typedContextQuery;
+    }
+
+    /**
+     * Maps an output context list to a map of context mapping names and their values
+     *
+     * see {@link org.elasticsearch.search.suggest.completion.context.ContextMappings.TypedContextField}
+     * @return a map of context names and their values
+     *
+     */
+    public Map<String, Set<CharSequence>> getNamedContexts(List<CharSequence> contexts) {
+        Map<String, Set<CharSequence>> contextMap = new HashMap<>(contexts.size());
+        for (CharSequence typedContext : contexts) {
+            int typeId = typedContext.charAt(0);
+            assert typeId < contextMappings.size() : "Returned context has invalid type";
+            ContextMapping mapping = contextMappings.get(typeId);
+            Set<CharSequence> contextEntries = contextMap.get(mapping.name());
+            if (contextEntries == null) {
+                contextEntries = new HashSet<>();
+                contextMap.put(mapping.name(), contextEntries);
+            }
+            contextEntries.add(typedContext.subSequence(1, typedContext.length()));
+        }
+        return contextMap;
+    }
+
+    /**
+     * Loads {@link ContextMappings} from configuration
+     *
+     * Expected configuration:
+     *  List of maps representing {@link ContextMapping}
+     *  [{"name": .., "type": .., ..}, {..}]
+     *
+     */
+    public static ContextMappings load(Object configuration, Version indexVersionCreated) throws ElasticsearchParseException {
+        final List<ContextMapping> contextMappings;
+        if (configuration instanceof List) {
+            contextMappings = new ArrayList<>();
+            List<Object> configurations = (List<Object>)configuration;
+            for (Object contextConfig : configurations) {
+                contextMappings.add(load((Map<String, Object>) contextConfig, indexVersionCreated));
+            }
+            if (contextMappings.size() == 0) {
+                throw new ElasticsearchParseException("expected at least one context mapping");
+            }
+        } else if (configuration instanceof Map) {
+            contextMappings = Collections.singletonList(load(((Map<String, Object>) configuration), indexVersionCreated));
+        } else {
+            throw new ElasticsearchParseException("expected a list or an entry of context mapping");
+        }
+        return new ContextMappings(contextMappings);
+    }
+
+    private static ContextMapping load(Map<String, Object> contextConfig, Version indexVersionCreated) {
+        String name = extractRequiredValue(contextConfig, FIELD_NAME);
+        String type = extractRequiredValue(contextConfig, FIELD_TYPE);
+        final ContextMapping contextMapping;
+        switch (Type.fromString(type)) {
+            case CATEGORY:
+                contextMapping = CategoryContextMapping.load(name, contextConfig);
+                break;
+            case GEO:
+                contextMapping = GeoContextMapping.load(name, contextConfig);
+                break;
+            default:
+                throw new ElasticsearchParseException("unknown context type[" + type + "]");
+        }
+        DocumentMapperParser.checkNoRemainingFields(name, contextConfig, indexVersionCreated);
+        return contextMapping;
+    }
+
+    private static String extractRequiredValue(Map<String, Object> contextConfig, String paramName) {
+        final Object paramValue = contextConfig.get(paramName);
+        if (paramValue == null) {
+            throw new ElasticsearchParseException("missing [" + paramName + "] in context mapping");
+        }
+        contextConfig.remove(paramName);
+        return paramValue.toString();
+    }
+
+    /**
+     * Writes a list of objects specified by the defined {@link ContextMapping}s
+     *
+     * see {@link ContextMapping#toXContent(XContentBuilder, Params)}
+     */
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        for (ContextMapping contextMapping : contextMappings) {
+            builder.startObject();
+            contextMapping.toXContent(builder, params);
+            builder.endObject();
+        }
+        return builder;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(contextMappings);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == null || (obj instanceof ContextMappings) == false) {
+            return false;
+        }
+        ContextMappings other = ((ContextMappings) obj);
+        return contextMappings.equals(other.contextMappings);
+    }
+}