construct suggestion context from query context

This commit is contained in:
Areek Zillur 2016-03-05 00:33:36 -05:00
parent 1264f37a1b
commit 5bb72dbcd2
22 changed files with 515 additions and 453 deletions

@ -716,13 +716,6 @@ public abstract class StreamInput extends InputStream {
return readNamedWriteable(SuggestionBuilder.class);
}
/**
* Reads a completion {@link QueryContext} from the current stream
*/
public QueryContext readCompletionSuggestionQueryContext() throws IOException {
return readNamedWriteable(QueryContext.class);
}
/**
* Reads a {@link org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder} from the current stream
*/

@ -740,10 +740,4 @@ public abstract class StreamOutput extends OutputStream {
writeNamedWriteable(suggestion);
}
/**
* Writes a completion {@link QueryContext} to the current stream
*/
public void writeCompletionSuggestionQueryContext(QueryContext queryContext) throws IOException {
writeNamedWriteable(queryContext);
}
}

@ -21,7 +21,6 @@ package org.elasticsearch.search.suggest.completion;
import org.apache.lucene.analysis.Analyzer;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
@ -34,14 +33,8 @@ import org.elasticsearch.index.query.RegexpFlag;
import org.elasticsearch.search.suggest.SuggestContextParser;
import org.elasticsearch.search.suggest.SuggestUtils.Fields;
import org.elasticsearch.search.suggest.SuggestionSearchContext;
import org.elasticsearch.search.suggest.completion.context.ContextMapping;
import org.elasticsearch.search.suggest.completion.context.ContextMappings;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
* Parses query options for {@link CompletionSuggester}
@ -74,27 +67,20 @@ import java.util.Map;
public class CompletionSuggestParser implements SuggestContextParser {
private static ObjectParser<CompletionSuggestionContext, ContextAndSuggest> TLP_PARSER = new ObjectParser<>(CompletionSuggestionBuilder.SUGGESTION_NAME, null);
private static ObjectParser<RegexOptions.Builder, ContextAndSuggest> REGEXP_PARSER = new ObjectParser<>(RegexOptions.REGEX_OPTIONS.getPreferredName(), RegexOptions.Builder::new);
private static ObjectParser<FuzzyOptions.Builder, ContextAndSuggest> FUZZY_PARSER = new ObjectParser<>(FuzzyOptions.FUZZY_OPTIONS.getPreferredName(), FuzzyOptions.Builder::new);
static {
FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyMinLength, FuzzyOptions.MIN_LENGTH_FIELD);
FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setMaxDeterminizedStates, FuzzyOptions.MAX_DETERMINIZED_STATES_FIELD);
FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setUnicodeAware, FuzzyOptions.UNICODE_AWARE_FIELD);
FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyPrefixLength, FuzzyOptions.PREFIX_LENGTH_FIELD);
FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setTranspositions, FuzzyOptions.TRANSPOSITION_FIELD);
FUZZY_PARSER.declareValue((a, b) -> {
try {
a.setFuzziness(Fuzziness.parse(b).asDistance());
} catch (IOException e) {
throw new ElasticsearchException(e);
}
}, Fuzziness.FIELD);
REGEXP_PARSER.declareInt(RegexOptions.Builder::setMaxDeterminizedStates, RegexOptions.MAX_DETERMINIZED_STATES);
REGEXP_PARSER.declareStringOrNull(RegexOptions.Builder::setFlags, RegexOptions.FLAGS_VALUE);
TLP_PARSER.declareStringArray(CompletionSuggestionContext::setPayloadFields, CompletionSuggestionBuilder.PAYLOAD_FIELD);
TLP_PARSER.declareObjectOrDefault(CompletionSuggestionContext::setFuzzyOptionsBuilder, FUZZY_PARSER, FuzzyOptions.Builder::new, FuzzyOptions.FUZZY_OPTIONS);
TLP_PARSER.declareObject(CompletionSuggestionContext::setRegexOptionsBuilder, REGEXP_PARSER, RegexOptions.REGEX_OPTIONS);
TLP_PARSER.declareField((parser, completionSuggestionContext, context) -> {
if (parser.currentToken() == XContentParser.Token.VALUE_BOOLEAN) {
if (parser.booleanValue()) {
completionSuggestionContext.setFuzzyOptions(new FuzzyOptions.Builder().build());
}
} else {
completionSuggestionContext.setFuzzyOptions(FuzzyOptions.parse(parser));
}
},
FuzzyOptions.FUZZY_OPTIONS, ObjectParser.ValueType.OBJECT_OR_BOOLEAN);
TLP_PARSER.declareField((parser, completionSuggestionContext, context) -> completionSuggestionContext.setRegexOptions(RegexOptions.parse(parser)),
RegexOptions.REGEX_OPTIONS, ObjectParser.ValueType.OBJECT);
TLP_PARSER.declareString(SuggestionSearchContext.SuggestionContext::setField, Fields.FIELD);
TLP_PARSER.declareField((p, v, c) -> {
String analyzerName = p.text();
@ -132,7 +118,7 @@ public class CompletionSuggestParser implements SuggestContextParser {
}
@Override
public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, QueryShardContext shardContext) throws IOException {
public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, QueryShardContext shardContext) throws IOException {
MapperService mapperService = shardContext.getMapperService();
final CompletionSuggestionContext suggestion = new CompletionSuggestionContext(shardContext);
final ContextAndSuggest contextAndSuggest = new ContextAndSuggest(mapperService);
@ -146,28 +132,12 @@ public class CompletionSuggestParser implements SuggestContextParser {
if (type.hasContextMappings() == false && contextParser != null) {
throw new IllegalArgumentException("suggester [" + type.name() + "] doesn't expect any context");
}
Map<String, List<ContextMapping.QueryContext>> queryContexts = Collections.emptyMap();
if (type.hasContextMappings() && contextParser != null) {
ContextMappings contextMappings = type.getContextMappings();
contextParser.nextToken();
queryContexts = new HashMap<>(contextMappings.size());
assert contextParser.currentToken() == XContentParser.Token.START_OBJECT;
XContentParser.Token currentToken;
String currentFieldName;
while ((currentToken = contextParser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (currentToken == XContentParser.Token.FIELD_NAME) {
currentFieldName = contextParser.currentName();
final ContextMapping mapping = contextMappings.get(currentFieldName);
queryContexts.put(currentFieldName, mapping.parseQueryContext(contextParser));
}
}
contextParser.close();
}
suggestion.setQueryContexts(CompletionSuggestionBuilder.parseQueryContexts(contextParser, type));
suggestion.setFieldType(type);
suggestion.setQueryContexts(queryContexts);
return suggestion;
} else {
throw new IllegalArgumentException("Field [" + suggestion.getField() + "] is not a completion suggest field");
}
}
}

@ -18,29 +18,37 @@
*/
package org.elasticsearch.search.suggest.completion;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentParser;
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.query.QueryParseContext;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.suggest.SuggestUtils;
import org.elasticsearch.search.suggest.SuggestionBuilder;
import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
import org.elasticsearch.search.suggest.completion.context.CategoryQueryContext;
import org.elasticsearch.search.suggest.completion.context.GeoQueryContext;
import org.elasticsearch.search.suggest.completion.context.ContextMapping;
import org.elasticsearch.search.suggest.completion.context.ContextMappings;
import org.elasticsearch.search.suggest.completion.context.QueryContext;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
/**
* Defines a suggest command based on a prefix, typically to provide "auto-complete" functionality
@ -55,10 +63,40 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
static final ParseField PAYLOAD_FIELD = new ParseField("payload");
static final ParseField CONTEXTS_FIELD = new ParseField("contexts", "context");
private static ObjectParser<CompletionSuggestionBuilder, Void> TLP_PARSER =
new ObjectParser<>(CompletionSuggestionBuilder.SUGGESTION_NAME, null);
static {
TLP_PARSER.declareStringArray(CompletionSuggestionBuilder::payload, CompletionSuggestionBuilder.PAYLOAD_FIELD);
TLP_PARSER.declareField((parser, completionSuggestionContext, context) -> {
if (parser.currentToken() == XContentParser.Token.VALUE_BOOLEAN) {
if (parser.booleanValue()) {
completionSuggestionContext.fuzzyOptions = new FuzzyOptions.Builder().build();
}
} else {
completionSuggestionContext.fuzzyOptions = FuzzyOptions.parse(parser);
}
},
FuzzyOptions.FUZZY_OPTIONS, ObjectParser.ValueType.OBJECT_OR_BOOLEAN);
TLP_PARSER.declareField((parser, completionSuggestionContext, context) ->
completionSuggestionContext.regexOptions = RegexOptions.parse(parser),
RegexOptions.REGEX_OPTIONS, ObjectParser.ValueType.OBJECT);
TLP_PARSER.declareString(CompletionSuggestionBuilder::field, SuggestUtils.Fields.FIELD);
TLP_PARSER.declareString(CompletionSuggestionBuilder::analyzer, SuggestUtils.Fields.ANALYZER);
TLP_PARSER.declareInt(CompletionSuggestionBuilder::size, SuggestUtils.Fields.SIZE);
TLP_PARSER.declareInt(CompletionSuggestionBuilder::shardSize, SuggestUtils.Fields.SHARD_SIZE);
TLP_PARSER.declareField((p, v, c) -> {
// Copy the current structure. We will parse, once the mapping is provided
XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
builder.copyCurrentStructure(p);
v.contextBytes = builder.bytes();
p.skipChildren();
}, CompletionSuggestionBuilder.CONTEXTS_FIELD, ObjectParser.ValueType.OBJECT); // context is deprecated
}
private FuzzyOptions fuzzyOptions;
private RegexOptions regexOptions;
private final Map<String, List<QueryContext>> queryContexts = new HashMap<>();
private final Set<String> payloadFields = new HashSet<>();
private BytesReference contextBytes = null;
private List<String> payloadFields = Collections.emptyList();
public CompletionSuggestionBuilder(String fieldname) {
super(fieldname);
@ -117,36 +155,33 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
* Note: Only doc values enabled fields are supported
*/
public CompletionSuggestionBuilder payload(List<String> fields) {
this.payloadFields.addAll(fields);
this.payloadFields = fields;
return this;
}
/**
* Sets query contexts for a category context
* @param name of the category context to execute on
* @param queryContexts a list of {@link CategoryQueryContext}
* Sets query contexts for completion
* @param queryContexts named query contexts
* see {@link org.elasticsearch.search.suggest.completion.context.CategoryQueryContext}
* and {@link org.elasticsearch.search.suggest.completion.context.GeoQueryContext}
*/
public CompletionSuggestionBuilder categoryContexts(String name, CategoryQueryContext... queryContexts) {
return contexts(name, queryContexts);
}
/**
* Sets query contexts for a geo context
* @param name of the geo context to execute on
* @param queryContexts a list of {@link GeoQueryContext}
*/
public CompletionSuggestionBuilder geoContexts(String name, GeoQueryContext... queryContexts) {
return contexts(name, queryContexts);
}
private CompletionSuggestionBuilder contexts(String name, QueryContext... queryContexts) {
List<QueryContext> contexts = this.queryContexts.get(name);
if (contexts == null) {
contexts = new ArrayList<>(2);
this.queryContexts.put(name, contexts);
public CompletionSuggestionBuilder contexts(Map<String, List<? extends QueryContext>> queryContexts) {
try {
XContentBuilder contentBuilder = XContentFactory.jsonBuilder();
contentBuilder.startObject();
for (Map.Entry<String, List<? extends QueryContext>> contextEntry : queryContexts.entrySet()) {
contentBuilder.startArray(contextEntry.getKey());
for (ToXContent queryContext : contextEntry.getValue()) {
queryContext.toXContent(contentBuilder, EMPTY_PARAMS);
}
contentBuilder.endArray();
}
contentBuilder.endObject();
contextBytes = contentBuilder.bytes();
return this;
} catch (IOException e) {
throw new IllegalArgumentException(e);
}
Collections.addAll(contexts, queryContexts);
return this;
}
@Override
@ -164,33 +199,44 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
if (regexOptions != null) {
regexOptions.toXContent(builder, params);
}
if (queryContexts.isEmpty() == false) {
builder.startObject(CONTEXTS_FIELD.getPreferredName());
for (Map.Entry<String, List<QueryContext>> entry : this.queryContexts.entrySet()) {
builder.startArray(entry.getKey());
for (ToXContent queryContext : entry.getValue()) {
queryContext.toXContent(builder, params);
}
builder.endArray();
}
builder.endObject();
if (contextBytes != null) {
XContentParser contextParser = XContentFactory.xContent(XContentType.JSON).createParser(contextBytes);
builder.field(CONTEXTS_FIELD.getPreferredName());
builder.copyCurrentStructure(contextParser);
}
return builder;
}
@Override
protected CompletionSuggestionBuilder innerFromXContent(QueryParseContext parseContext) throws IOException {
// NORELEASE implement parsing logic
throw new UnsupportedOperationException();
CompletionSuggestionBuilder builder = new CompletionSuggestionBuilder();
TLP_PARSER.parse(parseContext.parser(), builder);
return builder;
}
@Override
protected SuggestionContext innerBuild(QueryShardContext context) throws IOException {
CompletionSuggestionContext suggestionContext = new CompletionSuggestionContext(context);
// copy over common settings to each suggestion builder
populateCommonFields(context.getMapperService(), suggestionContext);
// NORELEASE
// still need to populate CompletionSuggestionContext's specific settings
final MapperService mapperService = context.getMapperService();
populateCommonFields(mapperService, suggestionContext);
suggestionContext.setPayloadFields(payloadFields);
suggestionContext.setFuzzyOptions(fuzzyOptions);
suggestionContext.setRegexOptions(regexOptions);
MappedFieldType mappedFieldType = mapperService.fullName(suggestionContext.getField());
if (mappedFieldType == null) {
throw new ElasticsearchException("Field [" + suggestionContext.getField() + "] is not a completion suggest field");
} else if (mappedFieldType instanceof CompletionFieldMapper.CompletionFieldType) {
CompletionFieldMapper.CompletionFieldType type = (CompletionFieldMapper.CompletionFieldType) mappedFieldType;
if (type.hasContextMappings() && contextBytes != null) {
XContentParser contextParser = XContentFactory.xContent(contextBytes).createParser(contextBytes);
suggestionContext.setQueryContexts(parseQueryContexts(contextParser, type));
} 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");
}
return suggestionContext;
}
@ -217,18 +263,10 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
if (regexOptions != null) {
regexOptions.writeTo(out);
}
boolean queryContextsExists = queryContexts.isEmpty() == false;
boolean queryContextsExists = contextBytes != null;
out.writeBoolean(queryContextsExists);
if (queryContextsExists) {
out.writeVInt(queryContexts.size());
for (Map.Entry<String, List<QueryContext>> namedQueryContexts : queryContexts.entrySet()) {
out.writeString(namedQueryContexts.getKey());
List<QueryContext> queryContexts = namedQueryContexts.getValue();
out.writeVInt(queryContexts.size());
for (QueryContext queryContext : queryContexts) {
out.writeCompletionSuggestionQueryContext(queryContext);
}
}
out.writeBytesReference(contextBytes);
}
}
@ -237,9 +275,11 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
CompletionSuggestionBuilder completionSuggestionBuilder = new CompletionSuggestionBuilder(fieldname);
if (in.readBoolean()) {
int numPayloadField = in.readVInt();
List<String> payloadFields = new ArrayList<>(numPayloadField);
for (int i = 0; i < numPayloadField; i++) {
completionSuggestionBuilder.payloadFields.add(in.readString());
payloadFields.add(in.readString());
}
completionSuggestionBuilder.payloadFields = payloadFields;
}
if (in.readBoolean()) {
completionSuggestionBuilder.fuzzyOptions = FuzzyOptions.readFuzzyOptions(in);
@ -248,30 +288,43 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
completionSuggestionBuilder.regexOptions = RegexOptions.readRegexOptions(in);
}
if (in.readBoolean()) {
int numNamedQueryContexts = in.readVInt();
for (int i = 0; i < numNamedQueryContexts; i++) {
String queryContextName = in.readString();
int numQueryContexts = in.readVInt();
List<QueryContext> queryContexts = new ArrayList<>(numQueryContexts);
for (int j = 0; j < numQueryContexts; j++) {
queryContexts.add(in.readCompletionSuggestionQueryContext());
}
completionSuggestionBuilder.queryContexts.put(queryContextName, queryContexts);
}
completionSuggestionBuilder.contextBytes = in.readBytesReference();
}
return completionSuggestionBuilder;
}
@Override
protected boolean doEquals(CompletionSuggestionBuilder other) {
return Objects.equals(payloadFields, other.payloadFields) &&
return Objects.equals(payloadFields, other.payloadFields) &&
Objects.equals(fuzzyOptions, other.fuzzyOptions) &&
Objects.equals(regexOptions, other.regexOptions) &&
Objects.equals(queryContexts, other.queryContexts);
Objects.equals(contextBytes, other.contextBytes);
}
@Override
protected int doHashCode() {
return Objects.hash(payloadFields, fuzzyOptions, regexOptions, queryContexts);
return Objects.hash(payloadFields, fuzzyOptions, regexOptions, contextBytes);
}
static Map<String, List<ContextMapping.InternalQueryContext>> parseQueryContexts(
XContentParser contextParser, CompletionFieldMapper.CompletionFieldType type) throws IOException {
Map<String, List<ContextMapping.InternalQueryContext>> queryContexts = Collections.emptyMap();
if (type.hasContextMappings() && contextParser != null) {
ContextMappings contextMappings = type.getContextMappings();
contextParser.nextToken();
queryContexts = new HashMap<>(contextMappings.size());
assert contextParser.currentToken() == XContentParser.Token.START_OBJECT;
XContentParser.Token currentToken;
String currentFieldName;
while ((currentToken = contextParser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (currentToken == XContentParser.Token.FIELD_NAME) {
currentFieldName = contextParser.currentName();
final ContextMapping mapping = contextMappings.get(currentFieldName);
queryContexts.put(currentFieldName, mapping.parseQueryContext(contextParser));
}
}
contextParser.close();
}
return queryContexts;
}
}

@ -44,7 +44,7 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
private CompletionFieldMapper.CompletionFieldType fieldType;
private FuzzyOptions fuzzyOptions;
private RegexOptions regexOptions;
private Map<String, List<ContextMapping.QueryContext>> queryContexts = Collections.emptyMap();
private Map<String, List<ContextMapping.InternalQueryContext>> queryContexts = Collections.emptyMap();
private Set<String> payloadFields = Collections.emptySet();
CompletionFieldMapper.CompletionFieldType getFieldType() {
@ -55,15 +55,15 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
this.fieldType = fieldType;
}
void setRegexOptionsBuilder(RegexOptions.Builder regexOptionsBuilder) {
this.regexOptions = regexOptionsBuilder.build();
void setRegexOptions(RegexOptions regexOptions) {
this.regexOptions = regexOptions;
}
void setFuzzyOptionsBuilder(FuzzyOptions.Builder fuzzyOptionsBuilder) {
this.fuzzyOptions = fuzzyOptionsBuilder.build();
void setFuzzyOptions(FuzzyOptions fuzzyOptions) {
this.fuzzyOptions = fuzzyOptions;
}
void setQueryContexts(Map<String, List<ContextMapping.QueryContext>> queryContexts) {
void setQueryContexts(Map<String, List<ContextMapping.InternalQueryContext>> queryContexts) {
this.queryContexts = queryContexts;
}
@ -79,6 +79,18 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
return payloadFields;
}
public FuzzyOptions getFuzzyOptions() {
return fuzzyOptions;
}
public RegexOptions getRegexOptions() {
return regexOptions;
}
public Map<String, List<ContextMapping.InternalQueryContext>> getQueryContexts() {
return queryContexts;
}
CompletionQuery toQuery() {
CompletionFieldMapper.CompletionFieldType fieldType = getFieldType();
final CompletionQuery query;

@ -21,13 +21,16 @@ package org.elasticsearch.search.suggest.completion;
import org.apache.lucene.search.suggest.document.FuzzyCompletionQuery;
import org.apache.lucene.util.automaton.Operations;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import java.io.IOException;
import java.util.Objects;
@ -43,6 +46,22 @@ public class FuzzyOptions implements ToXContent, Writeable<FuzzyOptions> {
static final ParseField UNICODE_AWARE_FIELD = new ParseField("unicode_aware");
static final ParseField MAX_DETERMINIZED_STATES_FIELD = new ParseField("max_determinized_states");
static ObjectParser<FuzzyOptions.Builder, Void> FUZZY_PARSER = new ObjectParser<>(FUZZY_OPTIONS.getPreferredName(), Builder::new);
static {
FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyMinLength, MIN_LENGTH_FIELD);
FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setMaxDeterminizedStates, MAX_DETERMINIZED_STATES_FIELD);
FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setUnicodeAware, UNICODE_AWARE_FIELD);
FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyPrefixLength, PREFIX_LENGTH_FIELD);
FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setTranspositions, TRANSPOSITION_FIELD);
FUZZY_PARSER.declareValue((a, b) -> {
try {
a.setFuzziness(Fuzziness.parse(b).asDistance());
} catch (IOException e) {
throw new ElasticsearchException(e);
}
}, Fuzziness.FIELD);
}
private int editDistance;
private boolean transpositions;
private int fuzzyMinLength;
@ -63,6 +82,10 @@ public class FuzzyOptions implements ToXContent, Writeable<FuzzyOptions> {
private FuzzyOptions() {
}
public static FuzzyOptions parse(XContentParser parser) throws IOException {
return FUZZY_PARSER.parse(parser).build();
}
public static Builder builder() {
return new Builder();
}

@ -21,12 +21,15 @@ package org.elasticsearch.search.suggest.completion;
import org.apache.lucene.util.automaton.Operations;
import org.apache.lucene.util.automaton.RegExp;
import org.elasticsearch.ElasticsearchParseException;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.RegexpFlag;
import java.io.IOException;
@ -39,6 +42,25 @@ public class RegexOptions implements ToXContent, Writeable<RegexOptions> {
static final ParseField REGEX_OPTIONS = new ParseField(NAME);
static final ParseField FLAGS_VALUE = new ParseField("flags", "flags_value");
static final ParseField MAX_DETERMINIZED_STATES = new ParseField("max_determinized_states");
private static ObjectParser<RegexOptions.Builder, Void> REGEXP_PARSER =
new ObjectParser<>(REGEX_OPTIONS.getPreferredName(), RegexOptions.Builder::new);
static {
REGEXP_PARSER.declareInt(RegexOptions.Builder::setMaxDeterminizedStates, MAX_DETERMINIZED_STATES);
REGEXP_PARSER.declareField((parser, builder, aVoid) -> {
if (parser.currentToken() == XContentParser.Token.VALUE_STRING) {
builder.setFlags(parser.text());
} else if (parser.currentToken() == XContentParser.Token.VALUE_NUMBER) {
builder.setFlagsValue(parser.intValue());
} else {
throw new ElasticsearchParseException(REGEX_OPTIONS.getPreferredName()
+ " " + FLAGS_VALUE.getPreferredName() + " supports string or number");
}
}, FLAGS_VALUE, ObjectParser.ValueType.VALUE);
REGEXP_PARSER.declareStringOrNull(RegexOptions.Builder::setFlags, FLAGS_VALUE);
}
private int flagsValue;
private int maxDeterminizedStates;
@ -69,6 +91,10 @@ public class RegexOptions implements ToXContent, Writeable<RegexOptions> {
return new Builder();
}
public static RegexOptions parse(XContentParser parser) throws IOException {
return REGEXP_PARSER.parse(parser).build();
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
@ -135,6 +161,11 @@ public class RegexOptions implements ToXContent, Writeable<RegexOptions> {
return this;
}
private Builder setFlagsValue(int flagsValue) {
this.flagsValue = flagsValue;
return this;
}
/**
* Sets the maximum automaton states allowed for the regular expression expansion
*/

@ -36,6 +36,7 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
/**
* A {@link ContextMapping} that uses a simple string as a criteria
@ -44,7 +45,7 @@ import java.util.Set;
* {@link CategoryQueryContext} defines options for constructing
* a unit of query context for this context type
*/
public class CategoryContextMapping extends ContextMapping {
public class CategoryContextMapping extends ContextMapping<CategoryQueryContext> {
private static final String FIELD_FIELDNAME = "path";
@ -137,6 +138,11 @@ public class CategoryContextMapping extends ContextMapping {
return (values == null) ? Collections.<CharSequence>emptySet() : values;
}
@Override
protected CategoryQueryContext prototype() {
return CategoryQueryContext.PROTOTYPE;
}
/**
* Parse a list of {@link CategoryQueryContext}
* using <code>parser</code>. A QueryContexts accepts one of the following forms:
@ -154,19 +160,13 @@ public class CategoryContextMapping extends ContextMapping {
* </ul>
*/
@Override
public List<QueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
List<QueryContext> queryContexts = new ArrayList<>();
Token token = parser.nextToken();
if (token == Token.START_OBJECT || token == Token.VALUE_STRING) {
CategoryQueryContext parse = CategoryQueryContext.PROTOTYPE.fromXContext(parser);
queryContexts.add(new QueryContext(parse.getCategory(), parse.getBoost(), parse.isPrefix()));
} else if (token == Token.START_ARRAY) {
while (parser.nextToken() != Token.END_ARRAY) {
CategoryQueryContext parse = CategoryQueryContext.PROTOTYPE.fromXContext(parser);
queryContexts.add(new QueryContext(parse.getCategory(), parse.getBoost(), parse.isPrefix()));
}
}
return queryContexts;
public List<InternalQueryContext> toInternalQueryContexts(List<CategoryQueryContext> queryContexts) {
List<InternalQueryContext> internalInternalQueryContexts = new ArrayList<>(queryContexts.size());
internalInternalQueryContexts.addAll(
queryContexts.stream()
.map(queryContext -> new InternalQueryContext(queryContext.getCategory(), queryContext.getBoost(), queryContext.isPrefix()))
.collect(Collectors.toList()));
return internalInternalQueryContexts;
}
@Override

@ -21,14 +21,11 @@ package org.elasticsearch.search.suggest.completion.context;
import org.elasticsearch.ElasticsearchParseException;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import java.io.IOException;
import java.util.Collections;
import java.util.Objects;
import static org.elasticsearch.search.suggest.completion.context.CategoryContextMapping.CONTEXT_BOOST;
@ -98,11 +95,6 @@ public final class CategoryQueryContext implements QueryContext {
return result;
}
@Override
public String getWriteableName() {
return NAME;
}
private static ObjectParser<Builder, Void> CATEGORY_PARSER = new ObjectParser<>(NAME, null);
static {
CATEGORY_PARSER.declareString(Builder::setCategory, new ParseField(CONTEXT_VALUE));
@ -134,22 +126,6 @@ public final class CategoryQueryContext implements QueryContext {
return builder;
}
@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeBoolean(isPrefix);
out.writeVInt(boost);
out.writeString(category);
}
@Override
public QueryContext readFrom(StreamInput in) throws IOException {
Builder builder = new Builder();
builder.isPrefix = in.readBoolean();
builder.boost = in.readVInt();
builder.category = in.readString();
return builder.build();
}
public static class Builder {
private String category;
private boolean isPrefix = false;

@ -23,11 +23,13 @@ import org.elasticsearch.ElasticsearchParseException;
import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.common.xcontent.XContentParser.Token;
import org.elasticsearch.common.xcontent.json.JsonXContent;
import org.elasticsearch.index.mapper.ParseContext;
import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Set;
@ -38,7 +40,7 @@ import java.util.Set;
*
* Implementations have to define how contexts are parsed at query/index time
*/
public abstract class ContextMapping implements ToXContent {
public abstract class ContextMapping<T extends QueryContext> implements ToXContent {
public static final String FIELD_TYPE = "type";
public static final String FIELD_NAME = "name";
@ -94,10 +96,25 @@ public abstract class ContextMapping implements ToXContent {
*/
protected abstract Set<CharSequence> parseContext(ParseContext.Document document);
protected abstract T prototype();
/**
* Parses query contexts for this mapper
*/
public abstract List<QueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException;
public List<InternalQueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
List<T> queryContexts = new ArrayList<>();
Token token = parser.nextToken();
if (token == Token.START_OBJECT || token == Token.VALUE_STRING) {
queryContexts.add((T) prototype().fromXContext(parser));
} else if (token == Token.START_ARRAY) {
while (parser.nextToken() != Token.END_ARRAY) {
queryContexts.add((T) prototype().fromXContext(parser));
}
}
return toInternalQueryContexts(queryContexts);
}
protected abstract List<InternalQueryContext> toInternalQueryContexts(List<T> queryContexts);
/**
* Implementations should add specific configurations
@ -136,17 +153,38 @@ public abstract class ContextMapping implements ToXContent {
}
}
public static class QueryContext {
public static class InternalQueryContext {
public final String context;
public final int boost;
public final boolean isPrefix;
public QueryContext(String context, int boost, boolean isPrefix) {
public InternalQueryContext(String context, int boost, boolean isPrefix) {
this.context = context;
this.boost = boost;
this.isPrefix = isPrefix;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
InternalQueryContext that = (InternalQueryContext) o;
if (boost != that.boost) return false;
if (isPrefix != that.isPrefix) return false;
return context != null ? context.equals(that.context) : that.context == null;
}
@Override
public int hashCode() {
int result = context != null ? context.hashCode() : 0;
result = 31 * result + boost;
result = 31 * result + (isPrefix ? 1 : 0);
return result;
}
@Override
public String toString() {
return "QueryContext{" +

@ -152,7 +152,7 @@ public class ContextMappings implements ToXContent {
* @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<ContextMapping.QueryContext>> queryContexts) {
public ContextQuery toContextQuery(CompletionQuery query, Map<String, List<ContextMapping.InternalQueryContext>> queryContexts) {
ContextQuery typedContextQuery = new ContextQuery(query);
if (queryContexts.isEmpty() == false) {
CharsRefBuilder scratch = new CharsRefBuilder();
@ -161,9 +161,9 @@ public class ContextMappings implements ToXContent {
scratch.setCharAt(0, (char) typeId);
scratch.setLength(1);
ContextMapping mapping = contextMappings.get(typeId);
List<ContextMapping.QueryContext> queryContext = queryContexts.get(mapping.name());
if (queryContext != null) {
for (ContextMapping.QueryContext context : queryContext) {
List<ContextMapping.InternalQueryContext> internalQueryContext = queryContexts.get(mapping.name());
if (internalQueryContext != null) {
for (ContextMapping.InternalQueryContext context : internalQueryContext) {
scratch.append(context.context);
typedContextQuery.addContext(scratch.toCharsRef(), context.boost, !context.isPrefix);
scratch.setLength(1);

@ -42,6 +42,7 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import static org.apache.lucene.spatial.util.GeoHashUtils.addNeighbors;
import static org.apache.lucene.spatial.util.GeoHashUtils.stringEncode;
@ -56,7 +57,7 @@ import static org.apache.lucene.spatial.util.GeoHashUtils.stringEncode;
* {@link GeoQueryContext} defines the options for constructing
* a unit of query context for this context type
*/
public class GeoContextMapping extends ContextMapping {
public class GeoContextMapping extends ContextMapping<GeoQueryContext> {
public static final String FIELD_PRECISION = "precision";
public static final String FIELD_FIELDNAME = "path";
@ -221,6 +222,11 @@ public class GeoContextMapping extends ContextMapping {
return locations;
}
@Override
protected GeoQueryContext prototype() {
return GeoQueryContext.PROTOTYPE;
}
/**
* Parse a list of {@link GeoQueryContext}
* using <code>parser</code>. A QueryContexts accepts one of the following forms:
@ -245,17 +251,8 @@ public class GeoContextMapping extends ContextMapping {
* see {@link GeoUtils#parseGeoPoint(String, GeoPoint)} for GEO POINT
*/
@Override
public List<QueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
List<GeoQueryContext> queryContexts = new ArrayList<>();
Token token = parser.nextToken();
if (token == Token.START_OBJECT || token == Token.VALUE_STRING) {
queryContexts.add(GeoQueryContext.PROTOTYPE.fromXContext(parser));
} else if (token == Token.START_ARRAY) {
while (parser.nextToken() != Token.END_ARRAY) {
queryContexts.add(GeoQueryContext.PROTOTYPE.fromXContext(parser));
}
}
List<QueryContext> queryContextList = new ArrayList<>();
public List<InternalQueryContext> toInternalQueryContexts(List<GeoQueryContext> queryContexts) {
List<InternalQueryContext> internalQueryContextList = new ArrayList<>();
for (GeoQueryContext queryContext : queryContexts) {
int minPrecision = Math.min(this.precision, queryContext.getPrecision());
GeoPoint point = queryContext.getGeoPoint();
@ -265,19 +262,20 @@ public class GeoContextMapping extends ContextMapping {
if (queryContext.getNeighbours().isEmpty() && geoHash.length() == this.precision) {
addNeighbors(geoHash, locations);
} else if (queryContext.getNeighbours().isEmpty() == false) {
for (Integer neighbourPrecision : queryContext.getNeighbours()) {
if (neighbourPrecision < geoHash.length()) {
queryContext.getNeighbours().stream()
.filter(neighbourPrecision -> neighbourPrecision < geoHash.length())
.forEach(neighbourPrecision -> {
String truncatedGeoHash = geoHash.substring(0, neighbourPrecision);
locations.add(truncatedGeoHash);
addNeighbors(truncatedGeoHash, locations);
}
}
}
for (String location : locations) {
queryContextList.add(new QueryContext(location, queryContext.getBoost(), location.length() < this.precision));
});
}
internalQueryContextList.addAll(
locations.stream()
.map(location -> new InternalQueryContext(location, queryContext.getBoost(), location.length() < this.precision))
.collect(Collectors.toList()));
}
return queryContextList;
return internalQueryContextList;
}
@Override
@ -301,7 +299,7 @@ public class GeoContextMapping extends ContextMapping {
private int precision = DEFAULT_PRECISION;
private String fieldName = null;
protected Builder(String name) {
public Builder(String name) {
super(name);
}

@ -23,14 +23,11 @@ import org.elasticsearch.ElasticsearchParseException;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.geo.GeoPoint;
import org.elasticsearch.common.geo.GeoUtils;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
@ -114,11 +111,6 @@ public final class GeoQueryContext implements QueryContext {
return new Builder();
}
@Override
public String getWriteableName() {
return NAME;
}
private static ObjectParser<GeoQueryContext.Builder, Void> GEO_CONTEXT_PARSER = new ObjectParser<>(NAME, null);
static {
GEO_CONTEXT_PARSER.declareField((parser, geoQueryContext, geoContextMapping) -> geoQueryContext.setGeoPoint(GeoUtils.parseGeoPoint(parser)), new ParseField(CONTEXT_VALUE), ObjectParser.ValueType.OBJECT);
@ -159,33 +151,6 @@ public final class GeoQueryContext implements QueryContext {
return builder;
}
@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeGeoPoint(geoPoint);
out.writeVInt(boost);
out.writeInt(precision);
out.writeVInt(neighbours.size());
for (Integer neighbour : neighbours) {
out.writeVInt(neighbour);
}
}
@Override
public QueryContext readFrom(StreamInput in) throws IOException {
Builder builder = new Builder();
builder.geoPoint = in.readGeoPoint();
builder.boost = in.readVInt();
builder.precision = in.readInt();
int nNeighbour = in.readVInt();
if (nNeighbour != 0) {
builder.neighbours = new ArrayList<>(nNeighbour);
for (int i = 0; i < nNeighbour; i++) {
builder.neighbours.add(in.readVInt());
}
}
return builder.build();
}
public static class Builder {
private GeoPoint geoPoint;
private int boost = 1;

@ -19,7 +19,6 @@
package org.elasticsearch.search.suggest.completion.context;
import org.elasticsearch.common.io.stream.NamedWriteable;
import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentParser;
@ -28,7 +27,7 @@ import java.io.IOException;
/**
* Interface for serializing/de-serializing completion query context
*/
public interface QueryContext extends ToXContent, NamedWriteable<QueryContext> {
public interface QueryContext extends ToXContent {
QueryContext fromXContext(XContentParser parser) throws IOException;
}

@ -21,6 +21,7 @@ package org.elasticsearch.search.suggest;
import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
import org.elasticsearch.common.ParseFieldMatcher;
import org.elasticsearch.common.collect.Tuple;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput;
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
@ -160,7 +161,8 @@ public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBui
assertTrue("suggestion builder is not equal to self", firstBuilder.equals(firstBuilder));
assertThat("same suggestion builder's hashcode returns different values if called multiple times", firstBuilder.hashCode(),
equalTo(firstBuilder.hashCode()));
assertThat("different suggestion builders should not be equal", mutate(firstBuilder), not(equalTo(firstBuilder)));
final SB mutate = mutate(firstBuilder);
assertThat("different suggestion builders should not be equal", mutate, not(equalTo(firstBuilder)));
SB secondBuilder = serializedCopy(firstBuilder);
assertTrue("suggestion builder is not equal to self", secondBuilder.equals(secondBuilder));
@ -211,6 +213,22 @@ public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBui
}
}
protected Tuple<MapperService, SB> mockMapperServiceAndSuggestionBuilder(
IndexSettings idxSettings, AnalysisService mockAnalysisService, SB suggestBuilder) {
final MapperService mapperService = new MapperService(idxSettings, mockAnalysisService, null,
new IndicesModule().getMapperRegistry(), null) {
@Override
public MappedFieldType fullName(String fullName) {
StringFieldType type = new StringFieldType();
if (randomBoolean()) {
type.setSearchAnalyzer(new NamedAnalyzer("foo", new WhitespaceAnalyzer()));
}
return type;
}
};
return new Tuple<>(mapperService, suggestBuilder);
}
/**
* parses random suggestion builder via old parseElement method and via
* build, comparing the results for equality
@ -226,30 +244,21 @@ public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBui
}
};
MapperService mockMapperService = new MapperService(idxSettings, mockAnalysisService, null,
new IndicesModule().getMapperRegistry(), null) {
@Override
public MappedFieldType fullName(String fullName) {
StringFieldType type = new StringFieldType();
if (randomBoolean()) {
type.setSearchAnalyzer(new NamedAnalyzer("foo", new WhitespaceAnalyzer()));
}
return type;
}
};
QueryShardContext mockShardContext = new QueryShardContext(idxSettings, null, null, mockMapperService, null, scriptService, null) {
@Override
public MappedFieldType fieldMapper(String name) {
StringFieldMapper.Builder builder = new StringFieldMapper.Builder(name);
return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
}
};
mockShardContext.setMapUnmappedFieldAsString(true);
for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
SuggestBuilder suggestBuilder = new SuggestBuilder();
SB suggestionBuilder = randomTestBuilder();
Tuple<MapperService, SB> mapperServiceSBTuple =
mockMapperServiceAndSuggestionBuilder(idxSettings, mockAnalysisService, suggestionBuilder);
suggestionBuilder = mapperServiceSBTuple.v2();
QueryShardContext mockShardContext = new QueryShardContext(idxSettings,
null, null, mapperServiceSBTuple.v1(), null, scriptService, null) {
@Override
public MappedFieldType fieldMapper(String name) {
StringFieldMapper.Builder builder = new StringFieldMapper.Builder(name);
return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
}
};
mockShardContext.setMapUnmappedFieldAsString(true);
suggestBuilder.addSuggestion(randomAsciiOfLength(10), suggestionBuilder);
if (suggestionBuilder.text() == null) {

@ -36,12 +36,14 @@ import org.elasticsearch.search.suggest.completion.context.ContextBuilder;
import org.elasticsearch.search.suggest.completion.context.ContextMapping;
import org.elasticsearch.search.suggest.completion.context.GeoContextMapping;
import org.elasticsearch.search.suggest.completion.context.GeoQueryContext;
import org.elasticsearch.search.suggest.completion.context.QueryContext;
import org.elasticsearch.test.ESIntegTestCase;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
@ -180,7 +182,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
indexRandom(true, indexRequestBuilders);
ensureYellow(INDEX);
CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat0").build());
.contexts(Collections.singletonMap("cat", Collections.singletonList(CategoryQueryContext.builder().setCategory("cat0").build())));
assertSuggestions("foo", prefix, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
}
@ -207,9 +209,9 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
indexRandom(true, indexRequestBuilders);
ensureYellow(INDEX);
CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
.categoryContexts("cat",
CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
CategoryQueryContext.builder().setCategory("cat1").build()
.contexts(Collections.singletonMap("cat",
Arrays.asList(CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
CategoryQueryContext.builder().setCategory("cat1").build()))
);
assertSuggestions("foo", prefix, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion2");
}
@ -267,24 +269,21 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
// filter only on context cat
CompletionSuggestionBuilder catFilterSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
catFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat0").build());
catFilterSuggest.contexts(Collections.singletonMap("cat", Collections.singletonList(CategoryQueryContext.builder().setCategory("cat0").build())));
assertSuggestions("foo", catFilterSuggest, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
// filter only on context type
CompletionSuggestionBuilder typeFilterSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
typeFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build(),
CategoryQueryContext.builder().setCategory("type1").build());
typeFilterSuggest.contexts(Collections.singletonMap("type", Arrays.asList(CategoryQueryContext.builder().setCategory("type2").build(),
CategoryQueryContext.builder().setCategory("type1").build())));
assertSuggestions("foo", typeFilterSuggest, "suggestion9", "suggestion6", "suggestion5", "suggestion2", "suggestion1");
CompletionSuggestionBuilder multiContextFilterSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
// query context order should never matter
if (randomBoolean()) {
multiContextFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build());
multiContextFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat2").build());
} else {
multiContextFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat2").build());
multiContextFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build());
}
Map<String, List<? extends QueryContext>> contextMap = new HashMap<>();
contextMap.put("type", Collections.singletonList(CategoryQueryContext.builder().setCategory("type2").build()));
contextMap.put("cat", Collections.singletonList(CategoryQueryContext.builder().setCategory("cat2").build()));
multiContextFilterSuggest.contexts(contextMap);
assertSuggestions("foo", multiContextFilterSuggest, "suggestion6", "suggestion2");
}
@ -315,36 +314,33 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
// boost only on context cat
CompletionSuggestionBuilder catBoostSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
catBoostSuggest.categoryContexts("cat",
CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
CategoryQueryContext.builder().setCategory("cat1").build());
catBoostSuggest.contexts(Collections.singletonMap("cat",
Arrays.asList(
CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
CategoryQueryContext.builder().setCategory("cat1").build())));
assertSuggestions("foo", catBoostSuggest, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion2");
// boost only on context type
CompletionSuggestionBuilder typeBoostSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
typeBoostSuggest.categoryContexts("type",
CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
typeBoostSuggest.contexts(Collections.singletonMap("type",
Arrays.asList(
CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
CategoryQueryContext.builder().setCategory("type1").setBoost(4).build())));
assertSuggestions("foo", typeBoostSuggest, "suggestion9", "suggestion5", "suggestion6", "suggestion1", "suggestion2");
// boost on both contexts
CompletionSuggestionBuilder multiContextBoostSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
// query context order should never matter
if (randomBoolean()) {
multiContextBoostSuggest.categoryContexts("type",
CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
multiContextBoostSuggest.categoryContexts("cat",
CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
CategoryQueryContext.builder().setCategory("cat1").build());
} else {
multiContextBoostSuggest.categoryContexts("cat",
CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
CategoryQueryContext.builder().setCategory("cat1").build());
multiContextBoostSuggest.categoryContexts("type",
CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
}
Map<String, List<? extends QueryContext>> contextMap = new HashMap<>();
contextMap.put("type", Arrays.asList(
CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
CategoryQueryContext.builder().setCategory("type1").setBoost(4).build())
);
contextMap.put("cat", Arrays.asList(
CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
CategoryQueryContext.builder().setCategory("cat1").build())
);
multiContextBoostSuggest.contexts(contextMap);
assertSuggestions("foo", multiContextBoostSuggest, "suggestion9", "suggestion6", "suggestion5", "suggestion2", "suggestion1");
}
@ -463,7 +459,8 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
CompletionSuggestionBuilder geoFilteringPrefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
.geoContexts("geo", GeoQueryContext.builder().setGeoPoint(new GeoPoint(geoPoints[0])).build());
.contexts(Collections.singletonMap("geo", Collections.singletonList(
GeoQueryContext.builder().setGeoPoint(new GeoPoint(geoPoints[0])).build())));
assertSuggestions("foo", geoFilteringPrefix, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
}
@ -497,7 +494,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
GeoQueryContext context1 = GeoQueryContext.builder().setGeoPoint(geoPoints[0]).setBoost(2).build();
GeoQueryContext context2 = GeoQueryContext.builder().setGeoPoint(geoPoints[1]).build();
CompletionSuggestionBuilder geoBoostingPrefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
.geoContexts("geo", context1, context2);
.contexts(Collections.singletonMap("geo", Arrays.asList(context1, context2)));
assertSuggestions("foo", geoBoostingPrefix, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion7");
}
@ -528,7 +525,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
indexRandom(true, indexRequestBuilders);
ensureYellow(INDEX);
CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
.geoContexts("geo", GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.2263, 4.543)).build());
.contexts(Collections.singletonMap("geo", Collections.singletonList(GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.2263, 4.543)).build())));
assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
}
@ -569,7 +566,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
CompletionSuggestionBuilder geoNeighbourPrefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
.geoContexts("geo", GeoQueryContext.builder().setGeoPoint(GeoPoint.fromGeohash(geohash)).build());
.contexts(Collections.singletonMap("geo", Collections.singletonList(GeoQueryContext.builder().setGeoPoint(GeoPoint.fromGeohash(geohash)).build())));
assertSuggestions("foo", geoNeighbourPrefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
}
@ -626,7 +623,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
String suggestionName = randomAsciiOfLength(10);
CompletionSuggestionBuilder context = SuggestBuilders.completionSuggestion(FIELD).text("h").size(10)
.geoContexts("st", GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.52, 13.4)).build());
.contexts(Collections.singletonMap("st", Collections.singletonList(GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.52, 13.4)).build())));
SuggestResponse suggestResponse = client().prepareSuggest(INDEX).addSuggestion(suggestionName, context).get();
assertEquals(suggestResponse.getSuggest().size(), 1);

@ -190,11 +190,11 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
XContentBuilder builder = jsonBuilder().value("context1");
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
CategoryContextMapping mapping = ContextBuilder.category("cat").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(1));
assertThat(queryContexts.get(0).context, equalTo("context1"));
assertThat(queryContexts.get(0).boost, equalTo(1));
assertThat(queryContexts.get(0).isPrefix, equalTo(false));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(1));
assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
assertThat(internalQueryContexts.get(0).boost, equalTo(1));
assertThat(internalQueryContexts.get(0).isPrefix, equalTo(false));
}
public void testQueryContextParsingArray() throws Exception {
@ -204,14 +204,14 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
.endArray();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
CategoryContextMapping mapping = ContextBuilder.category("cat").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(2));
assertThat(queryContexts.get(0).context, equalTo("context1"));
assertThat(queryContexts.get(0).boost, equalTo(1));
assertThat(queryContexts.get(0).isPrefix, equalTo(false));
assertThat(queryContexts.get(1).context, equalTo("context2"));
assertThat(queryContexts.get(1).boost, equalTo(1));
assertThat(queryContexts.get(1).isPrefix, equalTo(false));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(2));
assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
assertThat(internalQueryContexts.get(0).boost, equalTo(1));
assertThat(internalQueryContexts.get(0).isPrefix, equalTo(false));
assertThat(internalQueryContexts.get(1).context, equalTo("context2"));
assertThat(internalQueryContexts.get(1).boost, equalTo(1));
assertThat(internalQueryContexts.get(1).isPrefix, equalTo(false));
}
public void testQueryContextParsingObject() throws Exception {
@ -222,11 +222,11 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
.endObject();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
CategoryContextMapping mapping = ContextBuilder.category("cat").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(1));
assertThat(queryContexts.get(0).context, equalTo("context1"));
assertThat(queryContexts.get(0).boost, equalTo(10));
assertThat(queryContexts.get(0).isPrefix, equalTo(true));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(1));
assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
assertThat(internalQueryContexts.get(0).boost, equalTo(10));
assertThat(internalQueryContexts.get(0).isPrefix, equalTo(true));
}
@ -245,14 +245,14 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
.endArray();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
CategoryContextMapping mapping = ContextBuilder.category("cat").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(2));
assertThat(queryContexts.get(0).context, equalTo("context1"));
assertThat(queryContexts.get(0).boost, equalTo(2));
assertThat(queryContexts.get(0).isPrefix, equalTo(true));
assertThat(queryContexts.get(1).context, equalTo("context2"));
assertThat(queryContexts.get(1).boost, equalTo(3));
assertThat(queryContexts.get(1).isPrefix, equalTo(false));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(2));
assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
assertThat(internalQueryContexts.get(0).boost, equalTo(2));
assertThat(internalQueryContexts.get(0).isPrefix, equalTo(true));
assertThat(internalQueryContexts.get(1).context, equalTo("context2"));
assertThat(internalQueryContexts.get(1).boost, equalTo(3));
assertThat(internalQueryContexts.get(1).isPrefix, equalTo(false));
}
public void testQueryContextParsingMixed() throws Exception {
@ -266,14 +266,14 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
.endArray();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
CategoryContextMapping mapping = ContextBuilder.category("cat").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(2));
assertThat(queryContexts.get(0).context, equalTo("context1"));
assertThat(queryContexts.get(0).boost, equalTo(2));
assertThat(queryContexts.get(0).isPrefix, equalTo(true));
assertThat(queryContexts.get(1).context, equalTo("context2"));
assertThat(queryContexts.get(1).boost, equalTo(1));
assertThat(queryContexts.get(1).isPrefix, equalTo(false));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(2));
assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
assertThat(internalQueryContexts.get(0).boost, equalTo(2));
assertThat(internalQueryContexts.get(0).isPrefix, equalTo(true));
assertThat(internalQueryContexts.get(1).context, equalTo("context2"));
assertThat(internalQueryContexts.get(1).boost, equalTo(1));
assertThat(internalQueryContexts.get(1).isPrefix, equalTo(false));
}
public void testParsingContextFromDocument() throws Exception {

@ -38,25 +38,6 @@ public class CategoryQueryContextTests extends QueryContextTestCase<CategoryQuer
return randomCategoryQueryContext();
}
@Override
protected CategoryQueryContext createMutation(CategoryQueryContext original) throws IOException {
final CategoryQueryContext.Builder builder = CategoryQueryContext.builder();
builder.setCategory(original.getCategory()).setBoost(original.getBoost()).setPrefix(original.isPrefix());
switch (randomIntBetween(0, 2)) {
case 0:
builder.setCategory(randomValueOtherThan(original.getCategory(), () -> randomAsciiOfLength(10)));
break;
case 1:
builder.setBoost(randomValueOtherThan(original.getBoost(), () -> randomIntBetween(1, 5)));
break;
case 2:
builder.setPrefix(!original.isPrefix());
break;
}
return builder.build();
}
@Override
protected CategoryQueryContext prototype() {
return CategoryQueryContext.PROTOTYPE;

@ -19,10 +19,21 @@
package org.elasticsearch.search.suggest.completion;
import org.elasticsearch.common.collect.Tuple;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.index.IndexSettings;
import org.elasticsearch.index.analysis.AnalysisService;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.index.mapper.MapperService;
import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
import org.elasticsearch.indices.IndicesModule;
import org.elasticsearch.search.suggest.AbstractSuggestionBuilderTestCase;
import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
import org.elasticsearch.search.suggest.completion.context.CategoryContextMapping;
import org.elasticsearch.search.suggest.completion.context.CategoryQueryContext;
import org.elasticsearch.search.suggest.completion.context.ContextMapping;
import org.elasticsearch.search.suggest.completion.context.ContextMappings;
import org.elasticsearch.search.suggest.completion.context.GeoContextMapping;
import org.elasticsearch.search.suggest.completion.context.GeoQueryContext;
import org.elasticsearch.search.suggest.completion.context.QueryContext;
import org.junit.BeforeClass;
@ -30,18 +41,28 @@ import org.junit.BeforeClass;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import static org.hamcrest.core.IsInstanceOf.instanceOf;
public class CompletionSuggesterBuilderTests extends AbstractSuggestionBuilderTestCase<CompletionSuggestionBuilder> {
@BeforeClass
public static void initQueryContexts() {
namedWriteableRegistry.registerPrototype(QueryContext.class, CategoryQueryContext.PROTOTYPE);
namedWriteableRegistry.registerPrototype(QueryContext.class, GeoQueryContext.PROTOTYPE);
}
@Override
protected CompletionSuggestionBuilder randomSuggestionBuilder() {
return randomSuggestionBuilderWithContextInfo().builder;
}
private static class BuilderAndInfo {
CompletionSuggestionBuilder builder;
List<String> catContexts = new ArrayList<>();
List<String> geoContexts = new ArrayList<>();
}
private BuilderAndInfo randomSuggestionBuilderWithContextInfo() {
final BuilderAndInfo builderAndInfo = new BuilderAndInfo();
CompletionSuggestionBuilder testBuilder = new CompletionSuggestionBuilder(randomAsciiOfLengthBetween(2, 20));
switch (randomIntBetween(0, 3)) {
case 0:
@ -60,38 +81,66 @@ public class CompletionSuggesterBuilderTests extends AbstractSuggestionBuilderTe
List<String> payloads = new ArrayList<>();
Collections.addAll(payloads, generateRandomStringArray(5, 10, false, false));
maybeSet(testBuilder::payload, payloads);
Map<String, List<? extends QueryContext>> contextMap = new HashMap<>();
if (randomBoolean()) {
int numContext = randomIntBetween(1, 5);
CategoryQueryContext[] contexts = new CategoryQueryContext[numContext];
List<CategoryQueryContext> contexts = new ArrayList<>(numContext);
for (int i = 0; i < numContext; i++) {
contexts[i] = CategoryQueryContextTests.randomCategoryQueryContext();
contexts.add(CategoryQueryContextTests.randomCategoryQueryContext());
}
testBuilder.categoryContexts(randomAsciiOfLength(10), contexts);
String name = randomAsciiOfLength(10);
contextMap.put(name, contexts);
builderAndInfo.catContexts.add(name);
}
if (randomBoolean()) {
int numContext = randomIntBetween(1, 5);
GeoQueryContext[] contexts = new GeoQueryContext[numContext];
List<GeoQueryContext> contexts = new ArrayList<>(numContext);
for (int i = 0; i < numContext; i++) {
contexts[i] = GeoQueryContextTests.randomGeoQueryContext();
contexts.add(GeoQueryContextTests.randomGeoQueryContext());
}
testBuilder.geoContexts(randomAsciiOfLength(10), contexts);
String name = randomAsciiOfLength(10);
contextMap.put(name, contexts);
builderAndInfo.geoContexts.add(name);
}
return testBuilder;
testBuilder.contexts(contextMap);
builderAndInfo.builder = testBuilder;
return builderAndInfo;
}
@Override
protected void assertSuggestionContext(SuggestionContext oldSuggestion, SuggestionContext newSuggestion) {
assertThat(oldSuggestion, instanceOf(CompletionSuggestionContext.class));
assertThat(newSuggestion, instanceOf(CompletionSuggestionContext.class));
CompletionSuggestionContext oldCompletionSuggestion = (CompletionSuggestionContext) oldSuggestion;
CompletionSuggestionContext newCompletionSuggestion = (CompletionSuggestionContext) newSuggestion;
assertEquals(oldCompletionSuggestion.getPayloadFields(), newCompletionSuggestion.getPayloadFields());
assertEquals(oldCompletionSuggestion.getFuzzyOptions(), newCompletionSuggestion.getFuzzyOptions());
assertEquals(oldCompletionSuggestion.getRegexOptions(), newCompletionSuggestion.getRegexOptions());
assertEquals(oldCompletionSuggestion.getQueryContexts(), newCompletionSuggestion.getQueryContexts());
}
@Override
public void testBuild() throws IOException {
// skip for now
}
@Override
public void testFromXContent() throws IOException {
// skip for now
protected Tuple<MapperService, CompletionSuggestionBuilder> mockMapperServiceAndSuggestionBuilder(
IndexSettings idxSettings, AnalysisService mockAnalysisService, CompletionSuggestionBuilder suggestBuilder) {
final BuilderAndInfo builderAndInfo = randomSuggestionBuilderWithContextInfo();
final MapperService mapperService = new MapperService(idxSettings, mockAnalysisService, null,
new IndicesModule().getMapperRegistry(), null) {
@Override
public MappedFieldType fullName(String fullName) {
CompletionFieldMapper.CompletionFieldType type = new CompletionFieldMapper.CompletionFieldType();
List<ContextMapping> contextMappings = builderAndInfo.catContexts.stream()
.map(catContext -> new CategoryContextMapping.Builder(catContext).build())
.collect(Collectors.toList());
contextMappings.addAll(builderAndInfo.geoContexts.stream()
.map(geoContext -> new GeoContextMapping.Builder(geoContext).build())
.collect(Collectors.toList()));
type.setContextMappings(new ContextMappings(contextMappings));
return type;
}
};
final CompletionSuggestionBuilder builder = builderAndInfo.builder;
return new Tuple<>(mapperService, builder);
}
@Override
@ -103,20 +152,20 @@ public class CompletionSuggesterBuilderTests extends AbstractSuggestionBuilderTe
builder.payload(payloads);
break;
case 1:
int numCategoryContext = randomIntBetween(1, 5);
CategoryQueryContext[] categoryContexts = new CategoryQueryContext[numCategoryContext];
for (int i = 0; i < numCategoryContext; i++) {
categoryContexts[i] = CategoryQueryContextTests.randomCategoryQueryContext();
int nCatContext = randomIntBetween(1, 5);
List<CategoryQueryContext> contexts = new ArrayList<>(nCatContext);
for (int i = 0; i < nCatContext; i++) {
contexts.add(CategoryQueryContextTests.randomCategoryQueryContext());
}
builder.categoryContexts(randomAsciiOfLength(10), categoryContexts);
builder.contexts(Collections.singletonMap(randomAsciiOfLength(10), contexts));
break;
case 2:
int numGeoContext = randomIntBetween(1, 5);
GeoQueryContext[] geoContexts = new GeoQueryContext[numGeoContext];
for (int i = 0; i < numGeoContext; i++) {
geoContexts[i] = GeoQueryContextTests.randomGeoQueryContext();
int nGeoContext = randomIntBetween(1, 5);
List<GeoQueryContext> geoContexts = new ArrayList<>(nGeoContext);
for (int i = 0; i < nGeoContext; i++) {
geoContexts.add(GeoQueryContextTests.randomGeoQueryContext());
}
builder.geoContexts(randomAsciiOfLength(10), geoContexts);
builder.contexts(Collections.singletonMap(randomAsciiOfLength(10), geoContexts));
break;
case 3:
builder.prefix(randomAsciiOfLength(10), FuzzyOptionsTests.randomFuzzyOptions());

@ -202,15 +202,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
XContentBuilder builder = jsonBuilder().value("ezs42e44yx96");
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
GeoContextMapping mapping = ContextBuilder.geo("geo").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(1 + 8));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(1 + 8));
Collection<String> locations = new ArrayList<>();
locations.add("ezs42e");
addNeighbors("ezs42e", GeoContextMapping.DEFAULT_PRECISION, locations);
for (ContextMapping.QueryContext queryContext : queryContexts) {
assertThat(queryContext.context, isIn(locations));
assertThat(queryContext.boost, equalTo(1));
assertThat(queryContext.isPrefix, equalTo(false));
for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
assertThat(internalQueryContext.context, isIn(locations));
assertThat(internalQueryContext.boost, equalTo(1));
assertThat(internalQueryContext.isPrefix, equalTo(false));
}
}
@ -221,15 +221,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
.endObject();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
GeoContextMapping mapping = ContextBuilder.geo("geo").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(1 + 8));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(1 + 8));
Collection<String> locations = new ArrayList<>();
locations.add("wh0n94");
addNeighbors("wh0n94", GeoContextMapping.DEFAULT_PRECISION, locations);
for (ContextMapping.QueryContext queryContext : queryContexts) {
assertThat(queryContext.context, isIn(locations));
assertThat(queryContext.boost, equalTo(1));
assertThat(queryContext.isPrefix, equalTo(false));
for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
assertThat(internalQueryContext.context, isIn(locations));
assertThat(internalQueryContext.boost, equalTo(1));
assertThat(internalQueryContext.isPrefix, equalTo(false));
}
}
@ -244,8 +244,8 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
.endObject();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
GeoContextMapping mapping = ContextBuilder.geo("geo").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
Collection<String> locations = new ArrayList<>();
locations.add("wh0n94");
locations.add("w");
@ -254,10 +254,10 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
addNeighbors("wh", 2, locations);
locations.add("wh0");
addNeighbors("wh0", 3, locations);
for (ContextMapping.QueryContext queryContext : queryContexts) {
assertThat(queryContext.context, isIn(locations));
assertThat(queryContext.boost, equalTo(10));
assertThat(queryContext.isPrefix, equalTo(queryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
assertThat(internalQueryContext.context, isIn(locations));
assertThat(internalQueryContext.boost, equalTo(10));
assertThat(internalQueryContext.isPrefix, equalTo(internalQueryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
}
}
@ -282,8 +282,8 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
.endArray();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
GeoContextMapping mapping = ContextBuilder.geo("geo").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8 + 1 + 1 + 8));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8 + 1 + 1 + 8));
Collection<String> firstLocations = new ArrayList<>();
firstLocations.add("wh0n94");
firstLocations.add("w");
@ -296,15 +296,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
secondLocations.add("w5cx04");
secondLocations.add("w5cx0");
addNeighbors("w5cx0", 5, secondLocations);
for (ContextMapping.QueryContext queryContext : queryContexts) {
if (firstLocations.contains(queryContext.context)) {
assertThat(queryContext.boost, equalTo(10));
} else if (secondLocations.contains(queryContext.context)) {
assertThat(queryContext.boost, equalTo(2));
for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
if (firstLocations.contains(internalQueryContext.context)) {
assertThat(internalQueryContext.boost, equalTo(10));
} else if (secondLocations.contains(internalQueryContext.context)) {
assertThat(internalQueryContext.boost, equalTo(2));
} else {
fail(queryContext.context + " was not expected");
fail(internalQueryContext.context + " was not expected");
}
assertThat(queryContext.isPrefix, equalTo(queryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
assertThat(internalQueryContext.isPrefix, equalTo(internalQueryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
}
}
@ -325,8 +325,8 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
.endArray();
XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
GeoContextMapping mapping = ContextBuilder.geo("geo").build();
List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
assertThat(queryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
assertThat(internalQueryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
Collection<String> firstLocations = new ArrayList<>();
firstLocations.add("wh0n94");
firstLocations.add("w");
@ -336,15 +336,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
Collection<String> secondLocations = new ArrayList<>();
secondLocations.add("w5cx04");
addNeighbors("w5cx04", 6, secondLocations);
for (ContextMapping.QueryContext queryContext : queryContexts) {
if (firstLocations.contains(queryContext.context)) {
assertThat(queryContext.boost, equalTo(10));
} else if (secondLocations.contains(queryContext.context)) {
assertThat(queryContext.boost, equalTo(1));
for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
if (firstLocations.contains(internalQueryContext.context)) {
assertThat(internalQueryContext.boost, equalTo(10));
} else if (secondLocations.contains(internalQueryContext.context)) {
assertThat(internalQueryContext.boost, equalTo(1));
} else {
fail(queryContext.context + " was not expected");
fail(internalQueryContext.context + " was not expected");
}
assertThat(queryContext.isPrefix, equalTo(queryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
assertThat(internalQueryContext.isPrefix, equalTo(internalQueryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
}
}
}

@ -49,35 +49,6 @@ public class GeoQueryContextTests extends QueryContextTestCase<GeoQueryContext>
return randomGeoQueryContext();
}
@Override
protected GeoQueryContext createMutation(GeoQueryContext original) throws IOException {
final GeoQueryContext.Builder builder = GeoQueryContext.builder();
builder.setGeoPoint(original.getGeoPoint()).setBoost(original.getBoost())
.setNeighbours(original.getNeighbours()).setPrecision(original.getPrecision());
switch (randomIntBetween(0, 3)) {
case 0:
builder.setGeoPoint(randomValueOtherThan(original.getGeoPoint() ,() ->
new GeoPoint(randomDouble(), randomDouble())));
break;
case 1:
builder.setBoost(randomValueOtherThan(original.getBoost() ,() -> randomIntBetween(1, 5)));
break;
case 2:
builder.setPrecision(randomValueOtherThan(original.getPrecision() ,() -> randomIntBetween(1, 12)));
break;
case 3:
builder.setNeighbours(randomValueOtherThan(original.getNeighbours(), () -> {
List<Integer> newNeighbours = new ArrayList<>();
for (int i = 0; i < randomIntBetween(1, 12); i++) {
newNeighbours.add(randomIntBetween(1, 12));
}
return newNeighbours;
}));
break;
}
return builder.build();
}
@Override
protected GeoQueryContext prototype() {
return GeoQueryContext.PROTOTYPE;

@ -26,24 +26,27 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.search.suggest.completion.context.QueryContext;
import org.elasticsearch.test.ESTestCase;
import java.io.IOException;
import static junit.framework.TestCase.assertEquals;
public abstract class QueryContextTestCase<QC extends QueryContext> extends WritableTestCase<QC> {
public abstract class QueryContextTestCase<QC extends QueryContext> extends ESTestCase {
private static final int NUMBER_OF_RUNS = 20;
/**
* create random model that is put under test
*/
protected abstract QC createTestModel();
/**
* query context prototype to read serialized format
*/
protected abstract QC prototype();
@Override
protected QC readFrom(StreamInput in) throws IOException {
return (QC) prototype().readFrom(in);
}
public void testToXContext() throws IOException {
for (int i = 0; i < NUMBER_OF_RUNS; i++) {
QueryContext toXContent = createTestModel();