diff --git a/core/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java b/core/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java index 1557c266bd4..3bcc01daccb 100644 --- a/core/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java +++ b/core/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java @@ -32,6 +32,7 @@ import org.elasticsearch.search.aggregations.AbstractAggregationBuilder; import org.elasticsearch.search.builder.SearchSourceBuilder; import org.elasticsearch.search.fetch.innerhits.InnerHitsBuilder; import org.elasticsearch.search.highlight.HighlightBuilder; +import org.elasticsearch.search.rescore.RescoreBaseBuilder; import org.elasticsearch.search.rescore.RescoreBuilder; import org.elasticsearch.search.sort.SortBuilder; import org.elasticsearch.search.sort.SortOrder; @@ -391,25 +392,25 @@ public class SearchRequestBuilder extends ActionRequestBuilder(); diff --git a/core/src/main/java/org/elasticsearch/search/rescore/QueryRescorer.java b/core/src/main/java/org/elasticsearch/search/rescore/QueryRescorer.java index 7f95ff10824..319055639ac 100644 --- a/core/src/main/java/org/elasticsearch/search/rescore/QueryRescorer.java +++ b/core/src/main/java/org/elasticsearch/search/rescore/QueryRescorer.java @@ -27,7 +27,7 @@ import org.apache.lucene.search.TopDocs; import org.elasticsearch.common.ParseField; import org.elasticsearch.common.xcontent.ObjectParser; import org.elasticsearch.common.xcontent.XContentParser; -import org.elasticsearch.index.query.ParsedQuery; +import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.search.internal.ContextIndexSearcher; import org.elasticsearch.search.internal.SearchContext; @@ -120,17 +120,17 @@ public final class QueryRescorer implements Rescorer { } } - private static final ObjectParser RESCORE_PARSER = new ObjectParser<>("query", null); + private static final ObjectParser RESCORE_PARSER = new ObjectParser<>("query", null); static { - RESCORE_PARSER.declareObject(QueryRescoreContext::setParsedQuery, (p, c) -> c.indexShard().getQueryShardContext().parse(p), new ParseField("rescore_query")); + RESCORE_PARSER.declareObject(QueryRescoreContext::setQuery, (p, c) -> c.parse(p).query(), new ParseField("rescore_query")); RESCORE_PARSER.declareFloat(QueryRescoreContext::setQueryWeight, new ParseField("query_weight")); RESCORE_PARSER.declareFloat(QueryRescoreContext::setRescoreQueryWeight, new ParseField("rescore_query_weight")); RESCORE_PARSER.declareString(QueryRescoreContext::setScoreMode, new ParseField("score_mode")); } @Override - public RescoreSearchContext parse(XContentParser parser, SearchContext context) throws IOException { + public RescoreSearchContext parse(XContentParser parser, QueryShardContext context) throws IOException { return RESCORE_PARSER.parse(parser, new QueryRescoreContext(this), context); } @@ -178,22 +178,24 @@ public final class QueryRescorer implements Rescorer { public static class QueryRescoreContext extends RescoreSearchContext { + static final int DEFAULT_WINDOW_SIZE = 10; + public QueryRescoreContext(QueryRescorer rescorer) { - super(NAME, 10, rescorer); + super(NAME, DEFAULT_WINDOW_SIZE, rescorer); this.scoreMode = QueryRescoreMode.Total; } - private ParsedQuery parsedQuery; + private Query query; private float queryWeight = 1.0f; private float rescoreQueryWeight = 1.0f; private QueryRescoreMode scoreMode; - public void setParsedQuery(ParsedQuery parsedQuery) { - this.parsedQuery = parsedQuery; + public void setQuery(Query query) { + this.query = query; } public Query query() { - return parsedQuery.query(); + return query; } public float queryWeight() { diff --git a/core/src/main/java/org/elasticsearch/search/rescore/QueryRescorerBuilder.java b/core/src/main/java/org/elasticsearch/search/rescore/QueryRescorerBuilder.java new file mode 100644 index 00000000000..936353f7868 --- /dev/null +++ b/core/src/main/java/org/elasticsearch/search/rescore/QueryRescorerBuilder.java @@ -0,0 +1,239 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.rescore; + +import org.elasticsearch.common.ParseField; +import org.elasticsearch.common.ParsingException; +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.index.query.MatchAllQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryParseContext; +import org.elasticsearch.index.query.QueryShardContext; +import org.elasticsearch.search.rescore.QueryRescorer.QueryRescoreContext; + +import java.io.IOException; +import java.util.Locale; +import java.util.Objects; + +public class QueryRescorerBuilder implements RescoreBuilder { + + public static final String NAME = "query"; + + public static final QueryRescorerBuilder PROTOTYPE = new QueryRescorerBuilder(new MatchAllQueryBuilder()); + + public static final float DEFAULT_RESCORE_QUERYWEIGHT = 1.0f; + public static final float DEFAULT_QUERYWEIGHT = 1.0f; + public static final QueryRescoreMode DEFAULT_SCORE_MODE = QueryRescoreMode.Total; + private final QueryBuilder queryBuilder; + private float rescoreQueryWeight = DEFAULT_RESCORE_QUERYWEIGHT; + private float queryWeight = DEFAULT_QUERYWEIGHT; + private QueryRescoreMode scoreMode = DEFAULT_SCORE_MODE; + + private static ParseField RESCORE_QUERY_FIELD = new ParseField("rescore_query"); + private static ParseField QUERY_WEIGHT_FIELD = new ParseField("query_weight"); + private static ParseField RESCORE_QUERY_WEIGHT_FIELD = new ParseField("rescore_query_weight"); + private static ParseField SCORE_MODE_FIELD = new ParseField("score_mode"); + + private static final ObjectParser QUERY_RESCORE_PARSER = new ObjectParser<>(NAME, null); + + static { + QUERY_RESCORE_PARSER.declareObject(InnerBuilder::setQueryBuilder, (p, c) -> { + try { + return c.parseInnerQueryBuilder(); + } catch (IOException e) { + throw new ParsingException(p.getTokenLocation(), "Could not parse inner query", e); + } + } , RESCORE_QUERY_FIELD); + QUERY_RESCORE_PARSER.declareFloat(InnerBuilder::setQueryWeight, QUERY_WEIGHT_FIELD); + QUERY_RESCORE_PARSER.declareFloat(InnerBuilder::setRescoreQueryWeight, RESCORE_QUERY_WEIGHT_FIELD); + QUERY_RESCORE_PARSER.declareString((struct, value) -> struct.setScoreMode(QueryRescoreMode.fromString(value)), SCORE_MODE_FIELD); + } + + /** + * Creates a new {@link QueryRescorerBuilder} instance + * @param builder the query builder to build the rescore query from + */ + public QueryRescorerBuilder(QueryBuilder builder) { + this.queryBuilder = builder; + } + + /** + * @return the query used for this rescore query + */ + public QueryBuilder getRescoreQuery() { + return this.queryBuilder; + } + + /** + * Sets the original query weight for rescoring. The default is 1.0 + */ + public QueryRescorerBuilder setQueryWeight(float queryWeight) { + this.queryWeight = queryWeight; + return this; + } + + + /** + * Gets the original query weight for rescoring. The default is 1.0 + */ + public float getQueryWeight() { + return this.queryWeight; + } + + /** + * Sets the original query weight for rescoring. The default is 1.0 + */ + public QueryRescorerBuilder setRescoreQueryWeight(float rescoreQueryWeight) { + this.rescoreQueryWeight = rescoreQueryWeight; + return this; + } + + /** + * Gets the original query weight for rescoring. The default is 1.0 + */ + public float getRescoreQueryWeight() { + return this.rescoreQueryWeight; + } + + /** + * Sets the original query score mode. The default is {@link QueryRescoreMode#Total}. + */ + public QueryRescorerBuilder setScoreMode(QueryRescoreMode scoreMode) { + this.scoreMode = scoreMode; + return this; + } + + /** + * Gets the original query score mode. The default is total + */ + public QueryRescoreMode getScoreMode() { + return this.scoreMode; + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(NAME); + builder.field(RESCORE_QUERY_FIELD.getPreferredName(), queryBuilder); + builder.field(QUERY_WEIGHT_FIELD.getPreferredName(), queryWeight); + builder.field(RESCORE_QUERY_WEIGHT_FIELD.getPreferredName(), rescoreQueryWeight); + builder.field(SCORE_MODE_FIELD.getPreferredName(), scoreMode.name().toLowerCase(Locale.ROOT)); + builder.endObject(); + return builder; + } + + @Override + public QueryRescorerBuilder fromXContent(QueryParseContext parseContext) throws IOException { + InnerBuilder innerBuilder = QUERY_RESCORE_PARSER.parse(parseContext.parser(), new InnerBuilder(), parseContext); + return innerBuilder.build(); + } + + @Override + public QueryRescoreContext build(QueryShardContext context) throws IOException { + org.elasticsearch.search.rescore.QueryRescorer rescorer = new org.elasticsearch.search.rescore.QueryRescorer(); + QueryRescoreContext queryRescoreContext = new QueryRescoreContext(rescorer); + queryRescoreContext.setQuery(this.queryBuilder.toQuery(context)); + queryRescoreContext.setQueryWeight(this.queryWeight); + queryRescoreContext.setRescoreQueryWeight(this.rescoreQueryWeight); + queryRescoreContext.setScoreMode(this.scoreMode); + return queryRescoreContext; + } + + @Override + public final int hashCode() { + return Objects.hash(getClass(), scoreMode, queryWeight, rescoreQueryWeight, queryBuilder); + } + + @Override + public final boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (obj == null || getClass() != obj.getClass()) { + return false; + } + QueryRescorerBuilder other = (QueryRescorerBuilder) obj; + return Objects.equals(scoreMode, other.scoreMode) && + Objects.equals(queryWeight, other.queryWeight) && + Objects.equals(rescoreQueryWeight, other.rescoreQueryWeight) && + Objects.equals(queryBuilder, other.queryBuilder); + } + + @Override + public QueryRescorerBuilder readFrom(StreamInput in) throws IOException { + QueryRescorerBuilder rescorer = new QueryRescorerBuilder(in.readQuery()); + rescorer.setScoreMode(QueryRescoreMode.PROTOTYPE.readFrom(in)); + rescorer.setRescoreQueryWeight(in.readFloat()); + rescorer.setQueryWeight(in.readFloat()); + return rescorer; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeQuery(queryBuilder); + scoreMode.writeTo(out); + out.writeFloat(rescoreQueryWeight); + out.writeFloat(queryWeight); + } + + @Override + public String getWriteableName() { + return NAME; + } + + /** + * Helper to be able to use {@link ObjectParser}, since we need the inner query builder + * for the constructor of {@link QueryRescorerBuilder}, but {@link ObjectParser} only + * allows filling properties of an already constructed value. + */ + private class InnerBuilder { + + private QueryBuilder queryBuilder; + private float rescoreQueryWeight = DEFAULT_RESCORE_QUERYWEIGHT; + private float queryWeight = DEFAULT_QUERYWEIGHT; + private QueryRescoreMode scoreMode = DEFAULT_SCORE_MODE; + + void setQueryBuilder(QueryBuilder builder) { + this.queryBuilder = builder; + } + + QueryRescorerBuilder build() { + QueryRescorerBuilder queryRescoreBuilder = new QueryRescorerBuilder(queryBuilder); + queryRescoreBuilder.setQueryWeight(queryWeight); + queryRescoreBuilder.setRescoreQueryWeight(rescoreQueryWeight); + queryRescoreBuilder.setScoreMode(scoreMode); + return queryRescoreBuilder; + } + + void setQueryWeight(float queryWeight) { + this.queryWeight = queryWeight; + } + + void setRescoreQueryWeight(float rescoreQueryWeight) { + this.rescoreQueryWeight = rescoreQueryWeight; + } + + void setScoreMode(QueryRescoreMode scoreMode) { + this.scoreMode = scoreMode; + } + } +} \ No newline at end of file diff --git a/core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java b/core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java new file mode 100644 index 00000000000..ddcabf9ada5 --- /dev/null +++ b/core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java @@ -0,0 +1,173 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.rescore; + +import org.elasticsearch.ExceptionsHelper; +import org.elasticsearch.common.ParseField; +import org.elasticsearch.common.ParsingException; +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.ToXContent; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentFactory; +import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.index.query.MatchAllQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryParseContext; +import org.elasticsearch.index.query.QueryShardContext; + +import java.io.IOException; +import java.util.Objects; + +/** + * The base builder for rescorers. Wraps a conrete instance of {@link RescoreBuilder} and + * adds the ability to specify the optional `window_size` parameter + */ +public class RescoreBaseBuilder implements ToXContent, Writeable { + + private RescoreBuilder rescorer; + private Integer windowSize; + public static final RescoreBaseBuilder PROTOTYPE = new RescoreBaseBuilder(new QueryRescorerBuilder(new MatchAllQueryBuilder())); + + private static ParseField WINDOW_SIZE_FIELD = new ParseField("window_size"); + + public RescoreBaseBuilder(RescoreBuilder rescorer) { + if (rescorer == null) { + throw new IllegalArgumentException("rescorer cannot be null"); + } + this.rescorer = rescorer; + } + + public RescoreBuilder rescorer() { + return this.rescorer; + } + + public RescoreBaseBuilder windowSize(int windowSize) { + this.windowSize = windowSize; + return this; + } + + public Integer windowSize() { + return windowSize; + } + + public RescoreBaseBuilder fromXContent(QueryParseContext parseContext) throws IOException { + XContentParser parser = parseContext.parser(); + String fieldName = null; + RescoreBuilder rescorer = null; + Integer windowSize = null; + XContentParser.Token token; + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else if (token.isValue()) { + if (parseContext.parseFieldMatcher().match(fieldName, WINDOW_SIZE_FIELD)) { + windowSize = parser.intValue(); + } else { + throw new ParsingException(parser.getTokenLocation(), "rescore doesn't support [" + fieldName + "]"); + } + } else if (token == XContentParser.Token.START_OBJECT) { + // we only have QueryRescorer at this point + if (QueryRescorerBuilder.NAME.equals(fieldName)) { + rescorer = QueryRescorerBuilder.PROTOTYPE.fromXContent(parseContext); + } else { + throw new ParsingException(parser.getTokenLocation(), "rescore doesn't support rescorer with name [" + fieldName + "]"); + } + } else { + throw new ParsingException(parser.getTokenLocation(), "unexpected token [" + token + "] after [" + fieldName + "]"); + } + } + if (rescorer == null) { + throw new ParsingException(parser.getTokenLocation(), "missing rescore type"); + } + RescoreBaseBuilder rescoreBuilder = new RescoreBaseBuilder(rescorer); + if (windowSize != null) { + rescoreBuilder.windowSize(windowSize.intValue()); + } + return rescoreBuilder; + } + + public RescoreSearchContext build(QueryShardContext context) throws IOException { + RescoreSearchContext rescoreContext = this.rescorer.build(context); + if (windowSize != null) { + rescoreContext.setWindowSize(this.windowSize); + } + return rescoreContext; + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + if (windowSize != null) { + builder.field("window_size", windowSize); + } + rescorer.toXContent(builder, params); + return builder; + } + + public static QueryRescorerBuilder queryRescorer(QueryBuilder queryBuilder) { + return new QueryRescorerBuilder(queryBuilder); + } + + @Override + public final int hashCode() { + return Objects.hash(windowSize, rescorer); + } + + @Override + public final boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (obj == null || getClass() != obj.getClass()) { + return false; + } + RescoreBaseBuilder other = (RescoreBaseBuilder) obj; + return Objects.equals(windowSize, other.windowSize) && + Objects.equals(rescorer, other.rescorer); + } + + @Override + public RescoreBaseBuilder readFrom(StreamInput in) throws IOException { + RescoreBaseBuilder builder = new RescoreBaseBuilder(in.readRescorer()); + builder.windowSize = in.readOptionalVInt(); + return builder; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeRescorer(rescorer); + out.writeOptionalVInt(this.windowSize); + } + + @Override + public final String toString() { + try { + XContentBuilder builder = XContentFactory.jsonBuilder(); + builder.prettyPrint(); + builder.startObject(); + toXContent(builder, EMPTY_PARAMS); + builder.endObject(); + return builder.string(); + } catch (Exception e) { + return "{ \"error\" : \"" + ExceptionsHelper.detailedMessage(e) + "\"}"; + } + } +} diff --git a/core/src/main/java/org/elasticsearch/search/rescore/RescoreBuilder.java b/core/src/main/java/org/elasticsearch/search/rescore/RescoreBuilder.java index 7510d24f82d..eeefb1e5f5e 100644 --- a/core/src/main/java/org/elasticsearch/search/rescore/RescoreBuilder.java +++ b/core/src/main/java/org/elasticsearch/search/rescore/RescoreBuilder.java @@ -19,256 +19,16 @@ package org.elasticsearch.search.rescore; -import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.common.io.stream.NamedWriteable; -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.ToXContent; -import org.elasticsearch.common.xcontent.XContentBuilder; -import org.elasticsearch.common.xcontent.XContentFactory; -import org.elasticsearch.index.query.MatchAllQueryBuilder; -import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryParseContext; +import org.elasticsearch.index.query.QueryShardContext; import java.io.IOException; -import java.util.Locale; -import java.util.Objects; -public class RescoreBuilder implements ToXContent, Writeable { +public interface RescoreBuilder extends ToXContent, NamedWriteable { - private Rescorer rescorer; - private Integer windowSize; - public static final RescoreBuilder PROTOYPE = new RescoreBuilder(new QueryRescorer(new MatchAllQueryBuilder())); + RescoreSearchContext build(QueryShardContext context) throws IOException; - public RescoreBuilder(Rescorer rescorer) { - if (rescorer == null) { - throw new IllegalArgumentException("rescorer cannot be null"); - } - this.rescorer = rescorer; - } - - public Rescorer rescorer() { - return this.rescorer; - } - - public RescoreBuilder windowSize(int windowSize) { - this.windowSize = windowSize; - return this; - } - - public Integer windowSize() { - return windowSize; - } - - @Override - public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - if (windowSize != null) { - builder.field("window_size", windowSize); - } - rescorer.toXContent(builder, params); - return builder; - } - - public static QueryRescorer queryRescorer(QueryBuilder queryBuilder) { - return new QueryRescorer(queryBuilder); - } - - @Override - public final int hashCode() { - return Objects.hash(windowSize, rescorer); - } - - @Override - public final boolean equals(Object obj) { - if (this == obj) { - return true; - } - if (obj == null || getClass() != obj.getClass()) { - return false; - } - RescoreBuilder other = (RescoreBuilder) obj; - return Objects.equals(windowSize, other.windowSize) && - Objects.equals(rescorer, other.rescorer); - } - - @Override - public RescoreBuilder readFrom(StreamInput in) throws IOException { - RescoreBuilder builder = new RescoreBuilder(in.readRescorer()); - Integer windowSize = in.readOptionalVInt(); - if (windowSize != null) { - builder.windowSize(windowSize); - } - return builder; - } - - @Override - public void writeTo(StreamOutput out) throws IOException { - out.writeRescorer(rescorer); - out.writeOptionalVInt(this.windowSize); - } - - @Override - public final String toString() { - try { - XContentBuilder builder = XContentFactory.jsonBuilder(); - builder.prettyPrint(); - builder.startObject(); - toXContent(builder, EMPTY_PARAMS); - builder.endObject(); - return builder.string(); - } catch (Exception e) { - return "{ \"error\" : \"" + ExceptionsHelper.detailedMessage(e) + "\"}"; - } - } - - public static abstract class Rescorer implements ToXContent, NamedWriteable { - - private String name; - - public Rescorer(String name) { - this.name = name; - } - - @Override - public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - builder.startObject(name); - builder = innerToXContent(builder, params); - builder.endObject(); - return builder; - } - - protected abstract XContentBuilder innerToXContent(XContentBuilder builder, Params params) throws IOException; - - @Override - public abstract int hashCode(); - - @Override - public abstract boolean equals(Object obj); - } - - public static class QueryRescorer extends Rescorer { - - private static final String NAME = "query"; - public static final QueryRescorer PROTOTYPE = new QueryRescorer(new MatchAllQueryBuilder()); - public static final float DEFAULT_RESCORE_QUERYWEIGHT = 1.0f; - public static final float DEFAULT_QUERYWEIGHT = 1.0f; - public static final QueryRescoreMode DEFAULT_SCORE_MODE = QueryRescoreMode.Total; - private final QueryBuilder queryBuilder; - private float rescoreQueryWeight = DEFAULT_RESCORE_QUERYWEIGHT; - private float queryWeight = DEFAULT_QUERYWEIGHT; - private QueryRescoreMode scoreMode = DEFAULT_SCORE_MODE; - - /** - * Creates a new {@link QueryRescorer} instance - * @param builder the query builder to build the rescore query from - */ - public QueryRescorer(QueryBuilder builder) { - super(NAME); - this.queryBuilder = builder; - } - - /** - * @return the query used for this rescore query - */ - public QueryBuilder getRescoreQuery() { - return this.queryBuilder; - } - - /** - * Sets the original query weight for rescoring. The default is 1.0 - */ - public QueryRescorer setQueryWeight(float queryWeight) { - this.queryWeight = queryWeight; - return this; - } - - - /** - * Gets the original query weight for rescoring. The default is 1.0 - */ - public float getQueryWeight() { - return this.queryWeight; - } - - /** - * Sets the original query weight for rescoring. The default is 1.0 - */ - public QueryRescorer setRescoreQueryWeight(float rescoreQueryWeight) { - this.rescoreQueryWeight = rescoreQueryWeight; - return this; - } - - /** - * Gets the original query weight for rescoring. The default is 1.0 - */ - public float getRescoreQueryWeight() { - return this.rescoreQueryWeight; - } - - /** - * Sets the original query score mode. The default is {@link QueryRescoreMode#Total}. - */ - public QueryRescorer setScoreMode(QueryRescoreMode scoreMode) { - this.scoreMode = scoreMode; - return this; - } - - /** - * Gets the original query score mode. The default is total - */ - public QueryRescoreMode getScoreMode() { - return this.scoreMode; - } - - @Override - protected XContentBuilder innerToXContent(XContentBuilder builder, Params params) throws IOException { - builder.field("rescore_query", queryBuilder); - builder.field("query_weight", queryWeight); - builder.field("rescore_query_weight", rescoreQueryWeight); - builder.field("score_mode", scoreMode.name().toLowerCase(Locale.ROOT)); - return builder; - } - - @Override - public final int hashCode() { - return Objects.hash(getClass(), scoreMode, queryWeight, rescoreQueryWeight, queryBuilder); - } - - @Override - public final boolean equals(Object obj) { - if (this == obj) { - return true; - } - if (obj == null || getClass() != obj.getClass()) { - return false; - } - QueryRescorer other = (QueryRescorer) obj; - return Objects.equals(scoreMode, other.scoreMode) && - Objects.equals(queryWeight, other.queryWeight) && - Objects.equals(rescoreQueryWeight, other.rescoreQueryWeight) && - Objects.equals(queryBuilder, other.queryBuilder); - } - - @Override - public QueryRescorer readFrom(StreamInput in) throws IOException { - QueryRescorer rescorer = new QueryRescorer(in.readQuery()); - rescorer.setScoreMode(QueryRescoreMode.PROTOTYPE.readFrom(in)); - rescorer.setRescoreQueryWeight(in.readFloat()); - rescorer.setQueryWeight(in.readFloat()); - return rescorer; - } - - @Override - public void writeTo(StreamOutput out) throws IOException { - out.writeQuery(queryBuilder); - scoreMode.writeTo(out); - out.writeFloat(rescoreQueryWeight); - out.writeFloat(queryWeight); - } - - @Override - public String getWriteableName() { - return NAME; - } - } -} + RB fromXContent(QueryParseContext parseContext) throws IOException; +} \ No newline at end of file diff --git a/core/src/main/java/org/elasticsearch/search/rescore/RescoreParseElement.java b/core/src/main/java/org/elasticsearch/search/rescore/RescoreParseElement.java index 7f9f2725fbc..149db6cec2c 100644 --- a/core/src/main/java/org/elasticsearch/search/rescore/RescoreParseElement.java +++ b/core/src/main/java/org/elasticsearch/search/rescore/RescoreParseElement.java @@ -21,9 +21,12 @@ package org.elasticsearch.search.rescore; import org.elasticsearch.ElasticsearchParseException; import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.search.SearchParseElement; import org.elasticsearch.search.internal.SearchContext; +import java.io.IOException; + /** * */ @@ -33,14 +36,14 @@ public class RescoreParseElement implements SearchParseElement { public void parse(XContentParser parser, SearchContext context) throws Exception { if (parser.currentToken() == XContentParser.Token.START_ARRAY) { while (parser.nextToken() != XContentParser.Token.END_ARRAY) { - parseSingleRescoreContext(parser, context); + context.addRescore(parseSingleRescoreContext(parser, context.indexShard().getQueryShardContext())); } } else { - parseSingleRescoreContext(parser, context); + context.addRescore(parseSingleRescoreContext(parser, context.indexShard().getQueryShardContext())); } } - public void parseSingleRescoreContext(XContentParser parser, SearchContext context) throws Exception { + public RescoreSearchContext parseSingleRescoreContext(XContentParser parser, QueryShardContext context) throws ElasticsearchParseException, IOException { String fieldName = null; RescoreSearchContext rescoreContext = null; Integer windowSize = null; @@ -71,7 +74,7 @@ public class RescoreParseElement implements SearchParseElement { if (windowSize != null) { rescoreContext.setWindowSize(windowSize.intValue()); } - context.addRescore(rescoreContext); + return rescoreContext; } } diff --git a/core/src/main/java/org/elasticsearch/search/rescore/Rescorer.java b/core/src/main/java/org/elasticsearch/search/rescore/Rescorer.java index 3c90289fde5..e3465a4df8f 100644 --- a/core/src/main/java/org/elasticsearch/search/rescore/Rescorer.java +++ b/core/src/main/java/org/elasticsearch/search/rescore/Rescorer.java @@ -24,6 +24,7 @@ import org.apache.lucene.search.Explanation; import org.apache.lucene.search.TopDocs; import org.elasticsearch.action.search.SearchType; import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.search.internal.SearchContext; import java.io.IOException; @@ -68,11 +69,11 @@ public interface Rescorer { * Parses the {@link RescoreSearchContext} for this impelementation * * @param parser the parser to read the context from - * @param context the current search context + * @param context the current shard context * @return the parsed {@link RescoreSearchContext} * @throws IOException if an {@link IOException} occurs while parsing the context */ - public RescoreSearchContext parse(XContentParser parser, SearchContext context) throws IOException; + public RescoreSearchContext parse(XContentParser parser, QueryShardContext context) throws IOException; /** * Extracts all terms needed to exectue this {@link Rescorer}. This method @@ -81,7 +82,7 @@ public interface Rescorer { * {@link SearchType#DFS_QUERY_THEN_FETCH} */ public void extractTerms(SearchContext context, RescoreSearchContext rescoreContext, Set termsSet); - + /* * TODO: At this point we only have one implemenation which modifies the * TopDocs given. Future implemenations might return actual resutls that diff --git a/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java b/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java index 5a1b99fe05f..e992e33eeb6 100644 --- a/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java +++ b/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java @@ -57,7 +57,7 @@ import org.elasticsearch.search.fetch.innerhits.InnerHitsBuilder; import org.elasticsearch.search.fetch.innerhits.InnerHitsBuilder.InnerHit; import org.elasticsearch.search.fetch.source.FetchSourceContext; import org.elasticsearch.search.highlight.HighlightBuilderTests; -import org.elasticsearch.search.rescore.RescoreBuilder; +import org.elasticsearch.search.rescore.RescoreBaseBuilder; import org.elasticsearch.search.sort.SortBuilders; import org.elasticsearch.search.sort.SortOrder; import org.elasticsearch.search.suggest.SuggestBuilder; @@ -281,7 +281,7 @@ public class SearchSourceBuilderTests extends ESTestCase { int numRescores = randomIntBetween(1, 5); for (int i = 0; i < numRescores; i++) { // NORELEASE need a random rescore builder method - RescoreBuilder rescoreBuilder = new RescoreBuilder(RescoreBuilder.queryRescorer(QueryBuilders.termQuery(randomAsciiOfLengthBetween(5, 20), + RescoreBaseBuilder rescoreBuilder = new RescoreBaseBuilder(RescoreBaseBuilder.queryRescorer(QueryBuilders.termQuery(randomAsciiOfLengthBetween(5, 20), randomAsciiOfLengthBetween(5, 20)))); builder.addRescorer(rescoreBuilder); } diff --git a/core/src/test/java/org/elasticsearch/search/functionscore/QueryRescorerIT.java b/core/src/test/java/org/elasticsearch/search/functionscore/QueryRescorerIT.java index 5644f893603..c9b895b1995 100644 --- a/core/src/test/java/org/elasticsearch/search/functionscore/QueryRescorerIT.java +++ b/core/src/test/java/org/elasticsearch/search/functionscore/QueryRescorerIT.java @@ -38,8 +38,8 @@ import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders; import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.SearchHits; import org.elasticsearch.search.rescore.QueryRescoreMode; -import org.elasticsearch.search.rescore.RescoreBuilder; -import org.elasticsearch.search.rescore.RescoreBuilder.QueryRescorer; +import org.elasticsearch.search.rescore.QueryRescorerBuilder; +import org.elasticsearch.search.rescore.RescoreBaseBuilder; import org.elasticsearch.test.ESIntegTestCase; import java.util.Arrays; @@ -80,7 +80,7 @@ public class QueryRescorerIT extends ESIntegTestCase { for (int j = 0 ; j < iters; j++) { SearchResponse searchResponse = client().prepareSearch() .setQuery(QueryBuilders.matchAllQuery()) - .setRescorer(RescoreBuilder.queryRescorer( + .setRescorer(RescoreBaseBuilder.queryRescorer( QueryBuilders.functionScoreQuery(QueryBuilders.matchAllQuery(), ScoreFunctionBuilders.weightFactorFunction(100)).boostMode(CombineFunction.REPLACE)) .setQueryWeight(0.0f).setRescoreQueryWeight(1.0f), 1).setSize(randomIntBetween(2, 10)).execute() @@ -116,7 +116,7 @@ public class QueryRescorerIT extends ESIntegTestCase { SearchResponse searchResponse = client().prepareSearch() .setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR)) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "quick brown").slop(2).boost(4.0f)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "quick brown").slop(2).boost(4.0f)) .setRescoreQueryWeight(2), 5).execute().actionGet(); assertThat(searchResponse.getHits().totalHits(), equalTo(3l)); @@ -126,7 +126,7 @@ public class QueryRescorerIT extends ESIntegTestCase { searchResponse = client().prepareSearch() .setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR)) - .setRescorer(RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "the quick brown").slop(3)), 5) + .setRescorer(RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "the quick brown").slop(3)), 5) .execute().actionGet(); assertHitCount(searchResponse, 3); @@ -136,7 +136,7 @@ public class QueryRescorerIT extends ESIntegTestCase { searchResponse = client().prepareSearch() .setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR)) - .setRescorer(RescoreBuilder.queryRescorer((QueryBuilders.matchPhraseQuery("field1", "the quick brown"))), 5).execute() + .setRescorer(RescoreBaseBuilder.queryRescorer((QueryBuilders.matchPhraseQuery("field1", "the quick brown"))), 5).execute() .actionGet(); assertHitCount(searchResponse, 3); @@ -181,7 +181,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setFrom(0) .setSize(5) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) .setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 20).execute().actionGet(); assertThat(searchResponse.getHits().hits().length, equalTo(5)); @@ -197,7 +197,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setSize(5) .setSearchType(SearchType.DFS_QUERY_THEN_FETCH) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) .setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 20).execute().actionGet(); assertThat(searchResponse.getHits().hits().length, equalTo(5)); @@ -214,7 +214,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setSize(5) .setSearchType(SearchType.DFS_QUERY_THEN_FETCH) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) .setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 20).execute().actionGet(); assertThat(searchResponse.getHits().hits().length, equalTo(5)); @@ -263,7 +263,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setFrom(0) .setSize(5) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) .setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 2).execute().actionGet(); // Only top 2 hits were re-ordered: assertThat(searchResponse.getHits().hits().length, equalTo(4)); @@ -280,7 +280,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setFrom(0) .setSize(5) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) .setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 3).execute().actionGet(); // Only top 3 hits were re-ordered: @@ -333,7 +333,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setFrom(0) .setSize(5) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3)) .setQueryWeight(1.0f).setRescoreQueryWeight(-1f), 3).execute().actionGet(); // 6 and 1 got worse, and then the hit (2) outside the rescore window were sorted ahead: @@ -424,7 +424,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setFrom(0) .setSize(resultSize) .setRescorer( - RescoreBuilder + RescoreBaseBuilder .queryRescorer( QueryBuilders .constantScoreQuery(QueryBuilders.matchPhraseQuery("field1", intToEnglish).slop(3))) @@ -462,7 +462,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setFrom(0) .setSize(resultSize) .setRescorer( - RescoreBuilder + RescoreBaseBuilder .queryRescorer( QueryBuilders .constantScoreQuery(QueryBuilders.matchPhraseQuery("field1", "not in the index").slop(3))) @@ -480,7 +480,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setFrom(0) .setSize(resultSize) .setRescorer( - RescoreBuilder + RescoreBaseBuilder .queryRescorer( QueryBuilders.matchPhraseQuery("field1", intToEnglish).slop(0)) .setQueryWeight(1.0f).setRescoreQueryWeight(1.0f), 2 * rescoreWindow).execute().actionGet(); @@ -512,7 +512,7 @@ public class QueryRescorerIT extends ESIntegTestCase { .setSearchType(SearchType.DFS_QUERY_THEN_FETCH) .setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR)) .setRescorer( - RescoreBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "the quick brown").slop(2).boost(4.0f)) + RescoreBaseBuilder.queryRescorer(QueryBuilders.matchPhraseQuery("field1", "the quick brown").slop(2).boost(4.0f)) .setQueryWeight(0.5f).setRescoreQueryWeight(0.4f), 5).setExplain(true).execute() .actionGet(); assertHitCount(searchResponse, 3); @@ -538,7 +538,7 @@ public class QueryRescorerIT extends ESIntegTestCase { String[] scoreModes = new String[]{ "max", "min", "avg", "total", "multiply", "" }; String[] descriptionModes = new String[]{ "max of:", "min of:", "avg of:", "sum of:", "product of:", "sum of:" }; for (int innerMode = 0; innerMode < scoreModes.length; innerMode++) { - QueryRescorer innerRescoreQuery = RescoreBuilder.queryRescorer(QueryBuilders.matchQuery("field1", "the quick brown").boost(4.0f)) + QueryRescorerBuilder innerRescoreQuery = RescoreBaseBuilder.queryRescorer(QueryBuilders.matchQuery("field1", "the quick brown").boost(4.0f)) .setQueryWeight(0.5f).setRescoreQueryWeight(0.4f); if (!"".equals(scoreModes[innerMode])) { @@ -561,7 +561,7 @@ public class QueryRescorerIT extends ESIntegTestCase { } for (int outerMode = 0; outerMode < scoreModes.length; outerMode++) { - QueryRescorer outerRescoreQuery = RescoreBuilder.queryRescorer(QueryBuilders.matchQuery("field1", "the quick brown") + QueryRescorerBuilder outerRescoreQuery = RescoreBaseBuilder.queryRescorer(QueryBuilders.matchQuery("field1", "the quick brown") .boost(4.0f)).setQueryWeight(0.5f).setRescoreQueryWeight(0.4f); if (!"".equals(scoreModes[outerMode])) { @@ -599,7 +599,7 @@ public class QueryRescorerIT extends ESIntegTestCase { for (int i = 0; i < numDocs - 4; i++) { String[] intToEnglish = new String[] { English.intToEnglish(i), English.intToEnglish(i + 1), English.intToEnglish(i + 2), English.intToEnglish(i + 3) }; - QueryRescorer rescoreQuery = RescoreBuilder + QueryRescorerBuilder rescoreQuery = RescoreBaseBuilder .queryRescorer( QueryBuilders.boolQuery() .disableCoord(true) @@ -682,10 +682,10 @@ public class QueryRescorerIT extends ESIntegTestCase { public void testMultipleRescores() throws Exception { int numDocs = indexRandomNumbers("keyword", 1, true); - QueryRescorer eightIsGreat = RescoreBuilder.queryRescorer( + QueryRescorerBuilder eightIsGreat = RescoreBaseBuilder.queryRescorer( QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", English.intToEnglish(8)), ScoreFunctionBuilders.weightFactorFunction(1000.0f)).boostMode(CombineFunction.REPLACE)).setScoreMode(QueryRescoreMode.Total); - QueryRescorer sevenIsBetter = RescoreBuilder.queryRescorer( + QueryRescorerBuilder sevenIsBetter = RescoreBaseBuilder.queryRescorer( QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", English.intToEnglish(7)), ScoreFunctionBuilders.weightFactorFunction(10000.0f)).boostMode(CombineFunction.REPLACE)) .setScoreMode(QueryRescoreMode.Total); @@ -703,10 +703,10 @@ public class QueryRescorerIT extends ESIntegTestCase { // We have no idea what the second hit will be because we didn't get a chance to look for seven // Now use one rescore to drag the number we're looking for into the window of another - QueryRescorer ninetyIsGood = RescoreBuilder.queryRescorer( + QueryRescorerBuilder ninetyIsGood = RescoreBaseBuilder.queryRescorer( QueryBuilders.functionScoreQuery(QueryBuilders.queryStringQuery("*ninety*"), ScoreFunctionBuilders.weightFactorFunction(1000.0f)) .boostMode(CombineFunction.REPLACE)).setScoreMode(QueryRescoreMode.Total); - QueryRescorer oneToo = RescoreBuilder.queryRescorer( + QueryRescorerBuilder oneToo = RescoreBaseBuilder.queryRescorer( QueryBuilders.functionScoreQuery(QueryBuilders.queryStringQuery("*one*"), ScoreFunctionBuilders.weightFactorFunction(1000.0f)) .boostMode(CombineFunction.REPLACE)).setScoreMode(QueryRescoreMode.Total); request.clearRescorers().addRescorer(ninetyIsGood, numDocs).addRescorer(oneToo, 10); @@ -759,7 +759,7 @@ public class QueryRescorerIT extends ESIntegTestCase { request.setQuery(QueryBuilders.termQuery("text", "hello")); request.setFrom(1); request.setSize(4); - request.addRescorer(RescoreBuilder.queryRescorer(QueryBuilders.matchAllQuery()), 50); + request.addRescorer(RescoreBaseBuilder.queryRescorer(QueryBuilders.matchAllQuery()), 50); assertEquals(4, request.get().getHits().hits().length); } diff --git a/core/src/test/java/org/elasticsearch/search/rescore/QueryRescoreBuilderTests.java b/core/src/test/java/org/elasticsearch/search/rescore/QueryRescoreBuilderTests.java index 2aa55f8b626..a305e8ca5e0 100644 --- a/core/src/test/java/org/elasticsearch/search/rescore/QueryRescoreBuilderTests.java +++ b/core/src/test/java/org/elasticsearch/search/rescore/QueryRescoreBuilderTests.java @@ -19,19 +19,45 @@ package org.elasticsearch.search.rescore; +import org.elasticsearch.ElasticsearchParseException; +import org.elasticsearch.Version; +import org.elasticsearch.cluster.metadata.IndexMetaData; +import org.elasticsearch.common.ParseFieldMatcher; +import org.elasticsearch.common.ParsingException; import org.elasticsearch.common.io.stream.BytesStreamOutput; import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.xcontent.ToXContent; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentFactory; +import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.common.xcontent.XContentType; +import org.elasticsearch.index.Index; +import org.elasticsearch.index.IndexSettings; +import org.elasticsearch.index.mapper.ContentPath; +import org.elasticsearch.index.mapper.MappedFieldType; +import org.elasticsearch.index.mapper.Mapper; +import org.elasticsearch.index.mapper.MapperBuilders; +import org.elasticsearch.index.mapper.core.StringFieldMapper; import org.elasticsearch.index.query.MatchAllQueryBuilder; +import org.elasticsearch.index.query.MatchAllQueryParser; import org.elasticsearch.index.query.QueryBuilder; -import org.elasticsearch.search.rescore.RescoreBuilder.QueryRescorer; -import org.elasticsearch.search.rescore.RescoreBuilder.Rescorer; +import org.elasticsearch.index.query.QueryParseContext; +import org.elasticsearch.index.query.QueryParser; +import org.elasticsearch.index.query.QueryShardContext; +import org.elasticsearch.indices.query.IndicesQueriesRegistry; +import org.elasticsearch.search.rescore.QueryRescorer.QueryRescoreContext; import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.test.IndexSettingsModule; import org.junit.AfterClass; import org.junit.BeforeClass; import java.io.IOException; +import java.util.HashSet; +import java.util.Set; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.not; @@ -40,6 +66,7 @@ public class QueryRescoreBuilderTests extends ESTestCase { private static final int NUMBER_OF_TESTBUILDERS = 20; private static NamedWriteableRegistry namedWriteableRegistry; + private static IndicesQueriesRegistry indicesQueriesRegistry; /** * setup for the whole base test class @@ -47,13 +74,17 @@ public class QueryRescoreBuilderTests extends ESTestCase { @BeforeClass public static void init() { namedWriteableRegistry = new NamedWriteableRegistry(); - namedWriteableRegistry.registerPrototype(Rescorer.class, org.elasticsearch.search.rescore.RescoreBuilder.QueryRescorer.PROTOTYPE); - namedWriteableRegistry.registerPrototype(QueryBuilder.class, new MatchAllQueryBuilder()); + namedWriteableRegistry.registerPrototype(RescoreBuilder.class, org.elasticsearch.search.rescore.QueryRescorerBuilder.PROTOTYPE); + @SuppressWarnings("rawtypes") + Set injectedQueryParsers = new HashSet<>(); + injectedQueryParsers.add(new MatchAllQueryParser()); + indicesQueriesRegistry = new IndicesQueriesRegistry(Settings.settingsBuilder().build(), injectedQueryParsers, namedWriteableRegistry); } @AfterClass public static void afterClass() throws Exception { namedWriteableRegistry = null; + indicesQueriesRegistry = null; } /** @@ -61,8 +92,8 @@ public class QueryRescoreBuilderTests extends ESTestCase { */ public void testSerialization() throws IOException { for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) { - RescoreBuilder original = randomRescoreBuilder(); - RescoreBuilder deserialized = serializedCopy(original); + RescoreBaseBuilder original = randomRescoreBuilder(); + RescoreBaseBuilder deserialized = serializedCopy(original); assertEquals(deserialized, original); assertEquals(deserialized.hashCode(), original.hashCode()); assertNotSame(deserialized, original); @@ -74,7 +105,7 @@ public class QueryRescoreBuilderTests extends ESTestCase { */ public void testEqualsAndHashcode() throws IOException { for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) { - RescoreBuilder firstBuilder = randomRescoreBuilder(); + RescoreBaseBuilder firstBuilder = randomRescoreBuilder(); assertFalse("rescore builder is equal to null", firstBuilder.equals(null)); assertFalse("rescore builder is equal to incompatible type", firstBuilder.equals("")); assertTrue("rescore builder is not equal to self", firstBuilder.equals(firstBuilder)); @@ -82,13 +113,13 @@ public class QueryRescoreBuilderTests extends ESTestCase { equalTo(firstBuilder.hashCode())); assertThat("different rescore builder should not be equal", mutate(firstBuilder), not(equalTo(firstBuilder))); - RescoreBuilder secondBuilder = serializedCopy(firstBuilder); + RescoreBaseBuilder secondBuilder = serializedCopy(firstBuilder); assertTrue("rescore builder is not equal to self", secondBuilder.equals(secondBuilder)); assertTrue("rescore builder is not equal to its copy", firstBuilder.equals(secondBuilder)); assertTrue("equals is not symmetric", secondBuilder.equals(firstBuilder)); assertThat("rescore builder copy's hashcode is different from original hashcode", secondBuilder.hashCode(), equalTo(firstBuilder.hashCode())); - RescoreBuilder thirdBuilder = serializedCopy(secondBuilder); + RescoreBaseBuilder thirdBuilder = serializedCopy(secondBuilder); assertTrue("rescore builder is not equal to self", thirdBuilder.equals(thirdBuilder)); assertTrue("rescore builder is not equal to its copy", secondBuilder.equals(thirdBuilder)); assertThat("rescore builder copy's hashcode is different from original hashcode", secondBuilder.hashCode(), equalTo(thirdBuilder.hashCode())); @@ -99,8 +130,165 @@ public class QueryRescoreBuilderTests extends ESTestCase { } } - private RescoreBuilder mutate(RescoreBuilder original) throws IOException { - RescoreBuilder mutation = serializedCopy(original); + /** + * creates random rescorer, renders it to xContent and back to new instance that should be equal to original + */ + public void testFromXContent() throws IOException { + QueryParseContext context = new QueryParseContext(indicesQueriesRegistry); + context.parseFieldMatcher(new ParseFieldMatcher(Settings.EMPTY)); + for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) { + RescoreBaseBuilder rescoreBuilder = randomRescoreBuilder(); + + XContentParser parser = createParser(rescoreBuilder); + context.reset(parser); + parser.nextToken(); + RescoreBaseBuilder secondRescoreBuilder = RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + assertNotSame(rescoreBuilder, secondRescoreBuilder); + assertEquals(rescoreBuilder, secondRescoreBuilder); + assertEquals(rescoreBuilder.hashCode(), secondRescoreBuilder.hashCode()); + } + } + + private static XContentParser createParser(RescoreBaseBuilder rescoreBuilder) throws IOException { + XContentBuilder builder = XContentFactory.contentBuilder(randomFrom(XContentType.values())); + if (randomBoolean()) { + builder.prettyPrint(); + } + builder.startObject(); + rescoreBuilder.toXContent(builder, ToXContent.EMPTY_PARAMS); + builder.endObject(); + + return XContentHelper.createParser(builder.bytes()); + } + + /** + * test that build() outputs a {@link RescoreSearchContext} that is similar to the one + * we would get when parsing the xContent the test rescore builder is rendering out + */ + public void testBuildRescoreSearchContext() throws ElasticsearchParseException, IOException { + Settings indexSettings = Settings.settingsBuilder() + .put(IndexMetaData.SETTING_VERSION_CREATED, Version.CURRENT).build(); + Index index = new Index(randomAsciiOfLengthBetween(1, 10)); + IndexSettings idxSettings = IndexSettingsModule.newIndexSettings(index, indexSettings); + // shard context will only need indicesQueriesRegistry for building Query objects nested in query rescorer + QueryShardContext mockShardContext = new QueryShardContext(idxSettings, null, null, null, null, null, null, indicesQueriesRegistry) { + @Override + public MappedFieldType fieldMapper(String name) { + StringFieldMapper.Builder builder = MapperBuilders.stringField(name); + return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType(); + } + }; + + for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) { + RescoreBaseBuilder rescoreBuilder = randomRescoreBuilder(); + QueryRescoreContext rescoreContext = (QueryRescoreContext) rescoreBuilder.build(mockShardContext); + XContentParser parser = createParser(rescoreBuilder); + + QueryRescoreContext parsedRescoreContext = (QueryRescoreContext) new RescoreParseElement().parseSingleRescoreContext(parser, mockShardContext); + assertNotSame(rescoreContext, parsedRescoreContext); + assertEquals(rescoreContext.window(), parsedRescoreContext.window()); + assertEquals(rescoreContext.query(), parsedRescoreContext.query()); + assertEquals(rescoreContext.queryWeight(), parsedRescoreContext.queryWeight(), Float.MIN_VALUE); + assertEquals(rescoreContext.rescoreQueryWeight(), parsedRescoreContext.rescoreQueryWeight(), Float.MIN_VALUE); + assertEquals(rescoreContext.scoreMode(), parsedRescoreContext.scoreMode()); + } + } + + /** + * test parsing exceptions for incorrect rescorer syntax + */ + public void testUnknownFieldsExpection() throws IOException { + QueryParseContext context = new QueryParseContext(indicesQueriesRegistry); + context.parseFieldMatcher(new ParseFieldMatcher(Settings.EMPTY)); + + String rescoreElement = "{\n" + + " \"window_size\" : 20,\n" + + " \"bad_rescorer_name\" : { }\n" + + "}\n"; + prepareContext(context, rescoreElement); + try { + RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + fail("expected a parsing exception"); + } catch (ParsingException e) { + assertEquals("rescore doesn't support rescorer with name [bad_rescorer_name]", e.getMessage()); + } + + rescoreElement = "{\n" + + " \"bad_fieldName\" : 20\n" + + "}\n"; + prepareContext(context, rescoreElement); + try { + RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + fail("expected a parsing exception"); + } catch (ParsingException e) { + assertEquals("rescore doesn't support [bad_fieldName]", e.getMessage()); + } + + rescoreElement = "{\n" + + " \"window_size\" : 20,\n" + + " \"query\" : [ ]\n" + + "}\n"; + prepareContext(context, rescoreElement); + try { + RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + fail("expected a parsing exception"); + } catch (ParsingException e) { + assertEquals("unexpected token [START_ARRAY] after [query]", e.getMessage()); + } + + rescoreElement = "{ }"; + prepareContext(context, rescoreElement); + try { + RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + fail("expected a parsing exception"); + } catch (ParsingException e) { + assertEquals("missing rescore type", e.getMessage()); + } + + rescoreElement = "{\n" + + " \"window_size\" : 20,\n" + + " \"query\" : { \"bad_fieldname\" : 1.0 } \n" + + "}\n"; + prepareContext(context, rescoreElement); + try { + RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + fail("expected a parsing exception"); + } catch (IllegalArgumentException e) { + assertEquals("[query] unknown field [bad_fieldname], parser not found", e.getMessage()); + } + + rescoreElement = "{\n" + + " \"window_size\" : 20,\n" + + " \"query\" : { \"rescore_query\" : { \"unknown_queryname\" : { } } } \n" + + "}\n"; + prepareContext(context, rescoreElement); + try { + RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + fail("expected a parsing exception"); + } catch (ParsingException e) { + assertEquals("[query] failed to parse field [rescore_query]", e.getMessage()); + } + + rescoreElement = "{\n" + + " \"window_size\" : 20,\n" + + " \"query\" : { \"rescore_query\" : { \"match_all\" : { } } } \n" + + "}\n"; + prepareContext(context, rescoreElement); + RescoreBaseBuilder.PROTOTYPE.fromXContent(context); + } + + /** + * create a new parser from the rescorer string representation and reset context with it + */ + private static void prepareContext(QueryParseContext context, String rescoreElement) throws IOException { + XContentParser parser = XContentFactory.xContent(rescoreElement).createParser(rescoreElement); + context.reset(parser); + // move to first token, this is where the internal fromXContent + assertTrue(parser.nextToken() == XContentParser.Token.START_OBJECT); + } + + private static RescoreBaseBuilder mutate(RescoreBaseBuilder original) throws IOException { + RescoreBaseBuilder mutation = serializedCopy(original); if (randomBoolean()) { Integer windowSize = original.windowSize(); if (windowSize != null) { @@ -109,7 +297,7 @@ public class QueryRescoreBuilderTests extends ESTestCase { mutation.windowSize(randomIntBetween(0, 100)); } } else { - QueryRescorer queryRescorer = (QueryRescorer) mutation.rescorer(); + QueryRescorerBuilder queryRescorer = (QueryRescorerBuilder) mutation.rescorer(); switch (randomIntBetween(0, 3)) { case 0: queryRescorer.setQueryWeight(queryRescorer.getQueryWeight() + 0.1f); @@ -138,10 +326,10 @@ public class QueryRescoreBuilderTests extends ESTestCase { /** * create random shape that is put under test */ - private static RescoreBuilder randomRescoreBuilder() { + private static RescoreBaseBuilder randomRescoreBuilder() { QueryBuilder queryBuilder = new MatchAllQueryBuilder().boost(randomFloat()).queryName(randomAsciiOfLength(20)); - org.elasticsearch.search.rescore.RescoreBuilder.QueryRescorer rescorer = new - org.elasticsearch.search.rescore.RescoreBuilder.QueryRescorer(queryBuilder); + org.elasticsearch.search.rescore.QueryRescorerBuilder rescorer = new + org.elasticsearch.search.rescore.QueryRescorerBuilder(queryBuilder); if (randomBoolean()) { rescorer.setQueryWeight(randomFloat()); } @@ -151,18 +339,18 @@ public class QueryRescoreBuilderTests extends ESTestCase { if (randomBoolean()) { rescorer.setScoreMode(randomFrom(QueryRescoreMode.values())); } - RescoreBuilder builder = new RescoreBuilder(rescorer); + RescoreBaseBuilder builder = new RescoreBaseBuilder(rescorer); if (randomBoolean()) { builder.windowSize(randomIntBetween(0, 100)); } return builder; } - private static RescoreBuilder serializedCopy(RescoreBuilder original) throws IOException { + private static RescoreBaseBuilder serializedCopy(RescoreBaseBuilder original) throws IOException { try (BytesStreamOutput output = new BytesStreamOutput()) { original.writeTo(output); try (StreamInput in = new NamedWriteableAwareStreamInput(StreamInput.wrap(output.bytes()), namedWriteableRegistry)) { - return RescoreBuilder.PROTOYPE.readFrom(in); + return RescoreBaseBuilder.PROTOTYPE.readFrom(in); } } }