From ab7e870d5f82f6c0de236048bd7001e8e7d2a680 Mon Sep 17 00:00:00 2001 From: Rizwan Idrees Date: Thu, 12 Jun 2014 17:25:00 +0100 Subject: [PATCH] DATAES-96 - Add support for aggregation in ElasticsearchTemplate --- .../core/ElasticsearchOperations.java | 3 + .../core/ElasticsearchTemplate.java | 15 +- .../elasticsearch/core/ResultsExtractor.java | 23 ++ .../core/query/NativeSearchQuery.java | 25 +- .../core/query/NativeSearchQueryBuilder.java | 11 + .../elasticsearch/core/query/SearchQuery.java | 3 + ...ElasticsearchTemplateAggregationTests.java | 103 ++++++++ .../core/facet/ArticleEntity.java | 220 +++++++++--------- .../core/facet/ArticleEntityBuilder.java | 132 ++++++----- 9 files changed, 361 insertions(+), 174 deletions(-) create mode 100644 src/main/java/org/springframework/data/elasticsearch/core/ResultsExtractor.java create mode 100644 src/test/java/org/springframework/data/elasticsearch/core/aggregation/ElasticsearchTemplateAggregationTests.java diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java index 1b202c898..dc12b21e7 100755 --- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java +++ b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java @@ -471,4 +471,7 @@ public interface ElasticsearchOperations { * @return */ Set queryForAlias(String indexName); + + + T query(SearchQuery query, ResultsExtractor resultsExtractor); } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java index b61efe4df..251bd5cb4 100755 --- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java +++ b/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 query(SearchQuery query, ResultsExtractor resultsExtractor) { + SearchResponse response = doSearch(prepareSearch(query), query); + return resultsExtractor.extract(response); + } + + @Override public List queryForList(CriteriaQuery query, Class 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(); } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ResultsExtractor.java b/src/main/java/org/springframework/data/elasticsearch/core/ResultsExtractor.java new file mode 100644 index 000000000..90ae17ada --- /dev/null +++ b/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 extract(SearchResponse response); +} diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java index 55b2687b7..64d737391 100644 --- a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java +++ b/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 sorts; private List facets; + private List aggregations; private HighlightBuilder.Field[] highlightFields; @@ -94,4 +96,21 @@ public class NativeSearchQuery extends AbstractQuery implements SearchQuery { public List getFacets() { return facets; } + + @Override + public List getAggregations() { + return aggregations; + } + + + public void addAggregation(AbstractAggregationBuilder aggregationBuilder) { + if (aggregations == null) { + aggregations = new ArrayList(); + } + aggregations.add(aggregationBuilder); + } + + public void setAggregations(List aggregations) { + this.aggregations = aggregations; + } } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java index 00b0b7e9b..730dcbcd5 100755 --- a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java +++ b/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 sortBuilders = new ArrayList(); private List facetRequests = new ArrayList(); + private List aggregationBuilders = new ArrayList(); 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); } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/SearchQuery.java b/src/main/java/org/springframework/data/elasticsearch/core/query/SearchQuery.java index e96b99ed3..081ad7615 100644 --- a/src/main/java/org/springframework/data/elasticsearch/core/query/SearchQuery.java +++ b/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 getFacets(); + List getAggregations(); + HighlightBuilder.Field[] getHighlightFields(); } diff --git a/src/test/java/org/springframework/data/elasticsearch/core/aggregation/ElasticsearchTemplateAggregationTests.java b/src/test/java/org/springframework/data/elasticsearch/core/aggregation/ElasticsearchTemplateAggregationTests.java new file mode 100644 index 000000000..7fec7b9a8 --- /dev/null +++ b/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() { + @Override + public Aggregations extract(SearchResponse response) { + return response.getAggregations(); + } + }); + // then + assertThat(aggregations, is(notNullValue())); + assertThat(aggregations.asMap().get("subjects"), is(notNullValue())); + } + +} + + diff --git a/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntity.java b/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntity.java index 075a95a3d..8ed9d5bf0 100644 --- a/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntity.java +++ b/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntity.java @@ -1,106 +1,114 @@ -/* - * Copyright 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.facet; - -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.String; - -import java.util.ArrayList; -import java.util.List; - -import org.springframework.data.annotation.Id; -import org.springframework.data.elasticsearch.annotations.Document; -import org.springframework.data.elasticsearch.annotations.Field; -import org.springframework.data.elasticsearch.annotations.MultiField; -import org.springframework.data.elasticsearch.annotations.NestedField; - -/** - * Simple type to test facets - * - * @author Artur Konczak - * @author Mohsin Husen - */ -@Document(indexName = "articles", type = "article", shards = 1, replicas = 0, refreshInterval = "-1", indexStoreType = "memory") -public class ArticleEntity { - - @Id - private String id; - - private String title; - - @MultiField( - mainField = @Field(type = String, index = analyzed), - otherFields = { - @NestedField(dotSuffix = "untouched", type = String, store = true, index = not_analyzed), - @NestedField(dotSuffix = "sort", type = String, store = true, indexAnalyzer = "keyword") - } - ) - private List authors = new ArrayList(); - - @Field(type = Integer, store = true) - private List publishedYears = new ArrayList(); - - private int score; - - private ArticleEntity() { - - } - - public ArticleEntity(String id) { - this.id = id; - } - - public void setId(String id) { - this.id = id; - } - - public String getId() { - return id; - } - - public String getTitle() { - return title; - } - - public void setTitle(String title) { - this.title = title; - } - - public List getAuthors() { - return authors; - } - - public void setAuthors(List authors) { - this.authors = authors; - } - - public List getPublishedYears() { - return publishedYears; - } - - public void setPublishedYears(List publishedYears) { - this.publishedYears = publishedYears; - } - - public int getScore() { - return score; - } - - public void setScore(int score) { - this.score = score; - } -} +/* + * Copyright 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.facet; + +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.String; + +import java.util.ArrayList; +import java.util.List; + +import org.springframework.data.annotation.Id; +import org.springframework.data.elasticsearch.annotations.Document; +import org.springframework.data.elasticsearch.annotations.Field; +import org.springframework.data.elasticsearch.annotations.MultiField; +import org.springframework.data.elasticsearch.annotations.NestedField; + +/** + * Simple type to test facets + * + * @author Artur Konczak + * @author Mohsin Husen + */ +@Document(indexName = "articles", type = "article", shards = 1, replicas = 0, refreshInterval = "-1", indexStoreType = "memory") +public class ArticleEntity { + + @Id + private String id; + private String title; + private String subject; + + @MultiField( + mainField = @Field(type = String, index = analyzed), + otherFields = { + @NestedField(dotSuffix = "untouched", type = String, store = true, index = not_analyzed), + @NestedField(dotSuffix = "sort", type = String, store = true, indexAnalyzer = "keyword") + } + ) + private List authors = new ArrayList(); + + @Field(type = Integer, store = true) + private List publishedYears = new ArrayList(); + + private int score; + + private ArticleEntity() { + + } + + public ArticleEntity(String id) { + this.id = id; + } + + public void setId(String id) { + this.id = id; + } + + public String getId() { + return id; + } + + public String getTitle() { + return title; + } + + public void setTitle(String title) { + this.title = title; + } + + public String getSubject() { + return subject; + } + + public void setSubject(String subject) { + this.subject = subject; + } + + public List getAuthors() { + return authors; + } + + public void setAuthors(List authors) { + this.authors = authors; + } + + public List getPublishedYears() { + return publishedYears; + } + + public void setPublishedYears(List publishedYears) { + this.publishedYears = publishedYears; + } + + public int getScore() { + return score; + } + + public void setScore(int score) { + this.score = score; + } +} diff --git a/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntityBuilder.java b/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntityBuilder.java index 8b25fd008..caf49bae7 100644 --- a/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntityBuilder.java +++ b/src/test/java/org/springframework/data/elasticsearch/core/facet/ArticleEntityBuilder.java @@ -1,64 +1,68 @@ -/* - * Copyright 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.facet; - -import org.springframework.data.elasticsearch.core.query.IndexQuery; - -/** - * Simple type to test facets - * - * @author Artur Konczak - * @author Mohsin Husen - */ -public class ArticleEntityBuilder { - - private ArticleEntity result; - - public ArticleEntityBuilder(String id) { - result = new ArticleEntity(id); - } - - public ArticleEntityBuilder title(String title) { - result.setTitle(title); - return this; - } - - public ArticleEntityBuilder addAuthor(String author) { - result.getAuthors().add(author); - return this; - } - - public ArticleEntityBuilder addPublishedYear(Integer year) { - result.getPublishedYears().add(year); - return this; - } - - public ArticleEntityBuilder score(int score) { - result.setScore(score); - return this; - } - - public ArticleEntity build() { - return result; - } - - public IndexQuery buildIndex() { - IndexQuery indexQuery = new IndexQuery(); - indexQuery.setId(result.getId()); - indexQuery.setObject(result); - return indexQuery; - } -} +/* + * Copyright 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.facet; + +import org.springframework.data.elasticsearch.core.query.IndexQuery; + +/** + * Simple type to test facets + * + * @author Artur Konczak + * @author Mohsin Husen + */ +public class ArticleEntityBuilder { + + private ArticleEntity result; + + public ArticleEntityBuilder(String id) { + result = new ArticleEntity(id); + } + + public ArticleEntityBuilder title(String title) { + result.setTitle(title); + return this; + } + public ArticleEntityBuilder subject(String subject) { + result.setSubject(subject); + return this; + } + + public ArticleEntityBuilder addAuthor(String author) { + result.getAuthors().add(author); + return this; + } + + public ArticleEntityBuilder addPublishedYear(Integer year) { + result.getPublishedYears().add(year); + return this; + } + + public ArticleEntityBuilder score(int score) { + result.setScore(score); + return this; + } + + public ArticleEntity build() { + return result; + } + + public IndexQuery buildIndex() { + IndexQuery indexQuery = new IndexQuery(); + indexQuery.setId(result.getId()); + indexQuery.setObject(result); + return indexQuery; + } +}