DATAES-96 - Add support for aggregation in ElasticsearchTemplate

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.FilterBuilder;
 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 class ElasticsearchTemplate implements ElasticsearchOperations {
 		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 @@ public class ElasticsearchTemplate implements ElasticsearchOperations {
 			}
 		}
 
+        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 class NativeSearchQuery extends AbstractQuery implements SearchQuery {
 	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.apache.commons.collections.CollectionUtils;
 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 class NativeSearchQueryBuilder {
 		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 class NativeSearchQueryBuilder {
 			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 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;
@@ -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()));
+    }
+
+}
+
+

src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntity.java

@@ -1,106 +1,114 @@
-/*
+/*
- * Copyright 2014 the original author or authors.
+ * Copyright 2014 the original author or authors.
- *
+ *
- * Licensed under the Apache License, Version 2.0 (the "License");
+ * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
+ * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
+ * You may obtain a copy of the License at
- *
+ *
- *      http://www.apache.org/licenses/LICENSE-2.0
+ *      http://www.apache.org/licenses/LICENSE-2.0
- *
+ *
- * Unless required by applicable law or agreed to in writing, software
+ * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
+ * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
+ * See the License for the specific language governing permissions and
- * limitations under the License.
+ * limitations under the License.
- */
+ */
-package org.springframework.data.elasticsearch.core.facet;
+package org.springframework.data.elasticsearch.core.facet;
-
+
-import static org.springframework.data.elasticsearch.annotations.FieldIndex.*;
+import static org.springframework.data.elasticsearch.annotations.FieldIndex.*;
-import static org.springframework.data.elasticsearch.annotations.FieldType.Integer;
+import static org.springframework.data.elasticsearch.annotations.FieldType.Integer;
-import static org.springframework.data.elasticsearch.annotations.FieldType.String;
+import static org.springframework.data.elasticsearch.annotations.FieldType.String;
-
+
-import java.util.ArrayList;
+import java.util.ArrayList;
-import java.util.List;
+import java.util.List;
-
+
-import org.springframework.data.annotation.Id;
+import org.springframework.data.annotation.Id;
-import org.springframework.data.elasticsearch.annotations.Document;
+import org.springframework.data.elasticsearch.annotations.Document;
-import org.springframework.data.elasticsearch.annotations.Field;
+import org.springframework.data.elasticsearch.annotations.Field;
-import org.springframework.data.elasticsearch.annotations.MultiField;
+import org.springframework.data.elasticsearch.annotations.MultiField;
-import org.springframework.data.elasticsearch.annotations.NestedField;
+import org.springframework.data.elasticsearch.annotations.NestedField;
-
+
-/**
+/**
- * Simple type to test facets
+ * Simple type to test facets
- *
+ *
- * @author Artur Konczak
+ * @author Artur Konczak
- * @author Mohsin Husen
+ * @author Mohsin Husen
- */
+ */
-@Document(indexName = "articles", type = "article", shards = 1, replicas = 0, refreshInterval = "-1", indexStoreType = "memory")
+@Document(indexName = "articles", type = "article", shards = 1, replicas = 0, refreshInterval = "-1", indexStoreType = "memory")
-public class ArticleEntity {
+public class ArticleEntity {
-
+
-	@Id
+	@Id
-	private String id;
+	private String id;
-
+	private String title;
-	private String title;
+    private String subject;
-
+
-	@MultiField(
+	@MultiField(
-			mainField = @Field(type = String, index = analyzed),
+			mainField = @Field(type = String, index = analyzed),
-			otherFields = {
+			otherFields = {
-					@NestedField(dotSuffix = "untouched", type = String, store = true, index = not_analyzed),
+					@NestedField(dotSuffix = "untouched", type = String, store = true, index = not_analyzed),
-					@NestedField(dotSuffix = "sort", type = String, store = true, indexAnalyzer = "keyword")
+					@NestedField(dotSuffix = "sort", type = String, store = true, indexAnalyzer = "keyword")
-			}
+			}
-	)
+	)
-	private List<String> authors = new ArrayList<String>();
+	private List<String> authors = new ArrayList<String>();
-
+
-	@Field(type = Integer, store = true)
+	@Field(type = Integer, store = true)
-	private List<Integer> publishedYears = new ArrayList<Integer>();
+	private List<Integer> publishedYears = new ArrayList<Integer>();
-
+
-	private int score;
+	private int score;
-
+
-	private ArticleEntity() {
+	private ArticleEntity() {
-
+
-	}
+	}
-
+
-	public ArticleEntity(String id) {
+	public ArticleEntity(String id) {
-		this.id = id;
+		this.id = id;
-	}
+	}
-
+
-	public void setId(String id) {
+	public void setId(String id) {
-		this.id = id;
+		this.id = id;
-	}
+	}
-
+
-	public String getId() {
+	public String getId() {
-		return id;
+		return id;
-	}
+	}
-
+
-	public String getTitle() {
+	public String getTitle() {
-		return title;
+		return title;
-	}
+	}
-
+
-	public void setTitle(String title) {
+	public void setTitle(String title) {
-		this.title = title;
+		this.title = title;
-	}
+	}
-
+
-	public List<String> getAuthors() {
+    public String getSubject() {
-		return authors;
+        return subject;
-	}
+    }
-
+
-	public void setAuthors(List<String> authors) {
+    public void setSubject(String subject) {
-		this.authors = authors;
+        this.subject = subject;
-	}
+    }
-
+
-	public List<Integer> getPublishedYears() {
+    public List<String> getAuthors() {
-		return publishedYears;
+		return authors;
-	}
+	}
-
+
-	public void setPublishedYears(List<Integer> publishedYears) {
+	public void setAuthors(List<String> authors) {
-		this.publishedYears = publishedYears;
+		this.authors = authors;
-	}
+	}
-
+
-	public int getScore() {
+	public List<Integer> getPublishedYears() {
-		return score;
+		return publishedYears;
-	}
+	}
-
+
-	public void setScore(int score) {
+	public void setPublishedYears(List<Integer> publishedYears) {
-		this.score = score;
+		this.publishedYears = publishedYears;
-	}
+	}
-}
+
+	public int getScore() {
+		return score;
+	}
+
+	public void setScore(int score) {
+		this.score = score;
+	}
+}

