DATAES-536 - Add support for context suggester

Original pull request: #241

src/main/java/org/springframework/data/elasticsearch/annotations/CompletionContext.java

@@ -0,0 +1,29 @@
+package org.springframework.data.elasticsearch.annotations;
+
+import org.elasticsearch.search.suggest.completion.context.ContextMapping;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Based on reference doc - https://www.elastic.co/guide/en/elasticsearch/reference/current/suggester-context.html
+ *
+ * @author Robert Gruendler
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.FIELD)
+@Documented
+@Inherited
+public @interface CompletionContext {
+
+	String name();
+
+	ContextMapping.Type type();
+
+	String precision() default "";
+
+}

src/main/java/org/springframework/data/elasticsearch/annotations/CompletionContextType.java

@@ -0,0 +1,12 @@
+package org.springframework.data.elasticsearch.annotations;
+
+/**
+ * Based on reference doc - https://www.elastic.co/guide/en/elasticsearch/reference/current/suggester-context.html
+ *
+ * @author Robert Gruendler
+ */
+public enum CompletionContextType {
+
+	CATEGORY, GEO
+
+}

src/main/java/org/springframework/data/elasticsearch/annotations/CompletionField.java

@@ -15,12 +15,18 @@
  */
 package org.springframework.data.elasticsearch.annotations;
 
-import java.lang.annotation.*;
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
 
 /**
  * Based on the reference doc - http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-suggesters-completion.html
  *
  * @author Mewes Kochheim
+ * @author Robert Gruendler
  */
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.FIELD)
@@ -37,4 +43,6 @@ public @interface CompletionField {
 	boolean preservePositionIncrements() default true;
 
 	int maxInputLength() default 50;
+
+	CompletionContext[] contexts() default {};
 }

src/main/java/org/springframework/data/elasticsearch/core/MappingBuilder.java

@@ -24,6 +24,7 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.springframework.core.ResolvableType;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.data.annotation.Transient;
+import org.springframework.data.elasticsearch.annotations.CompletionContext;
 import org.springframework.data.elasticsearch.annotations.CompletionField;
 import org.springframework.data.elasticsearch.annotations.DateFormat;
 import org.springframework.data.elasticsearch.annotations.Field;
@@ -53,6 +54,7 @@ import static org.springframework.util.StringUtils.*;
  * @author Mark Paluch
  * @author Sascha Woo
  * @author Nordine Bittich
