HSEARCH-5426 Adjust the count aggregation DSL

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

@@ -87,7 +87,7 @@ public ElasticsearchIndexRootBuilder(ElasticsearchIndexFieldTypeFactoryProvider
 		this.defaultDynamicType = DynamicType.create( dynamicMapping );
 
 		this.typeBuilder.queryElementFactory( AggregationTypeKeys.COUNT_DOCUMENTS,
-				ElasticsearchCountDocumentAggregation.factory( false ) );
+				ElasticsearchCountDocumentAggregation.factory() );
 		this.addDefaultImplicitFields();
 	}
 

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

@@ -7,7 +7,7 @@
 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.ElasticsearchSearchIndexNodeContext;
 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;
@@ -28,19 +28,55 @@ public class ElasticsearchCountDocumentAggregation extends AbstractElasticsearch
 	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;
 	}
 
+	public static SearchQueryElementFactory<CountDocumentAggregationBuilder.TypeSelector,
+			ElasticsearchSearchIndexScope<?>,
+			ElasticsearchSearchIndexNodeContext> factory() {
+		return new Factory();
+	}
+
+	private static class Factory
+			implements SearchQueryElementFactory<CountDocumentAggregationBuilder.TypeSelector,
+					ElasticsearchSearchIndexScope<?>,
+					ElasticsearchSearchIndexNodeContext> {
+
+		@Override
+		public CountDocumentAggregationBuilder.TypeSelector create(ElasticsearchSearchIndexScope<?> scope,
+				ElasticsearchSearchIndexNodeContext node) {
+			return new TypeSelector( scope, node );
+		}
+
+		@Override
+		public void checkCompatibleWith(SearchQueryElementFactory<?, ?, ?> other) {
+			if ( !getClass().equals( other.getClass() ) ) {
+				throw QueryLog.INSTANCE.differentImplementationClassForQueryElement( getClass(), other.getClass() );
+			}
+		}
+	}
+
+	private static final class TypeSelector implements CountDocumentAggregationBuilder.TypeSelector {
+		private final ElasticsearchSearchIndexScope<?> scope;
+		private final ElasticsearchSearchIndexNodeContext node;
+
+		private TypeSelector(ElasticsearchSearchIndexScope<?> scope,
+				ElasticsearchSearchIndexNodeContext node) {
+			this.scope = scope;
+			this.node = node; // doesn't matter in this case
+		}
+
+		@Override
+		public CountDocumentAggregationBuilder builder() {
+			boolean isNested = node.isValueField() && !node.toValueField().nestedPathHierarchy().isEmpty();
+			return new ElasticsearchCountDocumentAggregation.Builder( scope, isNested );
+		}
+	}
+
 	@Override
 	public Extractor<Long> request(AggregationRequestContext context, AggregationKey<?> key, JsonObject jsonAggregations) {
 		return new CountDocumentsExtractor( key, isNested, context.isRootContext() );
@@ -73,40 +109,6 @@ public Long extract(JsonObject aggregationResult, AggregationExtractContext cont
 		}
 	}
 
-	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;

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

