DATAES-96 - Add support for aggregation in ElasticsearchTemplate

Author: Rizwan Idrees

src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java

@@ -471,4 +471,7 @@ public interface ElasticsearchOperations {
 	 * @return
 	 */
 	Set<String> queryForAlias(String indexName);
+
+
+    <T> T query(SearchQuery query, ResultsExtractor<T> resultsExtractor);
 }

src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java

@@ -64,6 +64,7 @@
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.facet.FacetBuilder;
 import org.elasticsearch.search.highlight.HighlightBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
@@ -247,7 +248,13 @@ public <T> FacetedPage<T> queryForPage(SearchQuery query, Class<T> clazz, Search
 		return mapper.mapResults(response, clazz, query.getPageable());
 	}
 
-	@Override
+    @Override
+    public <T> T query(SearchQuery query, ResultsExtractor<T> resultsExtractor) {
+        SearchResponse response = doSearch(prepareSearch(query), query);
+        return resultsExtractor.extract(response);
+    }
+
+    @Override
 	public <T> List<T> queryForList(CriteriaQuery query, Class<T> clazz) {
 		return queryForPage(query, clazz).getContent();
 	}
@@ -608,6 +615,12 @@ private SearchResponse doSearch(SearchRequestBuilder searchRequest, SearchQuery
 			}
 		}
 
+        if(CollectionUtils.isNotEmpty(searchQuery.getAggregations())){
+             for(AbstractAggregationBuilder aggregationBuilder : searchQuery.getAggregations()){
+                 searchRequest.addAggregation(aggregationBuilder);
+             }
+        }
+
 		return searchRequest.setQuery(searchQuery.getQuery()).execute().actionGet();
 	}
 

src/main/java/org/springframework/data/elasticsearch/core/ResultsExtractor.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2013-2014 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.data.elasticsearch.core;
+
+import org.elasticsearch.action.search.SearchResponse;
+
+public interface ResultsExtractor<T>  {
+    T extract(SearchResponse response);
+}

src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java

@@ -15,15 +15,16 @@
  */
 package org.springframework.data.elasticsearch.core.query;
 
-import java.util.ArrayList;
-import java.util.List;
-
 import org.elasticsearch.index.query.FilterBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.highlight.HighlightBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
 import org.springframework.data.elasticsearch.core.facet.FacetRequest;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * NativeSearchQuery
  *
