Skip to content

Commit 1e9f309

Browse files
marko-bekhtamarko-bekhta
committed
HSEARCH-5426 Adjust the count aggregation DSL
1 parent 273aa90 commit 1e9f309

File tree

68 files changed

+842
-957
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

68 files changed

+842
-957
lines changed

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
import org.hibernate.search.backend.elasticsearch.lowlevel.index.mapping.impl.RoutingType;
2727
import org.hibernate.search.backend.elasticsearch.lowlevel.index.settings.impl.IndexSettings;
2828
import org.hibernate.search.backend.elasticsearch.lowlevel.index.settings.impl.PropertyMappingIndexSettingsContributor;
29-
import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchCountDocumentAggregation;
29+
import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchCountAggregation;
3030
import org.hibernate.search.backend.elasticsearch.types.dsl.ElasticsearchIndexFieldTypeFactory;
3131
import org.hibernate.search.backend.elasticsearch.types.dsl.provider.impl.ElasticsearchIndexFieldTypeFactoryProvider;
3232
import org.hibernate.search.backend.elasticsearch.types.impl.ElasticsearchIndexCompositeNodeType;
@@ -86,8 +86,7 @@ public ElasticsearchIndexRootBuilder(ElasticsearchIndexFieldTypeFactoryProvider
8686
this.customIndexMapping = customIndexMapping;
8787
this.defaultDynamicType = DynamicType.create( dynamicMapping );
8888

89-
this.typeBuilder.queryElementFactory( AggregationTypeKeys.COUNT_DOCUMENTS,
90-
ElasticsearchCountDocumentAggregation.factory( false ) );
89+
this.typeBuilder.queryElementFactory( AggregationTypeKeys.COUNT, ElasticsearchCountAggregation.factory() );
9190
this.addDefaultImplicitFields();
9291
}
9392

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

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright Red Hat Inc. and Hibernate Authors
4+
*/
5+
package org.hibernate.search.backend.elasticsearch.search.aggregation.impl;
6+
7+
import java.util.List;
8+
9+
import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
10+
import org.hibernate.search.backend.elasticsearch.gson.impl.JsonElementTypes;
11+
import org.hibernate.search.backend.elasticsearch.logging.impl.ElasticsearchClientLog;
12+
import org.hibernate.search.backend.elasticsearch.logging.impl.QueryLog;
13+
import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexNodeContext;
14+
import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexScope;
15+
import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexValueFieldContext;
16+
import org.hibernate.search.backend.elasticsearch.search.predicate.impl.ElasticsearchSearchPredicate;
17+
import org.hibernate.search.backend.elasticsearch.search.query.impl.ElasticsearchSearchQueryExtractContext;
18+
import org.hibernate.search.engine.search.aggregation.AggregationKey;
19+
import org.hibernate.search.engine.search.aggregation.spi.CountAggregationBuilder;
20+
import org.hibernate.search.engine.search.aggregation.spi.CountDocumentAggregationBuilder;
21+
import org.hibernate.search.engine.search.aggregation.spi.CountValuesAggregationBuilder;
22+
import org.hibernate.search.engine.search.common.spi.SearchQueryElementFactory;
23+
24+
import com.google.gson.JsonElement;
25+
import com.google.gson.JsonObject;
26+
27+
public abstract class ElasticsearchCountAggregation {
28+
private ElasticsearchCountAggregation() {
29+
}
30+
31+
public static SearchQueryElementFactory<CountAggregationBuilder.TypeSelector,
32+
ElasticsearchSearchIndexScope<?>,
33+
ElasticsearchSearchIndexNodeContext> factory() {
34+
return new Factory();
35+
}
36+
37+
private static class Factory
38+
implements SearchQueryElementFactory<CountAggregationBuilder.TypeSelector,
39+
ElasticsearchSearchIndexScope<?>,
40+
ElasticsearchSearchIndexNodeContext> {
41+
42+
@Override
43+
public CountAggregationBuilder.TypeSelector create(ElasticsearchSearchIndexScope<?> scope,
44+
ElasticsearchSearchIndexNodeContext node) {
45+
return new ElasticsearchCountAggregation.TypeSelector( scope, node );
46+
}
47+
48+
@Override
49+
public void checkCompatibleWith(SearchQueryElementFactory<?, ?, ?> other) {
50+
if ( !getClass().equals( other.getClass() ) ) {
51+
throw QueryLog.INSTANCE.differentImplementationClassForQueryElement( getClass(), other.getClass() );
52+
}
53+
}
54+
}
55+
56+
private static final class TypeSelector implements CountAggregationBuilder.TypeSelector {
57+
private final ElasticsearchSearchIndexScope<?> scope;
58+
private final ElasticsearchSearchIndexNodeContext node;
59+
60+
private TypeSelector(ElasticsearchSearchIndexScope<?> scope,
61+
ElasticsearchSearchIndexNodeContext node) {
62+
this.scope = scope;
63+
this.node = node; // doesn't matter in this case
64+
}
65+
66+
@Override
67+
public CountValuesAggregationBuilder values() {
68+
return new ElasticsearchCountValuesAggregation.Builder( scope, node.toValueField() );
69+
}
70+
71+
@Override
72+
public CountDocumentAggregationBuilder documents() {
73+
boolean isNested = node.isValueField() && !node.toValueField().nestedPathHierarchy().isEmpty();
74+
return new ElasticsearchCountDocumentAggregation.Builder( scope, isNested );
75+
}
76+
}
77+
78+
private static class ElasticsearchCountDocumentAggregation extends AbstractElasticsearchAggregation<Long> {
79+
80+
private static final JsonAccessor<Long> TOTAL_HITS_VALUE_PROPERTY_ACCESSOR =
81+
JsonAccessor.root().property( "hits" )
82+
.property( "total" )
83+
.property( "value" ).asLong();
84+
85+
private static final JsonAccessor<Long> RESPONSE_DOC_COUNT_ACCESSOR =
86+
JsonAccessor.root().property( "doc_count" ).asLong();
87+
private static final JsonAccessor<Long> RESPONSE_ROOT_DOC_COUNT_ACCESSOR =
88+
JsonAccessor.root().property( "root_doc_count" ).property( "doc_count" ).asLong();
89+
90+
private final boolean isNested;
91+
92+
private ElasticsearchCountDocumentAggregation(Builder builder) {
93+
super( builder );
94+
this.isNested = builder.isNested;
95+
}
96+
97+
@Override
98+
public Extractor<Long> request(AggregationRequestContext context, AggregationKey<?> key, JsonObject jsonAggregations) {
99+
return new CountDocumentsExtractor( key, isNested, context.isRootContext() );
100+
}
101+
102+
private record CountDocumentsExtractor(AggregationKey<?> key, boolean isNested, boolean rootContext)
103+
implements Extractor<Long> {
104+
105+
@Override
106+
public Long extract(JsonObject aggregationResult, AggregationExtractContext context) {
107+
if ( rootContext ) {
108+
if ( context instanceof ElasticsearchSearchQueryExtractContext c ) {
109+
return TOTAL_HITS_VALUE_PROPERTY_ACCESSOR.get( c.getResponseBody() )
110+
.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
111+
}
112+
}
113+
else {
114+
if ( isNested ) {
115+
// We must return the number of root documents,
116+
// not the number of leaf documents that Elasticsearch returns by default.
117+
return RESPONSE_ROOT_DOC_COUNT_ACCESSOR.get( aggregationResult )
118+
.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
119+
}
120+
else {
121+
return RESPONSE_DOC_COUNT_ACCESSOR.get( aggregationResult )
122+
.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
123+
}
124+
}
125+
throw ElasticsearchClientLog.INSTANCE.elasticsearchResponseMissingData();
126+
}
127+
}
128+
129+
private static class Builder extends AbstractBuilder<Long>
130+
implements CountDocumentAggregationBuilder {
131+
private final boolean isNested;
132+
133+
private Builder(ElasticsearchSearchIndexScope<?> scope, boolean isNested) {
134+
super( scope );
135+
this.isNested = isNested;
136+
}
137+
138+
@Override
139+
public ElasticsearchCountDocumentAggregation build() {
140+
return new ElasticsearchCountDocumentAggregation( this );
141+
}
142+
}
143+
}
144+
145+
private static class ElasticsearchCountValuesAggregation extends AbstractElasticsearchNestableAggregation<Long> {
146+
147+
private static final JsonAccessor<JsonObject> COUNT_PROPERTY_ACCESSOR =
148+
JsonAccessor.root().property( "value_count" ).asObject();
149+
150+
private static final JsonAccessor<JsonObject> COUNT_DISTINCT_PROPERTY_ACCESSOR =
151+
JsonAccessor.root().property( "cardinality" ).asObject();
152+
153+
private static final JsonAccessor<String> FIELD_PROPERTY_ACCESSOR =
154+
JsonAccessor.root().property( "field" ).asString();
155+
156+
private final String absoluteFieldPath;
157+
private final JsonAccessor<JsonObject> operation;
158+
159+
private ElasticsearchCountValuesAggregation(Builder builder) {
160+
super( builder );
161+
this.absoluteFieldPath = builder.field.absolutePath();
162+
this.operation = builder.operation;
163+
}
164+
165+
@Override
166+
protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
167+
JsonObject outerObject = new JsonObject();
168+
JsonObject innerObject = new JsonObject();
169+
170+
operation.set( outerObject, innerObject );
171+
FIELD_PROPERTY_ACCESSOR.set( innerObject, absoluteFieldPath );
172+
return outerObject;
173+
}
174+
175+
@Override
176+
protected Extractor<Long> extractor(AggregationKey<?> key, AggregationRequestBuildingContextContext context) {
177+
return new MetricLongExtractor( key, nestedPathHierarchy, filter );
178+
}
179+
180+
private static class MetricLongExtractor extends AbstractExtractor<Long> {
181+
protected MetricLongExtractor(AggregationKey<?> key, List<String> nestedPathHierarchy,
182+
ElasticsearchSearchPredicate filter) {
183+
super( key, nestedPathHierarchy, filter );
184+
}
185+
186+
@Override
187+
protected Long doExtract(JsonObject aggregationResult, AggregationExtractContext context) {
188+
JsonElement value = aggregationResult.get( "value" );
189+
return JsonElementTypes.LONG.fromElement( value );
190+
}
191+
}
192+
193+
private static class Builder extends AbstractElasticsearchNestableAggregation.AbstractBuilder<Long>
194+
implements CountValuesAggregationBuilder {
195+
private JsonAccessor<JsonObject> operation;
196+
197+
private Builder(ElasticsearchSearchIndexScope<?> scope, ElasticsearchSearchIndexValueFieldContext<?> field) {
198+
super( scope, field );
199+
this.operation = COUNT_PROPERTY_ACCESSOR;
200+
}
201+
202+
@Override
203+
public ElasticsearchCountValuesAggregation build() {
204+
return new ElasticsearchCountValuesAggregation( this );
205+
}
206+
207+
@Override
208+
public void distinct(boolean distinct) {
209+
if ( distinct ) {
210+
operation = COUNT_DISTINCT_PROPERTY_ACCESSOR;
211+
}
212+
else {
213+
operation = COUNT_PROPERTY_ACCESSOR;
214+
}
215+
}
216+
}
217+
}
218+
}

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

Lines changed: 0 additions & 124 deletions
This file was deleted.

0 commit comments

Comments
 (0)