diff --git a/core/src/main/java/org/elasticsearch/search/SearchModule.java b/core/src/main/java/org/elasticsearch/search/SearchModule.java index c33471f4432..739b97034bf 100644 --- a/core/src/main/java/org/elasticsearch/search/SearchModule.java +++ b/core/src/main/java/org/elasticsearch/search/SearchModule.java @@ -19,14 +19,6 @@ package org.elasticsearch.search; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.function.Supplier; - import org.apache.lucene.search.BooleanQuery; import org.elasticsearch.common.geo.ShapesAvailability; import org.elasticsearch.common.geo.builders.CircleBuilder; @@ -227,9 +219,19 @@ import org.elasticsearch.search.highlight.HighlightPhase; import org.elasticsearch.search.highlight.Highlighter; import org.elasticsearch.search.highlight.Highlighters; import org.elasticsearch.search.query.QueryPhase; +import org.elasticsearch.search.rescore.QueryRescorerBuilder; +import org.elasticsearch.search.rescore.RescoreBuilder; import org.elasticsearch.search.suggest.Suggester; import org.elasticsearch.search.suggest.Suggesters; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Supplier; + /** * */ @@ -327,6 +329,7 @@ public class SearchModule extends AbstractModule { bind(IndicesQueriesRegistry.class).toInstance(buildQueryParserRegistry()); configureFetchSubPhase(); configureShapes(); + configureRescorers(); } protected void configureFetchSubPhase() { @@ -467,6 +470,10 @@ public class SearchModule extends AbstractModule { } } + private void configureRescorers() { + namedWriteableRegistry.registerPrototype(RescoreBuilder.class, QueryRescorerBuilder.PROTOTYPE); + } + private void registerBuiltinFunctionScoreParsers() { registerFunctionScoreParser(new ScriptScoreFunctionParser()); registerFunctionScoreParser(new GaussDecayFunctionParser()); diff --git a/core/src/main/java/org/elasticsearch/search/SearchService.java b/core/src/main/java/org/elasticsearch/search/SearchService.java index 6a84bb44ae7..a5511277bc5 100644 --- a/core/src/main/java/org/elasticsearch/search/SearchService.java +++ b/core/src/main/java/org/elasticsearch/search/SearchService.java @@ -23,6 +23,7 @@ import com.carrotsearch.hppc.ObjectFloatHashMap; import com.carrotsearch.hppc.ObjectHashSet; import com.carrotsearch.hppc.ObjectSet; import com.carrotsearch.hppc.cursors.ObjectCursor; + import org.apache.lucene.index.IndexOptions; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.NumericDocValues; @@ -100,6 +101,7 @@ import org.elasticsearch.search.query.QuerySearchRequest; import org.elasticsearch.search.query.QuerySearchResult; import org.elasticsearch.search.query.QuerySearchResultProvider; import org.elasticsearch.search.query.ScrollQuerySearchResult; +import org.elasticsearch.search.rescore.RescoreBaseBuilder; import org.elasticsearch.threadpool.ThreadPool; import java.io.IOException; @@ -772,33 +774,12 @@ public class SearchService extends AbstractLifecycleComponent imp } } if (source.rescores() != null) { - XContentParser completeRescoreParser = null; try { - XContentBuilder completeRescoreBuilder = XContentFactory.jsonBuilder(); - completeRescoreBuilder.startObject(); - completeRescoreBuilder.startArray("rescore"); - for (BytesReference rescore : source.rescores()) { - XContentParser parser = XContentFactory.xContent(rescore).createParser(rescore); - parser.nextToken(); - completeRescoreBuilder.copyCurrentStructure(parser); + for (RescoreBaseBuilder rescore : source.rescores()) { + context.addRescore(rescore.build(context.indexShard().getQueryShardContext())); } - completeRescoreBuilder.endArray(); - completeRescoreBuilder.endObject(); - BytesReference completeRescoreBytes = completeRescoreBuilder.bytes(); - completeRescoreParser = XContentFactory.xContent(completeRescoreBytes).createParser(completeRescoreBytes); - completeRescoreParser.nextToken(); - completeRescoreParser.nextToken(); - completeRescoreParser.nextToken(); - this.elementParsers.get("rescore").parse(completeRescoreParser, context); - } catch (Exception e) { - String sSource = "_na_"; - try { - sSource = source.toString(); - } catch (Throwable e1) { - // ignore - } - XContentLocation location = completeRescoreParser != null ? completeRescoreParser.getTokenLocation() : null; - throw new SearchParseException(context, "failed to parse rescore source [" + sSource + "]", location, e); + } catch (IOException e) { + throw new SearchContextException(context, "failed to create RescoreSearchContext", e); } } if (source.fields() != null) { diff --git a/core/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java b/core/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java index 777e32ee565..4db3a0583bb 100644 --- a/core/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java +++ b/core/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java @@ -21,6 +21,7 @@ package org.elasticsearch.search.builder; import com.carrotsearch.hppc.ObjectFloatHashMap; import com.carrotsearch.hppc.cursors.ObjectCursor; + import org.elasticsearch.Version; import org.elasticsearch.action.support.ToXContentToBytes; import org.elasticsearch.common.Nullable; @@ -151,7 +152,7 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ private BytesReference innerHitsBuilder; - private List rescoreBuilders; + private List rescoreBuilders; private ObjectFloatHashMap indexBoost = null; @@ -459,19 +460,11 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ } public SearchSourceBuilder addRescorer(RescoreBaseBuilder rescoreBuilder) { - try { if (rescoreBuilders == null) { rescoreBuilders = new ArrayList<>(); } - XContentBuilder builder = XContentFactory.jsonBuilder(); - builder.startObject(); - rescoreBuilder.toXContent(builder, EMPTY_PARAMS); - builder.endObject(); - rescoreBuilders.add(builder.bytes()); + rescoreBuilders.add(rescoreBuilder); return this; - } catch (IOException e) { - throw new RuntimeException(e); - } } public SearchSourceBuilder clearRescorers() { @@ -498,7 +491,7 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ /** * Gets the bytes representing the rescore builders for this request. */ - public List rescores() { + public List rescores() { return rescoreBuilders; } @@ -878,10 +871,9 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ } builder.sorts = sorts; } else if (context.parseFieldMatcher().match(currentFieldName, RESCORE_FIELD)) { - List rescoreBuilders = new ArrayList<>(); + List rescoreBuilders = new ArrayList<>(); while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) { - XContentBuilder xContentBuilder = XContentFactory.jsonBuilder().copyCurrentStructure(parser); - rescoreBuilders.add(xContentBuilder.bytes()); + rescoreBuilders.add(RescoreBaseBuilder.PROTOTYPE.fromXContent(context)); } builder.rescoreBuilders = rescoreBuilders; } else if (context.parseFieldMatcher().match(currentFieldName, STATS_FIELD)) { @@ -1048,10 +1040,8 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ if (rescoreBuilders != null) { builder.startArray(RESCORE_FIELD.getPreferredName()); - for (BytesReference rescoreBuilder : rescoreBuilders) { - XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(rescoreBuilder); - parser.nextToken(); - builder.copyCurrentStructure(parser); + for (RescoreBaseBuilder rescoreBuilder : rescoreBuilders) { + rescoreBuilder.toXContent(builder, params); } builder.endArray(); } @@ -1197,9 +1187,9 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ } if (in.readBoolean()) { int size = in.readVInt(); - List rescoreBuilders = new ArrayList<>(); + List rescoreBuilders = new ArrayList<>(); for (int i = 0; i < size; i++) { - rescoreBuilders.add(in.readBytesReference()); + rescoreBuilders.add(RescoreBaseBuilder.PROTOTYPE.readFrom(in)); } builder.rescoreBuilders = rescoreBuilders; } @@ -1313,8 +1303,8 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ out.writeBoolean(hasRescoreBuilders); if (hasRescoreBuilders) { out.writeVInt(rescoreBuilders.size()); - for (BytesReference rescoreBuilder : rescoreBuilders) { - out.writeBytesReference(rescoreBuilder); + for (RescoreBaseBuilder rescoreBuilder : rescoreBuilders) { + rescoreBuilder.writeTo(out); } } boolean hasScriptFields = scriptFields != null; diff --git a/core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java b/core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java index ddcabf9ada5..a0201ea8362 100644 --- a/core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java +++ b/core/src/main/java/org/elasticsearch/search/rescore/RescoreBaseBuilder.java @@ -115,10 +115,12 @@ public class RescoreBaseBuilder implements ToXContent, Writeable injectedQueryParsers = new HashSet<>(); - injectedQueryParsers.add(new MatchAllQueryParser()); - indicesQueriesRegistry = new IndicesQueriesRegistry(Settings.settingsBuilder().build(), injectedQueryParsers, namedWriteableRegistry); + namedWriteableRegistry.registerPrototype(RescoreBuilder.class, QueryRescorerBuilder.PROTOTYPE); + indicesQueriesRegistry = new SearchModule(Settings.EMPTY, namedWriteableRegistry).buildQueryParserRegistry(); } @AfterClass @@ -154,10 +148,7 @@ public class QueryRescoreBuilderTests extends ESTestCase { if (randomBoolean()) { builder.prettyPrint(); } - builder.startObject(); rescoreBuilder.toXContent(builder, ToXContent.EMPTY_PARAMS); - builder.endObject(); - return XContentHelper.createParser(builder.bytes()); } @@ -326,7 +317,7 @@ public class QueryRescoreBuilderTests extends ESTestCase { /** * create random shape that is put under test */ - private static RescoreBaseBuilder randomRescoreBuilder() { + public static RescoreBaseBuilder randomRescoreBuilder() { QueryBuilder queryBuilder = new MatchAllQueryBuilder().boost(randomFloat()).queryName(randomAsciiOfLength(20)); org.elasticsearch.search.rescore.QueryRescorerBuilder rescorer = new org.elasticsearch.search.rescore.QueryRescorerBuilder(queryBuilder);