Allow plugins to plug rescore implementations (#26368)

This allows plugins to plug rescore implementations into
Elasticsearch. While this is a fairly expert thing to do I've
done my best to point folks to the QueryRescorer as one that at
least documents the tradeoffs that it makes. I've attempted to
limit the API surface area by removing `SearchContext` from the
exposed interface, instead exposing just the IndexSearcher and
`QueryShardContext`. I also tried to make some of the class names
more consistent and do some general cleanup while I was there.

I entertained the notion of moving the `QueryRescorer` to module.
After all, it'd be a wonderful test to prove that you can plug
rescore implementation into Elasticsearch if the only built in
rescore implementation is in the module. But I decided against it
because the new module would require a client jar and it'd require
moving some more things around. I think if we really want to do
it, we should do it as a followup.

I did, on the other hand, create an "example" rescore plugin which
should both be a nice example for anyone wanting to plug in their
own rescore implementation and servers as a good integration test
to make sure that you can indeed plug one in.

Closes #26208
This commit is contained in:
Nik Everett 2017-08-25 13:46:57 -04:00 committed by GitHub
parent 74cd32942a
commit b3edd11aa0
36 changed files with 831 additions and 302 deletions

View File

@ -33,7 +33,7 @@ import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.PipelineAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.slice.SliceBuilder;
import org.elasticsearch.search.sort.SortBuilder;
import org.elasticsearch.search.sort.SortOrder;
@ -397,25 +397,25 @@ public class NoopSearchRequestBuilder extends ActionRequestBuilder<SearchRequest
/**
* Clears all rescorers on the builder and sets the first one. To use multiple rescore windows use
* {@link #addRescorer(org.elasticsearch.search.rescore.RescoreBuilder, int)}.
* {@link #addRescorer(org.elasticsearch.search.rescore.RescorerBuilder, int)}.
*
* @param rescorer rescorer configuration
* @return this for chaining
*/
public NoopSearchRequestBuilder setRescorer(RescoreBuilder<?> rescorer) {
public NoopSearchRequestBuilder setRescorer(RescorerBuilder<?> 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, int)}.
* {@link #addRescorer(org.elasticsearch.search.rescore.RescorerBuilder, int)}.
*
* @param rescorer rescorer configuration
* @param window rescore window
* @return this for chaining
*/
public NoopSearchRequestBuilder setRescorer(RescoreBuilder rescorer, int window) {
public NoopSearchRequestBuilder setRescorer(RescorerBuilder rescorer, int window) {
sourceBuilder().clearRescorers();
return addRescorer(rescorer.windowSize(window));
}
@ -426,7 +426,7 @@ public class NoopSearchRequestBuilder extends ActionRequestBuilder<SearchRequest
* @param rescorer rescorer configuration
* @return this for chaining
*/
public NoopSearchRequestBuilder addRescorer(RescoreBuilder<?> rescorer) {
public NoopSearchRequestBuilder addRescorer(RescorerBuilder<?> rescorer) {
sourceBuilder().addRescorer(rescorer);
return this;
}
@ -438,7 +438,7 @@ public class NoopSearchRequestBuilder extends ActionRequestBuilder<SearchRequest
* @param window rescore window
* @return this for chaining
*/
public NoopSearchRequestBuilder addRescorer(RescoreBuilder<?> rescorer, int window) {
public NoopSearchRequestBuilder addRescorer(RescorerBuilder<?> rescorer, int window) {
sourceBuilder().addRescorer(rescorer.windowSize(window));
return this;
}

View File

@ -40,7 +40,7 @@ import org.elasticsearch.search.SearchService;
import org.elasticsearch.search.internal.AliasFilter;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.internal.ShardSearchLocalRequest;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.rescore.Rescorer;
import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.transport.TransportService;
@ -105,9 +105,9 @@ public class TransportExplainAction extends TransportSingleShardAction<ExplainRe
context.preProcess(true);
int topLevelDocId = result.docIdAndVersion().docId + result.docIdAndVersion().context.docBase;
Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
for (RescoreSearchContext ctx : context.rescore()) {
for (RescoreContext ctx : context.rescore()) {
Rescorer rescorer = ctx.rescorer();
explanation = rescorer.explain(topLevelDocId, context, ctx, explanation);
explanation = rescorer.explain(topLevelDocId, context.searcher(), ctx, explanation);
}
if (request.storedFields() != null || (request.fetchSourceContext() != null && request.fetchSourceContext().fetchSource())) {
// Advantage is that we're not opening a second searcher to retrieve the _source. Also

View File

@ -32,7 +32,7 @@ import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.PipelineAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.slice.SliceBuilder;
import org.elasticsearch.search.sort.SortBuilder;
import org.elasticsearch.search.sort.SortOrder;
@ -415,25 +415,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, int)}.
* {@link #addRescorer(org.elasticsearch.search.rescore.RescorerBuilder, int)}.
*
* @param rescorer rescorer configuration
* @return this for chaining
*/
public SearchRequestBuilder setRescorer(RescoreBuilder<?> rescorer) {
public SearchRequestBuilder setRescorer(RescorerBuilder<?> 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, int)}.
* {@link #addRescorer(org.elasticsearch.search.rescore.RescorerBuilder, int)}.
*
* @param rescorer rescorer configuration
* @param window rescore window
* @return this for chaining
*/
public SearchRequestBuilder setRescorer(RescoreBuilder rescorer, int window) {
public SearchRequestBuilder setRescorer(RescorerBuilder rescorer, int window) {
sourceBuilder().clearRescorers();
return addRescorer(rescorer.windowSize(window));
}
@ -444,7 +444,7 @@ public class SearchRequestBuilder extends ActionRequestBuilder<SearchRequest, Se
* @param rescorer rescorer configuration
* @return this for chaining
*/
public SearchRequestBuilder addRescorer(RescoreBuilder<?> rescorer) {
public SearchRequestBuilder addRescorer(RescorerBuilder<?> rescorer) {
sourceBuilder().addRescorer(rescorer);
return this;
}
@ -456,7 +456,7 @@ public class SearchRequestBuilder extends ActionRequestBuilder<SearchRequest, Se
* @param window rescore window
* @return this for chaining
*/
public SearchRequestBuilder addRescorer(RescoreBuilder<?> rescorer, int window) {
public SearchRequestBuilder addRescorer(RescorerBuilder<?> rescorer, int window) {
sourceBuilder().addRescorer(rescorer.windowSize(window));
return this;
}

View File

@ -108,7 +108,7 @@ public final class ConstructingObjectParser<Value, Context> extends AbstractObje
* ObjectParser.
* @param builder A function that builds the object from an array of Objects. Declare this inline with the parser, casting the elements
* of the array to the arguments so they work with your favorite constructor. The objects in the array will be in the same order
* that you declared the {{@link #constructorArg()}s and none will be null. If any of the constructor arguments aren't defined in
* that you declared the {@link #constructorArg()}s and none will be null. If any of the constructor arguments aren't defined in
* the XContent then parsing will throw an error. We use an array here rather than a {@code Map<String, Object>} to save on
* allocations.
*/
@ -125,7 +125,7 @@ public final class ConstructingObjectParser<Value, Context> extends AbstractObje
* from external systems, never when parsing requests from users.
* @param builder A function that builds the object from an array of Objects. Declare this inline with the parser, casting the elements
* of the array to the arguments so they work with your favorite constructor. The objects in the array will be in the same order
* that you declared the {{@link #constructorArg()}s and none will be null. If any of the constructor arguments aren't defined in
* that you declared the {@link #constructorArg()}s and none will be null. If any of the constructor arguments aren't defined in
* the XContent then parsing will throw an error. We use an array here rather than a {@code Map<String, Object>} to save on
* allocations.
*/
@ -142,7 +142,7 @@ public final class ConstructingObjectParser<Value, Context> extends AbstractObje
* from external systems, never when parsing requests from users.
* @param builder A binary function that builds the object from an array of Objects and the parser context. Declare this inline with
* the parser, casting the elements of the array to the arguments so they work with your favorite constructor. The objects in
* the array will be in the same order that you declared the {{@link #constructorArg()}s and none will be null. The second
* the array will be in the same order that you declared the {@link #constructorArg()}s and none will be null. The second
* argument is the value of the context provided to the {@link #parse(XContentParser, Object) parse function}. If any of the
* constructor arguments aren't defined in the XContent then parsing will throw an error. We use an array here rather than a
* {@code Map<String, Object>} to save on allocations.

View File

@ -48,6 +48,8 @@ import org.elasticsearch.search.aggregations.pipeline.movavg.MovAvgPipelineAggre
import org.elasticsearch.search.aggregations.pipeline.movavg.models.MovAvgModel;
import org.elasticsearch.search.fetch.FetchSubPhase;
import org.elasticsearch.search.fetch.subphase.highlight.Highlighter;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.rescore.Rescorer;
import org.elasticsearch.search.suggest.Suggester;
import org.elasticsearch.search.suggest.SuggestionBuilder;
@ -126,6 +128,12 @@ public interface SearchPlugin {
default List<PipelineAggregationSpec> getPipelineAggregations() {
return emptyList();
}
/**
* The next {@link Rescorer}s added by this plugin.
*/
default List<RescorerSpec<?>> getRescorers() {
return emptyList();
}
/**
* The new search response listeners in the form of {@link BiConsumer}s added by this plugin.
* The listeners are invoked on the coordinating node, at the very end of the search request.
@ -360,6 +368,17 @@ public interface SearchPlugin {
}
}
class RescorerSpec<T extends RescorerBuilder<T>> extends SearchExtensionSpec<T, CheckedFunction<XContentParser, T, IOException>> {
public RescorerSpec(ParseField name, Writeable.Reader<? extends T> reader,
CheckedFunction<XContentParser, T, IOException> parser) {
super(name, reader, parser);
}
public RescorerSpec(String name, Writeable.Reader<? extends T> reader, CheckedFunction<XContentParser, T, IOException> parser) {
super(name, reader, parser);
}
}
/**
* Specification of search time behavior extension like a custom {@link MovAvgModel} or {@link ScoreFunction}.
*

View File

@ -68,7 +68,7 @@ import org.elasticsearch.search.internal.ShardSearchRequest;
import org.elasticsearch.search.profile.Profilers;
import org.elasticsearch.search.query.QueryPhaseExecutionException;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.slice.SliceBuilder;
import org.elasticsearch.search.sort.SortAndFormats;
import org.elasticsearch.search.suggest.SuggestionSearchContext;
@ -143,7 +143,7 @@ final class DefaultSearchContext extends SearchContext {
private SearchContextAggregations aggregations;
private SearchContextHighlight highlight;
private SuggestionSearchContext suggest;
private List<RescoreSearchContext> rescore;
private List<RescoreContext> rescore;
private volatile long keepAlive;
private final long originNanoTime = System.nanoTime();
private volatile long lastAccessTime = -1;
@ -213,11 +213,11 @@ final class DefaultSearchContext extends SearchContext {
}
if (rescore != null) {
int maxWindow = indexService.getIndexSettings().getMaxRescoreWindow();
for (RescoreSearchContext rescoreContext: rescore) {
if (rescoreContext.window() > maxWindow) {
throw new QueryPhaseExecutionException(this, "Rescore window [" + rescoreContext.window() + "] is too large. It must "
+ "be less than [" + maxWindow + "]. This prevents allocating massive heaps for storing the results to be "
+ "rescored. This limit can be set by changing the [" + IndexSettings.MAX_RESCORE_WINDOW_SETTING.getKey()
for (RescoreContext rescoreContext: rescore) {
if (rescoreContext.getWindowSize() > maxWindow) {
throw new QueryPhaseExecutionException(this, "Rescore window [" + rescoreContext.getWindowSize() + "] is too large. "
+ "It must be less than [" + maxWindow + "]. This prevents allocating massive heaps for storing the results "
+ "to be rescored. This limit can be set by changing the [" + IndexSettings.MAX_RESCORE_WINDOW_SETTING.getKey()
+ "] index level setting.");
}
@ -400,7 +400,7 @@ final class DefaultSearchContext extends SearchContext {
}
@Override
public List<RescoreSearchContext> rescore() {
public List<RescoreContext> rescore() {
if (rescore == null) {
return Collections.emptyList();
}
@ -408,7 +408,7 @@ final class DefaultSearchContext extends SearchContext {
}
@Override
public void addRescore(RescoreSearchContext rescore) {
public void addRescore(RescoreContext rescore) {
if (this.rescore == null) {
this.rescore = new ArrayList<>();
}

View File

@ -86,6 +86,7 @@ import org.elasticsearch.plugins.SearchPlugin.AggregationSpec;
import org.elasticsearch.plugins.SearchPlugin.FetchPhaseConstructionContext;
import org.elasticsearch.plugins.SearchPlugin.PipelineAggregationSpec;
import org.elasticsearch.plugins.SearchPlugin.QuerySpec;
import org.elasticsearch.plugins.SearchPlugin.RescorerSpec;
import org.elasticsearch.plugins.SearchPlugin.ScoreFunctionSpec;
import org.elasticsearch.plugins.SearchPlugin.SearchExtSpec;
import org.elasticsearch.plugins.SearchPlugin.SearchExtensionSpec;
@ -230,7 +231,7 @@ import org.elasticsearch.search.fetch.subphase.highlight.Highlighter;
import org.elasticsearch.search.fetch.subphase.highlight.PlainHighlighter;
import org.elasticsearch.search.fetch.subphase.highlight.UnifiedHighlighter;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.GeoDistanceSortBuilder;
import org.elasticsearch.search.sort.ScoreSortBuilder;
@ -281,7 +282,7 @@ public class SearchModule {
highlighters = setupHighlighters(settings, plugins);
registerScoreFunctions(plugins);
registerQueryParsers(plugins);
registerRescorers();
registerRescorers(plugins);
registerSorts();
registerValueFormats();
registerSignificanceHeuristics(plugins);
@ -526,8 +527,16 @@ public class SearchModule {
}
}
private void registerRescorers() {
namedWriteables.add(new NamedWriteableRegistry.Entry(RescoreBuilder.class, QueryRescorerBuilder.NAME, QueryRescorerBuilder::new));
private void registerRescorers(List<SearchPlugin> plugins) {
registerRescorer(new RescorerSpec<>(QueryRescorerBuilder.NAME, QueryRescorerBuilder::new, QueryRescorerBuilder::fromXContent));
registerFromPlugin(plugins, SearchPlugin::getRescorers, this::registerRescorer);
}
private void registerRescorer(RescorerSpec<?> spec) {
if (false == transportClient) {
namedXContents.add(new NamedXContentRegistry.Entry(RescorerBuilder.class, spec.getName(), (p, c) -> spec.getParser().apply(p)));
}
namedWriteables.add(new NamedWriteableRegistry.Entry(RescorerBuilder.class, spec.getName().getPreferredName(), spec.getReader()));
}
private void registerSorts() {

View File

@ -85,7 +85,7 @@ import org.elasticsearch.search.query.QueryPhase;
import org.elasticsearch.search.query.QuerySearchRequest;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.search.query.ScrollQuerySearchResult;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.searchafter.SearchAfterBuilder;
import org.elasticsearch.search.sort.SortAndFormats;
import org.elasticsearch.search.sort.SortBuilder;
@ -723,8 +723,8 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
}
if (source.rescores() != null) {
try {
for (RescoreBuilder<?> rescore : source.rescores()) {
context.addRescore(rescore.build(queryShardContext));
for (RescorerBuilder<?> rescore : source.rescores()) {
context.addRescore(rescore.buildContext(queryShardContext));
}
} catch (IOException e) {
throw new SearchContextException(context, "failed to create RescoreSearchContext", e);

View File

@ -45,7 +45,7 @@ import org.elasticsearch.search.fetch.FetchPhase;
import org.elasticsearch.search.fetch.FetchSearchResult;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.internal.SubSearchContext;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.sort.SortAndFormats;
import java.io.IOException;
@ -115,8 +115,8 @@ public class TopHitsAggregator extends MetricsAggregator {
SortAndFormats sort = subSearchContext.sort();
int topN = subSearchContext.from() + subSearchContext.size();
if (sort == null) {
for (RescoreSearchContext rescoreContext : context.rescore()) {
topN = Math.max(rescoreContext.window(), topN);
for (RescoreContext rescoreContext : context.rescore()) {
topN = Math.max(rescoreContext.getWindowSize(), topN);
}
}
// In the QueryPhase we don't need this protection, because it is build into the IndexSearcher,
@ -148,9 +148,9 @@ public class TopHitsAggregator extends MetricsAggregator {
} else {
TopDocs topDocs = topDocsCollector.topLevelCollector.topDocs();
if (subSearchContext.sort() == null) {
for (RescoreSearchContext ctx : context().rescore()) {
for (RescoreContext ctx : context().rescore()) {
try {
topDocs = ctx.rescorer().rescore(topDocs, context, ctx);
topDocs = ctx.rescorer().rescore(topDocs, context.searcher(), ctx);
} catch (IOException e) {
throw new ElasticsearchException("Rescore TopHits Failed", e);
}

View File

@ -50,7 +50,7 @@ import org.elasticsearch.search.fetch.StoredFieldsContext;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.searchafter.SearchAfterBuilder;
import org.elasticsearch.search.slice.SliceBuilder;
import org.elasticsearch.search.sort.ScoreSortBuilder;
@ -168,7 +168,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
private SuggestBuilder suggestBuilder;
private List<RescoreBuilder> rescoreBuilders;
private List<RescorerBuilder> rescoreBuilders;
private List<IndexBoost> indexBoosts = new ArrayList<>();
@ -202,7 +202,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
postQueryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
queryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
if (in.readBoolean()) {
rescoreBuilders = in.readNamedWriteableList(RescoreBuilder.class);
rescoreBuilders = in.readNamedWriteableList(RescorerBuilder.class);
}
if (in.readBoolean()) {
scriptFields = in.readList(ScriptField::new);
@ -621,7 +621,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
return suggestBuilder;
}
public SearchSourceBuilder addRescorer(RescoreBuilder<?> rescoreBuilder) {
public SearchSourceBuilder addRescorer(RescorerBuilder<?> rescoreBuilder) {
if (rescoreBuilders == null) {
rescoreBuilders = new ArrayList<>();
}
@ -653,7 +653,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
/**
* Gets the bytes representing the rescore builders for this request.
*/
public List<RescoreBuilder> rescores() {
public List<RescorerBuilder> rescores() {
return rescoreBuilders;
}
@ -891,7 +891,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
}
List<SortBuilder<?>> sorts = Rewriteable.rewrite(this.sorts, context);
List<RescoreBuilder> rescoreBuilders = Rewriteable.rewrite(this.rescoreBuilders, context);
List<RescorerBuilder> rescoreBuilders = Rewriteable.rewrite(this.rescoreBuilders, context);
HighlightBuilder highlightBuilder = this.highlightBuilder;
if (highlightBuilder != null) {
highlightBuilder = this.highlightBuilder.rewrite(context);
@ -919,7 +919,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
*/
private SearchSourceBuilder shallowCopy(QueryBuilder queryBuilder, QueryBuilder postQueryBuilder,
AggregatorFactories.Builder aggregations, SliceBuilder slice, List<SortBuilder<?>> sorts,
List<RescoreBuilder> rescoreBuilders, HighlightBuilder highlightBuilder) {
List<RescorerBuilder> rescoreBuilders, HighlightBuilder highlightBuilder) {
SearchSourceBuilder rewrittenBuilder = new SearchSourceBuilder();
rewrittenBuilder.aggregations = aggregations;
rewrittenBuilder.explain = explain;
@ -1034,7 +1034,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
sorts = new ArrayList<>(SortBuilder.fromXContent(parser));
} else if (RESCORE_FIELD.match(currentFieldName)) {
rescoreBuilders = new ArrayList<>();
rescoreBuilders.add(RescoreBuilder.parseFromXContent(parser));
rescoreBuilders.add(RescorerBuilder.parseFromXContent(parser));
} else if (EXT_FIELD.match(currentFieldName)) {
extBuilders = new ArrayList<>();
String extSectionName = null;
@ -1081,7 +1081,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
} else if (RESCORE_FIELD.match(currentFieldName)) {
rescoreBuilders = new ArrayList<>();
while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
rescoreBuilders.add(RescoreBuilder.parseFromXContent(parser));
rescoreBuilders.add(RescorerBuilder.parseFromXContent(parser));
}
} else if (STATS_FIELD.match(currentFieldName)) {
stats = new ArrayList<>();
@ -1222,7 +1222,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
if (rescoreBuilders != null) {
builder.startArray(RESCORE_FIELD.getPreferredName());
for (RescoreBuilder<?> rescoreBuilder : rescoreBuilders) {
for (RescorerBuilder<?> rescoreBuilder : rescoreBuilders) {
rescoreBuilder.toXContent(builder, params);
}
builder.endArray();

View File

@ -22,18 +22,19 @@ package org.elasticsearch.search.dfs;
import com.carrotsearch.hppc.ObjectHashSet;
import com.carrotsearch.hppc.ObjectObjectHashMap;
import com.carrotsearch.hppc.cursors.ObjectCursor;
import org.apache.lucene.index.IndexReaderContext;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermContext;
import org.apache.lucene.search.CollectionStatistics;
import org.apache.lucene.search.TermStatistics;
import org.elasticsearch.common.collect.HppcMaps;
import org.elasticsearch.search.SearchContextException;
import org.elasticsearch.search.SearchPhase;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.tasks.TaskCancelledException;
import java.io.IOException;
import java.util.AbstractSet;
import java.util.Collection;
import java.util.Iterator;
@ -53,8 +54,12 @@ public class DfsPhase implements SearchPhase {
final ObjectHashSet<Term> termsSet = new ObjectHashSet<>();
try {
context.searcher().createNormalizedWeight(context.query(), true).extractTerms(new DelegateSet(termsSet));
for (RescoreSearchContext rescoreContext : context.rescore()) {
rescoreContext.rescorer().extractTerms(context, rescoreContext, new DelegateSet(termsSet));
for (RescoreContext rescoreContext : context.rescore()) {
try {
rescoreContext.rescorer().extractTerms(context.searcher(), rescoreContext, new DelegateSet(termsSet));
} catch (IOException e) {
throw new IllegalStateException("Failed to extract terms", e);
}
}
Term[] terms = termsSet.toArray(Term.class);

View File

@ -22,7 +22,7 @@ import org.apache.lucene.search.Explanation;
import org.elasticsearch.search.fetch.FetchPhaseExecutionException;
import org.elasticsearch.search.fetch.FetchSubPhase;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import java.io.IOException;
@ -40,8 +40,8 @@ public final class ExplainFetchSubPhase implements FetchSubPhase {
final int topLevelDocId = hitContext.hit().docId();
Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
for (RescoreSearchContext rescore : context.rescore()) {
explanation = rescore.rescorer().explain(topLevelDocId, context, rescore, explanation);
for (RescoreContext rescore : context.rescore()) {
explanation = rescore.rescorer().explain(topLevelDocId, context.searcher(), rescore, explanation);
}
// we use the top level doc id, since we work with the top level searcher
hitContext.hit().explanation(explanation);

View File

@ -51,7 +51,7 @@ import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
import org.elasticsearch.search.lookup.SearchLookup;
import org.elasticsearch.search.profile.Profilers;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.sort.SortAndFormats;
import org.elasticsearch.search.suggest.SuggestionSearchContext;
@ -192,12 +192,12 @@ public abstract class FilteredSearchContext extends SearchContext {
}
@Override
public List<RescoreSearchContext> rescore() {
public List<RescoreContext> rescore() {
return in.rescore();
}
@Override
public void addRescore(RescoreSearchContext rescore) {
public void addRescore(RescoreContext rescore) {
in.addRescore(rescore);
}

View File

@ -58,7 +58,7 @@ import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
import org.elasticsearch.search.lookup.SearchLookup;
import org.elasticsearch.search.profile.Profilers;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.sort.SortAndFormats;
import org.elasticsearch.search.suggest.SuggestionSearchContext;
@ -175,9 +175,9 @@ public abstract class SearchContext extends AbstractRefCounted implements Releas
/**
* @return list of all rescore contexts. empty if there aren't any.
*/
public abstract List<RescoreSearchContext> rescore();
public abstract List<RescoreContext> rescore();
public abstract void addRescore(RescoreSearchContext rescore);
public abstract void addRescore(RescoreContext rescore);
public abstract boolean hasScriptFields();

View File

@ -31,7 +31,7 @@ import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.sort.SortAndFormats;
import org.elasticsearch.search.suggest.SuggestionSearchContext;
@ -111,7 +111,7 @@ public class SubSearchContext extends FilteredSearchContext {
}
@Override
public void addRescore(RescoreSearchContext rescore) {
public void addRescore(RescoreContext rescore) {
throw new UnsupportedOperationException("Not supported");
}

View File

@ -43,7 +43,7 @@ import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.collapse.CollapseContext;
import org.elasticsearch.search.internal.ScrollContext;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.sort.SortAndFormats;
import java.io.IOException;
@ -289,8 +289,8 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
} else {
int numDocs = Math.min(searchContext.from() + searchContext.size(), totalNumDocs);
final boolean rescore = searchContext.rescore().isEmpty() == false;
for (RescoreSearchContext rescoreContext : searchContext.rescore()) {
numDocs = Math.max(numDocs, rescoreContext.window());
for (RescoreContext rescoreContext : searchContext.rescore()) {
numDocs = Math.max(numDocs, rescoreContext.getWindowSize());
}
return new SimpleTopDocsCollectorContext(searchContext.sort(),
searchContext.searchAfter(),

View File

@ -21,11 +21,10 @@ package org.elasticsearch.search.rescore;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.elasticsearch.search.internal.ContextIndexSearcher;
import org.elasticsearch.search.internal.SearchContext;
import java.io.IOException;
import java.util.Arrays;
@ -35,15 +34,9 @@ import java.util.Set;
public final class QueryRescorer implements Rescorer {
public static final Rescorer INSTANCE = new QueryRescorer();
public static final String NAME = "query";
@Override
public String name() {
return NAME;
}
@Override
public TopDocs rescore(TopDocs topDocs, SearchContext context, RescoreSearchContext rescoreContext) throws IOException {
public TopDocs rescore(TopDocs topDocs, IndexSearcher searcher, RescoreContext rescoreContext) throws IOException {
assert rescoreContext != null;
if (topDocs == null || topDocs.totalHits == 0 || topDocs.scoreDocs.length == 0) {
@ -66,20 +59,19 @@ public final class QueryRescorer implements Rescorer {
};
// First take top slice of incoming docs, to be rescored:
TopDocs topNFirstPass = topN(topDocs, rescoreContext.window());
TopDocs topNFirstPass = topN(topDocs, rescoreContext.getWindowSize());
// Rescore them:
TopDocs rescored = rescorer.rescore(context.searcher(), topNFirstPass, rescoreContext.window());
TopDocs rescored = rescorer.rescore(searcher, topNFirstPass, rescoreContext.getWindowSize());
// Splice back to non-topN hits and resort all of them:
return combine(topDocs, rescored, (QueryRescoreContext) rescoreContext);
}
@Override
public Explanation explain(int topLevelDocId, SearchContext context, RescoreSearchContext rescoreContext,
public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreContext rescoreContext,
Explanation sourceExplanation) throws IOException {
QueryRescoreContext rescore = (QueryRescoreContext) rescoreContext;
ContextIndexSearcher searcher = context.searcher();
if (sourceExplanation == null) {
// this should not happen but just in case
return Explanation.noMatch("nothing matched");
@ -160,20 +152,17 @@ public final class QueryRescorer implements Rescorer {
return in;
}
public static class QueryRescoreContext extends RescoreSearchContext {
static final int DEFAULT_WINDOW_SIZE = 10;
public QueryRescoreContext(QueryRescorer rescorer) {
super(NAME, DEFAULT_WINDOW_SIZE, rescorer);
this.scoreMode = QueryRescoreMode.Total;
}
public static class QueryRescoreContext extends RescoreContext {
private Query query;
private float queryWeight = 1.0f;
private float rescoreQueryWeight = 1.0f;
private QueryRescoreMode scoreMode;
public QueryRescoreContext(int windowSize) {
super(windowSize, QueryRescorer.INSTANCE);
this.scoreMode = QueryRescoreMode.Total;
}
public void setQuery(Query query) {
this.query = query;
}
@ -212,12 +201,8 @@ public final class QueryRescorer implements Rescorer {
}
@Override
public void extractTerms(SearchContext context, RescoreSearchContext rescoreContext, Set<Term> termsSet) {
try {
context.searcher().createNormalizedWeight(((QueryRescoreContext) rescoreContext).query(), false).extractTerms(termsSet);
} catch (IOException e) {
throw new IllegalStateException("Failed to extract terms", e);
}
public void extractTerms(IndexSearcher searcher, RescoreContext rescoreContext, Set<Term> termsSet) throws IOException {
searcher.createNormalizedWeight(((QueryRescoreContext) rescoreContext).query(), false).extractTerms(termsSet);
}
}

View File

@ -29,7 +29,6 @@ import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryRewriteContext;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.index.query.Rewriteable;
import org.elasticsearch.search.rescore.QueryRescorer.QueryRescoreContext;
import java.io.IOException;
@ -38,25 +37,15 @@ import java.util.Objects;
import static org.elasticsearch.index.query.AbstractQueryBuilder.parseInnerQueryBuilder;
public class QueryRescorerBuilder extends RescoreBuilder<QueryRescorerBuilder> {
public class QueryRescorerBuilder extends RescorerBuilder<QueryRescorerBuilder> {
public static final String NAME = "query";
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 ParseField RESCORE_QUERY_FIELD = new ParseField("rescore_query");
private static final ParseField QUERY_WEIGHT_FIELD = new ParseField("query_weight");
private static final ParseField RESCORE_QUERY_WEIGHT_FIELD = new ParseField("rescore_query_weight");
private static final ParseField SCORE_MODE_FIELD = new ParseField("score_mode");
private static final ObjectParser<InnerBuilder, Void> QUERY_RESCORE_PARSER = new ObjectParser<>(NAME, null);
static {
QUERY_RESCORE_PARSER.declareObject(InnerBuilder::setQueryBuilder, (p, c) -> {
try {
@ -70,6 +59,14 @@ public class QueryRescorerBuilder extends RescoreBuilder<QueryRescorerBuilder> {
QUERY_RESCORE_PARSER.declareString((struct, value) -> struct.setScoreMode(QueryRescoreMode.fromString(value)), SCORE_MODE_FIELD);
}
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 QueryRescorerBuilder} instance
* @param builder the query builder to build the rescore query from
@ -100,6 +97,11 @@ public class QueryRescorerBuilder extends RescoreBuilder<QueryRescorerBuilder> {
out.writeFloat(queryWeight);
}
@Override
public String getWriteableName() {
return NAME;
}
/**
* @return the query used for this rescore query
*/
@ -169,17 +171,13 @@ public class QueryRescorerBuilder extends RescoreBuilder<QueryRescorerBuilder> {
}
@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);
public QueryRescoreContext innerBuildContext(int windowSize, QueryShardContext context) throws IOException {
QueryRescoreContext queryRescoreContext = new QueryRescoreContext(windowSize);
// query is rewritten at this point already
queryRescoreContext.setQuery(queryBuilder.toQuery(context));
queryRescoreContext.setQueryWeight(this.queryWeight);
queryRescoreContext.setRescoreQueryWeight(this.rescoreQueryWeight);
queryRescoreContext.setScoreMode(this.scoreMode);
if (this.windowSize != null) {
queryRescoreContext.setWindowSize(this.windowSize);
}
return queryRescoreContext;
}
@ -205,11 +203,6 @@ public class QueryRescorerBuilder extends RescoreBuilder<QueryRescorerBuilder> {
Objects.equals(queryBuilder, other.queryBuilder);
}
@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
@ -248,7 +241,7 @@ public class QueryRescorerBuilder extends RescoreBuilder<QueryRescorerBuilder> {
}
@Override
public RescoreBuilder rewrite(QueryRewriteContext ctx) throws IOException {
public QueryRescorerBuilder rewrite(QueryRewriteContext ctx) throws IOException {
QueryBuilder rewrite = queryBuilder.rewrite(ctx);
if (rewrite == queryBuilder) {
return this;

View File

@ -19,37 +19,35 @@
package org.elasticsearch.search.rescore;
public class RescoreSearchContext {
private int windowSize;
private final String type;
/**
* Context available to the rescore while it is running. Rescore
* implementations should extend this with any additional resources that
* they will need while rescoring.
*/
public class RescoreContext {
private final int windowSize;
private final Rescorer rescorer;
public RescoreSearchContext(String type, int windowSize, Rescorer rescorer) {
super();
this.type = type;
/**
* Build the context.
* @param rescorer the rescorer actually performing the rescore.
*/
public RescoreContext(int windowSize, Rescorer rescorer) {
this.windowSize = windowSize;
this.rescorer = rescorer;
}
/**
* The rescorer to actually apply.
*/
public Rescorer rescorer() {
return rescorer;
}
public String getType() {
return type;
}
public void setWindowSize(int windowSize) {
this.windowSize = windowSize;
}
public int window() {
/**
* Size of the window to rescore.
*/
public int getWindowSize() {
return windowSize;
}
}

View File

@ -45,8 +45,8 @@ public class RescorePhase extends AbstractComponent implements SearchPhase {
public void execute(SearchContext context) {
try {
TopDocs topDocs = context.queryResult().topDocs();
for (RescoreSearchContext ctx : context.rescore()) {
topDocs = ctx.rescorer().rescore(topDocs, context, ctx);
for (RescoreContext ctx : context.rescore()) {
topDocs = ctx.rescorer().rescore(topDocs, context.searcher(), ctx);
}
context.queryResult().topDocs(topDocs, context.queryResult().sortValueFormats());
} catch (IOException e) {

View File

@ -21,9 +21,9 @@ package org.elasticsearch.search.rescore;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.TopDocs;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.search.internal.SearchContext;
import java.io.IOException;
import java.util.Set;
@ -31,36 +31,34 @@ import java.util.Set;
/**
* A query rescorer interface used to re-rank the Top-K results of a previously
* executed search.
*
* Subclasses should borrow heavily from {@link QueryRescorer} because it is
* fairly well behaved and documents that tradeoffs that it is making. There
* is also an {@code ExampleRescorer} that is worth looking at.
*/
public interface Rescorer {
/**
* Returns the name of this rescorer
*/
String name();
/**
* Modifies the result of the previously executed search ({@link TopDocs})
* in place based on the given {@link RescoreSearchContext}.
* in place based on the given {@link RescoreContext}.
*
* @param topDocs the result of the previously executed search
* @param context the current {@link SearchContext}. This will never be <code>null</code>.
* @param rescoreContext the {@link RescoreSearchContext}. This will never be <code>null</code>
* @param searcher the searcher used for this search. This will never be <code>null</code>.
* @param rescoreContext the {@link RescoreContext}. This will never be <code>null</code>
* @throws IOException if an {@link IOException} occurs during rescoring
*/
TopDocs rescore(TopDocs topDocs, SearchContext context, RescoreSearchContext rescoreContext) throws IOException;
TopDocs rescore(TopDocs topDocs, IndexSearcher searcher, RescoreContext rescoreContext) throws IOException;
/**
* Executes an {@link Explanation} phase on the rescorer.
*
* @param topLevelDocId the global / top-level document ID to explain
* @param context the explanation for the results being fed to this rescorer
* @param searcher the searcher used for this search. This will never be <code>null</code>.
* @param rescoreContext context for this rescorer
* @param sourceExplanation explanation of the source of the documents being fed into this rescore
* @return the explain for the given top level document ID.
* @throws IOException if an {@link IOException} occurs
*/
Explanation explain(int topLevelDocId, SearchContext context, RescoreSearchContext rescoreContext,
Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreContext rescoreContext,
Explanation sourceExplanation) throws IOException;
/**
@ -68,17 +66,5 @@ public interface Rescorer {
* is executed in a distributed frequency collection roundtrip for
* {@link SearchType#DFS_QUERY_THEN_FETCH}
*/
void extractTerms(SearchContext context, RescoreSearchContext rescoreContext, Set<Term> termsSet);
/*
* TODO: At this point we only have one implementation which modifies the
* TopDocs given. Future implementations might return actual results that
* contain information about the rescore context. For example a pair wise
* reranker might return the feature vector for the top N window in order to
* merge results on the callers side. For now we don't have a return type at
* all since something like this requires a more general refactoring how
* documents are merged since in such a case we don't really have a score
* per document rather a "X is more relevant than Y" relation
*/
void extractTerms(IndexSearcher searcher, RescoreContext rescoreContext, Set<Term> termsSet) throws IOException;
}

View File

@ -28,34 +28,33 @@ import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.ToXContentObject;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.index.query.Rewriteable;
import org.elasticsearch.search.rescore.QueryRescorer.QueryRescoreContext;
import java.io.IOException;
import java.util.Objects;
/**
* The abstract base builder for instances of {@link RescoreBuilder}.
* The abstract base builder for instances of {@link RescorerBuilder}.
*/
public abstract class RescoreBuilder<RB extends RescoreBuilder<RB>>
implements NamedWriteable, ToXContentObject, Rewriteable<RescoreBuilder<RB>> {
public abstract class RescorerBuilder<RB extends RescorerBuilder<RB>>
implements NamedWriteable, ToXContentObject, Rewriteable<RescorerBuilder<RB>> {
public static final int DEFAULT_WINDOW_SIZE = 10;
protected Integer windowSize;
private static ParseField WINDOW_SIZE_FIELD = new ParseField("window_size");
private static final ParseField WINDOW_SIZE_FIELD = new ParseField("window_size");
/**
* Construct an empty RescoreBuilder.
*/
public RescoreBuilder() {
public RescorerBuilder() {
}
/**
* Read from a stream.
*/
protected RescoreBuilder(StreamInput in) throws IOException {
protected RescorerBuilder(StreamInput in) throws IOException {
windowSize = in.readOptionalVInt();
}
@ -77,9 +76,9 @@ public abstract class RescoreBuilder<RB extends RescoreBuilder<RB>>
return windowSize;
}
public static RescoreBuilder<?> parseFromXContent(XContentParser parser) throws IOException {
public static RescorerBuilder<?> parseFromXContent(XContentParser parser) throws IOException {
String fieldName = null;
RescoreBuilder<?> rescorer = null;
RescorerBuilder<?> rescorer = null;
Integer windowSize = null;
XContentParser.Token token;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
@ -92,12 +91,7 @@ public abstract class RescoreBuilder<RB extends RescoreBuilder<RB>>
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.fromXContent(parser);
} else {
throw new ParsingException(parser.getTokenLocation(), "rescore doesn't support rescorer with name [" + fieldName + "]");
}
rescorer = parser.namedObject(RescorerBuilder.class, fieldName, null);
} else {
throw new ParsingException(parser.getTokenLocation(), "unexpected token [" + token + "] after [" + fieldName + "]");
}
@ -124,12 +118,21 @@ public abstract class RescoreBuilder<RB extends RescoreBuilder<RB>>
protected abstract void doXContent(XContentBuilder builder, Params params) throws IOException;
public abstract QueryRescoreContext build(QueryShardContext context) throws IOException;
public static QueryRescorerBuilder queryRescorer(QueryBuilder queryBuilder) {
return new QueryRescorerBuilder(queryBuilder);
/**
* Build the {@linkplain RescoreContext} that will be used to actually
* execute the rescore against a particular shard.
*/
public final RescoreContext buildContext(QueryShardContext context) throws IOException {
int finalWindowSize = windowSize == null ? DEFAULT_WINDOW_SIZE : windowSize;
RescoreContext rescoreContext = innerBuildContext(finalWindowSize, context);
return rescoreContext;
}
/**
* Extensions override this to build the context that they need for rescoring.
*/
protected abstract RescoreContext innerBuildContext(int windowSize, QueryShardContext context) throws IOException;
@Override
public int hashCode() {
return Objects.hash(windowSize);
@ -143,8 +146,7 @@ public abstract class RescoreBuilder<RB extends RescoreBuilder<RB>>
if (obj == null || getClass() != obj.getClass()) {
return false;
}
@SuppressWarnings("rawtypes")
RescoreBuilder other = (RescoreBuilder) obj;
RescorerBuilder<?> other = (RescorerBuilder<?>) obj;
return Objects.equals(windowSize, other.windowSize);
}

View File

@ -35,7 +35,7 @@ import org.elasticsearch.plugins.SearchPlugin;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.collapse.CollapseBuilderTests;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilderTests;
import org.elasticsearch.search.rescore.QueryRescoreBuilderTests;
import org.elasticsearch.search.rescore.QueryRescorerBuilderTests;
import org.elasticsearch.search.suggest.SuggestBuilderTests;
import org.elasticsearch.test.ESTestCase;
@ -90,7 +90,7 @@ public abstract class AbstractSearchTestCase extends ESTestCase {
return RandomSearchRequestGenerator.randomSearchSourceBuilder(
HighlightBuilderTests::randomHighlighterBuilder,
SuggestBuilderTests::randomSuggestBuilder,
QueryRescoreBuilderTests::randomRescoreBuilder,
QueryRescorerBuilderTests::randomRescoreBuilder,
randomExtBuilders,
CollapseBuilderTests::randomCollapseBuilder);
}

View File

@ -26,6 +26,8 @@ import org.elasticsearch.common.xcontent.NamedXContentRegistry;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryRewriteContext;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.index.query.functionscore.GaussDecayFunctionBuilder;
import org.elasticsearch.plugins.SearchPlugin;
@ -58,6 +60,9 @@ import org.elasticsearch.search.fetch.subphase.highlight.Highlighter;
import org.elasticsearch.search.fetch.subphase.highlight.PlainHighlighter;
import org.elasticsearch.search.fetch.subphase.highlight.UnifiedHighlighter;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.suggest.CustomSuggesterSearchIT.CustomSuggestionBuilder;
import org.elasticsearch.search.suggest.SuggestionBuilder;
import org.elasticsearch.search.suggest.term.TermSuggestionBuilder;
@ -87,8 +92,7 @@ public class SearchModuleTests extends ModuleTestCase {
return singletonMap("plain", new PlainHighlighter());
}
};
expectThrows(IllegalArgumentException.class,
() -> new SearchModule(Settings.EMPTY, false, singletonList(registersDupeHighlighter)));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeHighlighter));
SearchPlugin registersDupeSuggester = new SearchPlugin() {
@Override
@ -96,8 +100,7 @@ public class SearchModuleTests extends ModuleTestCase {
return singletonList(new SuggesterSpec<>("term", TermSuggestionBuilder::new, TermSuggestionBuilder::fromXContent));
}
};
expectThrows(IllegalArgumentException.class, () -> new NamedXContentRegistry(
new SearchModule(Settings.EMPTY, false, singletonList(registersDupeSuggester)).getNamedXContents()));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeSuggester));
SearchPlugin registersDupeScoreFunction = new SearchPlugin() {
@Override
@ -106,8 +109,7 @@ public class SearchModuleTests extends ModuleTestCase {
GaussDecayFunctionBuilder.PARSER));
}
};
expectThrows(IllegalArgumentException.class, () -> new NamedXContentRegistry(
new SearchModule(Settings.EMPTY, false, singletonList(registersDupeScoreFunction)).getNamedXContents()));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeScoreFunction));
SearchPlugin registersDupeSignificanceHeuristic = new SearchPlugin() {
@Override
@ -115,8 +117,7 @@ public class SearchModuleTests extends ModuleTestCase {
return singletonList(new SearchExtensionSpec<>(ChiSquare.NAME, ChiSquare::new, ChiSquare.PARSER));
}
};
expectThrows(IllegalArgumentException.class, () -> new SearchModule(Settings.EMPTY, false,
singletonList(registersDupeSignificanceHeuristic)));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeSignificanceHeuristic));
SearchPlugin registersDupeMovAvgModel = new SearchPlugin() {
@Override
@ -124,8 +125,7 @@ public class SearchModuleTests extends ModuleTestCase {
return singletonList(new SearchExtensionSpec<>(SimpleModel.NAME, SimpleModel::new, SimpleModel.PARSER));
}
};
expectThrows(IllegalArgumentException.class, () -> new SearchModule(Settings.EMPTY, false,
singletonList(registersDupeMovAvgModel)));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeMovAvgModel));
SearchPlugin registersDupeFetchSubPhase = new SearchPlugin() {
@Override
@ -133,8 +133,7 @@ public class SearchModuleTests extends ModuleTestCase {
return singletonList(new ExplainFetchSubPhase());
}
};
expectThrows(IllegalArgumentException.class, () -> new SearchModule(Settings.EMPTY, false,
singletonList(registersDupeFetchSubPhase)));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeFetchSubPhase));
SearchPlugin registersDupeQuery = new SearchPlugin() {
@Override
@ -142,8 +141,7 @@ public class SearchModuleTests extends ModuleTestCase {
return singletonList(new QuerySpec<>(TermQueryBuilder.NAME, TermQueryBuilder::new, TermQueryBuilder::fromXContent));
}
};
expectThrows(IllegalArgumentException.class, () -> new NamedXContentRegistry(
new SearchModule(Settings.EMPTY, false, singletonList(registersDupeQuery)).getNamedXContents()));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeQuery));
SearchPlugin registersDupeAggregation = new SearchPlugin() {
@Override
@ -152,8 +150,7 @@ public class SearchModuleTests extends ModuleTestCase {
TermsAggregationBuilder::parse));
}
};
expectThrows(IllegalArgumentException.class, () -> new NamedXContentRegistry(new SearchModule(Settings.EMPTY, false,
singletonList(registersDupeAggregation)).getNamedXContents()));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeAggregation));
SearchPlugin registersDupePipelineAggregation = new SearchPlugin() {
@Override
@ -166,8 +163,19 @@ public class SearchModuleTests extends ModuleTestCase {
.addResultReader(InternalDerivative::new));
}
};
expectThrows(IllegalArgumentException.class, () -> new NamedXContentRegistry(new SearchModule(Settings.EMPTY, false,
singletonList(registersDupePipelineAggregation)).getNamedXContents()));
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupePipelineAggregation));
SearchPlugin registersDupeRescorer = new SearchPlugin() {
public List<RescorerSpec<?>> getRescorers() {
return singletonList(
new RescorerSpec<>(QueryRescorerBuilder.NAME, QueryRescorerBuilder::new, QueryRescorerBuilder::fromXContent));
}
};
expectThrows(IllegalArgumentException.class, registryForPlugin(registersDupeRescorer));
}
private ThrowingRunnable registryForPlugin(SearchPlugin plugin) {
return () -> new NamedXContentRegistry(new SearchModule(Settings.EMPTY, false, singletonList(plugin)).getNamedXContents());
}
public void testRegisterSuggester() {
@ -262,6 +270,20 @@ public class SearchModuleTests extends ModuleTestCase {
hasSize(1));
}
public void testRegisterRescorer() {
SearchModule module = new SearchModule(Settings.EMPTY, false, singletonList(new SearchPlugin() {
@Override
public List<RescorerSpec<?>> getRescorers() {
return singletonList(new RescorerSpec<>("test", TestRescorerBuilder::new, TestRescorerBuilder::fromXContent));
}
}));
assertThat(
module.getNamedXContents().stream()
.filter(entry -> entry.categoryClass.equals(RescorerBuilder.class) && entry.name.match("test"))
.collect(toList()),
hasSize(1));
}
private static final String[] NON_DEPRECATED_QUERIES = new String[] {
"bool",
"boosting",
@ -424,4 +446,37 @@ public class SearchModuleTests extends ModuleTestCase {
return null;
}
}
private static class TestRescorerBuilder extends RescorerBuilder<TestRescorerBuilder> {
public static TestRescorerBuilder fromXContent(XContentParser parser) {
return null;
}
TestRescorerBuilder(StreamInput in) throws IOException {
super(in);
}
@Override
public String getWriteableName() {
return "test";
}
@Override
public RescorerBuilder<TestRescorerBuilder> rewrite(QueryRewriteContext ctx) throws IOException {
return this;
}
@Override
protected void doWriteTo(StreamOutput out) throws IOException {
}
@Override
protected void doXContent(XContentBuilder builder, Params params) throws IOException {
}
@Override
public RescoreContext innerBuildContext(int windowSize, QueryShardContext context) throws IOException {
return null;
}
}
}

View File

@ -49,7 +49,7 @@ import org.elasticsearch.search.aggregations.metrics.max.Max;
import org.elasticsearch.search.aggregations.metrics.tophits.TopHits;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
@ -1054,7 +1054,7 @@ public class TopHitsIT extends ESIntegTestCase {
SearchResponse response = client()
.prepareSearch("idx")
.addRescorer(
RescoreBuilder.queryRescorer(new MatchAllQueryBuilder().boost(3.0f))
new QueryRescorerBuilder(new MatchAllQueryBuilder().boost(3.0f))
)
.addAggregation(terms("terms")
.field(TERMS_AGGS_FIELD)
@ -1076,7 +1076,7 @@ public class TopHitsIT extends ESIntegTestCase {
SearchResponse response = client()
.prepareSearch("idx")
.addRescorer(
RescoreBuilder.queryRescorer(new MatchAllQueryBuilder().boost(3.0f))
new QueryRescorerBuilder(new MatchAllQueryBuilder().boost(3.0f))
)
.addAggregation(terms("terms")
.field(TERMS_AGGS_FIELD)
@ -1099,7 +1099,7 @@ public class TopHitsIT extends ESIntegTestCase {
SearchResponse response = client()
.prepareSearch("idx")
.addRescorer(
RescoreBuilder.queryRescorer(new MatchAllQueryBuilder().boost(3.0f))
new QueryRescorerBuilder(new MatchAllQueryBuilder().boost(3.0f))
)
.addAggregation(terms("terms")
.field(TERMS_AGGS_FIELD)
@ -1121,7 +1121,7 @@ public class TopHitsIT extends ESIntegTestCase {
SearchResponse response = client()
.prepareSearch("idx")
.addRescorer(
RescoreBuilder.queryRescorer(new MatchAllQueryBuilder().boost(3.0f))
new QueryRescorerBuilder(new MatchAllQueryBuilder().boost(3.0f))
)
.addAggregation(terms("terms")
.field(TERMS_AGGS_FIELD)

View File

@ -38,7 +38,6 @@ import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.rescore.QueryRescoreMode;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.test.ESIntegTestCase;
import java.util.Arrays;
@ -51,10 +50,12 @@ import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
import static org.elasticsearch.index.query.QueryBuilders.constantScoreQuery;
import static org.elasticsearch.index.query.QueryBuilders.functionScoreQuery;
import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
import static org.elasticsearch.index.query.QueryBuilders.matchPhraseQuery;
import static org.elasticsearch.index.query.QueryBuilders.matchQuery;
import static org.elasticsearch.index.query.QueryBuilders.queryStringQuery;
import static org.elasticsearch.index.query.QueryBuilders.termQuery;
import static org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders.weightFactorFunction;
import static org.elasticsearch.search.rescore.RescoreBuilder.queryRescorer;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertFirstHit;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertFourthHit;
@ -84,8 +85,8 @@ public class QueryRescorerIT extends ESIntegTestCase {
for (int j = 0 ; j < iters; j++) {
SearchResponse searchResponse = client().prepareSearch()
.setQuery(QueryBuilders.matchAllQuery())
.setRescorer(queryRescorer(
QueryBuilders.functionScoreQuery(QueryBuilders.matchAllQuery(),
.setRescorer(new QueryRescorerBuilder(
functionScoreQuery(matchAllQuery(),
ScoreFunctionBuilders.weightFactorFunction(100)).boostMode(CombineFunction.REPLACE))
.setQueryWeight(0.0f).setRescoreQueryWeight(1.0f), 1).setSize(randomIntBetween(2, 10)).execute()
.actionGet();
@ -120,7 +121,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
SearchResponse searchResponse = client().prepareSearch()
.setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR))
.setRescorer(
queryRescorer(QueryBuilders.matchPhraseQuery("field1", "quick brown").slop(2).boost(4.0f))
new QueryRescorerBuilder(matchPhraseQuery("field1", "quick brown").slop(2).boost(4.0f))
.setRescoreQueryWeight(2), 5).execute().actionGet();
assertThat(searchResponse.getHits().getTotalHits(), equalTo(3L));
@ -131,7 +132,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
searchResponse = client().prepareSearch()
.setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR))
.setRescorer(queryRescorer(QueryBuilders.matchPhraseQuery("field1", "the quick brown").slop(3)), 5)
.setRescorer(new QueryRescorerBuilder(matchPhraseQuery("field1", "the quick brown").slop(3)), 5)
.execute().actionGet();
assertHitCount(searchResponse, 3);
@ -141,7 +142,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
searchResponse = client().prepareSearch()
.setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR))
.setRescorer(queryRescorer((QueryBuilders.matchPhraseQuery("field1", "the quick brown"))), 5).execute()
.setRescorer(new QueryRescorerBuilder(matchPhraseQuery("field1", "the quick brown")), 5).execute()
.actionGet();
assertHitCount(searchResponse, 3);
@ -187,7 +188,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setQuery(QueryBuilders.matchQuery("field1", "lexington avenue massachusetts").operator(Operator.OR))
.setFrom(0)
.setSize(5)
.setRescorer(queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setRescorer(new QueryRescorerBuilder(matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 20).execute().actionGet();
assertThat(searchResponse.getHits().getHits().length, equalTo(5));
@ -202,7 +203,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setFrom(0)
.setSize(5)
.setSearchType(SearchType.DFS_QUERY_THEN_FETCH)
.setRescorer(queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setRescorer(new QueryRescorerBuilder(matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 20).execute().actionGet();
assertThat(searchResponse.getHits().getHits().length, equalTo(5));
@ -219,7 +220,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setFrom(2)
.setSize(5)
.setSearchType(SearchType.DFS_QUERY_THEN_FETCH)
.setRescorer(queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setRescorer(new QueryRescorerBuilder(matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 20).execute().actionGet();
assertThat(searchResponse.getHits().getHits().length, equalTo(5));
@ -270,7 +271,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setQuery(QueryBuilders.matchQuery("field1", "massachusetts"))
.setFrom(0)
.setSize(5)
.setRescorer(queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setRescorer(new QueryRescorerBuilder(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().getHits().length, equalTo(4));
@ -287,7 +288,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setQuery(QueryBuilders.matchQuery("field1", "massachusetts"))
.setFrom(0)
.setSize(5)
.setRescorer(queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setRescorer(new QueryRescorerBuilder(matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setQueryWeight(0.6f).setRescoreQueryWeight(2.0f), 3).execute().actionGet();
// Only top 3 hits were re-ordered:
@ -342,7 +343,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setQuery(QueryBuilders.matchQuery("field1", "massachusetts").operator(Operator.OR))
.setFrom(0)
.setSize(5)
.setRescorer(queryRescorer(QueryBuilders.matchPhraseQuery("field1", "lexington avenue massachusetts").slop(3))
.setRescorer(new QueryRescorerBuilder(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:
@ -410,7 +411,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setQuery(QueryBuilders.matchQuery("field1", query).operator(Operator.OR))
.setFrom(0)
.setSize(resultSize)
.setRescorer(queryRescorer(constantScoreQuery(QueryBuilders.matchPhraseQuery("field1", intToEnglish).slop(3)))
.setRescorer(new QueryRescorerBuilder(constantScoreQuery(matchPhraseQuery("field1", intToEnglish).slop(3)))
.setQueryWeight(1.0f)
// no weight - so we basically use the same score as the actual query
.setRescoreQueryWeight(0.0f), rescoreWindow)
@ -432,7 +433,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.setQuery(QueryBuilders.matchQuery("field1", query).operator(Operator.OR))
.setFrom(0)
.setSize(resultSize)
.setRescorer(queryRescorer(constantScoreQuery(matchPhraseQuery("field1", "not in the index").slop(3)))
.setRescorer(new QueryRescorerBuilder(constantScoreQuery(matchPhraseQuery("field1", "not in the index").slop(3)))
.setQueryWeight(1.0f)
.setRescoreQueryWeight(1.0f), rescoreWindow).execute()
.actionGet();
@ -462,7 +463,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.prepareSearch()
.setSearchType(SearchType.DFS_QUERY_THEN_FETCH)
.setQuery(QueryBuilders.matchQuery("field1", "the quick brown").operator(Operator.OR))
.setRescorer(queryRescorer(matchPhraseQuery("field1", "the quick brown").slop(2).boost(4.0f))
.setRescorer(new QueryRescorerBuilder(matchPhraseQuery("field1", "the quick brown").slop(2).boost(4.0f))
.setQueryWeight(0.5f).setRescoreQueryWeight(0.4f), 5).setExplain(true).execute()
.actionGet();
assertHitCount(searchResponse, 3);
@ -490,7 +491,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++) {
QueryRescorerBuilder innerRescoreQuery = queryRescorer(QueryBuilders.matchQuery("field1", "the quick brown").boost(4.0f))
QueryRescorerBuilder innerRescoreQuery = new QueryRescorerBuilder(matchQuery("field1", "the quick brown").boost(4.0f))
.setQueryWeight(0.5f).setRescoreQueryWeight(0.4f);
if (!"".equals(scoreModes[innerMode])) {
@ -513,7 +514,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
}
for (int outerMode = 0; outerMode < scoreModes.length; outerMode++) {
QueryRescorerBuilder outerRescoreQuery = queryRescorer(QueryBuilders.matchQuery("field1", "the quick brown").boost(4.0f))
QueryRescorerBuilder outerRescoreQuery = new QueryRescorerBuilder(matchQuery("field1", "the quick brown").boost(4.0f))
.setQueryWeight(0.5f).setRescoreQueryWeight(0.4f);
if (!"".equals(scoreModes[outerMode])) {
@ -557,7 +558,7 @@ public class QueryRescorerIT extends ESIntegTestCase {
.should(functionScoreQuery(termQuery("field1", intToEnglish[1]), weightFactorFunction(3.0f)).boostMode(REPLACE))
.should(functionScoreQuery(termQuery("field1", intToEnglish[2]), weightFactorFunction(5.0f)).boostMode(REPLACE))
.should(functionScoreQuery(termQuery("field1", intToEnglish[3]), weightFactorFunction(0.2f)).boostMode(REPLACE));
QueryRescorerBuilder rescoreQuery = queryRescorer(boolQuery()
QueryRescorerBuilder rescoreQuery = new QueryRescorerBuilder(boolQuery()
.should(functionScoreQuery(termQuery("field1", intToEnglish[0]), weightFactorFunction(5.0f)).boostMode(REPLACE))
.should(functionScoreQuery(termQuery("field1", intToEnglish[1]), weightFactorFunction(7.0f)).boostMode(REPLACE))
.should(functionScoreQuery(termQuery("field1", intToEnglish[3]), weightFactorFunction(0.0f)).boostMode(REPLACE)));
@ -621,12 +622,12 @@ public class QueryRescorerIT extends ESIntegTestCase {
public void testMultipleRescores() throws Exception {
int numDocs = indexRandomNumbers("keyword", 1, true);
QueryRescorerBuilder eightIsGreat = RescoreBuilder
.queryRescorer(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", English.intToEnglish(8)),
QueryRescorerBuilder eightIsGreat = new QueryRescorerBuilder(functionScoreQuery(
termQuery("field1", English.intToEnglish(8)),
ScoreFunctionBuilders.weightFactorFunction(1000.0f)).boostMode(CombineFunction.REPLACE))
.setScoreMode(QueryRescoreMode.Total);
QueryRescorerBuilder sevenIsBetter = RescoreBuilder
.queryRescorer(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", English.intToEnglish(7)),
QueryRescorerBuilder sevenIsBetter = new QueryRescorerBuilder(functionScoreQuery(
termQuery("field1", English.intToEnglish(7)),
ScoreFunctionBuilders.weightFactorFunction(10000.0f)).boostMode(CombineFunction.REPLACE))
.setScoreMode(QueryRescoreMode.Total);
@ -643,11 +644,11 @@ 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
QueryRescorerBuilder ninetyIsGood = RescoreBuilder.queryRescorer(QueryBuilders
.functionScoreQuery(QueryBuilders.queryStringQuery("*ninety*"), ScoreFunctionBuilders.weightFactorFunction(1000.0f))
QueryRescorerBuilder ninetyIsGood = new QueryRescorerBuilder(functionScoreQuery(
queryStringQuery("*ninety*"), ScoreFunctionBuilders.weightFactorFunction(1000.0f))
.boostMode(CombineFunction.REPLACE)).setScoreMode(QueryRescoreMode.Total);
QueryRescorerBuilder oneToo = RescoreBuilder.queryRescorer(QueryBuilders
.functionScoreQuery(QueryBuilders.queryStringQuery("*one*"), ScoreFunctionBuilders.weightFactorFunction(1000.0f))
QueryRescorerBuilder oneToo = new QueryRescorerBuilder(functionScoreQuery(
queryStringQuery("*one*"), ScoreFunctionBuilders.weightFactorFunction(1000.0f))
.boostMode(CombineFunction.REPLACE)).setScoreMode(QueryRescoreMode.Total);
request.clearRescorers().addRescorer(ninetyIsGood, numDocs).addRescorer(oneToo, 10);
response = request.setSize(2).get();
@ -700,7 +701,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(new QueryRescorerBuilder(matchAllQuery()), 50);
assertEquals(4, request.get().getHits().getHits().length);
}

View File

@ -53,8 +53,9 @@ import java.io.IOException;
import static java.util.Collections.emptyList;
import static org.elasticsearch.test.EqualsHashCodeTestUtils.checkEqualsAndHashCode;
import static org.hamcrest.Matchers.containsString;
public class QueryRescoreBuilderTests extends ESTestCase {
public class QueryRescorerBuilderTests extends ESTestCase {
private static final int NUMBER_OF_TESTBUILDERS = 20;
private static NamedWriteableRegistry namedWriteableRegistry;
@ -81,8 +82,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 = copy(original);
RescorerBuilder<?> original = randomRescoreBuilder();
RescorerBuilder<?> deserialized = copy(original);
assertEquals(deserialized, original);
assertEquals(deserialized.hashCode(), original.hashCode());
assertNotSame(deserialized, original);
@ -94,13 +95,13 @@ public class QueryRescoreBuilderTests extends ESTestCase {
*/
public void testEqualsAndHashcode() throws IOException {
for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
checkEqualsAndHashCode(randomRescoreBuilder(), this::copy, QueryRescoreBuilderTests::mutate);
checkEqualsAndHashCode(randomRescoreBuilder(), this::copy, QueryRescorerBuilderTests::mutate);
}
}
private RescoreBuilder<?> copy(RescoreBuilder<?> original) throws IOException {
private RescorerBuilder<?> copy(RescorerBuilder<?> original) throws IOException {
return copyWriteable(original, namedWriteableRegistry,
namedWriteableRegistry.getReader(RescoreBuilder.class, original.getWriteableName()));
namedWriteableRegistry.getReader(RescorerBuilder.class, original.getWriteableName()));
}
/**
@ -108,7 +109,7 @@ public class QueryRescoreBuilderTests extends ESTestCase {
*/
public void testFromXContent() throws IOException {
for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
RescoreBuilder<?> rescoreBuilder = randomRescoreBuilder();
RescorerBuilder<?> rescoreBuilder = randomRescoreBuilder();
XContentBuilder builder = XContentFactory.contentBuilder(randomFrom(XContentType.values()));
if (randomBoolean()) {
builder.prettyPrint();
@ -119,7 +120,7 @@ public class QueryRescoreBuilderTests extends ESTestCase {
XContentParser parser = createParser(shuffled);
parser.nextToken();
RescoreBuilder<?> secondRescoreBuilder = RescoreBuilder.parseFromXContent(parser);
RescorerBuilder<?> secondRescoreBuilder = RescorerBuilder.parseFromXContent(parser);
assertNotSame(rescoreBuilder, secondRescoreBuilder);
assertEquals(rescoreBuilder, secondRescoreBuilder);
assertEquals(rescoreBuilder.hashCode(), secondRescoreBuilder.hashCode());
@ -127,7 +128,7 @@ public class QueryRescoreBuilderTests extends ESTestCase {
}
/**
* test that build() outputs a {@link RescoreSearchContext} that has the same properties
* test that build() outputs a {@link RescoreContext} that has the same properties
* than the test builder
*/
public void testBuildRescoreSearchContext() throws ElasticsearchParseException, IOException {
@ -147,10 +148,10 @@ public class QueryRescoreBuilderTests extends ESTestCase {
for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
QueryRescorerBuilder rescoreBuilder = randomRescoreBuilder();
QueryRescoreContext rescoreContext = rescoreBuilder.build(mockShardContext);
int expectedWindowSize = rescoreBuilder.windowSize() == null ? QueryRescoreContext.DEFAULT_WINDOW_SIZE :
QueryRescoreContext rescoreContext = (QueryRescoreContext) rescoreBuilder.buildContext(mockShardContext);
int expectedWindowSize = rescoreBuilder.windowSize() == null ? RescorerBuilder.DEFAULT_WINDOW_SIZE :
rescoreBuilder.windowSize().intValue();
assertEquals(expectedWindowSize, rescoreContext.window());
assertEquals(expectedWindowSize, rescoreContext.getWindowSize());
Query expectedQuery = Rewriteable.rewrite(rescoreBuilder.getRescoreQuery(), mockShardContext).toQuery(mockShardContext);
assertEquals(expectedQuery, rescoreContext.query());
assertEquals(rescoreBuilder.getQueryWeight(), rescoreContext.queryWeight(), Float.MIN_VALUE);
@ -173,22 +174,18 @@ public class QueryRescoreBuilderTests extends ESTestCase {
" \"window_size\" : 20,\n" +
" \"bad_rescorer_name\" : { }\n" +
"}\n";
XContentParser parser = createParser(rescoreElement);
try {
RescoreBuilder.parseFromXContent(parser);
fail("expected a parsing exception");
} catch (ParsingException e) {
assertEquals("rescore doesn't support rescorer with name [bad_rescorer_name]", e.getMessage());
{
XContentParser parser = createParser(rescoreElement);
Exception e = expectThrows(ParsingException.class, () -> RescorerBuilder.parseFromXContent(parser));
assertEquals("Unknown RescorerBuilder [bad_rescorer_name]", e.getMessage());
}
rescoreElement = "{\n" +
" \"bad_fieldName\" : 20\n" +
"}\n";
parser = createParser(rescoreElement);
try {
RescoreBuilder.parseFromXContent(parser);
fail("expected a parsing exception");
} catch (ParsingException e) {
{
XContentParser parser = createParser(rescoreElement);
Exception e = expectThrows(ParsingException.class, () -> RescorerBuilder.parseFromXContent(parser));
assertEquals("rescore doesn't support [bad_fieldName]", e.getMessage());
}
@ -196,20 +193,16 @@ public class QueryRescoreBuilderTests extends ESTestCase {
" \"window_size\" : 20,\n" +
" \"query\" : [ ]\n" +
"}\n";
parser = createParser(rescoreElement);
try {
RescoreBuilder.parseFromXContent(parser);
fail("expected a parsing exception");
} catch (ParsingException e) {
{
XContentParser parser = createParser(rescoreElement);
Exception e = expectThrows(ParsingException.class, () -> RescorerBuilder.parseFromXContent(parser));
assertEquals("unexpected token [START_ARRAY] after [query]", e.getMessage());
}
rescoreElement = "{ }";
parser = createParser(rescoreElement);
try {
RescoreBuilder.parseFromXContent(parser);
fail("expected a parsing exception");
} catch (ParsingException e) {
{
XContentParser parser = createParser(rescoreElement);
Exception e = expectThrows(ParsingException.class, () -> RescorerBuilder.parseFromXContent(parser));
assertEquals("missing rescore type", e.getMessage());
}
@ -217,11 +210,9 @@ public class QueryRescoreBuilderTests extends ESTestCase {
" \"window_size\" : 20,\n" +
" \"query\" : { \"bad_fieldname\" : 1.0 } \n" +
"}\n";
parser = createParser(rescoreElement);
try {
RescoreBuilder.parseFromXContent(parser);
fail("expected a parsing exception");
} catch (IllegalArgumentException e) {
{
XContentParser parser = createParser(rescoreElement);
Exception e = expectThrows(IllegalArgumentException.class, () -> RescorerBuilder.parseFromXContent(parser));
assertEquals("[query] unknown field [bad_fieldname], parser not found", e.getMessage());
}
@ -229,11 +220,9 @@ public class QueryRescoreBuilderTests extends ESTestCase {
" \"window_size\" : 20,\n" +
" \"query\" : { \"rescore_query\" : { \"unknown_queryname\" : { } } } \n" +
"}\n";
parser = createParser(rescoreElement);
try {
RescoreBuilder.parseFromXContent(parser);
fail("expected a parsing exception");
} catch (ParsingException e) {
{
XContentParser parser = createParser(rescoreElement);
Exception e = expectThrows(ParsingException.class, () -> RescorerBuilder.parseFromXContent(parser));
assertEquals("[query] failed to parse field [rescore_query]", e.getMessage());
}
@ -241,8 +230,8 @@ public class QueryRescoreBuilderTests extends ESTestCase {
" \"window_size\" : 20,\n" +
" \"query\" : { \"rescore_query\" : { \"match_all\" : { } } } \n"
+ "}\n";
parser = createParser(rescoreElement);
RescoreBuilder.parseFromXContent(parser);
XContentParser parser = createParser(rescoreElement);
RescorerBuilder.parseFromXContent(parser);
}
/**
@ -260,8 +249,8 @@ public class QueryRescoreBuilderTests extends ESTestCase {
return xContentRegistry;
}
private static RescoreBuilder<?> mutate(RescoreBuilder<?> original) throws IOException {
RescoreBuilder<?> mutation = ESTestCase.copyWriteable(original, namedWriteableRegistry, QueryRescorerBuilder::new);
private static RescorerBuilder<?> mutate(RescorerBuilder<?> original) throws IOException {
RescorerBuilder<?> mutation = ESTestCase.copyWriteable(original, namedWriteableRegistry, QueryRescorerBuilder::new);
if (randomBoolean()) {
Integer windowSize = original.windowSize();
if (windowSize != null) {

View File

@ -0,0 +1,26 @@
/*
* 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.
*/
apply plugin: 'elasticsearch.esplugin'
esplugin {
name 'example-rescore'
description 'An example plugin implementing rescore and verifying that plugins *can* implement rescore'
classname 'org.elasticsearch.example.rescore.ExampleRescorePlugin'
}

View File

@ -0,0 +1,232 @@
/*
* 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.example.rescore;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.elasticsearch.common.Nullable;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.fielddata.AtomicFieldData;
import org.elasticsearch.index.fielddata.AtomicNumericFieldData;
import org.elasticsearch.index.fielddata.IndexFieldData;
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
import org.elasticsearch.index.query.QueryRewriteContext;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.rescore.Rescorer;
import org.elasticsearch.search.rescore.RescorerBuilder;
import java.io.IOException;
import java.util.Arrays;
import java.util.Iterator;
import java.util.Objects;
import java.util.Set;
import static java.util.Collections.singletonList;
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;
/**
* Example rescorer that multiplies the score of the hit by some factor and doesn't resort them.
*/
public class ExampleRescoreBuilder extends RescorerBuilder<ExampleRescoreBuilder> {
public static final String NAME = "example";
private final float factor;
private final String factorField;
public ExampleRescoreBuilder(float factor, @Nullable String factorField) {
this.factor = factor;
this.factorField = factorField;
}
ExampleRescoreBuilder(StreamInput in) throws IOException {
super(in);
factor = in.readFloat();
factorField = in.readOptionalString();
}
@Override
protected void doWriteTo(StreamOutput out) throws IOException {
out.writeFloat(factor);
out.writeOptionalString(factorField);
}
@Override
public String getWriteableName() {
return NAME;
}
@Override
public RescorerBuilder<ExampleRescoreBuilder> rewrite(QueryRewriteContext ctx) throws IOException {
return this;
}
private static final ParseField FACTOR = new ParseField("factor");
private static final ParseField FACTOR_FIELD = new ParseField("factor_field");
@Override
protected void doXContent(XContentBuilder builder, Params params) throws IOException {
builder.field(FACTOR.getPreferredName(), factor);
if (factorField != null) {
builder.field(FACTOR_FIELD.getPreferredName(), factorField);
}
}
private static final ConstructingObjectParser<ExampleRescoreBuilder, Void> PARSER = new ConstructingObjectParser<>(NAME,
args -> new ExampleRescoreBuilder((float) args[0], (String) args[1]));
static {
PARSER.declareFloat(constructorArg(), FACTOR);
PARSER.declareString(optionalConstructorArg(), FACTOR_FIELD);
}
public static ExampleRescoreBuilder fromXContent(XContentParser parser) {
return PARSER.apply(parser, null);
}
@Override
public RescoreContext innerBuildContext(int windowSize, QueryShardContext context) throws IOException {
IndexFieldData<?> factorField =
this.factorField == null ? null : context.getForField(context.fieldMapper(this.factorField));
return new ExampleRescoreContext(windowSize, factor, factorField);
}
@Override
public boolean equals(Object obj) {
if (false == super.equals(obj)) {
return false;
}
ExampleRescoreBuilder other = (ExampleRescoreBuilder) obj;
return factor == other.factor
&& Objects.equals(factorField, other.factorField);
}
@Override
public int hashCode() {
return Objects.hash(super.hashCode(), factor, factorField);
}
float factor() {
return factor;
}
@Nullable
String factorField() {
return factorField;
}
private static class ExampleRescoreContext extends RescoreContext {
private final float factor;
@Nullable
private final IndexFieldData<?> factorField;
ExampleRescoreContext(int windowSize, float factor, @Nullable IndexFieldData<?> factorField) {
super(windowSize, ExampleRescorer.INSTANCE);
this.factor = factor;
this.factorField = factorField;
}
}
private static class ExampleRescorer implements Rescorer {
private static final ExampleRescorer INSTANCE = new ExampleRescorer();
@Override
public TopDocs rescore(TopDocs topDocs, IndexSearcher searcher, RescoreContext rescoreContext) throws IOException {
ExampleRescoreContext context = (ExampleRescoreContext) rescoreContext;
int end = Math.min(topDocs.scoreDocs.length, rescoreContext.getWindowSize());
for (int i = 0; i < end; i++) {
topDocs.scoreDocs[i].score *= context.factor;
}
if (context.factorField != null) {
/*
* Since this example looks up a single field value it should
* access them in docId order because that is the order in
* which they are stored on disk and we want reads to be
* forwards and close together if possible.
*
* If accessing multiple fields we'd be better off accessing
* them in (reader, field, docId) order because that is the
* order they are on disk.
*/
ScoreDoc[] sortedByDocId = new ScoreDoc[topDocs.scoreDocs.length];
System.arraycopy(topDocs.scoreDocs, 0, sortedByDocId, 0, topDocs.scoreDocs.length);
Arrays.sort(sortedByDocId, (a, b) -> a.doc - b.doc); // Safe because doc ids >= 0
Iterator<LeafReaderContext> leaves = searcher.getIndexReader().leaves().iterator();
LeafReaderContext leaf = null;
SortedNumericDoubleValues data = null;
int endDoc = 0;
for (int i = 0; i < end; i++) {
if (topDocs.scoreDocs[i].doc >= endDoc) {
do {
leaf = leaves.next();
endDoc = leaf.docBase + leaf.reader().maxDoc();
} while (topDocs.scoreDocs[i].doc >= endDoc);
AtomicFieldData fd = context.factorField.load(leaf);
if (false == (fd instanceof AtomicNumericFieldData)) {
throw new IllegalArgumentException("[" + context.factorField.getFieldName() + "] is not a number");
}
data = ((AtomicNumericFieldData) fd).getDoubleValues();
}
if (false == data.advanceExact(topDocs.scoreDocs[i].doc)) {
throw new IllegalArgumentException("document [" + topDocs.scoreDocs[i].doc
+ "] does not have the field [" + context.factorField.getFieldName() + "]");
}
if (data.docValueCount() > 1) {
throw new IllegalArgumentException("document [" + topDocs.scoreDocs[i].doc
+ "] has more than one value for [" + context.factorField.getFieldName() + "]");
}
topDocs.scoreDocs[i].score *= data.nextValue();
}
}
// Sort by score descending, then docID ascending, just like lucene's QueryRescorer
Arrays.sort(topDocs.scoreDocs, (a, b) -> {
if (a.score > b.score) {
return -1;
}
if (a.score < b.score) {
return 1;
}
// Safe because doc ids >= 0
return a.doc - b.doc;
});
return topDocs;
}
@Override
public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreContext rescoreContext,
Explanation sourceExplanation) throws IOException {
ExampleRescoreContext context = (ExampleRescoreContext) rescoreContext;
// Note that this is inaccurate because it ignores factor field
return Explanation.match(context.factor, "test", singletonList(sourceExplanation));
}
@Override
public void extractTerms(IndexSearcher searcher, RescoreContext rescoreContext, Set<Term> termsSet) {
// Since we don't use queries there are no terms to extract.
}
}
}

View File

@ -0,0 +1,35 @@
/*
* 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.example.rescore;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.plugins.SearchPlugin;
import java.util.List;
import static java.util.Collections.singletonList;
public class ExampleRescorePlugin extends Plugin implements SearchPlugin {
@Override
public List<RescorerSpec<?>> getRescorers() {
return singletonList(
new RescorerSpec<>(ExampleRescoreBuilder.NAME, ExampleRescoreBuilder::new, ExampleRescoreBuilder::fromXContent));
}
}

View File

@ -0,0 +1,80 @@
/*
* 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.example.rescore;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.elasticsearch.common.io.stream.Writeable.Reader;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.test.AbstractWireSerializingTestCase;
import org.elasticsearch.test.ESTestCase;
import java.io.IOException;
import java.util.function.Supplier;
public class ExampleRescoreBuilderTests extends AbstractWireSerializingTestCase<ExampleRescoreBuilder> {
@Override
protected ExampleRescoreBuilder createTestInstance() {
String factorField = randomBoolean() ? null : randomAlphaOfLength(5);
return new ExampleRescoreBuilder(randomFloat(), factorField).windowSize(between(0, Integer.MAX_VALUE));
}
@Override
protected Reader<ExampleRescoreBuilder> instanceReader() {
return ExampleRescoreBuilder::new;
}
@Override
protected ExampleRescoreBuilder mutateInstance(ExampleRescoreBuilder instance) throws IOException {
@SuppressWarnings("unchecked")
Supplier<ExampleRescoreBuilder> supplier = randomFrom(
() -> new ExampleRescoreBuilder(instance.factor(), instance.factorField())
.windowSize(randomValueOtherThan(instance.windowSize(), () -> between(0, Integer.MAX_VALUE))),
() -> new ExampleRescoreBuilder(randomValueOtherThan(instance.factor(), ESTestCase::randomFloat), instance.factorField())
.windowSize(instance.windowSize()),
() -> new ExampleRescoreBuilder(
instance.factor(), randomValueOtherThan(instance.factorField(), () -> randomAlphaOfLength(5)))
.windowSize(instance.windowSize()));
return supplier.get();
}
public void testRewrite() throws IOException {
ExampleRescoreBuilder builder = createTestInstance();
assertSame(builder, builder.rewrite(null));
}
public void testRescore() throws IOException {
// Always use a factor > 1 so rescored fields are sorted in front of the unrescored fields.
float factor = (float) randomDoubleBetween(1.0d, Float.MAX_VALUE, false);
// Skipping factorField because it is much harder to mock. We'll catch it in an integration test.
String fieldFactor = null;
ExampleRescoreBuilder builder = new ExampleRescoreBuilder(factor, fieldFactor).windowSize(2);
RescoreContext context = builder.buildContext(null);
TopDocs docs = new TopDocs(10, new ScoreDoc[3], 0);
docs.scoreDocs[0] = new ScoreDoc(0, 1.0f);
docs.scoreDocs[1] = new ScoreDoc(1, 1.0f);
docs.scoreDocs[2] = new ScoreDoc(2, 1.0f);
context.rescorer().rescore(docs, null, context);
assertEquals(factor, docs.scoreDocs[0].score, 0.0f);
assertEquals(factor, docs.scoreDocs[1].score, 0.0f);
assertEquals(1.0f, docs.scoreDocs[2].score, 0.0f);
}
}

View File

@ -0,0 +1,38 @@
/*
* 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.example.rescore;
import com.carrotsearch.randomizedtesting.annotations.Name;
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
public class ExampleRescoreClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
public ExampleRescoreClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) {
super(testCandidate);
}
@ParametersFactory
public static Iterable<Object[]> parameters() throws Exception {
return ESClientYamlSuiteTestCase.createParameters();
}
}

View File

@ -0,0 +1,13 @@
# Integration tests for the expert scoring script example plugin
#
"Plugin loaded":
- do:
cluster.state: {}
# Get master node id
- set: { master_node: master }
- do:
nodes.info: {}
- match: { nodes.$master.plugins.0.name: example-rescore }

View File

@ -0,0 +1,63 @@
---
setup:
- do:
indices.create:
index: test
body:
settings:
number_of_shards: 1
number_of_replicas: 1
- do:
index:
index: test
type: test
id: 1
body: { "test": 1 }
- do:
index:
index: test
type: test
id: 2
body: { "test": 2 }
- do:
indices.refresh: {}
---
"just factor":
- do:
search:
index: test
body:
rescore:
example:
factor: 0
- length: { hits.hits: 2 }
- match: { hits.hits.0._score: 0 }
- match: { hits.hits.1._score: 0 }
- do:
search:
index: test
body:
rescore:
window_size: 1
example:
factor: 0
- length: { hits.hits: 2 }
- match: { hits.hits.0._score: 1 }
- match: { hits.hits.1._score: 0 }
---
"with factor field":
- do:
search:
index: test
body:
rescore:
example:
factor: 1
factor_field: test
- length: { hits.hits: 2 }
- match: { hits.hits.0._score: 2 }
- match: { hits.hits.1._score: 1 }

View File

@ -37,7 +37,7 @@ import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.collapse.CollapseBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.searchafter.SearchAfterBuilder;
import org.elasticsearch.search.slice.SliceBuilder;
import org.elasticsearch.search.sort.ScriptSortBuilder;
@ -115,7 +115,7 @@ public class RandomSearchRequestGenerator {
public static SearchSourceBuilder randomSearchSourceBuilder(
Supplier<HighlightBuilder> randomHighlightBuilder,
Supplier<SuggestBuilder> randomSuggestBuilder,
Supplier<RescoreBuilder<?>> randomRescoreBuilder,
Supplier<RescorerBuilder<?>> randomRescoreBuilder,
Supplier<List<SearchExtBuilder>> randomExtBuilders,
Supplier<CollapseBuilder> randomCollapseBuilder) {
SearchSourceBuilder builder = new SearchSourceBuilder();

View File

@ -58,7 +58,7 @@ import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.internal.ShardSearchRequest;
import org.elasticsearch.search.profile.Profilers;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.search.rescore.RescoreSearchContext;
import org.elasticsearch.search.rescore.RescoreContext;
import org.elasticsearch.search.sort.SortAndFormats;
import org.elasticsearch.search.suggest.SuggestionSearchContext;
import org.elasticsearch.threadpool.ThreadPool;
@ -219,12 +219,12 @@ public class TestSearchContext extends SearchContext {
}
@Override
public List<RescoreSearchContext> rescore() {
public List<RescoreContext> rescore() {
return Collections.emptyList();
}
@Override
public void addRescore(RescoreSearchContext rescore) {
public void addRescore(RescoreContext rescore) {
}
@Override