src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntityBuilder.java

@@ -1,64 +1,68 @@
-/*
+/*
- * Copyright 2014 the original author or authors.
+ * Copyright 2014 the original author or authors.
- *
+ *
- * Licensed under the Apache License, Version 2.0 (the "License");
+ * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
+ * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
+ * You may obtain a copy of the License at
- *
+ *
- *      http://www.apache.org/licenses/LICENSE-2.0
+ *      http://www.apache.org/licenses/LICENSE-2.0
- *
+ *
- * Unless required by applicable law or agreed to in writing, software
+ * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
+ * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
+ * See the License for the specific language governing permissions and
- * limitations under the License.
+ * limitations under the License.
- */
+ */
-package org.springframework.data.elasticsearch.core.facet;
+package org.springframework.data.elasticsearch.core.facet;
-
+
-import org.springframework.data.elasticsearch.core.query.IndexQuery;
+import org.springframework.data.elasticsearch.core.query.IndexQuery;
-
+
-/**
+/**
- * Simple type to test facets
+ * Simple type to test facets
- *
+ *
- * @author Artur Konczak
+ * @author Artur Konczak
- * @author Mohsin Husen
+ * @author Mohsin Husen
- */
+ */
-public class ArticleEntityBuilder {
+public class ArticleEntityBuilder {
-
+
-	private ArticleEntity result;
+	private ArticleEntity result;
-
+
-	public ArticleEntityBuilder(String id) {
+	public ArticleEntityBuilder(String id) {
-		result = new ArticleEntity(id);
+		result = new ArticleEntity(id);
-	}
+	}
-
+
-	public ArticleEntityBuilder title(String title) {
+	public ArticleEntityBuilder title(String title) {
-		result.setTitle(title);
+		result.setTitle(title);
-		return this;
+		return this;
-	}
+	}
-
+    public ArticleEntityBuilder subject(String subject) {
-	public ArticleEntityBuilder addAuthor(String author) {
+        result.setSubject(subject);
-		result.getAuthors().add(author);
+        return this;
-		return this;
+    }
-	}
+
-
+	public ArticleEntityBuilder addAuthor(String author) {
-	public ArticleEntityBuilder addPublishedYear(Integer year) {
+		result.getAuthors().add(author);
-		result.getPublishedYears().add(year);
+		return this;
-		return this;
+	}
-	}
+
-
+	public ArticleEntityBuilder addPublishedYear(Integer year) {
-	public ArticleEntityBuilder score(int score) {
+		result.getPublishedYears().add(year);
-		result.setScore(score);
+		return this;
-		return this;
+	}
-	}
+
-
+	public ArticleEntityBuilder score(int score) {
-	public ArticleEntity build() {
+		result.setScore(score);
-		return result;
+		return this;
-	}
+	}
-
+
-	public IndexQuery buildIndex() {
+	public ArticleEntity build() {
-		IndexQuery indexQuery = new IndexQuery();
+		return result;
-		indexQuery.setId(result.getId());
+	}
-		indexQuery.setObject(result);
+
-		return indexQuery;
+	public IndexQuery buildIndex() {
-	}
+		IndexQuery indexQuery = new IndexQuery();
-}
+		indexQuery.setId(result.getId());
+		indexQuery.setObject(result);
+		return indexQuery;
+	}
+}