@@ -37,6 +38,7 @@ public class NativeSearchQuery extends AbstractQuery implements SearchQuery {
 	private FilterBuilder filter;
 	private List<SortBuilder> sorts;
 	private List<FacetRequest> facets;
+    private List<AbstractAggregationBuilder> aggregations;
 	private HighlightBuilder.Field[] highlightFields;
 
 
@@ -94,4 +96,21 @@ public void setFacets(List<FacetRequest> facets) {
 	public List<FacetRequest> getFacets() {
 		return facets;
 	}
+
+    @Override
+    public List<AbstractAggregationBuilder> getAggregations() {
+        return aggregations;
+    }
+
+
+    public void addAggregation(AbstractAggregationBuilder aggregationBuilder) {
+        if (aggregations == null) {
+            aggregations = new ArrayList<AbstractAggregationBuilder>();
+        }
+        aggregations.add(aggregationBuilder);
+    }
+
+    public void setAggregations(List<AbstractAggregationBuilder> aggregations) {
+        this.aggregations = aggregations;
+    }
 }

src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.index.query.FilterBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.highlight.HighlightBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
 import org.springframework.data.domain.Pageable;
@@ -42,6 +43,7 @@ public class NativeSearchQueryBuilder {
 	private FilterBuilder filterBuilder;
 	private List<SortBuilder> sortBuilders = new ArrayList<SortBuilder>();
 	private List<FacetRequest> facetRequests = new ArrayList<FacetRequest>();
+    private List<AbstractAggregationBuilder> aggregationBuilders = new ArrayList<AbstractAggregationBuilder>();
 	private HighlightBuilder.Field[] highlightFields;
 	private Pageable pageable;
 	private String[] indices;
@@ -67,6 +69,11 @@ public NativeSearchQueryBuilder withSort(SortBuilder sortBuilder) {
 		return this;
 	}
 
+    public NativeSearchQueryBuilder addAggregation(AbstractAggregationBuilder aggregationBuilder){
+        this.aggregationBuilders.add(aggregationBuilder);
+        return this;
+    }
+
 	public NativeSearchQueryBuilder withFacet(FacetRequest facetRequest) {
 		facetRequests.add(facetRequest);
 		return this;
@@ -139,6 +146,10 @@ public NativeSearchQuery build() {
 			nativeSearchQuery.setFacets(facetRequests);
 		}
 
+        if (CollectionUtils.isNotEmpty(aggregationBuilders)) {
+            nativeSearchQuery.setAggregations(aggregationBuilders);
+        }
+
 		if (minScore > 0) {
 			nativeSearchQuery.setMinScore(minScore);
 		}

src/main/java/org/springframework/data/elasticsearch/core/query/SearchQuery.java

@@ -19,6 +19,7 @@
 
 import org.elasticsearch.index.query.FilterBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.highlight.HighlightBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
 import org.springframework.data.elasticsearch.core.facet.FacetRequest;
@@ -40,5 +41,7 @@ public interface SearchQuery extends Query {
 
 	List<FacetRequest> getFacets();
 
+    List<AbstractAggregationBuilder> getAggregations();
+
 	HighlightBuilder.Field[] getHighlightFields();
 }

src/test/java/org/springframework/data/elasticsearch/core/aggregation/ElasticsearchTemplateAggregationTests.java

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2013 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.elasticsearch.core.aggregation;
+
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.search.aggregations.Aggregations;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.data.elasticsearch.core.ElasticsearchTemplate;
+import org.springframework.data.elasticsearch.core.ResultsExtractor;
+import org.springframework.data.elasticsearch.core.facet.ArticleEntity;
+import org.springframework.data.elasticsearch.core.facet.ArticleEntityBuilder;
+import org.springframework.data.elasticsearch.core.query.IndexQuery;
+import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
+import org.springframework.data.elasticsearch.core.query.SearchQuery;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import static org.elasticsearch.action.search.SearchType.COUNT;
+import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
+import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.junit.Assert.assertThat;
+
+/**
+ * @author Rizwan Idrees
+ * @author Mohsin Husen
+ * @author Jonathan Yan
+ * @author Artur Konczak
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration("classpath:elasticsearch-template-test.xml")
+public class ElasticsearchTemplateAggregationTests {
+
+    public static final String RIZWAN_IDREES = "Rizwan Idrees";
+    public static final String MOHSIN_HUSEN = "Mohsin Husen";
+    public static final String JONATHAN_YAN = "Jonathan Yan";
+    public static final String ARTUR_KONCZAK = "Artur Konczak";
+    public static final int YEAR_2002 = 2002;
+    public static final int YEAR_2001 = 2001;
+    public static final int YEAR_2000 = 2000;
+    @Autowired
+    private ElasticsearchTemplate elasticsearchTemplate;
+
+    @Before
+    public void before() {
+        elasticsearchTemplate.deleteIndex(ArticleEntity.class);
+        elasticsearchTemplate.createIndex(ArticleEntity.class);
+        elasticsearchTemplate.putMapping(ArticleEntity.class);
+        elasticsearchTemplate.refresh(ArticleEntity.class, true);
+
+        IndexQuery article1 = new ArticleEntityBuilder("1").title("article four").subject("computing").addAuthor(RIZWAN_IDREES).addAuthor(ARTUR_KONCZAK).addAuthor(MOHSIN_HUSEN).addAuthor(JONATHAN_YAN).score(10).buildIndex();
+        IndexQuery article2 = new ArticleEntityBuilder("2").title("article three").subject("computing").addAuthor(RIZWAN_IDREES).addAuthor(ARTUR_KONCZAK).addAuthor(MOHSIN_HUSEN).addPublishedYear(YEAR_2000).score(20).buildIndex();
+        IndexQuery article3 = new ArticleEntityBuilder("3").title("article two").subject("computing").addAuthor(RIZWAN_IDREES).addAuthor(ARTUR_KONCZAK).addPublishedYear(YEAR_2001).addPublishedYear(YEAR_2000).score(30).buildIndex();
+        IndexQuery article4 = new ArticleEntityBuilder("4").title("article one").subject("accounting").addAuthor(RIZWAN_IDREES).addPublishedYear(YEAR_2002).addPublishedYear(YEAR_2001).addPublishedYear(YEAR_2000).score(40).buildIndex();
+
+        elasticsearchTemplate.index(article1);
+        elasticsearchTemplate.index(article2);
+        elasticsearchTemplate.index(article3);
+        elasticsearchTemplate.index(article4);
+        elasticsearchTemplate.refresh(ArticleEntity.class, true);
+    }
+
+    @Test
+    public void shouldReturnAggregatedResponseForGivenSearchQuery() {
+        // given
+        SearchQuery searchQuery = new NativeSearchQueryBuilder()
+                .withQuery(matchAllQuery())
+                .withSearchType(COUNT)
+                .withIndices("articles").withTypes("article")
+                .addAggregation(terms("subjects").field("subject"))
+                .build();
+        // when
+        Aggregations aggregations = elasticsearchTemplate.query(searchQuery, new ResultsExtractor<Aggregations>() {
+            @Override
+            public Aggregations extract(SearchResponse response) {
+                return response.getAggregations();
+            }
+        });
+        // then
+        assertThat(aggregations, is(notNullValue()));
+        assertThat(aggregations.asMap().get("subjects"), is(notNullValue()));
+    }
+
+}
+
+