Skip to content

Commit 409642c

Browse files
marko-bekhtamarko-bekhta
committed
HSEARCH-3661 Adjust how extractors are built for bucket Elasticsearch aggregations
so that the state is not "stored" in the aggregation but passed around through a context.
1 parent 7043f3c commit 409642c

File tree

7 files changed

+118
-34
lines changed

7 files changed

+118
-34
lines changed

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,15 @@ public abstract class AbstractElasticsearchBucketAggregation<K, V>
2929
private static final JsonAccessor<JsonObject> REQUEST_AGGREGATIONS_ROOT_DOC_COUNT_ACCESSOR =
3030
JsonAccessor.root().property( "aggregations" ).property( ROOT_DOC_COUNT_NAME ).asObject();
3131

32+
protected static final String INNER_EXTRACTOR_KEY = "innerExtractorKey";
33+
protected static final String INNER_EXTRACTOR = "innerExtractor";
34+
3235
AbstractElasticsearchBucketAggregation(AbstractBuilder<K, V> builder) {
3336
super( builder );
3437
}
3538

3639
@Override
37-
protected final JsonObject doRequest(AggregationRequestContext context) {
40+
protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
3841
JsonObject outerObject = new JsonObject();
3942
JsonObject innerObject = new JsonObject();
4043

@@ -51,7 +54,8 @@ protected final JsonObject doRequest(AggregationRequestContext context) {
5154
return outerObject;
5255
}
5356

54-
protected abstract void doRequest(JsonObject outerObject, JsonObject innerObject, AggregationRequestContext context);
57+
protected abstract void doRequest(JsonObject outerObject, JsonObject innerObject,
58+
AggregationRequestBuildingContextContext context);
5559

5660
protected abstract class AbstractBucketExtractor<A, B> extends AbstractExtractor<Map<A, B>> {
5761

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

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,11 +50,12 @@ public abstract class AbstractElasticsearchNestableAggregation<A> extends Abstra
5050

5151
@Override
5252
public final Extractor<A> request(AggregationRequestContext context, AggregationKey<?> key, JsonObject jsonAggregations) {
53-
jsonAggregations.add( key.name(), request( context ) );
54-
return extractor( context );
53+
AggregationRequestBuildingContextContext buildingContext = new AggregationRequestBuildingContextContext( context );
54+
jsonAggregations.add( key.name(), request( buildingContext ) );
55+
return extractor( buildingContext );
5556
}
5657

57-
private JsonObject request(AggregationRequestContext context) {
58+
private JsonObject request(AggregationRequestBuildingContextContext context) {
5859
JsonObject result = doRequest( context );
5960

6061
if ( nestedPathHierarchy.isEmpty() ) {
@@ -90,9 +91,9 @@ private JsonObject request(AggregationRequestContext context) {
9091
return result;
9192
}
9293

93-
protected abstract JsonObject doRequest(AggregationRequestContext context);
94+
protected abstract JsonObject doRequest(AggregationRequestBuildingContextContext context);
9495

95-
protected abstract Extractor<A> extractor(AggregationRequestContext context);
96+
protected abstract Extractor<A> extractor(AggregationRequestBuildingContextContext context);
9697

9798
protected abstract static class AbstractExtractor<T> implements Extractor<T> {
9899

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
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.HashMap;
8+
import java.util.Map;
9+
import java.util.Objects;
10+
11+
import org.hibernate.search.backend.elasticsearch.search.predicate.impl.PredicateRequestContext;
12+
import org.hibernate.search.util.common.annotation.Incubating;
13+
14+
/**
15+
* Sometimes we need to pass something we created while building up the json in one of the "doRequest" methods
16+
* in the aggregation build up to the "later" steps e.g. to when we create the extractor.
17+
*/
18+
@Incubating
19+
public final class AggregationRequestBuildingContextContext implements AggregationRequestContext {
20+
private final AggregationRequestContext aggregationRequestContext;
21+
private final Map<Key<?>, Object> buildingContext = new HashMap<>();
22+
23+
public AggregationRequestBuildingContextContext(AggregationRequestContext aggregationRequestContext) {
24+
this.aggregationRequestContext = aggregationRequestContext;
25+
}
26+
27+
public <T> T get(Key<T> key) {
28+
Object value = buildingContext.get( key );
29+
return key.cast( value );
30+
}
31+
32+
public void add(Key<?> key, Object value) {
33+
buildingContext.put( key, value );
34+
}
35+
36+
public AggregationRequestContext rootAggregationRequestContext() {
37+
return aggregationRequestContext;
38+
}
39+
40+
@Override
41+
public PredicateRequestContext getRootPredicateContext() {
42+
return aggregationRequestContext.getRootPredicateContext();
43+
}
44+
45+
public static <V> Key<V> buildingContextKey(String name) {
46+
return new Key<>( name );
47+
}
48+
49+
public static class Key<V> {
50+
51+
private final String name;
52+
53+
private Key(String name) {
54+
this.name = name;
55+
}
56+
57+
@SuppressWarnings("unchecked")
58+
private V cast(Object value) {
59+
return (V) value;
60+
}
61+
62+
@Override
63+
public boolean equals(Object o) {
64+
if ( o == null || getClass() != o.getClass() ) {
65+
return false;
66+
}
67+
Key<?> key = (Key<?>) o;
68+
return Objects.equals( name, key.name );
69+
}
70+
71+
@Override
72+
public int hashCode() {
73+
return Objects.hashCode( name );
74+
}
75+
}
76+
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ private ElasticsearchMetricFieldAggregation(Builder<F, K> builder) {
7474
}
7575

7676
@Override
77-
protected final JsonObject doRequest(AggregationRequestContext context) {
77+
protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
7878
JsonObject outerObject = new JsonObject();
7979
JsonObject innerObject = new JsonObject();
8080

@@ -84,7 +84,7 @@ protected final JsonObject doRequest(AggregationRequestContext context) {
8484
}
8585

8686
@Override
87-
protected Extractor<K> extractor(AggregationRequestContext context) {
87+
protected Extractor<K> extractor(AggregationRequestBuildingContextContext context) {
8888
return metricFieldExtractorCreator.extractor( filter );
8989
}
9090

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ private ElasticsearchMetricLongAggregation(Builder builder) {
4747
}
4848

4949
@Override
50-
protected final JsonObject doRequest(AggregationRequestContext context) {
50+
protected final JsonObject doRequest(AggregationRequestBuildingContextContext context) {
5151
JsonObject outerObject = new JsonObject();
5252
JsonObject innerObject = new JsonObject();
5353

@@ -57,7 +57,7 @@ protected final JsonObject doRequest(AggregationRequestContext context) {
5757
}
5858

5959
@Override
60-
protected Extractor<Long> extractor(AggregationRequestContext context) {
60+
protected Extractor<Long> extractor(AggregationRequestBuildingContextContext context) {
6161
return new MetricLongExtractor( nestedPathHierarchy, filter );
6262
}
6363

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

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
*/
55
package org.hibernate.search.backend.elasticsearch.search.aggregation.impl;
66

7+
import static org.hibernate.search.backend.elasticsearch.search.aggregation.impl.AggregationRequestBuildingContextContext.buildingContextKey;
8+
79
import java.util.ArrayList;
810
import java.util.List;
911
import java.util.Map;
@@ -44,10 +46,6 @@ public class ElasticsearchRangeAggregation<F, K, V>
4446

4547
private final ElasticsearchSearchAggregation<V> aggregation;
4648

47-
// TODO: do not store these two here:
48-
private Extractor<V> innerExtractor;
49-
private AggregationKey<V> innerExtractorKey;
50-
5149
private ElasticsearchRangeAggregation(Builder<F, K, V> builder) {
5250
super( builder );
5351
this.absoluteFieldPath = builder.field.absolutePath();
@@ -57,27 +55,27 @@ private ElasticsearchRangeAggregation(Builder<F, K, V> builder) {
5755
}
5856

5957
@Override
60-
protected void doRequest(JsonObject outerObject, JsonObject innerObject, AggregationRequestContext context) {
58+
protected void doRequest(JsonObject outerObject, JsonObject innerObject, AggregationRequestBuildingContextContext context) {
6159
outerObject.add( "range", innerObject );
6260
innerObject.addProperty( "field", absoluteFieldPath );
6361
innerObject.addProperty( "keyed", true );
6462
innerObject.add( "ranges", rangesJson );
6563

6664
JsonObject subOuterObject = new JsonObject();
67-
// this is just a "random name" so we can get the aggregation back from the response.
68-
// once we switch to the "composite aggregation" where we compute multiple aggregations for a range,
69-
// 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
71-
innerExtractorKey = AggregationKey.of( "agg" );
72-
innerExtractor = aggregation.request( context, innerExtractorKey, subOuterObject );
65+
AggregationKey<V> innerExtractorKey = AggregationKey.of( "agg" );
66+
context.add( buildingContextKey( INNER_EXTRACTOR_KEY ), innerExtractorKey );
67+
context.add( buildingContextKey( INNER_EXTRACTOR ), aggregation.request( context, innerExtractorKey, subOuterObject ) );
68+
7369
if ( !subOuterObject.isEmpty() ) {
7470
outerObject.add( "aggs", subOuterObject );
7571
}
7672
}
7773

7874
@Override
79-
protected Extractor<Map<Range<K>, V>> extractor(AggregationRequestContext context) {
80-
return new RangeBucketExtractor( nestedPathHierarchy, filter, rangesInOrder );
75+
protected Extractor<Map<Range<K>, V>> extractor(AggregationRequestBuildingContextContext context) {
76+
AggregationKey<V> innerExtractorKey = context.get( buildingContextKey( INNER_EXTRACTOR_KEY ) );
77+
Extractor<V> innerExtractor = context.get( buildingContextKey( INNER_EXTRACTOR ) );
78+
return new RangeBucketExtractor( nestedPathHierarchy, filter, rangesInOrder, innerExtractorKey, innerExtractor );
8179
}
8280

8381
public static class Factory<F>
@@ -112,11 +110,15 @@ public <T> Builder<F, T, Long> type(Class<T> expectedType, ValueModel valueModel
112110

113111
protected class RangeBucketExtractor extends AbstractBucketExtractor<Range<K>, V> {
114112
private final List<Range<K>> rangesInOrder;
113+
private final Extractor<V> innerExtractor;
114+
private final AggregationKey<V> innerExtractorKey;
115115

116116
protected RangeBucketExtractor(List<String> nestedPathHierarchy, ElasticsearchSearchPredicate filter,
117-
List<Range<K>> rangesInOrder) {
117+
List<Range<K>> rangesInOrder, AggregationKey<V> innerExtractorKey, Extractor<V> innerExtractor) {
118118
super( nestedPathHierarchy, filter );
119119
this.rangesInOrder = rangesInOrder;
120+
this.innerExtractorKey = innerExtractorKey;
121+
this.innerExtractor = innerExtractor;
120122
}
121123

122124

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

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
*/
55
package org.hibernate.search.backend.elasticsearch.search.aggregation.impl;
66

7+
import static org.hibernate.search.backend.elasticsearch.search.aggregation.impl.AggregationRequestBuildingContextContext.buildingContextKey;
8+
79
import java.util.List;
810
import java.util.Map;
911
import java.util.function.BiFunction;
@@ -44,10 +46,6 @@ public class ElasticsearchTermsAggregation<F, K, T, V>
4446
private final int size;
4547
private final int minDocCount;
4648

47-
// TODO: do not store these two here:
48-
private Extractor<V> innerExtractor;
49-
private AggregationKey<V> innerExtractorKey;
50-
5149
private ElasticsearchTermsAggregation(Builder<F, K, T, V> builder) {
5250
super( builder );
5351
this.absoluteFieldPath = builder.field.absolutePath();
@@ -60,7 +58,7 @@ private ElasticsearchTermsAggregation(Builder<F, K, T, V> builder) {
6058
}
6159

6260
@Override
63-
protected void doRequest(JsonObject outerObject, JsonObject innerObject, AggregationRequestContext context) {
61+
protected void doRequest(JsonObject outerObject, JsonObject innerObject, AggregationRequestBuildingContextContext context) {
6462
outerObject.add( "terms", innerObject );
6563
innerObject.addProperty( "field", absoluteFieldPath );
6664
if ( order != null ) {
@@ -69,17 +67,20 @@ protected void doRequest(JsonObject outerObject, JsonObject innerObject, Aggrega
6967
innerObject.addProperty( "size", size );
7068
innerObject.addProperty( "min_doc_count", minDocCount );
7169

72-
// TODO: not really good that we have state saved into aggregation within the request, we should pass it up instead
7370
JsonObject subOuterObject = new JsonObject();
74-
innerExtractorKey = AggregationKey.of( "agg" );
75-
innerExtractor = aggregation.request( context, innerExtractorKey, subOuterObject );
71+
AggregationKey<V> innerExtractorKey = AggregationKey.of( "agg" );
72+
context.add( buildingContextKey( INNER_EXTRACTOR_KEY ), innerExtractorKey );
73+
context.add( buildingContextKey( INNER_EXTRACTOR ), aggregation.request( context, innerExtractorKey, subOuterObject ) );
74+
7675
if ( !subOuterObject.isEmpty() ) {
7776
outerObject.add( "aggs", subOuterObject );
7877
}
7978
}
8079

8180
@Override
82-
protected Extractor<Map<K, V>> extractor(AggregationRequestContext context) {
81+
protected Extractor<Map<K, V>> extractor(AggregationRequestBuildingContextContext context) {
82+
AggregationKey<V> innerExtractorKey = context.get( buildingContextKey( INNER_EXTRACTOR_KEY ) );
83+
Extractor<V> innerExtractor = context.get( buildingContextKey( INNER_EXTRACTOR ) );
8384
return new TermsBucketExtractor( nestedPathHierarchy, filter, innerExtractorKey, innerExtractor );
8485
}
8586

0 commit comments

Comments
 (0)