fix searchbox-io#626: Support composite aggregation. (searchbox-io#683)

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

@@ -31,7 +31,8 @@ public enum AggregationField {
     AVG("avg"),
     SUM("sum"),
     DOC_COUNT_ERROR_UPPER_BOUND("doc_count_error_upper_bound"),
-    SUM_OTHER_DOC_COUNT("sum_other_doc_count");
+    SUM_OTHER_DOC_COUNT("sum_other_doc_count"),
+    AFTER_KEY("after_key");
 
     private final String field;
 

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

@@ -0,0 +1,70 @@
+package io.searchbox.core.search.aggregation;
+
+import com.google.gson.JsonArray;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+import static io.searchbox.core.search.aggregation.AggregationField.*;
+
+public class CompositeAggregation extends BucketAggregation{
+
+    private List<Entry> buckets;
+    private JsonObject afterKey;
+
+    public CompositeAggregation(String name, JsonObject agg) {
+        super(name, agg);
+        if (agg.has(AFTER_KEY.toString()) && agg.get(AFTER_KEY.toString()).isJsonObject()) {
+            afterKey = agg.get(AFTER_KEY.toString()).getAsJsonObject();
+        }
+        if (agg.has(BUCKETS.toString()) && agg.get(BUCKETS.toString()).isJsonArray()) {
+            final JsonArray bucketsArray = agg.get(BUCKETS.toString()).getAsJsonArray();
+            buckets = StreamSupport.stream(bucketsArray.spliterator(), false)
+                    .map(JsonElement::getAsJsonObject)
+                    .map(bucket ->   new Entry(bucket,
+                                        bucket.get(KEY.toString()).getAsJsonObject(),
+                                        bucket.get(DOC_COUNT.toString()).getAsLong()))
+                    .collect(Collectors.toList());
+        }
+    }
+
+    @Override
+    public List<Entry> getBuckets() {
+        return buckets;
+    }
+
+    public JsonObject getAfterKey() {
+        return afterKey;
+    }
+
+    public class Entry extends Bucket {
+        private JsonObject key;
+
+        public Entry(JsonObject bucketRoot,  JsonObject key, Long count) {
+            super(bucketRoot, count);
+            this.key = key;
+        }
+
+        public JsonObject getKey() {
+            return key;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            if (!super.equals(o)) return false;
+            Entry entry = (Entry) o;
+            return Objects.equals(key, entry.key);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(super.hashCode(), key);
+        }
+    }
+}

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

@@ -73,6 +73,14 @@ public CardinalityAggregation getCardinalityAggregation(String aggName) {
         return jsonRoot.has(aggName) ? new CardinalityAggregation(aggName, jsonRoot.getAsJsonObject(aggName)) : null;
     }
 
+    /**
+     * @param aggName Name of the CompositeAggregation
+     * @return a new CompositeAggregation object if aggName is found within sub-aggregations of current aggregation level or null if not found
+     */
+    public CompositeAggregation getCompositeAggregation(String aggName) {
+        return jsonRoot.has(aggName) ? new CompositeAggregation(aggName, jsonRoot.getAsJsonObject(aggName)) : null;
+    }
+
     /**
      * @param aggName Name of the DateHistogramAggregation
      * @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

@@ -0,0 +1,65 @@
+package io.searchbox.core.search.aggregation;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonObject;
+
+import java.util.Arrays;
+import java.util.List;
+
+import org.junit.Test;
+import static org.junit.Assert.*;
+
+public class CompositeAggregationTest {
+
+    final String afterKey = "'after_key':{'field':'valB'}";
+    final String bucketsArr = "'buckets':[" +
+            "{  'key':{'field':'val'}," +
+            "   'doc_count':1," +
+            "   'sub_agg':{}" +
+            "}," +
+            "{  'key':{'field':'valA'}," +
+            "   'doc_count':2," +
+            "   'sub_agg':{}" +
+            "}" + "]";
+    final Gson gson = new Gson();
+
+    @Test
+    public void testParseBuckets() {
+        for (String input : Arrays.asList("{" + afterKey + ", " + bucketsArr + "}", "{" + bucketsArr + "}")) {
+            JsonObject compositeAggregationJson = gson.fromJson(
+                    input.replace('\'', '\"'), JsonObject.class);
+            CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
+            List<CompositeAggregation.Entry> buckets = compositeAggregation.getBuckets();
+            assertNotNull(buckets);
+            assertEquals(2, buckets.size());
+
+            assertEquals("val", buckets.get(0).getKey().get("field").getAsString());
+            assertEquals(Long.valueOf(1L), buckets.get(0).getCount());
+            assertTrue(buckets.get(0).getTermsAggregation("sub_agg").getBuckets().isEmpty());
+
+            assertEquals("valA", buckets.get(1).getKey().get("field").getAsString());
+            assertEquals(Long.valueOf(2L), buckets.get(1).getCount());
+            assertTrue(buckets.get(1).getTermsAggregation("sub_agg").getBuckets().isEmpty());
+        }
+    }
+
+    @Test
+    public void testParseAfterKey() {
+        for (String input : Arrays.asList("{" + afterKey + ", " + bucketsArr + "}", "{" + afterKey + "}")) {
+            JsonObject compositeAggregationJson = gson.fromJson(
+                    input.replace('\'', '\"'), JsonObject.class);
+            CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
+            assertEquals("valB", compositeAggregation.getAfterKey().get("field").getAsString());
+        }
+    }
+
+    @Test
+    public void testNoAfterKey() {
+        for (String input : Arrays.asList("{"  + bucketsArr + "}", "{" + "}")) {
+            JsonObject compositeAggregationJson = gson.fromJson(
+                    input.replace('\'', '\"'), JsonObject.class);
+            CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
+            assertNull(compositeAggregation.getAfterKey());
+        }
+    }
+}

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

@@ -0,0 +1,114 @@
+package io.searchbox.core.search.aggregation;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonObject;
+import io.searchbox.common.AbstractIntegrationTest;
+import io.searchbox.core.Search;
+import io.searchbox.core.SearchResult;
+import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
+import org.elasticsearch.action.admin.indices.mapping.put.PutMappingResponse;
+import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.test.ESIntegTestCase;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@ESIntegTestCase.ClusterScope (scope = ESIntegTestCase.Scope.TEST, numDataNodes = 1)
+public class CompositeAggregationIntegrationTest extends AbstractIntegrationTest {
+
+    public static final Gson gson = new Gson();
+    private final String INDEX = "composite_aggregation";
+    private final String TYPE = "document";
+
+    @Test
+    public void testGetCompositeAggregation()
+            throws IOException {
+        createIndex(INDEX);
+        PutMappingResponse putMappingResponse = client().admin().indices().putMapping(new PutMappingRequest(INDEX)
+                .type(TYPE)
+                .source("{\"document\":{\"properties\":{\"gender\":{\"store\":true,\"type\":\"keyword\"}}}}", XContentType.JSON)
+        ).actionGet();
+
+        assertTrue(putMappingResponse.isAcknowledged());
+
+        index(INDEX, TYPE, null, "{\"gender\":\"male\"}");
+        index(INDEX, TYPE, null, "{\"gender\":\"male\"}");
+        index(INDEX, TYPE, null, "{\"gender\":\"female\"}");
+        refresh();
+        ensureSearchable(INDEX);
+        final JsonObject afterKey;
+        {
+            String query = "{\n" +
+                    "    \"query\" : {\n" +
+                    "        \"match_all\" : {}\n" +
+                    "    },\n" +
+                    "    \"aggs\" : {\n" +
+                    "        \"composite1\" : {\n" +
+                    "            \"composite\" : {\n" +
+                    "            \"size\" : 10,\n" +
+                    "                \"sources\":[{\"term1\": { \"terms\":{\"field\" : \"gender\"}}}]\n" +
+                    "            }\n" +
+                    "        }\n" +
+                    "    }\n" +
+                    "}";
+            Search search = new Search.Builder(query)
+                    .addIndex(INDEX)
+                    .addType(TYPE)
+                    .build();
+            SearchResult result = client.execute(search);
+            assertTrue(result.getErrorMessage(), result.isSucceeded());
+
+            CompositeAggregation composite = result.getAggregations().getCompositeAggregation("composite1");
+            assertEquals("composite1", composite.getName());
+            assertEquals(2, composite.getBuckets().size());
+            assertTrue(2L == composite.getBuckets().get(1).getCount());
+            assertEquals(gson.fromJson("{\"term1\":\"male\"}", JsonObject.class), composite.getBuckets().get(1).getKey());
+            assertTrue(1L == composite.getBuckets().get(0).getCount());
+            assertEquals(gson.fromJson("{\"term1\":\"female\"}", JsonObject.class), composite.getBuckets().get(0).getKey());
+
+            Aggregation aggregation = result.getAggregations().getAggregation("composite1", CompositeAggregation.class);
+            assertTrue(aggregation instanceof CompositeAggregation);
+            CompositeAggregation compositeByType = (CompositeAggregation) aggregation;
+            assertEquals(composite, compositeByType);
+
+            Map<String, Class> nameToTypeMap = new HashMap<String, Class>();
+            nameToTypeMap.put("composite1", CompositeAggregation.class);
+            List<Aggregation> aggregations = result.getAggregations().getAggregations(nameToTypeMap);
+            assertEquals(1, aggregations.size());
+            assertTrue(aggregations.get(0) instanceof CompositeAggregation);
+            CompositeAggregation termsWithMap = (CompositeAggregation) aggregations.get(0);
+            assertEquals(termsWithMap, compositeByType);
+
+            afterKey = composite.getAfterKey();
+            assertNotNull("ain't sure why", afterKey);
+        }
+
+        String query2 = "{\n" +
+                "    \"query\" : {\n" +
+                "        \"match_all\" : {}\n" +
+                "    },\n" +
+                "    \"aggs\" : {\n" +
+                "        \"composite1\" : {\n" +
+                "            \"composite\" : {\n" +
+                "            \"after\" : "+ gson.toJson(afterKey)+",\n" +
+                "                \"sources\":[{\"term1\": { \"terms\":{\"field\" : \"gender\"}}}]\n" +
+                "            }\n" +
+                "        }\n" +
+                "    }\n" +
+                "}";
+        Search search2 = new Search.Builder(query2)
+                .addIndex(INDEX)
+                .addType(TYPE)
+                .build();
+        SearchResult result2 = client.execute(search2);
+        assertTrue(result2.getErrorMessage(), result2.isSucceeded());
+
+        CompositeAggregation composite2 = result2.getAggregations().getCompositeAggregation("composite1");
+        assertEquals("composite1", composite2.getName());
+        assertEquals(0, composite2.getBuckets().size());
+        assertNull(composite2.getAfterKey());
+    }
+}