HSEARCH-3661 Setting the value of bucket aggregations to something more than just the document count

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/document/model/dsl/impl/ElasticsearchIndexRootBuilder.java

@@ -26,6 +26,7 @@
 import org.hibernate.search.backend.elasticsearch.lowlevel.index.mapping.impl.RoutingType;
 import org.hibernate.search.backend.elasticsearch.lowlevel.index.settings.impl.IndexSettings;
 import org.hibernate.search.backend.elasticsearch.lowlevel.index.settings.impl.PropertyMappingIndexSettingsContributor;
+import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchCountDocumentAggregation;
 import org.hibernate.search.backend.elasticsearch.types.dsl.ElasticsearchIndexFieldTypeFactory;
 import org.hibernate.search.backend.elasticsearch.types.dsl.provider.impl.ElasticsearchIndexFieldTypeFactoryProvider;
 import org.hibernate.search.backend.elasticsearch.types.impl.ElasticsearchIndexCompositeNodeType;
@@ -45,6 +46,7 @@
 import org.hibernate.search.engine.common.tree.spi.TreeNodeInclusion;
 import org.hibernate.search.engine.mapper.mapping.building.spi.IndexFieldTypeDefaultsProvider;
 import org.hibernate.search.engine.reporting.spi.EventContexts;
+import org.hibernate.search.engine.search.aggregation.spi.AggregationTypeKeys;
 import org.hibernate.search.util.common.reporting.EventContext;
 
 public class ElasticsearchIndexRootBuilder extends AbstractElasticsearchIndexCompositeNodeBuilder
@@ -84,6 +86,8 @@ public ElasticsearchIndexRootBuilder(ElasticsearchIndexFieldTypeFactoryProvider
 		this.customIndexMapping = customIndexMapping;
 		this.defaultDynamicType = DynamicType.create( dynamicMapping );
 
+		this.typeBuilder.queryElementFactory( AggregationTypeKeys.COUNT_DOCUMENTS,
+				ElasticsearchCountDocumentAggregation.factory( false ) );
 		this.addDefaultImplicitFields();
 	}
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/impl/AbstractElasticsearchBucketAggregation.java

@@ -8,7 +8,6 @@
 import java.util.Map;
 
 import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
-import org.hibernate.search.backend.elasticsearch.logging.impl.ElasticsearchClientLog;
 import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexScope;
 import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexValueFieldContext;
 import org.hibernate.search.backend.elasticsearch.search.predicate.impl.ElasticsearchSearchPredicate;
@@ -29,21 +28,20 @@ public abstract class AbstractElasticsearchBucketAggregation<K, V>
 	private static final String ROOT_DOC_COUNT_NAME = "root_doc_count";
 	private static final JsonAccessor<JsonObject> REQUEST_AGGREGATIONS_ROOT_DOC_COUNT_ACCESSOR =
 			JsonAccessor.root().property( "aggregations" ).property( ROOT_DOC_COUNT_NAME ).asObject();
-	private static final JsonAccessor<Long> RESPONSE_DOC_COUNT_ACCESSOR =
-			JsonAccessor.root().property( "doc_count" ).asLong();
-	private static final JsonAccessor<Long> RESPONSE_ROOT_DOC_COUNT_ACCESSOR =
-			JsonAccessor.root().property( ROOT_DOC_COUNT_NAME ).property( "doc_count" ).asLong();
+
+	protected static final String INNER_EXTRACTOR_KEY = "innerExtractorKey";
+	protected static final String INNER_EXTRACTOR = "innerExtractor";
 
 	AbstractElasticsearchBucketAggregation(AbstractBuilder<K, V> builder) {
 		super( builder );
 	}
 
 	@Override
