Add Rescore functionality

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

@@ -24,6 +24,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 
 import org.elasticsearch.action.DocWriteRequest;
@@ -37,7 +38,6 @@
 import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequestBuilder;
 import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
 import org.elasticsearch.action.admin.indices.settings.get.GetSettingsRequest;
-import org.elasticsearch.action.admin.indices.settings.get.GetSettingsResponse;
 import org.elasticsearch.action.admin.indices.template.delete.DeleteIndexTemplateRequest;
 import org.elasticsearch.action.bulk.BulkRequest;
 import org.elasticsearch.action.bulk.BulkRequestBuilder;
@@ -65,7 +65,6 @@
 import org.elasticsearch.client.indices.PutIndexTemplateRequest;
 import org.elasticsearch.client.indices.PutMappingRequest;
 import org.elasticsearch.common.geo.GeoDistance;
-import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.DistanceUnit;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.index.VersionType;
@@ -83,6 +82,8 @@
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 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.GeoDistanceSortBuilder;
 import org.elasticsearch.search.sort.ScoreSortBuilder;
@@ -106,6 +107,7 @@
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
 import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.elasticsearch.core.query.*;
+import org.springframework.data.elasticsearch.core.query.RescorerQuery.ScoreMode;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -119,6 +121,7 @@
  * @author Roman Puchkovskiy
  * @author Subhobrata Dey
  * @author Farid Faoudi
+ * @author Peer Mueller
  * @since 4.0
  */
 class RequestFactory {
@@ -1050,6 +1053,9 @@ private SearchRequest prepareSearchRequest(Query query, @Nullable Class<?> clazz
 			sourceBuilder.searchAfter(query.getSearchAfter().toArray());
 		}
 
+		query.getRescorerQueries().forEach(rescorer -> sourceBuilder.addRescorer(
+				getQueryRescorerBuilder(rescorer)));
+
 		request.source(sourceBuilder);
 		return request;
 	}
@@ -1136,6 +1142,9 @@ private SearchRequestBuilder prepareSearchRequestBuilder(Query query, Client cli
 			searchRequestBuilder.searchAfter(query.getSearchAfter().toArray());
 		}
 
+		query.getRescorerQueries().forEach(rescorer -> searchRequestBuilder.addRescorer(
+				getQueryRescorerBuilder(rescorer)));
+
 		return searchRequestBuilder;
 	}
 
@@ -1260,6 +1269,30 @@ private SortBuilder<?> getSortBuilder(Sort.Order order, @Nullable ElasticsearchP
 			}
 		}
 	}
+
+	private QueryRescorerBuilder getQueryRescorerBuilder(RescorerQuery rescorerQuery) {
+
+		QueryRescorerBuilder builder = new QueryRescorerBuilder(Objects.requireNonNull(getQuery(rescorerQuery.getQuery())));
+
+		if (rescorerQuery.getScoreMode() != ScoreMode.Default) {
+			builder.setScoreMode(QueryRescoreMode.valueOf(rescorerQuery.getScoreMode().name()));
+		}
+
+		if (rescorerQuery.getQueryWeight() != null) {
+			builder.setQueryWeight(rescorerQuery.getQueryWeight());
+		}
+
+		if (rescorerQuery.getRescoreQueryWeight() != null) {
+			builder.setRescoreQueryWeight(rescorerQuery.getRescoreQueryWeight());
+		}
+
+		if (rescorerQuery.getWindowSize() != null) {
+			builder.windowSize(rescorerQuery.getWindowSize());
+		}
+
+		return builder;
+
+	}
 	// endregion
 
 	// region update

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 {
 
@@ -63,6 +64,7 @@ abstract class AbstractQuery implements Query {
 	@Nullable private TimeValue timeout;
 	private boolean explain = false;
 	@Nullable private List<Object> searchAfter;
+	protected List<RescorerQuery> rescorerQueries = new ArrayList<>();
 
 	@Override
 	@Nullable
@@ -295,4 +297,20 @@ public void setSearchAfter(@Nullable List<Object> searchAfter) {
 	public List<Object> getSearchAfter() {
 		return searchAfter;
 	}
+
+	@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<RescorerQuery> rescorerQueries = 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(RescorerQuery rescorerQuery) {
+		this.rescorerQueries.add(rescorerQuery);
+		return this;
+	}
+
 	public NativeSearchQuery build() {
 
 		NativeSearchQuery nativeSearchQuery = new NativeSearchQuery(queryBuilder, filterBuilder, sortBuilders,
@@ -250,11 +259,16 @@ public NativeSearchQuery build() {
 		}
 
 		nativeSearchQuery.setTrackTotalHits(trackTotalHits);
-		
+
 		if (timeout != null) {
 			nativeSearchQuery.setTimeout(timeout);
 		}
 
+		if (!isEmpty(rescorerQueries)) {
+			nativeSearchQuery.setRescorerQueries(
+					rescorerQueries);
+		}
+
 		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;
 
@@ -41,6 +42,7 @@
  * @author Christoph Strobl
  * @author Farid Azaza
  * @author Peter-Josef Meisch
+ * @author Peer Mueller
  */
 public interface Query {
 
@@ -297,7 +299,7 @@ default boolean getExplain() {
 
 	/**
 	 * Sets the setSearchAfter objects for this query.
-	 * 
+	 *
 	 * @param searchAfter the setSearchAfter objects. These are obtained with {@link SearchHit#getSortValues()} from a
 	 *          search result.
 	 * @since 4.2
@@ -310,4 +312,24 @@ default boolean getExplain() {
 	 */
 	@Nullable
 	List<Object> getSearchAfter();
+
+	/**
+	 * 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,94 @@
+/*
+ * 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.index.query.QueryBuilder;
+import org.elasticsearch.search.rescore.QueryRescorerBuilder;
+import org.springframework.lang.Nullable;
+
+/**
+ * Implementation of RescorerQuery to be used for rescoring filtered search results.
+ *
+ * @author Peer Mueller
+ * @since 4.2
+ */
+public class RescorerQuery {
+
+	private final Query query;
+	private ScoreMode scoreMode = ScoreMode.Default;
+	@Nullable private Integer windowSize;
+	@Nullable private Float queryWeight;
+	@Nullable private Float rescoreQueryWeight;
+
+	public RescorerQuery(Query query) {
+		this.query = query;
+	}
+
+	public Query getQuery() {
+		return query;
+	}
+
+	public ScoreMode getScoreMode() {
+		return scoreMode;
+	}
+
+	@Nullable
+	public Integer getWindowSize() {
+		return windowSize;
+	}
+
+	@Nullable
+	public Float getQueryWeight() {
+		return queryWeight;
+	}
+
+	@Nullable
+	public Float getRescoreQueryWeight() {
+		return rescoreQueryWeight;
+	}
+
+	public RescorerQuery withScoreMode(ScoreMode scoreMode) {
+		this.scoreMode = scoreMode;
+		return this;
+	}
+
+	public RescorerQuery withWindowSize(int windowSize) {
+		this.windowSize = windowSize;
+		return this;
+	}
+
+	public RescorerQuery withQueryWeight(float queryWeight) {
+		this.queryWeight = queryWeight;
+		return this;
+	}
+
+	public RescorerQuery withRescoreQueryWeight(float rescoreQueryWeight) {
+		this.rescoreQueryWeight = rescoreQueryWeight;
+		return this;
+	}
+
+
+
+	public enum ScoreMode {
+		Default,
+		Avg,
+		Max,
+		Min,
+		Total,
+		Multiply
+	}
+
+}