Skip to content

Commit 916fd41

Browse files
marko-bekhtamarko-bekhta
committed
HSEARCH-3661 Make Elasticsearch's terms aggregation accept "value"
1 parent 876fbf3 commit 916fd41

File tree

6 files changed

+122
-52
lines changed

6 files changed

+122
-52
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,8 @@ public ElasticsearchIndexRootBuilder(ElasticsearchIndexFieldTypeFactoryProvider
8686
this.customIndexMapping = customIndexMapping;
8787
this.defaultDynamicType = DynamicType.create( dynamicMapping );
8888

89-
this.typeBuilder.queryElementFactory( AggregationTypeKeys.COUNT_DOCUMENTS, ElasticsearchCountDocumentAggregation.factory( false ) );
89+
this.typeBuilder.queryElementFactory( AggregationTypeKeys.COUNT_DOCUMENTS,
90+
ElasticsearchCountDocumentAggregation.factory( false ) );
9091
this.addDefaultImplicitFields();
9192
}
9293

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

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import java.util.Map;
99

1010
import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
11-
import org.hibernate.search.backend.elasticsearch.logging.impl.ElasticsearchClientLog;
1211
import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexScope;
1312
import org.hibernate.search.backend.elasticsearch.search.common.impl.ElasticsearchSearchIndexValueFieldContext;
1413
import org.hibernate.search.backend.elasticsearch.search.predicate.impl.ElasticsearchSearchPredicate;
@@ -29,10 +28,6 @@ public abstract class AbstractElasticsearchBucketAggregation<K, V>
2928
private static final String ROOT_DOC_COUNT_NAME = "root_doc_count";
3029
private static final JsonAccessor<JsonObject> REQUEST_AGGREGATIONS_ROOT_DOC_COUNT_ACCESSOR =
3130
JsonAccessor.root().property( "aggregations" ).property( ROOT_DOC_COUNT_NAME ).asObject();
32-
private static final JsonAccessor<Long> RESPONSE_DOC_COUNT_ACCESSOR =
33-
JsonAccessor.root().property( "doc_count" ).asLong();
34-
private static final JsonAccessor<Long> RESPONSE_ROOT_DOC_COUNT_ACCESSOR =
35-
JsonAccessor.root().property( ROOT_DOC_COUNT_NAME ).property( "doc_count" ).asLong();
3631

