honeymoose
/
OpenSearch
mirror of https://github.com/honeymoose/OpenSearch.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

RescoreBuilder: Add parsing and creating of RescoreSearchContext 

Adding the ability to parse from xContent to the rescore builder.
Also making RescoreBuilder an interface and renaming the current
base builder that encapsulates the `window_size` setting the the
concrete rescorer implementation to RescoreBaseBuilder.
This commit is contained in:
Christoph Büscher 2016-01-13 16:38:38 +01:00
parent b9dc5acf2c
commit 428043700d
13 changed files with 690 additions and 324 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
core/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java
View File

@ -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<SearchRequest, Se
/**
* Clears all rescorers on the builder and sets the first one. To use multiple rescore windows use
* {@link #addRescorer(org.elasticsearch.search.rescore.RescoreBuilder.Rescorer, int)}.
* {@link #addRescorer(org.elasticsearch.search.rescore.RescoreBuilder, int)}.
*
* @param rescorer rescorer configuration
* @return this for chaining
*/
public SearchRequestBuilder setRescorer(RescoreBuilder.Rescorer rescorer) {
public SearchRequestBuilder setRescorer(RescoreBuilder rescorer) {
sourceBuilder().clearRescorers();
return addRescorer(rescorer);
}
/**
* Clears all rescorers on the builder and sets the first one. To use multiple rescore windows use
* {@link #addRescorer(org.elasticsearch.search.rescore.RescoreBuilder.Rescorer, int)}.
* {@link #addRescorer(org.elasticsearch.search.rescore.RescoreBuilder, int)}.
*
* @param rescorer rescorer configuration
* @param window rescore window
* @return this for chaining
*/
public SearchRequestBuilder setRescorer(RescoreBuilder.Rescorer rescorer, int window) {
public SearchRequestBuilder setRescorer(RescoreBuilder rescorer, int window) {
sourceBuilder().clearRescorers();
return addRescorer(rescorer, window);
}
@ -420,8 +421,8 @@ public class SearchRequestBuilder extends ActionRequestBuilder<SearchRequest, Se
* @param rescorer rescorer configuration
* @return this for chaining
*/
public SearchRequestBuilder addRescorer(RescoreBuilder.Rescorer rescorer) {
sourceBuilder().addRescorer(new RescoreBuilder(rescorer));
public SearchRequestBuilder addRescorer(RescoreBuilder rescorer) {
sourceBuilder().addRescorer(new RescoreBaseBuilder(rescorer));
return this;
}
@ -432,8 +433,8 @@ public class SearchRequestBuilder extends ActionRequestBuilder<SearchRequest, Se
* @param window rescore window
* @return this for chaining
*/
public SearchRequestBuilder addRescorer(RescoreBuilder.Rescorer rescorer, int window) {
sourceBuilder().addRescorer(new RescoreBuilder(rescorer).windowSize(window));
public SearchRequestBuilder addRescorer(RescoreBuilder rescorer, int window) {
sourceBuilder().addRescorer(new RescoreBaseBuilder(rescorer).windowSize(window));
return this;
}

9
core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
View File

@ -37,7 +37,7 @@ import org.elasticsearch.common.geo.builders.ShapeBuilder;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder;
import org.elasticsearch.search.rescore.RescoreBuilder.Rescorer;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
@ -61,7 +61,6 @@ import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;
import static org.elasticsearch.ElasticsearchException.readException;
@ -678,10 +677,10 @@ public abstract class StreamInput extends InputStream {
}
/**
* Reads a {@link QueryBuilder} from the current stream
* Reads a {@link RescoreBuilder} from the current stream
*/
public Rescorer readRescorer() throws IOException {
return readNamedWriteable(Rescorer.class);
public RescoreBuilder readRescorer() throws IOException {
return readNamedWriteable(RescoreBuilder.class);
}
/**

6
core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
View File

@ -36,7 +36,7 @@ import org.elasticsearch.common.geo.builders.ShapeBuilder;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder;
import org.elasticsearch.search.rescore.RescoreBuilder.Rescorer;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.joda.time.ReadableInstant;
import java.io.EOFException;
@ -679,9 +679,9 @@ public abstract class StreamOutput extends OutputStream {
}
/**
* Writes a {@link Rescorer} to the current stream
* Writes a {@link RescoreBuilder} to the current stream
*/
public void writeRescorer(Rescorer rescorer) throws IOException {
public void writeRescorer(RescoreBuilder rescorer) throws IOException {
writeNamedWriteable(rescorer);
}
}

4
core/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java
View File

@ -45,7 +45,7 @@ import org.elasticsearch.search.fetch.innerhits.InnerHitsBuilder;
import org.elasticsearch.search.fetch.source.FetchSourceContext;
import org.elasticsearch.search.highlight.HighlightBuilder;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.RescoreBaseBuilder;
import org.elasticsearch.search.sort.SortBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
@ -458,7 +458,7 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ
return suggestBuilder;
}
public SearchSourceBuilder addRescorer(RescoreBuilder rescoreBuilder) {
public SearchSourceBuilder addRescorer(RescoreBaseBuilder rescoreBuilder) {
try {
if (rescoreBuilders == null) {
rescoreBuilders = new ArrayList<>();

20
core/src/main/java/org/elasticsearch/search/rescore/QueryRescorer.java
View File

@ -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<QueryRescoreContext, SearchContext> RESCORE_PARSER = new ObjectParser<>("query", null);
private static final ObjectParser<QueryRescoreContext, QueryShardContext> 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() {

239
core/src/main/java/org/elasticsearch/search/rescore/QueryRescorerBuilder.java Normal file
View File

@ -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<QueryRescorerBuilder> {
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<InnerBuilder, QueryParseContext> 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 <tt>1.0</tt>
*/
public QueryRescorerBuilder setQueryWeight(float queryWeight) {
this.queryWeight = queryWeight;
return this;
}
/**
* Gets the original query weight for rescoring. The default is <tt>1.0</tt>
*/
public float getQueryWeight() {
return this.queryWeight;
}
/**
* Sets the original query weight for rescoring. The default is <tt>1.0</tt>
*/
public QueryRescorerBuilder setRescoreQueryWeight(float rescoreQueryWeight) {
this.rescoreQueryWeight = rescoreQueryWeight;
return this;
}
/**
* Gets the original query weight for rescoring. The default is <tt>1.0</tt>
*/
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 <tt>total</tt>
*/
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;
}
}
}

173
core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java Normal file
View File

@ -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<RescoreBaseBuilder> {
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) + "\"}";
}
}
}

252
core/src/main/java/org/elasticsearch/search/rescore/RescoreBuilder.java
View File

@ -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<RescoreBuilder> {
public interface RescoreBuilder<RB extends RescoreBuilder> extends ToXContent, NamedWriteable<RB> {
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<Rescorer> {
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 <tt>1.0</tt>
*/
public QueryRescorer setQueryWeight(float queryWeight) {
this.queryWeight = queryWeight;
return this;
}
/**
* Gets the original query weight for rescoring. The default is <tt>1.0</tt>
*/
public float getQueryWeight() {
return this.queryWeight;
}
/**
* Sets the original query weight for rescoring. The default is <tt>1.0</tt>
*/
public QueryRescorer setRescoreQueryWeight(float rescoreQueryWeight) {
this.rescoreQueryWeight = rescoreQueryWeight;
return this;
}
/**
* Gets the original query weight for rescoring. The default is <tt>1.0</tt>
*/
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 <tt>total</tt>
*/
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;
}

11
core/src/main/java/org/elasticsearch/search/rescore/RescoreParseElement.java
View File

@ -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;
}
}

7
core/src/main/java/org/elasticsearch/search/rescore/Rescorer.java
View File

@ -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<Term> termsSet);
/*
* TODO: At this point we only have one implemenation which modifies the
* TopDocs given. Future implemenations might return actual resutls that

4
core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java
View File

@ -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);
}

48
core/src/test/java/org/elasticsearch/search/functionscore/QueryRescorerIT.java
View File

@ -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);
}

224
core/src/test/java/org/elasticsearch/search/rescore/QueryRescoreBuilderTests.java
View File

@ -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<QueryParser> 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<MatchAllQueryBuilder> 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);
}
}
}