-	protected final JsonObject doRequest(AggregationRequestContext context) {
+	protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
 		JsonObject outerObject = new JsonObject();
 		JsonObject innerObject = new JsonObject();
 
-		doRequest( outerObject, innerObject );
+		doRequest( outerObject, innerObject, context );
 
 		if ( isNested() ) {
 			JsonObject rootDocCountSubAggregationOuterObject = new JsonObject();
@@ -56,20 +54,8 @@ protected final JsonObject doRequest(AggregationRequestContext context) {
 		return outerObject;
 	}
 
-	protected abstract void doRequest(JsonObject outerObject, JsonObject innerObject);
-
-	protected final long getBucketDocCount(JsonObject bucket) {
-		if ( isNested() ) {
-			// We must return the number of root documents,
-			// not the number of leaf documents that Elasticsearch returns by default.
-			return RESPONSE_ROOT_DOC_COUNT_ACCESSOR.get( bucket )
-					.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
-		}
-		else {
-			return RESPONSE_DOC_COUNT_ACCESSOR.get( bucket )
-					.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
-		}
-	}
+	protected abstract void doRequest(JsonObject outerObject, JsonObject innerObject,
+			AggregationRequestBuildingContextContext context);
 
 	protected abstract class AbstractBucketExtractor<A, B> extends AbstractExtractor<Map<A, B>> {
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/impl/AbstractElasticsearchNestableAggregation.java

@@ -50,11 +50,12 @@ public abstract class AbstractElasticsearchNestableAggregation<A> extends Abstra
 
 	@Override
 	public final Extractor<A> request(AggregationRequestContext context, AggregationKey<?> key, JsonObject jsonAggregations) {
-		jsonAggregations.add( key.name(), request( context ) );
-		return extractor( context );
+		AggregationRequestBuildingContextContext buildingContext = new AggregationRequestBuildingContextContext( context );
+		jsonAggregations.add( key.name(), request( buildingContext ) );
+		return extractor( buildingContext );
 	}
 
-	private JsonObject request(AggregationRequestContext context) {
+	private JsonObject request(AggregationRequestBuildingContextContext context) {
 		JsonObject result = doRequest( context );
 
 		if ( nestedPathHierarchy.isEmpty() ) {
@@ -90,9 +91,9 @@ private JsonObject request(AggregationRequestContext context) {
 		return result;
 	}
 
-	protected abstract JsonObject doRequest(AggregationRequestContext context);
+	protected abstract JsonObject doRequest(AggregationRequestBuildingContextContext context);
 
-	protected abstract Extractor<A> extractor(AggregationRequestContext context);
+	protected abstract Extractor<A> extractor(AggregationRequestBuildingContextContext context);
 
 	protected abstract static class AbstractExtractor<T> implements Extractor<T> {
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/impl/AggregationRequestBuildingContextContext.java

@@ -0,0 +1,76 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.search.backend.elasticsearch.search.aggregation.impl;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+import org.hibernate.search.backend.elasticsearch.search.predicate.impl.PredicateRequestContext;
+import org.hibernate.search.util.common.annotation.Incubating;
+
+/**
+ * Sometimes we need to pass something we created while building up the json in one of the "doRequest" methods
+ * in the aggregation build up to the "later" steps e.g. to when we create the extractor.
+ */
+@Incubating
+public final class AggregationRequestBuildingContextContext implements AggregationRequestContext {
+	private final AggregationRequestContext aggregationRequestContext;
+	private final Map<Key<?>, Object> buildingContext = new HashMap<>();
+
+	public AggregationRequestBuildingContextContext(AggregationRequestContext aggregationRequestContext) {
+		this.aggregationRequestContext = aggregationRequestContext;
+	}
+
+	public <T> T get(Key<T> key) {
+		Object value = buildingContext.get( key );
+		return key.cast( value );
+	}
+
+	public void add(Key<?> key, Object value) {
+		buildingContext.put( key, value );
+	}
+
+	public AggregationRequestContext rootAggregationRequestContext() {
+		return aggregationRequestContext;
+	}
+
+	@Override
+	public PredicateRequestContext getRootPredicateContext() {
+		return aggregationRequestContext.getRootPredicateContext();
+	}
+
+	public static <V> Key<V> buildingContextKey(String name) {
+		return new Key<>( name );
+	}
+
+	public static class Key<V> {
+
+		private final String name;
+
+		private Key(String name) {
+			this.name = name;
+		}
+
+		@SuppressWarnings("unchecked")
+		private V cast(Object value) {
+			return (V) value;
+		}
+
+		@Override
+		public boolean equals(Object o) {
+			if ( o == null || getClass() != o.getClass() ) {
+				return false;
+			}
+			Key<?> key = (Key<?>) o;
+			return Objects.equals( name, key.name );
+		}
+
+		@Override
+		public int hashCode() {
+			return Objects.hashCode( name );
+		}
+	}
+}

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/impl/ElasticsearchCountDocumentAggregation.java

@@ -0,0 +1,121 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.search.backend.elasticsearch.search.aggregation.impl;
+
+import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
+import org.hibernate.search.backend.elasticsearch.logging.impl.ElasticsearchClientLog;
+import org.hibernate.search.backend.elasticsearch.logging.impl.QueryLog;
+import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexCompositeNodeContext;
+import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexScope;
+import org.hibernate.search.backend.elasticsearch.search.query.impl.ElasticsearchSearchQueryExtractContext;
+import org.hibernate.search.engine.search.aggregation.AggregationKey;
+import org.hibernate.search.engine.search.aggregation.spi.CountDocumentAggregationBuilder;
+import org.hibernate.search.engine.search.common.spi.SearchQueryElementFactory;
+
+import com.google.gson.JsonObject;
+
+public class ElasticsearchCountDocumentAggregation extends AbstractElasticsearchAggregation<Long> {
+
+	private static final JsonAccessor<Long> TOTAL_HITS_VALUE_PROPERTY_ACCESSOR =
+			JsonAccessor.root().property( "hits" )
+					.property( "total" )
+					.property( "value" ).asLong();
+
+	private static final JsonAccessor<Long> RESPONSE_DOC_COUNT_ACCESSOR =
+			JsonAccessor.root().property( "doc_count" ).asLong();
+	private static final JsonAccessor<Long> RESPONSE_ROOT_DOC_COUNT_ACCESSOR =
+			JsonAccessor.root().property( "root_doc_count" ).property( "doc_count" ).asLong();
+
+	public static SearchQueryElementFactory<CountDocumentAggregationBuilder.TypeSelector,
+			ElasticsearchSearchIndexScope<?>,
+			ElasticsearchSearchIndexCompositeNodeContext> factory(boolean isNested) {
+		return new ElasticsearchCountDocumentAggregation.Factory( isNested );
+	}
+
+	private final boolean isNested;
+
+	private ElasticsearchCountDocumentAggregation(Builder builder) {
+		super( builder );
+		this.isNested = builder.isNested;
+	}
+
+	@Override
+	public Extractor<Long> request(AggregationRequestContext context, AggregationKey<?> key, JsonObject jsonAggregations) {
+		return new CountDocumentsExtractor( isNested );
+	}
+
+	private record CountDocumentsExtractor(boolean isNested) implements Extractor<Long> {
+
+		@Override
+		public Long extract(JsonObject aggregationResult, AggregationExtractContext context) {
+			if ( aggregationResult != null ) {
+				if ( isNested ) {
+					// We must return the number of root documents,
+					// not the number of leaf documents that Elasticsearch returns by default.
+					return RESPONSE_ROOT_DOC_COUNT_ACCESSOR.get( aggregationResult )
+							.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
+				}
+				else {
+					return RESPONSE_DOC_COUNT_ACCESSOR.get( aggregationResult )
+							.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
+				}
+			}
+			else if ( context instanceof ElasticsearchSearchQueryExtractContext c ) {
+				return TOTAL_HITS_VALUE_PROPERTY_ACCESSOR.get( c.getResponseBody() )
+						.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
+			}
+			throw ElasticsearchClientLog.INSTANCE.elasticsearchResponseMissingData();
+		}
+	}
+
+	private static class Factory
+			implements
+			SearchQueryElementFactory<CountDocumentAggregationBuilder.TypeSelector,
+					ElasticsearchSearchIndexScope<?>,
+					ElasticsearchSearchIndexCompositeNodeContext> {
+		private final boolean isNested;
+
+		public Factory(boolean isNested) {
+			this.isNested = isNested;
+		}
+
+		@Override
+		public CountDocumentAggregationBuilder.TypeSelector create(ElasticsearchSearchIndexScope<?> scope,
+				ElasticsearchSearchIndexCompositeNodeContext node) {
+			return new ElasticsearchCountDocumentAggregation.TypeSelector( scope, isNested );
+		}
+
+		@Override
+		public void checkCompatibleWith(SearchQueryElementFactory<?, ?, ?> other) {
+			if ( !getClass().equals( other.getClass() ) ) {
+				throw QueryLog.INSTANCE.differentImplementationClassForQueryElement( getClass(), other.getClass() );
+			}
+		}
+	}
+
+	private record TypeSelector(ElasticsearchSearchIndexScope<?> scope, boolean isNested)
+			implements CountDocumentAggregationBuilder.TypeSelector {
+
+		@Override
+		public CountDocumentAggregationBuilder type() {
+			return new Builder( scope, isNested );
+		}
+	}
+
+	private static class Builder extends AbstractBuilder<Long>
+			implements CountDocumentAggregationBuilder {
+		private final boolean isNested;
+
+		private Builder(ElasticsearchSearchIndexScope<?> scope, boolean isNested) {
+			super( scope );
+			this.isNested = isNested;
+		}
+
+		@Override
+		public ElasticsearchCountDocumentAggregation build() {
+			return new ElasticsearchCountDocumentAggregation( this );
+		}
+	}
+}

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/impl/ElasticsearchMetricFieldAggregation.java

@@ -74,7 +74,7 @@ private ElasticsearchMetricFieldAggregation(Builder<F, K> builder) {
 	}
 
 	@Override
-	protected final JsonObject doRequest(AggregationRequestContext context) {
+	protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
 		JsonObject outerObject = new JsonObject();
 		JsonObject innerObject = new JsonObject();
 
@@ -84,7 +84,7 @@ protected final JsonObject doRequest(AggregationRequestContext context) {
 	}
 
 	@Override
-	protected Extractor<K> extractor(AggregationRequestContext context) {
+	protected Extractor<K> extractor(AggregationRequestBuildingContextContext context) {
 		return metricFieldExtractorCreator.extractor( filter );
 	}
 

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/impl/ElasticsearchMetricLongAggregation.java

@@ -47,7 +47,7 @@ private ElasticsearchMetricLongAggregation(Builder builder) {
 	}
 
 	@Override
-	protected final JsonObject doRequest(AggregationRequestContext context) {
+	protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
 		JsonObject outerObject = new JsonObject();
 		JsonObject innerObject = new JsonObject();
 
@@ -57,7 +57,7 @@ protected final JsonObject doRequest(AggregationRequestContext context) {
 	}
 
 	@Override
-	protected Extractor<Long> extractor(AggregationRequestContext context) {
+	protected Extractor<Long> extractor(AggregationRequestBuildingContextContext context) {
 		return new MetricLongExtractor( nestedPathHierarchy, filter );
 	}