3732
AbstractElasticsearchBucketAggregation(AbstractBuilder<K, V> builder) {
3833
super( builder );
@@ -58,19 +53,6 @@ protected final JsonObject doRequest(AggregationRequestContext context) {
5853

5954
protected abstract void doRequest(JsonObject outerObject, JsonObject innerObject, AggregationRequestContext context);
6055

61-
protected final long getBucketDocCount(JsonObject bucket) {
62-
if ( isNested() ) {
63-
// We must return the number of root documents,
64-
// not the number of leaf documents that Elasticsearch returns by default.
65-
return RESPONSE_ROOT_DOC_COUNT_ACCESSOR.get( bucket )
66-
.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
67-
}
68-
else {
69-
return RESPONSE_DOC_COUNT_ACCESSOR.get( bucket )
70-
.orElseThrow( ElasticsearchClientLog.INSTANCE::elasticsearchResponseMissingData );
71-
}
72-
}
73-
7456
protected abstract class AbstractBucketExtractor<A, B> extends AbstractExtractor<Map<A, B>> {
7557

7658
protected AbstractBucketExtractor(List<String> nestedPathHierarchy,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ private ElasticsearchCountDocumentAggregation(Builder builder) {
4343

4444
@Override
4545
public Extractor<Long> request(AggregationRequestContext context, AggregationKey<?> key, JsonObject jsonAggregations) {
46-
return new CountDocumentsExtractor(isNested);
46+
return new CountDocumentsExtractor( isNested );
4747
}
4848

4949
private record CountDocumentsExtractor(boolean isNested) implements Extractor<Long> {

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public class ElasticsearchRangeAggregation<F, K, V>
4444

4545
private final ElasticsearchSearchAggregation<V> aggregation;
4646

47+
// TODO: do not store these two here:
4748
private Extractor<V> innerExtractor;
4849
private AggregationKey<V> innerExtractorKey;
4950

@@ -66,6 +67,7 @@ protected void doRequest(JsonObject outerObject, JsonObject innerObject, Aggrega
6667
// this is just a "random name" so we can get the aggregation back from the response.
6768
// once we switch to the "composite aggregation" where we compute multiple aggregations for a range,
6869
// this should be moved into a new "aggregation" that would handle all the logic for adding and then extracting 0-n aggregations.
70+
// TODO: not really good that we have state saved into aggregation within the request, we should pass it up instead
6971
innerExtractorKey = AggregationKey.of( "agg" );
7072
innerExtractor = aggregation.request( context, innerExtractorKey, subOuterObject );
7173
if ( !subOuterObject.isEmpty() ) {
@@ -103,7 +105,8 @@ private TypeSelector(ElasticsearchSearchIndexScope<?> scope,
103105

104106
@Override
105107
public <T> Builder<F, T, Long> type(Class<T> expectedType, ValueModel valueModel) {
106-
return new CountBuilder<>( scope, field, field.encodingContext().encoder( scope, field, expectedType, valueModel ) );
108+
return new CountBuilder<>( scope, field,
109+
field.encodingContext().encoder( scope, field, expectedType, valueModel ) );
107110
}
108111
}
109112

@@ -125,22 +128,23 @@ protected Map<Range<K>, V> doExtract(AggregationExtractContext context, JsonElem
125128
JsonObject bucket = bucketMap.get( String.valueOf( i ) ).getAsJsonObject();
126129
Range<K> range = rangesInOrder.get( i );
127130
if ( bucket.has( innerExtractorKey.name() ) ) {
128-
bucket = bucket.getAsJsonObject( innerExtractorKey.name() );
131+
bucket = bucket.getAsJsonObject( innerExtractorKey.name() );
129132
}
130133
result.put( range, innerExtractor.extract( bucket, context ) );
131134
}
132135
return result;
133136
}
134137
}
135138

136-
public static class CountBuilder<F, K> extends Builder<F, K, Long> {
139+
private static class CountBuilder<F, K> extends Builder<F, K, Long> {
137140

138141
protected CountBuilder(ElasticsearchSearchIndexScope<?> scope,
139142
ElasticsearchSearchIndexValueFieldContext<?> field,
140143
Function<? super K, JsonElement> encoder) {
141144
super( scope, field, encoder, new ArrayList<>(), new JsonArray(),
142145
ElasticsearchSearchAggregation.from( scope,
143-
ElasticsearchCountDocumentAggregation.factory(field.nestedPathHierarchy().isEmpty()).create( scope, null ).type().build() ) );
146+
ElasticsearchCountDocumentAggregation.factory( !field.nestedPathHierarchy().isEmpty() )
147+
.create( scope, null ).type().build() ) );
144148
}
145149
}
146150

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

Lines changed: 82 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
import org.hibernate.search.backend.elasticsearch.types.codec.impl.ElasticsearchFieldCodec;
1616
import org.hibernate.search.engine.backend.types.converter.runtime.FromDocumentValueConvertContext;
1717
import org.hibernate.search.engine.backend.types.converter.spi.ProjectionConverter;
18+
import org.hibernate.search.engine.search.aggregation.AggregationKey;
19+
import org.hibernate.search.engine.search.aggregation.SearchAggregation;
1820
import org.hibernate.search.engine.search.aggregation.spi.TermsAggregationBuilder;
1921
import org.hibernate.search.engine.search.common.ValueModel;
2022
import org.hibernate.search.util.common.impl.CollectionHelper;
@@ -28,26 +30,33 @@
2830
* @param <K> The type of keys in the returned map. It can be {@code F}
2931
* or a different type if value converters are used.
3032
*/
31-
public class ElasticsearchTermsAggregation<F, K, T>
32-
extends AbstractElasticsearchBucketAggregation<K, Long> {
33+
public class ElasticsearchTermsAggregation<F, K, T, V>
34+
extends AbstractElasticsearchBucketAggregation<K, V> {
3335

3436
private final String absoluteFieldPath;
3537

3638
private final ProjectionConverter<T, ? extends K> fromFieldValueConverter;
3739
private final BiFunction<JsonElement, JsonElement, T> decodeFunction;
3840

41+
private final ElasticsearchSearchAggregation<V> aggregation;
42+
3943
private final JsonObject order;
4044
private final int size;
4145
private final int minDocCount;
4246

43-
private ElasticsearchTermsAggregation(Builder<F, K, T> builder) {
47+
// TODO: do not store these two here:
48+
private Extractor<V> innerExtractor;
49+
private AggregationKey<V> innerExtractorKey;
50+
51+
private ElasticsearchTermsAggregation(Builder<F, K, T, V> builder) {
4452
super( builder );
4553
this.absoluteFieldPath = builder.field.absolutePath();
4654
this.fromFieldValueConverter = builder.fromFieldValueConverter;
4755
this.decodeFunction = builder.decodeFunction;
4856
this.order = builder.order;
4957
this.size = builder.size;
5058
this.minDocCount = builder.minDocCount;
59+
this.aggregation = builder.aggregation;
5160
}
5261

5362
@Override
@@ -59,11 +68,19 @@ protected void doRequest(JsonObject outerObject, JsonObject innerObject, Aggrega
5968
}
6069
innerObject.addProperty( "size", size );
6170
innerObject.addProperty( "min_doc_count", minDocCount );
71+
72+
// TODO: not really good that we have state saved into aggregation within the request, we should pass it up instead
73+
JsonObject subOuterObject = new JsonObject();
74+
innerExtractorKey = AggregationKey.of( "agg" );
75+
innerExtractor = aggregation.request( context, innerExtractorKey, subOuterObject );
76+
if ( !subOuterObject.isEmpty() ) {
77+
outerObject.add( "aggs", subOuterObject );
78+
}
6279
}
6380

6481
@Override
65-
protected Extractor<Map<K, Long>> extractor(AggregationRequestContext context) {
66-
return new TermsBucketExtractor( nestedPathHierarchy, filter );
82+
protected Extractor<Map<K, V>> extractor(AggregationRequestContext context) {
83+
return new TermsBucketExtractor( nestedPathHierarchy, filter, innerExtractorKey, innerExtractor );
6784
}
6885

6986
public static class Factory<F>
@@ -93,34 +110,40 @@ private TypeSelector(ElasticsearchFieldCodec<F> codec,
93110

94111
@SuppressWarnings("unchecked")
95112
@Override
96-
public <T> Builder<F, T, ?> type(Class<T> expectedType, ValueModel valueModel) {
113+
public <T> Builder<F, T, ?, Long> type(Class<T> expectedType, ValueModel valueModel) {
97114
if ( ValueModel.RAW.equals( valueModel ) ) {
98-
return new Builder<>(
99-
(key, string) -> string != null && !string.isJsonNull() ? string : key,
115+
return new CountBuilder<>(
100116
scope, field,
117+
(key, string) -> string != null && !string.isJsonNull() ? string : key,
101118
// unchecked cast to make eclipse-compiler happy
102119
// we know that Elasticsearch projection converters work with the String
103120
( (ProjectionConverter<JsonElement, ?>) field.type().rawProjectionConverter() )
104121
.withConvertedType( expectedType, field )
105122
);
106123
}
107124
else {
108-
return new Builder<>( codec::decodeAggregationKey, scope, field,
125+
return new CountBuilder<>( scope, field, codec::decodeAggregationKey,
109126
field.type().projectionConverter( valueModel ).withConvertedType( expectedType, field ) );
110127
}
111128
}
112129
}
113130

114-
protected class TermsBucketExtractor extends AbstractBucketExtractor<K, Long> {
131+
protected class TermsBucketExtractor extends AbstractBucketExtractor<K, V> {
132+
private final AggregationKey<V> innerExtractorKey;
133+
private final Extractor<V> innerExtractor;
134+
115135
protected TermsBucketExtractor(List<String> nestedPathHierarchy,
116-
ElasticsearchSearchPredicate filter) {
136+
ElasticsearchSearchPredicate filter, AggregationKey<V> innerExtractorKey, Extractor<V> innerExtractor
137+
) {
117138
super( nestedPathHierarchy, filter );
139+
this.innerExtractorKey = innerExtractorKey;
140+
this.innerExtractor = innerExtractor;
118141
}
119142

120143
@Override
121-
protected Map<K, Long> doExtract(AggregationExtractContext context, JsonElement buckets) {
144+
protected Map<K, V> doExtract(AggregationExtractContext context, JsonElement buckets) {
122145
JsonArray bucketArray = buckets.getAsJsonArray();
123-
Map<K, Long> result = CollectionHelper.newLinkedHashMap( bucketArray.size() );
146+
Map<K, V> result = CollectionHelper.newLinkedHashMap( bucketArray.size() );
124147
FromDocumentValueConvertContext convertContext = context.fromDocumentValueConvertContext();
125148
for ( JsonElement bucketElement : bucketArray ) {
126149
JsonObject bucket = bucketElement.getAsJsonObject();
@@ -130,29 +153,60 @@ protected Map<K, Long> doExtract(AggregationExtractContext context, JsonElement
130153
decodeFunction.apply( keyJson, keyAsStringJson ),
131154
convertContext
132155
);
133-
long documentCount = getBucketDocCount( bucket );
134-
result.put( key, documentCount );
156+
157+
if ( bucket.has( innerExtractorKey.name() ) ) {
158+
bucket = bucket.getAsJsonObject( innerExtractorKey.name() );
159+
}
160+
result.put( key, innerExtractor.extract( bucket, context ) );
135161
}
136162
return result;
137163
}
138164
}
139165

140-
private static class Builder<F, K, T> extends AbstractBuilder<K, Long>
141-
implements TermsAggregationBuilder<K> {
166+
private static class CountBuilder<F, K, T> extends Builder<F, K, T, Long> {
167+
168+
protected CountBuilder(ElasticsearchSearchIndexScope<?> scope,
169+
ElasticsearchSearchIndexValueFieldContext<F> field,
170+
BiFunction<JsonElement, JsonElement, T> decodeFunction,
171+
ProjectionConverter<T, ? extends K> fromFieldValueConverter) {
172+
super( scope, field, decodeFunction, fromFieldValueConverter,
173+
ElasticsearchSearchAggregation.from( scope,
174+
ElasticsearchCountDocumentAggregation.factory( !field.nestedPathHierarchy().isEmpty() )
175+
.create( scope, null ).type().build() ) );
176+
}
177+
}
178+
179+
private static class Builder<F, K, T, V> extends AbstractBuilder<K, V>
180+
implements TermsAggregationBuilder<K, V> {
142181

143182
private final BiFunction<JsonElement, JsonElement, T> decodeFunction;
144183
private final ProjectionConverter<T, ? extends K> fromFieldValueConverter;
184+
private final ElasticsearchSearchAggregation<V> aggregation;
145185

146186
private JsonObject order;
147-
private int minDocCount = 1;
148-
private int size = 100;
187+
private int minDocCount;
188+
private int size;
149189

150-
private Builder(BiFunction<JsonElement, JsonElement, T> decodeFunction, ElasticsearchSearchIndexScope<?> scope,
190+
private Builder(ElasticsearchSearchIndexScope<?> scope,
151191
ElasticsearchSearchIndexValueFieldContext<F> field,
152-
ProjectionConverter<T, ? extends K> fromFieldValueConverter) {
192+
BiFunction<JsonElement, JsonElement, T> decodeFunction,
193+
ProjectionConverter<T, ? extends K> fromFieldValueConverter,
194+
ElasticsearchSearchAggregation<V> aggregation) {
195+
this( scope, field, decodeFunction, fromFieldValueConverter, aggregation, null, 1, 100 );
196+
}
197+
198+
public Builder(ElasticsearchSearchIndexScope<?> scope, ElasticsearchSearchIndexValueFieldContext<?> field,
199+
BiFunction<JsonElement, JsonElement, T> decodeFunction,
200+
ProjectionConverter<T, ? extends K> fromFieldValueConverter,
201+
ElasticsearchSearchAggregation<V> aggregation,
202+
JsonObject order, int minDocCount, int size) {
153203
super( scope, field );
204+
this.order = order;
154205
this.decodeFunction = decodeFunction;
155206
this.fromFieldValueConverter = fromFieldValueConverter;
207+
this.aggregation = aggregation;
208+
this.minDocCount = minDocCount;
209+
this.size = size;
156210
}
157211

158212
@Override
@@ -186,7 +240,13 @@ public void maxTermCount(int maxTermCount) {
186240
}
187241

188242
@Override
189-
public ElasticsearchTermsAggregation<F, K, T> build() {
243+
public <R> TermsAggregationBuilder<K, R> withValue(SearchAggregation<R> aggregation) {
244+
return new Builder<>( scope, field, decodeFunction, fromFieldValueConverter,
245+
ElasticsearchSearchAggregation.from( scope, aggregation ), order, minDocCount, size );
246+
}
247+
248+
@Override
249+
public ElasticsearchTermsAggregation<F, K, T, V> build() {
190250
return new ElasticsearchTermsAggregation<>( this );
191251
}
192252

‎documentation/src/test/java/org/hibernate/search/documentation/search/aggregation/AggregationDslIT.java

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -285,9 +285,10 @@ void terms_value() {
285285
// end::terms-sum[]
286286
assertThat( sumByPrice )
287287
.containsExactly(
288-
entry( 10.0, 7.99 ),
289-
entry( 20.0 , 35.98 ),
290-
entry( null , 24.99 )
288+
entry( 7.99, 7.99 ),
289+
entry( 15.99, 15.99 ),
290+
entry( 19.99, 19.99 ),
291+
entry( 24.99, 24.99 )
291292
);
292293
} );
293294

@@ -306,9 +307,31 @@ void terms_value() {
306307
// end::terms-count[]
307308
assertThat( countsByPrice )
308309
.containsExactly(
309-
entry( Range.canonical( 0.0, 10.0 ), 1L ),
310-
entry( Range.canonical( 10.0, 20.0 ), 2L ),
311-
entry( Range.canonical( 20.0, null ), 1L )
310+
entry( 7.99, 1L ),
311+
entry( 15.99, 1L ),
312+
entry( 19.99, 1L ),
313+
entry( 24.99, 1L )
314+
);
315+
} );
316+
317+
withinSearchSession( searchSession -> {
318+
// tag::terms-count-implicit[]
319+
AggregationKey<Map<Double, Long>> countsByPriceKey = AggregationKey.of( "countsByPrice" );
320+
SearchResult<Book> result = searchSession.search( Book.class )
321+
.where( f -> f.matchAll() )
322+
.aggregation(
323+
countsByPriceKey, f -> f.terms()
324+
.field( "price", Double.class ) // <1>
325+
)
326+
.fetch( 20 );
327+
Map<Double, Long> countsByPrice = result.aggregation( countsByPriceKey );
328+
// end::terms-count-implicit[]
329+
assertThat( countsByPrice )
330+
.containsExactly(
331+
entry( 7.99, 1L ),
332+
entry( 15.99, 1L ),
333+
entry( 19.99, 1L ),
334+
entry( 24.99, 1L )
312335
);
313336
} );
314337
}

0 commit comments

Comments
 (0)