Add Rescore functionality

src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java

@@ -120,6 +120,7 @@
  * @author Roman Puchkovskiy
  * @author Subhobrata Dey
  * @author Farid Faoudi
+ * @author Peer Mueller
  * @since 4.0
  */
 class RequestFactory {
@@ -1147,6 +1148,9 @@ private SearchRequest prepareSearchRequest(Query query, @Nullable Class<?> clazz
 
 		sourceBuilder.explain(query.getExplain());
 
+		query.getRescorerQueries().stream().map(RescorerQuery::getRescorerBuilder)
+				.forEach(sourceBuilder::addRescorer);
+
 		request.source(sourceBuilder);
 		return request;
 	}
@@ -1229,6 +1233,9 @@ private SearchRequestBuilder prepareSearchRequestBuilder(Query query, Client cli
 
 		searchRequestBuilder.setExplain(query.getExplain());
 
+		query.getRescorerQueries().stream().map(RescorerQuery::getRescorerBuilder)
+				.forEach(searchRequestBuilder::addRescorer);
+
 		return searchRequestBuilder;
 	}
 

src/main/java/org/springframework/data/elasticsearch/core/query/AbstractQuery.java

@@ -41,6 +41,7 @@
  * @author Sascha Woo
  * @author Farid Azaza
  * @author Peter-Josef Meisch
+ * @author Peer Mueller
  */
 abstract class AbstractQuery implements Query {
 
@@ -62,6 +63,7 @@ abstract class AbstractQuery implements Query {
 	@Nullable private Duration scrollTime;
 	@Nullable private TimeValue timeout;
 	private boolean explain = false;
+	protected List<RescorerQuery> rescorerQueries = new ArrayList<>();
 
 	@Override
 	@Nullable
@@ -283,4 +285,20 @@ public boolean getExplain() {
 	public void setExplain(boolean explain) {
 		this.explain = explain;
 	}
+
+	@Override
+	public void addRescorerQuery(RescorerQuery rescorerQuery) {
+		this.rescorerQueries.add(rescorerQuery);
+	}
+
+	@Override
+	public void setRescorerQueries(List<RescorerQuery> rescorerQueryList) {
+		this.rescorerQueries.clear();
+		this.rescorerQueries.addAll(rescorerQueryList);
+	}
+
+	@Override
+	public List<RescorerQuery> getRescorerQueries() {
+		return rescorerQueries;
+	}
 }

src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java

@@ -20,6 +20,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.action.support.IndicesOptions;
@@ -28,6 +29,7 @@
 import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.collapse.CollapseBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
+import org.elasticsearch.search.rescore.QueryRescorerBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
 import org.springframework.data.domain.Pageable;
 import org.springframework.lang.Nullable;
@@ -45,6 +47,7 @@
  * @author Martin Choraine
  * @author Farid Azaza
  * @author Peter-Josef Meisch
+ * @author Peer Mueller
  */
 public class NativeSearchQueryBuilder {
 
@@ -70,6 +73,7 @@ public class NativeSearchQueryBuilder {
 	@Nullable private Integer maxResults;
 	@Nullable private Boolean trackTotalHits;
 	@Nullable private TimeValue timeout;
+	private final List<QueryRescorerBuilder> queryRescorerBuilders = new ArrayList<>();
 
 	public NativeSearchQueryBuilder withQuery(QueryBuilder queryBuilder) {
 		this.queryBuilder = queryBuilder;
@@ -183,12 +187,17 @@ public NativeSearchQueryBuilder withTrackTotalHits(Boolean trackTotalHits) {
 		this.trackTotalHits = trackTotalHits;
 		return this;
 	}
-	
-	public NativeSearchQueryBuilder withTimeout(TimeValue timeout) {   
+
+	public NativeSearchQueryBuilder withTimeout(TimeValue timeout) {
 		this.timeout = timeout;
 		return this;
 	}
 
+	public NativeSearchQueryBuilder withRescorerQuery(QueryRescorerBuilder queryRescorerBuilder) {
+		this.queryRescorerBuilders.add(queryRescorerBuilder);
+		return this;
+	}
+
 	public NativeSearchQuery build() {
 
 		NativeSearchQuery nativeSearchQuery = new NativeSearchQuery(queryBuilder, filterBuilder, sortBuilders,
@@ -250,11 +259,18 @@ public NativeSearchQuery build() {
 		}
 
 		nativeSearchQuery.setTrackTotalHits(trackTotalHits);
-		
+
 		if (timeout != null) {
 			nativeSearchQuery.setTimeout(timeout);
 		}
 
+		if (!isEmpty(queryRescorerBuilders)) {
+			nativeSearchQuery.setRescorerQueries(
+					queryRescorerBuilders.stream()
+							.map(RescorerQuery::new)
+							.collect(Collectors.toList()));
+		}
+
 		return nativeSearchQuery;
 	}
 }

src/main/java/org/springframework/data/elasticsearch/core/query/Query.java

@@ -17,6 +17,7 @@
 
 import java.time.Duration;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
@@ -40,6 +41,7 @@
  * @author Christoph Strobl
  * @author Farid Azaza
  * @author Peter-Josef Meisch
+ * @author Peer Mueller
  */
 public interface Query {
 
@@ -293,4 +295,24 @@ default boolean hasScrollTime() {
 	default boolean getExplain() {
 		return false;
 	}
+
+	/**
+	 * Sets the {@link RescorerQuery}.
+	 *
+	 * @param rescorerQuery the query to add to the list of rescorer queries
+	 * @since 4.2
+	 */
+	void addRescorerQuery(RescorerQuery rescorerQuery);
+
+	/**
+	 * Sets the {@link RescorerQuery}.
+	 *
+	 * @param rescorerQueryList list of rescorer queries set
+	 * @since 4.2
+	 */
+	void setRescorerQueries(List<RescorerQuery> rescorerQueryList);
+
+	default List<RescorerQuery> getRescorerQueries() {
+		return Collections.emptyList();
+	}
 }

src/main/java/org/springframework/data/elasticsearch/core/query/RescorerQuery.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2020-2021 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.data.elasticsearch.core.query;
+
+import org.elasticsearch.search.rescore.QueryRescorerBuilder;
+
+/**
+ * Encapsulates an Elasticsearch {@link QueryRescorerBuilder} to prevent leaking of Elasticsearch
+ * classes into the query API.
+ *
+ * @author Peer Mueller
+ * @since 4.2
+ */
+public class RescorerQuery {
+
+	private final QueryRescorerBuilder queryRescorerBuilder;
+
+	public RescorerQuery(QueryRescorerBuilder queryRescorerBuilder) {
+		this.queryRescorerBuilder = queryRescorerBuilder;
+	}
+
+	public QueryRescorerBuilder getRescorerBuilder() {
+		return queryRescorerBuilder;
+	}
+
+
+}

src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java

@@ -53,12 +53,19 @@
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.action.update.UpdateRequest;
 import org.elasticsearch.cluster.metadata.AliasMetadata;
+import org.elasticsearch.common.lucene.search.function.CombineFunction;
+import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
 import org.elasticsearch.index.VersionType;
+import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
+import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder.FilterFunctionBuilder;
+import org.elasticsearch.index.query.functionscore.GaussDecayFunctionBuilder;
 import org.elasticsearch.join.query.ParentIdQueryBuilder;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
+import org.elasticsearch.search.rescore.QueryRescoreMode;
+import org.elasticsearch.search.rescore.QueryRescorerBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.search.sort.SortBuilders;
 import org.elasticsearch.search.sort.SortOrder;
@@ -121,6 +128,7 @@
  * @author Roman Puchkovskiy
  * @author Subhobrata Dey
  * @author Farid Faoudi
+ * @author Peer Mueller
  */
 @SpringIntegrationTest
 public abstract class ElasticsearchTemplateTests {
@@ -3108,7 +3116,63 @@ void shouldReturnHighlightFieldsInSearchHit() {
 		assertThat(highlightField.get(1)).contains("<em>message</em>");
 	}
 
-	@Test // DATAES-738
+	@Test
+	void shouldRunRescoreQueryInSearchQuery() {
+		IndexCoordinates index = IndexCoordinates.of("test-index-rescore-entity-template");
+
+		// matches main query better
+		SampleEntity entity = SampleEntity.builder() //
+				.id("1") //
+				.message("some message") //
+				.rate(java.lang.Integer.MAX_VALUE)
+				.version(System.currentTimeMillis()) //
+				.build();
+
+		// high score from rescore query
+		SampleEntity entity2 = SampleEntity.builder() //
+				.id("2") //
+				.message("nothing") //
+				.rate(1)
+				.version(System.currentTimeMillis()) //
+				.build();
+
+		List<IndexQuery> indexQueries = getIndexQueries(Arrays.asList(entity, entity2));
+
+		operations.bulkIndex(indexQueries, index);
+		indexOperations.refresh();
+
+		NativeSearchQuery query = new NativeSearchQueryBuilder() //
+				.withQuery(
+						boolQuery().filter(existsQuery("rate")).should(termQuery("message", "message"))) //
+				.withRescorerQuery(new QueryRescorerBuilder(
+						new FunctionScoreQueryBuilder(
+								new FunctionScoreQueryBuilder.FilterFunctionBuilder[]{
+										new FilterFunctionBuilder(
+												new GaussDecayFunctionBuilder("rate", 0, 10, null, 0.5)
+														.setWeight(1f)),
+										new FilterFunctionBuilder(
+												new GaussDecayFunctionBuilder("rate", 0, 10, null, 0.5)
+														.setWeight(100f))})
+								.scoreMode(FunctionScoreQuery.ScoreMode.SUM)
+								.maxBoost(80f)
+								.boostMode(CombineFunction.REPLACE))
+						.setScoreMode(QueryRescoreMode.Max)
+						.windowSize(100))
+				.build();
+
+		SearchHits<SampleEntity> searchHits = operations.search(query, SampleEntity.class, index);
+
+		assertThat(searchHits).isNotNull();
+		assertThat(searchHits.getSearchHits()).hasSize(2);
+
+		SearchHit<SampleEntity> searchHit = searchHits.getSearchHit(0);
+		assertThat(searchHit.getContent().getMessage()).isEqualTo("nothing");
+		//score capped to 80
+		assertThat(searchHit.getScore()).isEqualTo(80f);
+	}
+
+	@Test
+		// DATAES-738
 	void shouldSaveEntityWithIndexCoordinates() {
 		String id = "42";
 		SampleEntity entity = new SampleEntity();