Skip to content
This repository was archived by the owner on Jun 20, 2023. It is now read-only.

Commit fbe146e

Browse files
mkhludnevmkhludnev
authored
fix #626: Support composite aggregation. (#683)
1 parent debfca8 commit fbe146e

File tree

5 files changed

+259
-1
lines changed

5 files changed

+259
-1
lines changed

‎jest-common/src/main/java/io/searchbox/core/search/aggregation/AggregationField.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@ public enum AggregationField {
3131
AVG("avg"),
3232
SUM("sum"),
3333
DOC_COUNT_ERROR_UPPER_BOUND("doc_count_error_upper_bound"),
34-
SUM_OTHER_DOC_COUNT("sum_other_doc_count");
34+
SUM_OTHER_DOC_COUNT("sum_other_doc_count"),
35+
AFTER_KEY("after_key");
3536

3637
private final String field;
3738

‎jest-common/src/main/java/io/searchbox/core/search/aggregation/CompositeAggregation.java

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package io.searchbox.core.search.aggregation;
2+
3+
import com.google.gson.JsonArray;
4+
import com.google.gson.JsonElement;
5+
import com.google.gson.JsonObject;
6+
7+
import java.util.List;
8+
import java.util.Objects;
9+
import java.util.stream.Collectors;
10+
import java.util.stream.StreamSupport;
11+
12+
import static io.searchbox.core.search.aggregation.AggregationField.*;
13+
14+
public class CompositeAggregation extends BucketAggregation{
15+
16+
private List<Entry> buckets;
17+
private JsonObject afterKey;
18+
19+
public CompositeAggregation(String name, JsonObject agg) {
20+
super(name, agg);
21+
if (agg.has(AFTER_KEY.toString()) && agg.get(AFTER_KEY.toString()).isJsonObject()) {
22+
afterKey = agg.get(AFTER_KEY.toString()).getAsJsonObject();
23+
}
24+
if (agg.has(BUCKETS.toString()) && agg.get(BUCKETS.toString()).isJsonArray()) {
25+
final JsonArray bucketsArray = agg.get(BUCKETS.toString()).getAsJsonArray();
26+
buckets = StreamSupport.stream(bucketsArray.spliterator(), false)
27+
.map(JsonElement::getAsJsonObject)
28+
.map(bucket -> new Entry(bucket,
29+
bucket.get(KEY.toString()).getAsJsonObject(),
30+
bucket.get(DOC_COUNT.toString()).getAsLong()))
31+
.collect(Collectors.toList());
32+
}
33+
}
34+
35+
@Override
36+
public List<Entry> getBuckets() {
37+
return buckets;
38+
}
39+
40+
public JsonObject getAfterKey() {
41+
return afterKey;
42+
}
43+
44+
public class Entry extends Bucket {
45+
private JsonObject key;
46+
47+
public Entry(JsonObject bucketRoot, JsonObject key, Long count) {
48+
super(bucketRoot, count);
49+
this.key = key;
50+
}
51+
52+
public JsonObject getKey() {
53+
return key;
54+
}
55+
56+
@Override
57+
public boolean equals(Object o) {
58+
if (this == o) return true;
59+
if (o == null || getClass() != o.getClass()) return false;
60+
if (!super.equals(o)) return false;
61+
Entry entry = (Entry) o;
62+
return Objects.equals(key, entry.key);
63+
}
64+
65+
@Override
66+
public int hashCode() {
67+
return Objects.hash(super.hashCode(), key);
68+
}
69+
}
70+
}

‎jest-common/src/main/java/io/searchbox/core/search/aggregation/MetricAggregation.java

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,14 @@ public CardinalityAggregation getCardinalityAggregation(String aggName) {
7373
return jsonRoot.has(aggName) ? new CardinalityAggregation(aggName, jsonRoot.getAsJsonObject(aggName)) : null;
7474
}
7575

76+
/**
77+
* @param aggName Name of the CompositeAggregation
78+
* @return a new CompositeAggregation object if aggName is found within sub-aggregations of current aggregation level or null if not found
79+
*/
80+
public CompositeAggregation getCompositeAggregation(String aggName) {
81+
return jsonRoot.has(aggName) ? new CompositeAggregation(aggName, jsonRoot.getAsJsonObject(aggName)) : null;
82+
}
83+
7684
/**
7785
* @param aggName Name of the DateHistogramAggregation
7886
* @return a new DateHistogramAggregation object if aggName is found within sub-aggregations of current aggregation level or null if not found

‎jest-common/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationTest.java

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
package io.searchbox.core.search.aggregation;
2+
3+
import com.google.gson.Gson;
4+
import com.google.gson.JsonObject;
5+
6+
import java.util.Arrays;
7+
import java.util.List;
8+
9+
import org.junit.Test;
10+
import static org.junit.Assert.*;
11+
12+
public class CompositeAggregationTest {
13+
14+
final String afterKey = "'after_key':{'field':'valB'}";
15+
final String bucketsArr = "'buckets':[" +
16+
"{ 'key':{'field':'val'}," +
17+
" 'doc_count':1," +
18+
" 'sub_agg':{}" +
19+
"}," +
20+
"{ 'key':{'field':'valA'}," +
21+
" 'doc_count':2," +
22+
" 'sub_agg':{}" +
23+
"}" + "]";
24+
final Gson gson = new Gson();
25+
26+
@Test
27+
public void testParseBuckets() {
28+
for (String input : Arrays.asList("{" + afterKey + ", " + bucketsArr + "}", "{" + bucketsArr + "}")) {
29+
JsonObject compositeAggregationJson = gson.fromJson(
30+
input.replace('\'', '\"'), JsonObject.class);
31+
CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
32+
List<CompositeAggregation.Entry> buckets = compositeAggregation.getBuckets();
33+
assertNotNull(buckets);
34+
assertEquals(2, buckets.size());
35+
36+
assertEquals("val", buckets.get(0).getKey().get("field").getAsString());
37+
assertEquals(Long.valueOf(1L), buckets.get(0).getCount());
38+
assertTrue(buckets.get(0).getTermsAggregation("sub_agg").getBuckets().isEmpty());
39+
40+
assertEquals("valA", buckets.get(1).getKey().get("field").getAsString());
41+
assertEquals(Long.valueOf(2L), buckets.get(1).getCount());
42+
assertTrue(buckets.get(1).getTermsAggregation("sub_agg").getBuckets().isEmpty());
43+
}
44+
}
45+
46+
@Test
47+
public void testParseAfterKey() {
48+
for (String input : Arrays.asList("{" + afterKey + ", " + bucketsArr + "}", "{" + afterKey + "}")) {
49+
JsonObject compositeAggregationJson = gson.fromJson(
50+
input.replace('\'', '\"'), JsonObject.class);
51+
CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
52+
assertEquals("valB", compositeAggregation.getAfterKey().get("field").getAsString());
53+
}
54+
}
55+
56+
@Test
57+
public void testNoAfterKey() {
58+
for (String input : Arrays.asList("{" + bucketsArr + "}", "{" + "}")) {
59+
JsonObject compositeAggregationJson = gson.fromJson(
60+
input.replace('\'', '\"'), JsonObject.class);
61+
CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
62+
assertNull(compositeAggregation.getAfterKey());
63+
}
64+
}
65+
}

‎jest/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationIntegrationTest.java

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
package io.searchbox.core.search.aggregation;
2+
3+
import com.google.gson.Gson;
4+
import com.google.gson.JsonObject;
5+
import io.searchbox.common.AbstractIntegrationTest;
6+
import io.searchbox.core.Search;
7+
import io.searchbox.core.SearchResult;
8+
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
9+
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingResponse;
10+
import org.elasticsearch.common.xcontent.XContentType;
11+
import org.elasticsearch.test.ESIntegTestCase;
12+
import org.junit.Test;
13+
14+
import java.io.IOException;
15+
import java.util.HashMap;
16+
import java.util.List;
17+
import java.util.Map;
18+
19+
@ESIntegTestCase.ClusterScope (scope = ESIntegTestCase.Scope.TEST, numDataNodes = 1)
20+
public class CompositeAggregationIntegrationTest extends AbstractIntegrationTest {
21+
22+
public static final Gson gson = new Gson();
23+
private final String INDEX = "composite_aggregation";
24+
private final String TYPE = "document";
25+
26+
@Test
27+
public void testGetCompositeAggregation()
28+
throws IOException {
29+
createIndex(INDEX);
30+
PutMappingResponse putMappingResponse = client().admin().indices().putMapping(new PutMappingRequest(INDEX)
31+
.type(TYPE)
32+
.source("{\"document\":{\"properties\":{\"gender\":{\"store\":true,\"type\":\"keyword\"}}}}", XContentType.JSON)
33+
).actionGet();
34+
35+
assertTrue(putMappingResponse.isAcknowledged());
36+
37+
index(INDEX, TYPE, null, "{\"gender\":\"male\"}");
38+
index(INDEX, TYPE, null, "{\"gender\":\"male\"}");
39+
index(INDEX, TYPE, null, "{\"gender\":\"female\"}");
40+
refresh();
41+
ensureSearchable(INDEX);
42+
final JsonObject afterKey;
43+
{
44+
String query = "{\n" +
45+
" \"query\" : {\n" +
46+
" \"match_all\" : {}\n" +
47+
" },\n" +
48+
" \"aggs\" : {\n" +
49+
" \"composite1\" : {\n" +
50+
" \"composite\" : {\n" +
51+
" \"size\" : 10,\n" +
52+
" \"sources\":[{\"term1\": { \"terms\":{\"field\" : \"gender\"}}}]\n" +
53+
" }\n" +
54+
" }\n" +
55+
" }\n" +
56+
"}";
57+
Search search = new Search.Builder(query)
58+
.addIndex(INDEX)
59+
.addType(TYPE)
60+
.build();
61+
SearchResult result = client.execute(search);
62+
assertTrue(result.getErrorMessage(), result.isSucceeded());
63+
64+
CompositeAggregation composite = result.getAggregations().getCompositeAggregation("composite1");
65+
assertEquals("composite1", composite.getName());
66+
assertEquals(2, composite.getBuckets().size());
67+
assertTrue(2L == composite.getBuckets().get(1).getCount());
68+
assertEquals(gson.fromJson("{\"term1\":\"male\"}", JsonObject.class), composite.getBuckets().get(1).getKey());
69+
assertTrue(1L == composite.getBuckets().get(0).getCount());
70+
assertEquals(gson.fromJson("{\"term1\":\"female\"}", JsonObject.class), composite.getBuckets().get(0).getKey());
71+
72+
Aggregation aggregation = result.getAggregations().getAggregation("composite1", CompositeAggregation.class);
73+
assertTrue(aggregation instanceof CompositeAggregation);
74+
CompositeAggregation compositeByType = (CompositeAggregation) aggregation;
75+
assertEquals(composite, compositeByType);
76+
77+
Map<String, Class> nameToTypeMap = new HashMap<String, Class>();
78+
nameToTypeMap.put("composite1", CompositeAggregation.class);
79+
List<Aggregation> aggregations = result.getAggregations().getAggregations(nameToTypeMap);
80+
assertEquals(1, aggregations.size());
81+
assertTrue(aggregations.get(0) instanceof CompositeAggregation);
82+
CompositeAggregation termsWithMap = (CompositeAggregation) aggregations.get(0);
83+
assertEquals(termsWithMap, compositeByType);
84+
85+
afterKey = composite.getAfterKey();
86+
assertNotNull("ain't sure why", afterKey);
87+
}
88+
89+
String query2 = "{\n" +
90+
" \"query\" : {\n" +
91+
" \"match_all\" : {}\n" +
92+
" },\n" +
93+
" \"aggs\" : {\n" +
94+
" \"composite1\" : {\n" +
95+
" \"composite\" : {\n" +
96+
" \"after\" : "+ gson.toJson(afterKey)+",\n" +
97+
" \"sources\":[{\"term1\": { \"terms\":{\"field\" : \"gender\"}}}]\n" +
98+
" }\n" +
99+
" }\n" +
100+
" }\n" +
101+
"}";
102+
Search search2 = new Search.Builder(query2)
103+
.addIndex(INDEX)
104+
.addType(TYPE)
105+
.build();
106+
SearchResult result2 = client.execute(search2);
107+
assertTrue(result2.getErrorMessage(), result2.isSucceeded());
108+
109+
CompositeAggregation composite2 = result2.getAggregations().getCompositeAggregation("composite1");
110+
assertEquals("composite1", composite2.getName());
111+
assertEquals(0, composite2.getBuckets().size());
112+
assertNull(composite2.getAfterKey());
113+
}
114+
}

0 commit comments

Comments
 (0)