+ * @author Robert Gruendler
  */
 class MappingBuilder {
 
@@ -67,10 +69,14 @@ class MappingBuilder {
 	public static final String FIELD_PROPERTIES = "properties";
 	public static final String FIELD_PARENT = "_parent";
 	public static final String FIELD_COPY_TO = "copy_to";
+	public static final String FIELD_CONTEXT_NAME = "name";
+	public static final String FIELD_CONTEXT_TYPE = "type";
+	public static final String FIELD_CONTEXT_PRECISION = "precision";
 
 	public static final String COMPLETION_PRESERVE_SEPARATORS = "preserve_separators";
 	public static final String COMPLETION_PRESERVE_POSITION_INCREMENTS = "preserve_position_increments";
 	public static final String COMPLETION_MAX_INPUT_LENGTH = "max_input_length";
+	public static final String COMPLETION_CONTEXTS = "contexts";
 
 	public static final String TYPE_VALUE_KEYWORD = "keyword";
 	public static final String TYPE_VALUE_GEO_POINT = "geo_point";
@@ -79,6 +85,7 @@ class MappingBuilder {
 	public static final String TYPE_VALUE_GEO_HASH_PRECISION = "geohash_precision";
 
 	private static SimpleTypeHolder SIMPLE_TYPE_HOLDER = SimpleTypeHolder.DEFAULT;
+	private XContentBuilder xContentBuilder;
 
 	static XContentBuilder buildMapping(Class clazz, String indexType, String idFieldName, String parentType) throws IOException {
 
@@ -212,6 +219,20 @@ class MappingBuilder {
 			if (!StringUtils.isEmpty(annotation.analyzer())) {
 				xContentBuilder.field(FIELD_INDEX_ANALYZER, annotation.analyzer());
 			}
+			if (annotation.contexts().length > 0) {
+				xContentBuilder.startArray(COMPLETION_CONTEXTS);
+				for (CompletionContext context : annotation.contexts()) {
+					xContentBuilder.startObject();
+					xContentBuilder.field(FIELD_CONTEXT_NAME, context.name());
+					xContentBuilder.field(FIELD_CONTEXT_TYPE, context.type().name().toLowerCase());
+					if (context.precision().length() > 0) {
+						xContentBuilder.field(FIELD_CONTEXT_PRECISION, context.precision());
+					}
+					xContentBuilder.endObject();
+				}
+				xContentBuilder.endArray();
+			}
+
 		}
 		xContentBuilder.endObject();
 	}

src/main/java/org/springframework/data/elasticsearch/core/completion/Completion.java

@@ -2,19 +2,25 @@ package org.springframework.data.elasticsearch.core.completion;
 
 import com.fasterxml.jackson.annotation.JsonInclude;
 
+import java.util.List;
+import java.util.Map;
+
 /**
- * Based on the reference doc - http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-suggesters-completion.html
+ * Based on the reference doc -
+ * http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-suggesters-completion.html
  *
  * @author Mewes Kochheim
+ * @author Robert Gruendler
  */
 @JsonInclude(value = JsonInclude.Include.NON_NULL)
 public class Completion {
 
 	private String[] input;
+	private Map<String, List<String>> contexts;
 	private Integer weight;
 
 	private Completion() {
-		//required by mapper to instantiate object
+		// required by mapper to instantiate object
 	}
 
 	public Completion(String[] input) {
@@ -36,4 +42,13 @@ public class Completion {
 	public void setWeight(Integer weight) {
 		this.weight = weight;
 	}
+
+	public Map<String, List<String>> getContexts() {
+		return contexts;
+	}
+
+	public void setContexts(Map<String, List<String>> contexts) {
+		this.contexts = contexts;
+	}
+
 }

src/test/java/org/springframework/data/elasticsearch/core/completion/ContextCompletionEntity.java

@@ -0,0 +1,57 @@
+package org.springframework.data.elasticsearch.core.completion;
+
+import org.elasticsearch.search.suggest.completion.context.ContextMapping;
+import org.springframework.data.annotation.Id;
+import org.springframework.data.elasticsearch.annotations.CompletionContext;
+import org.springframework.data.elasticsearch.annotations.CompletionContextType;
+import org.springframework.data.elasticsearch.annotations.CompletionField;
+import org.springframework.data.elasticsearch.annotations.Document;
+
+/**
+ * @author Mewes Kochheim
+ * @author Robert Gruendler
+ */
+@Document(indexName = "test-index-context-completion", type = "context-completion-type", shards = 1, replicas = 0, refreshInterval = "-1")
+public class ContextCompletionEntity {
+
+	public static final String LANGUAGE_CATEGORY = "language";
+	@Id
+	private String id;
+	private String name;
+
+	@CompletionField(maxInputLength = 100, contexts = {
+			@CompletionContext(name = LANGUAGE_CATEGORY, type = ContextMapping.Type.CATEGORY)
+	})
+	private Completion suggest;
+
+	private ContextCompletionEntity() {
+	}
+
+	public ContextCompletionEntity(String id) {
+		this.id = id;
+	}
+
+	public String getId() {
+		return id;
+	}
+
+	public void setId(String id) {
+		this.id = id;
+	}
+
+	public String getName() {
+		return name;
+	}
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+	public Completion getSuggest() {
+		return suggest;
+	}
+
+	public void setSuggest(Completion suggest) {
+		this.suggest = suggest;
+	}
+}

src/test/java/org/springframework/data/elasticsearch/core/completion/ContextCompletionEntityBuilder.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2013-2019 the original author or authors.
+ *
+ * Licensed 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.springframework.data.elasticsearch.core.completion;
+
+import org.springframework.data.elasticsearch.core.query.IndexQuery;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * @author Robert Gruendler
+ */
+public class ContextCompletionEntityBuilder {
+
+	private ContextCompletionEntity result;
+
+	public ContextCompletionEntityBuilder(String id) {
+		result = new ContextCompletionEntity(id);
+	}
+
+	public ContextCompletionEntityBuilder name(String name) {
+		result.setName(name);
+		return this;
+	}
+
+	public ContextCompletionEntityBuilder suggest(String[] input, Map<String, List<String>> contexts) {
+		Completion suggest = new Completion(input);
+		suggest.setContexts(contexts);
+
+		result.setSuggest(suggest);
+		return this;
+	}
+
+	public ContextCompletionEntity build() {
+		return result;
+	}
+
+	public IndexQuery buildIndex() {
+		IndexQuery indexQuery = new IndexQuery();
+		indexQuery.setId(result.getId());
+		indexQuery.setObject(result);
+		return indexQuery;
+	}
+}

src/test/java/org/springframework/data/elasticsearch/core/completion/ElasticsearchTemplateCompletionWithContextsTests.java

@@ -0,0 +1,182 @@
+/*
+ * Copyright 2013-2019 the original author or authors.
+ *
+ * Licensed 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.springframework.data.elasticsearch.core.completion;
+
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.common.unit.Fuzziness;
+import org.elasticsearch.common.xcontent.ToXContent;
+import org.elasticsearch.search.suggest.SuggestBuilder;
+import org.elasticsearch.search.suggest.SuggestBuilders;
+import org.elasticsearch.search.suggest.SuggestionBuilder;
+import org.elasticsearch.search.suggest.completion.CompletionSuggestion;
+import org.elasticsearch.search.suggest.completion.CompletionSuggestionBuilder;
+import org.elasticsearch.search.suggest.completion.context.CategoryQueryContext;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.data.elasticsearch.core.ElasticsearchTemplate;
+import org.springframework.data.elasticsearch.core.query.IndexQuery;
+import org.springframework.data.elasticsearch.entities.NonDocumentEntity;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.*;
+import static org.junit.Assert.*;
+
+/**
+ * @author Robert Gruendler
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration("classpath:elasticsearch-template-test.xml")
+public class ElasticsearchTemplateCompletionWithContextsTests {
+
+	@Autowired
+	private ElasticsearchTemplate elasticsearchTemplate;
+
+	private void loadContextCompletionObjectEntities() {
+		elasticsearchTemplate.deleteIndex(ContextCompletionEntity.class);
+		elasticsearchTemplate.createIndex(ContextCompletionEntity.class);
+		elasticsearchTemplate.refresh(ContextCompletionEntity.class);
+		elasticsearchTemplate.putMapping(ContextCompletionEntity.class);
+
+		NonDocumentEntity nonDocumentEntity = new NonDocumentEntity();
+		nonDocumentEntity.setSomeField1("foo");
+		nonDocumentEntity.setSomeField2("bar");
+
+		List<IndexQuery> indexQueries = new ArrayList<>();
+
+		Map<String, List<String>> context1 = new HashMap<>();
+		context1.put(ContextCompletionEntity.LANGUAGE_CATEGORY, Arrays.asList("java", "elastic"));
+		indexQueries.add(new ContextCompletionEntityBuilder("1").name("Rizwan Idrees").suggest(new String[]{"Rizwan Idrees"}, context1).buildIndex());
+
+		Map<String, List<String>> context2 = new HashMap<>();
+		context2.put(ContextCompletionEntity.LANGUAGE_CATEGORY, Arrays.asList("kotlin", "mongo"));
+		indexQueries.add(new ContextCompletionEntityBuilder("2").name("Franck Marchand").suggest(new String[]{"Franck", "Marchand"}, context2).buildIndex());
+
+		Map<String, List<String>> context3 = new HashMap<>();
+		context3.put(ContextCompletionEntity.LANGUAGE_CATEGORY, Arrays.asList("kotlin", "elastic"));
+		indexQueries.add(new ContextCompletionEntityBuilder("3").name("Mohsin Husen").suggest(new String[]{"Mohsin", "Husen"}, context3).buildIndex());
+
+		Map<String, List<String>> context4 = new HashMap<>();
+		context4.put(ContextCompletionEntity.LANGUAGE_CATEGORY, Arrays.asList("java", "kotlin", "redis"));
+		indexQueries.add(new ContextCompletionEntityBuilder("4").name("Artur Konczak").suggest(new String[]{"Artur", "Konczak"}, context4).buildIndex());
+
+		elasticsearchTemplate.bulkIndex(indexQueries);
+		elasticsearchTemplate.refresh(ContextCompletionEntity.class);
+	}
+
+	@Test
+	public void shouldPutMappingForGivenEntity() throws Exception {
+		//given
+		Class entity = ContextCompletionEntity.class;
+		elasticsearchTemplate.createIndex(entity);
+
+		//when
+		assertThat(elasticsearchTemplate.putMapping(entity), is(true));
+	}
+
+	@Test // DATAES-536
+	public void shouldFindSuggestionsForGivenCriteriaQueryUsingContextCompletionEntityOfMongo() {
+		//given
+		loadContextCompletionObjectEntities();
+		SuggestionBuilder completionSuggestionFuzzyBuilder = SuggestBuilders.completionSuggestion("suggest").prefix("m", Fuzziness.AUTO);
+
+		Map<String, List<? extends ToXContent>> contextMap = new HashMap<>();
+		List<CategoryQueryContext> contexts = new ArrayList<>(1);
+
+		final CategoryQueryContext.Builder builder = CategoryQueryContext.builder();
+		builder.setCategory("mongo");
+		CategoryQueryContext queryContext = builder.build();
+		contexts.add(queryContext);
+		contextMap.put(ContextCompletionEntity.LANGUAGE_CATEGORY, contexts);
+
+		((CompletionSuggestionBuilder) completionSuggestionFuzzyBuilder).contexts(contextMap);
+
+		//when
+		final SearchResponse suggestResponse  = elasticsearchTemplate.suggest(new SuggestBuilder().addSuggestion("test-suggest", completionSuggestionFuzzyBuilder), ContextCompletionEntity.class);
+		assertNotNull(suggestResponse.getSuggest());
+		CompletionSuggestion completionSuggestion = suggestResponse.getSuggest().getSuggestion("test-suggest");
+		List<CompletionSuggestion.Entry.Option> options = completionSuggestion.getEntries().get(0).getOptions();
+
+		//then
+		assertThat(options.size(), is(1));
+		assertThat(options.get(0).getText().string(), isOneOf("Marchand"));
+	}
+
+	@Test // DATAES-536
+	public void shouldFindSuggestionsForGivenCriteriaQueryUsingContextCompletionEntityOfElastic() {
+		//given
+		loadContextCompletionObjectEntities();
+		SuggestionBuilder completionSuggestionFuzzyBuilder = SuggestBuilders.completionSuggestion("suggest").prefix("m", Fuzziness.AUTO);
+
+		Map<String, List<? extends ToXContent>> contextMap = new HashMap<>();
+		List<CategoryQueryContext> contexts = new ArrayList<>(1);
+
+		final CategoryQueryContext.Builder builder = CategoryQueryContext.builder();
+		builder.setCategory("elastic");
+		CategoryQueryContext queryContext = builder.build();
+		contexts.add(queryContext);
+		contextMap.put(ContextCompletionEntity.LANGUAGE_CATEGORY, contexts);
+
+		((CompletionSuggestionBuilder) completionSuggestionFuzzyBuilder).contexts(contextMap);
+
+		//when
+		final SearchResponse suggestResponse  = elasticsearchTemplate.suggest(new SuggestBuilder().addSuggestion("test-suggest", completionSuggestionFuzzyBuilder), ContextCompletionEntity.class);
+		assertNotNull(suggestResponse.getSuggest());
+		CompletionSuggestion completionSuggestion = suggestResponse.getSuggest().getSuggestion("test-suggest");
+		List<CompletionSuggestion.Entry.Option> options = completionSuggestion.getEntries().get(0).getOptions();
+
+		//then
+		assertThat(options.size(), is(1));
+		assertThat(options.get(0).getText().string(), isOneOf( "Mohsin"));
+	}
+
+	@Test // DATAES-536
+	public void shouldFindSuggestionsForGivenCriteriaQueryUsingContextCompletionEntityOfKotlin() {
+		//given
+		loadContextCompletionObjectEntities();
+		SuggestionBuilder completionSuggestionFuzzyBuilder = SuggestBuilders.completionSuggestion("suggest").prefix("m", Fuzziness.AUTO);
+
+		Map<String, List<? extends ToXContent>> contextMap = new HashMap<>();
+		List<CategoryQueryContext> contexts = new ArrayList<>(1);
+
+		final CategoryQueryContext.Builder builder = CategoryQueryContext.builder();
+		builder.setCategory("kotlin");
+		CategoryQueryContext queryContext = builder.build();
+		contexts.add(queryContext);
+		contextMap.put(ContextCompletionEntity.LANGUAGE_CATEGORY, contexts);
+
+		((CompletionSuggestionBuilder) completionSuggestionFuzzyBuilder).contexts(contextMap);
+
+		//when
+		final SearchResponse suggestResponse  = elasticsearchTemplate.suggest(new SuggestBuilder().addSuggestion("test-suggest", completionSuggestionFuzzyBuilder), ContextCompletionEntity.class);
+		assertNotNull(suggestResponse.getSuggest());
+		CompletionSuggestion completionSuggestion = suggestResponse.getSuggest().getSuggestion("test-suggest");
+		List<CompletionSuggestion.Entry.Option> options = completionSuggestion.getEntries().get(0).getOptions();
+
+		//then
+		assertThat(options.size(), is(2));
+		assertThat(options.get(0).getText().string(), isOneOf("Marchand", "Mohsin"));
+		assertThat(options.get(1).getText().string(), isOneOf("Marchand", "Mohsin"));
+	}
+}