HSEARCH-3649 Use a separate context for Elasticsearch projection requ…

…est and extraction

To be consistent with what we already have in Lucene, and with what we
will soon have for aggregations in Elasticsearch.

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/AbstractElasticsearchCompositeProjection.java

@@ -35,10 +35,10 @@ public String toString() {
 	}
 
 	@Override
-	public final void contributeRequest(JsonObject requestBody,
-			SearchProjectionExtractContext context) {
+	public final void request(JsonObject requestBody,
+			SearchProjectionRequestContext context) {
 		for ( ElasticsearchSearchProjection<?, ?> child : children ) {
-			child.contributeRequest( requestBody, context );
+			child.request( requestBody, context );
 		}
 	}
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/DistanceSortKey.java

@@ -0,0 +1,43 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.backend.elasticsearch.search.projection.impl;
+
+import java.util.Objects;
+
+import org.hibernate.search.engine.spatial.GeoPoint;
+
+public class DistanceSortKey {
+
+	private final String absoluteFieldPath;
+
+	private final GeoPoint location;
+
+	public DistanceSortKey(String absoluteFieldPath, GeoPoint location) {
+		this.absoluteFieldPath = absoluteFieldPath;
+		this.location = location;
+	}
+
+	@Override
+	public boolean equals(Object obj) {
+		if ( obj == this ) {
+			return true;
+		}
+		if ( !(obj instanceof DistanceSortKey ) ) {
+			return false;
+		}
+
+		DistanceSortKey other = (DistanceSortKey) obj;
+
+		return Objects.equals( this.absoluteFieldPath, other.absoluteFieldPath )
+				&& Objects.equals( this.location, other.location );
+	}
+
+	@Override
+	public int hashCode() {
+		return Objects.hash( absoluteFieldPath, location );
+	}
+}

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/DocumentReferenceExtractorHelper.java

@@ -32,7 +32,7 @@ public DocumentReferenceExtractorHelper(Function<String, String> indexNameConver
 		this.multiTenancyStrategy = multiTenancyStrategy;
 	}
 
-	public void contributeRequest(JsonObject requestBody) {
+	public void requestDocumentReference(JsonObject requestBody) {
 		multiTenancyStrategy.contributeToSearchRequest( requestBody );
 	}
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchDistanceToFieldProjection.java

@@ -67,7 +67,7 @@ class ElasticsearchDistanceToFieldProjection implements ElasticsearchSearchProje
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
 		if ( context.getDistanceSortIndex( absoluteFieldPath, center ) == null ) {
 			// we rely on a script to compute the distance
 			SCRIPT_FIELDS_ACCESSOR

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchDocumentReferenceProjection.java

@@ -26,8 +26,8 @@ class ElasticsearchDocumentReferenceProjection
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
-		helper.contributeRequest( requestBody );
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
+		helper.requestDocumentReference( requestBody );
 	}
 
 	@Override

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchEntityProjection.java

@@ -24,8 +24,8 @@ public ElasticsearchEntityProjection(Set<String> indexNames, DocumentReferenceEx
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
-		helper.contributeRequest( requestBody );
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
+		helper.requestDocumentReference( requestBody );
 	}
 
 	@Override

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchEntityReferenceProjection.java

@@ -24,8 +24,8 @@ public ElasticsearchEntityReferenceProjection(Set<String> indexNames, DocumentRe
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
-		helper.contributeRequest( requestBody );
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
+		helper.requestDocumentReference( requestBody );
 	}
 
 	@SuppressWarnings("unchecked")

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchExplanationProjection.java

@@ -30,7 +30,7 @@ class ElasticsearchExplanationProjection implements ElasticsearchSearchProjectio
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
 		REQUEST_EXPLAIN_ACCESSOR.set( requestBody, true );
 	}
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchFieldProjection.java

@@ -47,7 +47,7 @@ class ElasticsearchFieldProjection<F, V> implements ElasticsearchSearchProjectio
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
 		JsonArray source = REQUEST_SOURCE_ACCESSOR.getOrCreate( requestBody, JsonArray::new );
 		JsonPrimitive fieldPathJson = new JsonPrimitive( absoluteFieldPath );
 		if ( !source.contains( fieldPathJson ) ) {

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchScoreProjection.java

@@ -26,7 +26,7 @@ class ElasticsearchScoreProjection implements ElasticsearchSearchProjection<Floa
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
 		TRACK_SCORES_ACCESSOR.set( requestBody, true );
 	}
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchSearchProjection.java

@@ -19,10 +19,10 @@ public interface ElasticsearchSearchProjection<E, P> extends SearchProjection<P>
 	/**
 	 * Contribute to the request, making sure that the requirements for this projection are met.
 	 * @param requestBody The request body.
-	 * @param context An execution context that will also be passed to
+	 * @param context An execution context that will share state with the context passed to
 	 * {@link #extract(ProjectionHitMapper, JsonObject, JsonObject, SearchProjectionExtractContext)}.
 	 */
-	void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context);
+	void request(JsonObject requestBody, SearchProjectionRequestContext context);
 
 	/**
 	 * Perform hit extraction.

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/ElasticsearchSourceProjection.java

@@ -35,7 +35,7 @@ class ElasticsearchSourceProjection implements ElasticsearchSearchProjection<Str
 	}
 
 	@Override
-	public void contributeRequest(JsonObject requestBody, SearchProjectionExtractContext context) {
+	public void request(JsonObject requestBody, SearchProjectionRequestContext context) {
 		JsonArray source = REQUEST_SOURCE_ACCESSOR.getOrCreate( requestBody, JsonArray::new );
 		if ( !source.contains( WILDCARD_ALL ) ) {
 			source.add( WILDCARD_ALL );

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/SearchProjectionExtractContext.java

@@ -6,57 +6,18 @@
  */
 package org.hibernate.search.backend.elasticsearch.search.projection.impl;
 