@@ -0,0 +1,138 @@
+/*
+ * 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.List;
+
+import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
+import org.hibernate.search.backend.elasticsearch.gson.impl.JsonElementTypes;
+import org.hibernate.search.backend.elasticsearch.logging.impl.QueryLog;
+import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexNodeContext;
+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;
+import org.hibernate.search.engine.search.aggregation.AggregationKey;
+import org.hibernate.search.engine.search.aggregation.spi.CountValuesAggregationBuilder;
+import org.hibernate.search.engine.search.common.spi.SearchQueryElementFactory;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+
+public class ElasticsearchCountValuesAggregation extends AbstractElasticsearchNestableAggregation<Long> {
+
+	private static final JsonAccessor<JsonObject> COUNT_PROPERTY_ACCESSOR =
+			JsonAccessor.root().property( "value_count" ).asObject();
+
+	private static final JsonAccessor<JsonObject> COUNT_DISTINCT_PROPERTY_ACCESSOR =
+			JsonAccessor.root().property( "cardinality" ).asObject();
+
+	private static final JsonAccessor<String> FIELD_PROPERTY_ACCESSOR =
+			JsonAccessor.root().property( "field" ).asString();
+
+	private final String absoluteFieldPath;
+	private final JsonAccessor<JsonObject> operation;
+
+	private ElasticsearchCountValuesAggregation(Builder builder) {
+		super( builder );
+		this.absoluteFieldPath = builder.field.absolutePath();
+		this.operation = builder.operation;
+	}
+
+	public static SearchQueryElementFactory<CountValuesAggregationBuilder.TypeSelector,
+			ElasticsearchSearchIndexScope<?>,
+			ElasticsearchSearchIndexNodeContext> factory() {
+		return new Factory();
+	}
+
+	private static class Factory
+			implements SearchQueryElementFactory<CountValuesAggregationBuilder.TypeSelector,
+					ElasticsearchSearchIndexScope<?>,
+					ElasticsearchSearchIndexNodeContext> {
+
+		@Override
+		public CountValuesAggregationBuilder.TypeSelector create(ElasticsearchSearchIndexScope<?> scope,
+				ElasticsearchSearchIndexNodeContext node) {
+			return new TypeSelector( scope, node );
+		}
+
+		@Override
+		public void checkCompatibleWith(SearchQueryElementFactory<?, ?, ?> other) {
+			if ( !getClass().equals( other.getClass() ) ) {
+				throw QueryLog.INSTANCE.differentImplementationClassForQueryElement( getClass(), other.getClass() );
+			}
+		}
+	}
+
+	private static final class TypeSelector implements CountValuesAggregationBuilder.TypeSelector {
+		private final ElasticsearchSearchIndexScope<?> scope;
+		private final ElasticsearchSearchIndexNodeContext node;
+
+		private TypeSelector(ElasticsearchSearchIndexScope<?> scope,
+				ElasticsearchSearchIndexNodeContext node) {
+			this.scope = scope;
+			this.node = node; // doesn't matter in this case
+		}
+
+		@Override
+		public CountValuesAggregationBuilder builder() {
+			return new ElasticsearchCountValuesAggregation.Builder( scope, node.toValueField() );
+		}
+	}
+
+	@Override
+	protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
+		JsonObject outerObject = new JsonObject();
+		JsonObject innerObject = new JsonObject();
+
+		operation.set( outerObject, innerObject );
+		FIELD_PROPERTY_ACCESSOR.set( innerObject, absoluteFieldPath );
+		return outerObject;
+	}
+
+	@Override
+	protected Extractor<Long> extractor(AggregationKey<?> key, AggregationRequestBuildingContextContext context) {
+		return new MetricLongExtractor( key, nestedPathHierarchy, filter );
+	}
+
+	private static class MetricLongExtractor extends AbstractExtractor<Long> {
+		protected MetricLongExtractor(
+				AggregationKey<?> key, List<String> nestedPathHierarchy,
+				ElasticsearchSearchPredicate filter
+		) {
+			super( key, nestedPathHierarchy, filter );
+		}
+
+		@Override
+		protected Long doExtract(JsonObject aggregationResult, AggregationExtractContext context) {
+			JsonElement value = aggregationResult.get( "value" );
+			return JsonElementTypes.LONG.fromElement( value );
+		}
+	}
+
+	private static class Builder extends AbstractBuilder<Long>
+			implements CountValuesAggregationBuilder {
+		private JsonAccessor<JsonObject> operation;
+
+		private Builder(ElasticsearchSearchIndexScope<?> scope, ElasticsearchSearchIndexValueFieldContext<?> field) {
+			super( scope, field );
+			this.operation = COUNT_PROPERTY_ACCESSOR;
+		}
+
+		@Override
+		public ElasticsearchCountValuesAggregation build() {
+			return new ElasticsearchCountValuesAggregation( this );
+		}
+
+		@Override
+		public void distinct(boolean distinct) {
+			if ( distinct ) {
+				operation = COUNT_DISTINCT_PROPERTY_ACCESSOR;
+			}
+			else {
+				operation = COUNT_PROPERTY_ACCESSOR;
+			}
+		}
+	}
+}

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

@@ -138,8 +138,7 @@ protected CountBuilder(ElasticsearchSearchIndexScope<?> scope,
 				Function<? super K, JsonElement> encoder) {
 			super( scope, field, encoder, new ArrayList<>(), new JsonArray(),
 					ElasticsearchSearchAggregation.from( scope,
-							ElasticsearchCountDocumentAggregation.factory( !field.nestedPathHierarchy().isEmpty() )
-									.create( scope, null ).type().build() ) );
+							ElasticsearchCountDocumentAggregation.factory().create( scope, field ).builder().build() ) );
 		}
 	}
 

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

@@ -163,8 +163,7 @@ protected CountBuilder(ElasticsearchSearchIndexScope<?> scope,
 				ProjectionConverter<T, ? extends K> fromFieldValueConverter) {
 			super( scope, field, decodeFunction, fromFieldValueConverter,
 					ElasticsearchSearchAggregation.from( scope,
-							ElasticsearchCountDocumentAggregation.factory( !field.nestedPathHierarchy().isEmpty() )
-									.create( scope, null ).type().build() ) );
+							ElasticsearchCountDocumentAggregation.factory().create( scope, field ).builder().build() ) );
 		}
 	}