support highlight_query (#2793)

* support highlight_query

* implement highlight query with spring data elasticsearch query

* highight query by StringQuery

* split highligh fields assertion into different parts

src/main/antora/modules/ROOT/pages/migration-guides/migration-guide-4.4-5.0.adoc

@@ -49,7 +49,7 @@ Also the reactive implementation that was provided up to now has been moved here
 If you are using `ElasticsearchRestTemplate` directly and not the `ElasticsearchOperations` interface you'll need to adjust your imports as well.
 
 When working with the `NativeSearchQuery` class, you'll need to switch to the `NativeQuery` class, which can take a
-`Query` instance comign from the new Elasticsearch client libraries.
+`Query` instance coming from the new Elasticsearch client libraries.
 You'll find plenty of examples in the test code.
 
 [[elasticsearch-migration-guide-4.4-5.0.breaking-changes-records]]

src/main/java/org/springframework/data/elasticsearch/client/elc/HighlightQueryBuilder.java

@@ -35,14 +35,17 @@ import org.springframework.util.StringUtils;
  * {@link co.elastic.clients.elasticsearch.core.search.Highlight}.
  *
  * @author Peter-Josef Meisch
+ * @author Haibo Liu
  * @since 4.4
  */
 class HighlightQueryBuilder {
 	private final MappingContext<? extends ElasticsearchPersistentEntity<?>, ElasticsearchPersistentProperty> mappingContext;
+	private final RequestConverter requestConverter;
 
 	HighlightQueryBuilder(
-			MappingContext<? extends ElasticsearchPersistentEntity<?>, ElasticsearchPersistentProperty> mappingContext) {
+			MappingContext<? extends ElasticsearchPersistentEntity<?>, ElasticsearchPersistentProperty> mappingContext, RequestConverter requestConverter) {
 		this.mappingContext = mappingContext;
+		this.requestConverter = requestConverter;
 	}
 
 	public co.elastic.clients.elasticsearch.core.search.Highlight getHighlight(Highlight highlight,
@@ -52,7 +55,7 @@ class HighlightQueryBuilder {
 
 		// in the old implementation we could use one addParameters method, but in the new Elasticsearch client
 		// the builder for highlight and highlightfield share no code
-		addParameters(highlight.getParameters(), highlightBuilder);
+		addParameters(highlight.getParameters(), highlightBuilder, type);
 
 		for (HighlightField highlightField : highlight.getFields()) {
 			String mappedName = mapFieldName(highlightField.getName(), type);
@@ -69,7 +72,7 @@ class HighlightQueryBuilder {
 	 * the builder for highlight and highlight fields don't share code, so we have these two methods here that basically are almost copies
 	 */
 	private void addParameters(HighlightParameters parameters,
-			co.elastic.clients.elasticsearch.core.search.Highlight.Builder builder) {
+			co.elastic.clients.elasticsearch.core.search.Highlight.Builder builder, @Nullable Class<?> type) {
 
 		if (StringUtils.hasLength(parameters.getBoundaryChars())) {
 			builder.boundaryChars(parameters.getBoundaryChars());
@@ -103,6 +106,10 @@ class HighlightQueryBuilder {
 			builder.numberOfFragments(parameters.getNumberOfFragments());
 		}
 
+		if (parameters.getHighlightQuery() != null) {
+			builder.highlightQuery(requestConverter.getQuery(parameters.getHighlightQuery(), type));
+		}
+
 		if (StringUtils.hasLength(parameters.getOrder())) {
 			builder.order(highlighterOrder(parameters.getOrder()));
 		}
@@ -174,6 +181,10 @@ class HighlightQueryBuilder {
 			builder.numberOfFragments(parameters.getNumberOfFragments());
 		}
 
+		if (parameters.getHighlightQuery() != null) {
+			builder.highlightQuery(requestConverter.getQuery(parameters.getHighlightQuery(), type));
+		}
+
 		if (StringUtils.hasLength(parameters.getOrder())) {
 			builder.order(highlighterOrder(parameters.getOrder()));
 		}

src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java

@@ -105,6 +105,7 @@ import org.springframework.util.StringUtils;
  * @author Sascha Woo
  * @author cdalxndr
  * @author scoobyzhang
+ * @author Haibo Liu
  * @since 4.4
  */
 @SuppressWarnings("ClassCanBeRecord")
@@ -1494,7 +1495,7 @@ class RequestConverter {
 	private void addHighlight(Query query, SearchRequest.Builder builder) {
 
 		Highlight highlight = query.getHighlightQuery()
-				.map(highlightQuery -> new HighlightQueryBuilder(elasticsearchConverter.getMappingContext())
+				.map(highlightQuery -> new HighlightQueryBuilder(elasticsearchConverter.getMappingContext(), this)
 						.getHighlight(highlightQuery.getHighlight(), highlightQuery.getType()))
 				.orElse(null);
 
@@ -1504,7 +1505,7 @@ class RequestConverter {
 	private void addHighlight(Query query, MultisearchBody.Builder builder) {
 
 		Highlight highlight = query.getHighlightQuery()
-				.map(highlightQuery -> new HighlightQueryBuilder(elasticsearchConverter.getMappingContext())
+				.map(highlightQuery -> new HighlightQueryBuilder(elasticsearchConverter.getMappingContext(), this)
 						.getHighlight(highlightQuery.getHighlight(), highlightQuery.getType()))
 				.orElse(null);
 
@@ -1646,7 +1647,7 @@ class RequestConverter {
 	}
 
 	@Nullable
-	private co.elastic.clients.elasticsearch._types.query_dsl.Query getQuery(@Nullable Query query,
+	co.elastic.clients.elasticsearch._types.query_dsl.Query getQuery(@Nullable Query query,
 			@Nullable Class<?> clazz) {
 
 		if (query == null) {

src/main/java/org/springframework/data/elasticsearch/core/query/highlight/HighlightCommonParameters.java

@@ -15,10 +15,13 @@
  */
 package org.springframework.data.elasticsearch.core.query.highlight;
 
+import org.springframework.data.elasticsearch.core.query.Query;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
 /**
  * @author Peter-Josef Meisch
+ * @author Haibo Liu
  * @since 4.3
  */
 public abstract class HighlightCommonParameters {
@@ -31,6 +34,7 @@ public abstract class HighlightCommonParameters {
 	private final int fragmentSize;
 	private final int noMatchSize;
 	private final int numberOfFragments;
+	@Nullable private final Query highlightQuery;
 	private final String order;
 	private final int phraseLimit;
 	private final String[] preTags;
@@ -51,6 +55,7 @@ public abstract class HighlightCommonParameters {
 		fragmentSize = builder.fragmentSize;
 		noMatchSize = builder.noMatchSize;
 		numberOfFragments = builder.numberOfFragments;
+		highlightQuery = builder.highlightQuery;
 		order = builder.order;
 		phraseLimit = builder.phraseLimit;
 		preTags = builder.preTags;
@@ -95,6 +100,11 @@ public abstract class HighlightCommonParameters {
 		return numberOfFragments;
 	}
 
+	@Nullable
+	public Query getHighlightQuery() {
+		return highlightQuery;
+	}
+
 	public String getOrder() {
 		return order;
 	}
@@ -130,6 +140,11 @@ public abstract class HighlightCommonParameters {
 		private int fragmentSize = -1;
 		private int noMatchSize = -1;
 		private int numberOfFragments = -1;
+		/**
+		 * Only the search query part of the {@link Query} takes effect,
+		 * others are just ignored.
+		 */
+		@Nullable private Query highlightQuery = null;
 		private String order = "";
 		private int phraseLimit = -1;
 		private String[] preTags = new String[0];
@@ -184,6 +199,11 @@ public abstract class HighlightCommonParameters {
 			return (SELF) this;
 		}
 
+		public SELF withHighlightQuery(Query highlightQuery) {
+			this.highlightQuery = highlightQuery;
+			return (SELF) this;
+		}
+
 		public SELF withOrder(String order) {
 			this.order = order;
 			return (SELF) this;

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

@@ -66,6 +66,8 @@ import org.springframework.data.elasticsearch.core.query.*;
 import org.springframework.data.elasticsearch.core.query.Query;
 import org.springframework.data.elasticsearch.core.query.highlight.Highlight;
 import org.springframework.data.elasticsearch.core.query.highlight.HighlightField;
+import org.springframework.data.elasticsearch.core.query.highlight.HighlightFieldParameters;
+import org.springframework.data.elasticsearch.core.query.highlight.HighlightParameters;
 import org.springframework.data.elasticsearch.junit.jupiter.SpringIntegrationTest;
 import org.springframework.data.elasticsearch.utils.IndexNameProvider;
 import org.springframework.data.util.StreamUtils;
@@ -2964,6 +2966,128 @@ public abstract class ElasticsearchIntegrationTests {
 		assertThat(highlightField.get(1)).contains("<em>message</em>");
 	}
 
+	@Test // #2636
+	void shouldReturnHighlightFieldsWithHighlightQueryInSearchHit() {
+		IndexCoordinates index = createIndexCoordinatesWithHighlightMessage();
+
+		// a highlight query equals to the search query
+		var sameHighlightQuery = HighlightFieldParameters.builder()
+				.withHighlightQuery(getBuilderWithTermQuery("message", "message").build())
+				.build();
+		Query query = getBuilderWithTermQuery("message", "message") //
+				.withHighlightQuery(
+						new HighlightQuery(new Highlight(singletonList(new HighlightField("message", sameHighlightQuery))), HighlightEntity.class)
+				)
+				.build();
+		SearchHits<HighlightEntity> searchHits = operations.search(query, HighlightEntity.class, index);
+
+		assertThat(searchHits).isNotNull();
+		assertThat(searchHits.getSearchHits()).hasSize(1);
+
+		SearchHit<HighlightEntity> searchHit = searchHits.getSearchHit(0);
+		List<String> highlightField = searchHit.getHighlightField("message");
+		assertThat(highlightField).hasSize(2);
+		assertThat(highlightField.get(0)).contains("<em>message</em>");
+		assertThat(highlightField.get(1)).contains("<em>message</em>");
+	}
+
+	@Test // #2636
+	void shouldReturnDifferentHighlightFieldsWithDifferentHighlightQueryInSearchHit() {
+		IndexCoordinates index = createIndexCoordinatesWithHighlightMessage();
+
+		// a different highlight query from the search query
+		var differentHighlightQueryInField = HighlightFieldParameters.builder()
+				.withHighlightQuery(getBuilderWithTermQuery("message", "initial").build())
+				.build();
+		// highlight_query in field
+		Query highlightQueryInField = getBuilderWithTermQuery("message", "message") //
+				.withHighlightQuery(
+						new HighlightQuery(new Highlight(singletonList(new HighlightField("message", differentHighlightQueryInField))), HighlightEntity.class)
+				)
+				.build();
+		assertThatHighlightFieldsIsDifferentFromHighlightQuery(highlightQueryInField, index);
+	}
+
+	@Test // #2636
+	void shouldReturnDifferentHighlightFieldsWithDifferentParamHighlightQueryInSearchHit() {
+		IndexCoordinates index = createIndexCoordinatesWithHighlightMessage();
+
+		// a different highlight query from the search query and used in highlight param rather than field
+		var differentHighlightQueryInParam = HighlightParameters.builder()
+				.withHighlightQuery(getBuilderWithTermQuery("message", "initial").build())
+				.build();
+		// highlight_query in param
+		Query highlightQueryInParam = getBuilderWithTermQuery("message", "message") //
+				.withHighlightQuery(
+						new HighlightQuery(new Highlight(differentHighlightQueryInParam, singletonList(new HighlightField("message"))), HighlightEntity.class)
+				)
+				.build();
+		assertThatHighlightFieldsIsDifferentFromHighlightQuery(highlightQueryInParam, index);
+	}
+
+	@Test // #2636
+	void shouldReturnDifferentHighlightFieldsWithDifferentHighlightCriteriaQueryInSearchHit() {
+		IndexCoordinates index = createIndexCoordinatesWithHighlightMessage();
+		// a different highlight query from the search query, written by CriteriaQuery rather than NativeQuery
+		var criteriaHighlightQueryInParam = HighlightParameters.builder()
+				.withHighlightQuery(new CriteriaQuery(new Criteria("message").is("initial")))
+				.build();
+		// highlight_query in param
+		Query differentHighlightQueryUsingCriteria = getBuilderWithTermQuery("message", "message") //
+				.withHighlightQuery(
+						new HighlightQuery(new Highlight(criteriaHighlightQueryInParam, singletonList(new HighlightField("message"))), HighlightEntity.class)
+				)
+				.build();
+		assertThatHighlightFieldsIsDifferentFromHighlightQuery(differentHighlightQueryUsingCriteria, index);
+	}
+
+	@Test // #2636
+	void shouldReturnDifferentHighlightFieldsWithDifferentHighlightStringQueryInSearchHit() {
+		IndexCoordinates index = createIndexCoordinatesWithHighlightMessage();
+		// a different highlight query from the search query, written by StringQuery
+		var stringHighlightQueryInParam = HighlightParameters.builder()
+				.withHighlightQuery(new StringQuery(
+						"""
+								{
+								    "term": {
+								      "message": {
+								        "value": "initial"
+								      }
+								    }
+								}
+								"""
+				))
+				.build();
+		// highlight_query in param
+		Query differentHighlightQueryUsingStringQuery = getBuilderWithTermQuery("message", "message") //
+				.withHighlightQuery(
+						new HighlightQuery(new Highlight(stringHighlightQueryInParam, singletonList(new HighlightField("message"))), HighlightEntity.class)
+				)
+				.build();
+		assertThatHighlightFieldsIsDifferentFromHighlightQuery(differentHighlightQueryUsingStringQuery, index);
+	}
+
+	private IndexCoordinates createIndexCoordinatesWithHighlightMessage() {
+		IndexCoordinates index = IndexCoordinates.of("test-index-highlight-entity-template");
+		HighlightEntity entity = new HighlightEntity("1",
+				"This message is a long text which contains the word to search for "
+						+ "in two places, the first being near the beginning and the second near the end of the message. "
+						+ "However, i'll use a different highlight query from the initial search query");
+		IndexQuery indexQuery = new IndexQueryBuilder().withId(entity.getId()).withObject(entity).build();
+		operations.index(indexQuery, index);
+		operations.indexOps(index).refresh();
+		return index;
+	}
+
+	private void assertThatHighlightFieldsIsDifferentFromHighlightQuery(Query query, IndexCoordinates index) {
+		SearchHits<HighlightEntity> searchHits = operations.search(query, HighlightEntity.class, index);
+
+		SearchHit<HighlightEntity> searchHit = searchHits.getSearchHit(0);
+		List<String> highlightField = searchHit.getHighlightField("message");
+		assertThat(highlightField).hasSize(1);
+		assertThat(highlightField.get(0)).contains("<em>initial</em>");
+	}
+
 	@Test // #1686
 	void shouldRunRescoreQueryInSearchQuery() {
 		IndexCoordinates index = IndexCoordinates.of(indexNameProvider.getPrefix() + "rescore-entity");