-import java.util.Collections;
-import java.util.Map;
-import java.util.Objects;
-
 import org.hibernate.search.engine.spatial.GeoPoint;
 
 public class SearchProjectionExtractContext {
 
-	private final Map<DistanceSortKey, Integer> distanceSorts;
+	private final SearchProjectionRequestContext requestContext;
 
-	public SearchProjectionExtractContext(Map<DistanceSortKey, Integer> distanceSorts) {
-		this.distanceSorts = distanceSorts != null ? Collections.unmodifiableMap( distanceSorts ) : null;
+	public SearchProjectionExtractContext(SearchProjectionRequestContext requestContext) {
+		this.requestContext = requestContext;
 	}
 
 	Integer getDistanceSortIndex(String absoluteFieldPath, GeoPoint location) {
-		if ( distanceSorts == null ) {
-			return null;
-		}
-
-		return distanceSorts.get( new DistanceSortKey( absoluteFieldPath, location ) );
+		return requestContext.getDistanceSortIndex( absoluteFieldPath, location );
 	}
 
-	public static class DistanceSortKey {
-
-		private final String absoluteFieldPath;
-
-		private final GeoPoint location;
-
-		public DistanceSortKey(String absoluteFieldPath, GeoPoint location) {
-			this.absoluteFieldPath = absoluteFieldPath;
-			this.location = location;
-		}
-
-		@Override
-		public boolean equals(Object obj) {
-			if ( obj == this ) {
-				return true;
-			}
-			if ( !(obj instanceof DistanceSortKey) ) {
-				return false;
-			}
-
-			DistanceSortKey other = (DistanceSortKey) obj;
-
-			return Objects.equals( this.absoluteFieldPath, other.absoluteFieldPath )
-					&& Objects.equals( this.location, other.location );
-		}
-
-		@Override
-		public int hashCode() {
-			return Objects.hash( absoluteFieldPath, location );
-		}
-	}
 }

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/projection/impl/SearchProjectionRequestContext.java

@@ -0,0 +1,15 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.backend.elasticsearch.search.projection.impl;
+
+import org.hibernate.search.engine.spatial.GeoPoint;
+
+public interface SearchProjectionRequestContext {
+
+	Integer getDistanceSortIndex(String absoluteFieldPath, GeoPoint location);
+
+}

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/query/impl/ElasticsearchSearchQueryBuilder.java

@@ -15,8 +15,8 @@
 import org.hibernate.search.backend.elasticsearch.search.impl.ElasticsearchSearchContext;
 import org.hibernate.search.backend.elasticsearch.search.impl.ElasticsearchSearchQueryElementCollector;
 import org.hibernate.search.backend.elasticsearch.search.predicate.impl.ElasticsearchSearchPredicateContext;
+import org.hibernate.search.backend.elasticsearch.search.projection.impl.DistanceSortKey;
 import org.hibernate.search.backend.elasticsearch.search.projection.impl.ElasticsearchSearchProjection;
-import org.hibernate.search.backend.elasticsearch.search.projection.impl.SearchProjectionExtractContext;
 import org.hibernate.search.backend.elasticsearch.search.query.ElasticsearchSearchQuery;
 import org.hibernate.search.backend.elasticsearch.work.builder.factory.impl.ElasticsearchWorkBuilderFactory;
 import org.hibernate.search.backend.elasticsearch.work.impl.ElasticsearchSearchResultExtractor;
@@ -50,7 +50,7 @@ public class ElasticsearchSearchQueryBuilder<H>
 	private final Set<String> routingKeys;
 	private JsonObject jsonPredicate;
 	private JsonArray jsonSort;
-	private Map<SearchProjectionExtractContext.DistanceSortKey, Integer> distanceSorts;
+	private Map<DistanceSortKey, Integer> distanceSorts;
 
 	public ElasticsearchSearchQueryBuilder(
 			ElasticsearchWorkBuilderFactory workFactory,
@@ -112,7 +112,7 @@ public void collectDistanceSort(JsonElement sort, String absoluteFieldPath, GeoP
 			distanceSorts = CollectionHelper.newHashMap( 3 );
 		}
 
-		distanceSorts.put( new SearchProjectionExtractContext.DistanceSortKey( absoluteFieldPath, center ), index );
+		distanceSorts.put( new DistanceSortKey( absoluteFieldPath, center ), index );
 	}
 
 	@Override
@@ -130,17 +130,14 @@ public ElasticsearchSearchQuery<H> build() {
 			payload.add( "sort", jsonSort );
 		}
 
-		SearchProjectionExtractContext searchProjectionExecutionContext =
-				new SearchProjectionExtractContext( distanceSorts );
-
-		rootProjection.contributeRequest( payload, searchProjectionExecutionContext );
-
 		LoadingContext<?, ?> loadingContext = loadingContextBuilder.build();
 
 		ElasticsearchSearchQueryRequestContext requestContext = new ElasticsearchSearchQueryRequestContext(
 				sessionContext, loadingContext, distanceSorts
 		);
 
+		rootProjection.request( payload, requestContext );
+
 		ElasticsearchSearchResultExtractor<ElasticsearchLoadableSearchResult<H>> searchResultExtractor =
 				searchResultExtractorFactory.createResultExtractor(
 						requestContext,

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/query/impl/ElasticsearchSearchQueryExtractContext.java

@@ -6,10 +6,7 @@
  */
 package org.hibernate.search.backend.elasticsearch.search.query.impl;
 
-import java.util.Map;
-
 import org.hibernate.search.backend.elasticsearch.search.projection.impl.SearchProjectionExtractContext;
-import org.hibernate.search.backend.elasticsearch.search.projection.impl.SearchProjectionExtractContext.DistanceSortKey;
 import org.hibernate.search.backend.elasticsearch.search.projection.impl.SearchProjectionTransformContext;
 import org.hibernate.search.engine.backend.types.converter.runtime.FromDocumentFieldValueConvertContext;
 import org.hibernate.search.engine.backend.types.converter.runtime.spi.FromDocumentFieldValueConvertContextImpl;
@@ -24,18 +21,18 @@
  */
 class ElasticsearchSearchQueryExtractContext {
 
-	private final Map<DistanceSortKey, Integer> distanceSorts;
+	private final ElasticsearchSearchQueryRequestContext requestContext;
 	private final ProjectionHitMapper<?, ?> projectionHitMapper;
 	private final FromDocumentFieldValueConvertContext convertContext;
 
 	private final JsonObject responseBody;
 
-	ElasticsearchSearchQueryExtractContext(SessionContextImplementor sessionContext,
+	ElasticsearchSearchQueryExtractContext(ElasticsearchSearchQueryRequestContext requestContext,
+			SessionContextImplementor sessionContext,
 			ProjectionHitMapper<?, ?> projectionHitMapper,
-			Map<DistanceSortKey, Integer> distanceSorts,
 			JsonObject responseBody) {
+		this.requestContext = requestContext;
 		this.projectionHitMapper = projectionHitMapper;
-		this.distanceSorts = distanceSorts;
 		this.convertContext = new FromDocumentFieldValueConvertContextImpl( sessionContext );
 		this.responseBody = responseBody;
 	}
@@ -49,7 +46,7 @@ JsonObject getResponseBody() {
 	}
 
 	SearchProjectionExtractContext createProjectionExtractContext() {
-		return new SearchProjectionExtractContext( distanceSorts );
+		return new SearchProjectionExtractContext( requestContext );
 	}
 
 	SearchProjectionTransformContext createProjectionTransformContext() {

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/query/impl/ElasticsearchSearchQueryRequestContext.java

@@ -6,40 +6,57 @@
  */
 package org.hibernate.search.backend.elasticsearch.search.query.impl;
 
+import java.util.Collections;
 import java.util.Map;
 
-import org.hibernate.search.backend.elasticsearch.search.projection.impl.SearchProjectionExtractContext;
+import org.hibernate.search.backend.elasticsearch.search.projection.impl.DistanceSortKey;
+import org.hibernate.search.backend.elasticsearch.search.projection.impl.SearchProjectionRequestContext;
 import org.hibernate.search.engine.mapper.session.context.spi.SessionContextImplementor;
 import org.hibernate.search.engine.search.loading.context.spi.LoadingContext;
+import org.hibernate.search.engine.spatial.GeoPoint;
 
 import com.google.gson.JsonObject;
 
 /**
  * The context holding all the useful information pertaining to the Elasticsearch search query,
- * to be used when extracting data from the response,
- * to get an "extract" context linked to the session/loading context
- * ({@link #createExtractContext(JsonObject)}.
+ * to be used:
+ * <ul>
+ *     <li>When building later parts of the query, to get information on more basic parts of the query.
+ *     For example distance projections need to inspect distance sorts (if any) for optimization purposes.
+ *     ({@link #getDistanceSortIndex(String, GeoPoint)}</li>
+ *     <li>When extracting data from the response, to get an "extract" context linked to the session/loading context
+ *     ({@link #createExtractContext(JsonObject)}</li>
+ * </ul>
  */
-class ElasticsearchSearchQueryRequestContext {
+class ElasticsearchSearchQueryRequestContext implements SearchProjectionRequestContext {
 
 	private final SessionContextImplementor sessionContext;
 	private final LoadingContext<?, ?> loadingContext;
-	private final Map<SearchProjectionExtractContext.DistanceSortKey, Integer> distanceSorts;
+	private final Map<DistanceSortKey, Integer> distanceSorts;
 
 	ElasticsearchSearchQueryRequestContext(
 			SessionContextImplementor sessionContext,
 			LoadingContext<?, ?> loadingContext,
-			Map<SearchProjectionExtractContext.DistanceSortKey, Integer> distanceSorts) {
+			Map<DistanceSortKey, Integer> distanceSorts) {
 		this.sessionContext = sessionContext;
 		this.loadingContext = loadingContext;
-		this.distanceSorts = distanceSorts;
+		this.distanceSorts = distanceSorts != null ? Collections.unmodifiableMap( distanceSorts ) : null;
+	}
+
+	@Override
+	public Integer getDistanceSortIndex(String absoluteFieldPath, GeoPoint location) {
+		if ( distanceSorts == null ) {
+			return null;
+		}
+
+		return distanceSorts.get( new DistanceSortKey( absoluteFieldPath, location ) );
 	}
 
 	ElasticsearchSearchQueryExtractContext createExtractContext(JsonObject responseBody) {
 		return new ElasticsearchSearchQueryExtractContext(
+				this,
 				sessionContext,
 				loadingContext.getProjectionHitMapper(),
-				distanceSorts,
 				responseBody
 		);
 	}