Skip to content

Commit e264bd7

Browse files
zzt93akonczak
zzt93
authored and
akonczak
committed
DATAES-487 - Support for multi search API.
1 parent 5ff6238 commit e264bd7

File tree

4 files changed

+262
-8
lines changed

4 files changed

+262
-8
lines changed

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

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
package org.springframework.data.elasticsearch.core;
1717

1818
import org.elasticsearch.action.update.UpdateResponse;
19-
import org.elasticsearch.client.Client;
2019
import org.elasticsearch.cluster.metadata.AliasMetaData;
2120
import org.elasticsearch.common.Nullable;
2221
import org.springframework.data.domain.Page;
@@ -28,13 +27,15 @@
2827
import java.util.LinkedList;
2928
import java.util.List;
3029
import java.util.Map;
30+
import java.util.stream.Collectors;
3131

3232
/**
3333
* ElasticsearchOperations
3434
*
3535
* @author Rizwan Idrees
3636
* @author Mohsin Husen
3737
* @author Kevin Leturc
38+
* @author Zetang Zeng
3839
*/
3940
public interface ElasticsearchOperations {
4041

@@ -186,6 +187,42 @@ public interface ElasticsearchOperations {
186187
*/
187188
<T> Page<T> queryForPage(SearchQuery query, Class<T> clazz, SearchResultMapper mapper);
188189

190+
/**
191+
* Execute the multi-search against elasticsearch and return result as {@link List} of {@link Page}
192+
*
193+
* @param queries
194+
* @param clazz
195+
* @return
196+
*/
197+
<T> List<Page<T>> queryForPage(List<SearchQuery> queries, Class<T> clazz);
198+
199+
/**
200+
* Execute the multi-search against elasticsearch and return result as {@link List} of {@link Page} using custom mapper
201+
*
202+
* @param queries
203+
* @param clazz
204+
* @return
205+
*/
206+
<T> List<Page<T>> queryForPage(List<SearchQuery> queries, Class<T> clazz, SearchResultMapper mapper);
207+
208+
/**
209+
* Execute the multi-search against elasticsearch and return result as {@link List} of {@link Page}
210+
*
211+
* @param queries
212+
* @param classes
213+
* @return
214+
*/
215+
List<Page<?>> queryForPage(List<SearchQuery> queries, List<Class<?>> classes);
216+
217+
/**
218+
* Execute the multi-search against elasticsearch and return result as {@link List} of {@link Page} using custom mapper
219+
*
220+
* @param queries
221+
* @param classes
222+
* @return
223+
*/
224+
List<Page<?>> queryForPage(List<SearchQuery> queries, List<Class<?>> classes, SearchResultMapper mapper);
225+
189226
/**
190227
* Execute the query against elasticsearch and return result as {@link Page}
191228
*
@@ -283,6 +320,29 @@ public interface ElasticsearchOperations {
283320
*/
284321
<T> List<T> queryForList(SearchQuery query, Class<T> clazz);
285322

323+
/**
324+
* Execute the multi search query against elasticsearch and return result as {@link List}
325+
*
326+
* @param queries
327+
* @param clazz
328+
* @param <T>
329+
* @return
330+
*/
331+
default <T> List<List<T>> queryForList(List<SearchQuery> queries, Class<T> clazz) {
332+
return queryForPage(queries, clazz).stream().map(Page::getContent).collect(Collectors.toList());
333+
}
334+
335+
/**
336+
* Execute the multi search query against elasticsearch and return result as {@link List}
337+
*
338+
* @param queries
339+
* @param classes
340+
* @return
341+
*/
342+
default List<List<?>> queryForList(List<SearchQuery> queries, List<Class<?>> classes) {
343+
return queryForPage(queries, classes).stream().map(Page::getContent).collect(Collectors.toList());
344+
}
345+
286346
/**
287347
* Execute the query against elasticsearch and return ids
288348
*

‎src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java

Lines changed: 76 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@
4747
import org.elasticsearch.action.index.IndexRequest;
4848
import org.elasticsearch.action.search.ClearScrollRequest;
4949
import org.elasticsearch.action.search.ClearScrollResponse;
50+
import org.elasticsearch.action.search.MultiSearchRequest;
51+
import org.elasticsearch.action.search.MultiSearchResponse;
5052
import org.elasticsearch.action.search.SearchRequest;
5153
import org.elasticsearch.action.search.SearchResponse;
5254
import org.elasticsearch.action.search.SearchScrollRequest;
@@ -127,6 +129,7 @@
127129
* @author Sascha Woo
128130
* @author Ted Liang
129131
* @author Don Wellington
132+
* @author Zetang Zeng
130133
*/
131134
public class ElasticsearchRestTemplate
132135
implements ElasticsearchOperations, EsClient<RestHighLevelClient>, ApplicationContextAware {
@@ -335,6 +338,68 @@ public <T> AggregatedPage<T> queryForPage(SearchQuery query, Class<T> clazz, Sea
335338
return mapper.mapResults(response, clazz, query.getPageable());
336339
}
337340

341+
@Override
342+
public <T> List<Page<T>> queryForPage(List<SearchQuery> queries, Class<T> clazz) {
343+
return queryForPage(queries, clazz, resultsMapper);
344+
}
345+
346+
private <T> List<Page<T>> doMultiSearch(List<SearchQuery> queries, Class<T> clazz, MultiSearchRequest request, SearchResultMapper resultsMapper) {
347+
MultiSearchResponse.Item[] items = getMultiSearchResult(request);
348+
List<Page<T>> res = new ArrayList<>(queries.size());
349+
int c = 0;
350+
for (SearchQuery query : queries) {
351+
res.add(resultsMapper.mapResults(items[c++].getResponse(), clazz, query.getPageable()));
352+
}
353+
return res;
354+
}
355+
356+
private List<Page<?>> doMultiSearch(List<SearchQuery> queries, List<Class<?>> classes, MultiSearchRequest request, SearchResultMapper resultsMapper) {
357+
MultiSearchResponse.Item[] items = getMultiSearchResult(request);
358+
List<Page<?>> res = new ArrayList<>(queries.size());
359+
int c = 0;
360+
Iterator<Class<?>> it = classes.iterator();
361+
for (SearchQuery query : queries) {
362+
res.add(resultsMapper.mapResults(items[c++].getResponse(), it.next(), query.getPageable()));
363+
}
364+
return res;
365+
}
366+
367+
private MultiSearchResponse.Item[] getMultiSearchResult(MultiSearchRequest request) {
368+
MultiSearchResponse response;
369+
try {
370+
response = client.multiSearch(request);
371+
} catch (IOException e) {
372+
throw new ElasticsearchException("Error for search request: " + request.toString(), e);
373+
}
374+
MultiSearchResponse.Item[] items = response.getResponses();
375+
Assert.isTrue(items.length == request.requests().size(), "Response should has same length with queries");
376+
return items;
377+
}
378+
379+
@Override
380+
public <T> List<Page<T>> queryForPage(List<SearchQuery> queries, Class<T> clazz, SearchResultMapper mapper) {
381+
MultiSearchRequest request = new MultiSearchRequest();
382+
for (SearchQuery query : queries) {
383+
request.add(prepareSearch(prepareSearch(query, clazz), query));
384+
}
385+
return doMultiSearch(queries, clazz, request, mapper);
386+
}
387+
388+
@Override
389+
public List<Page<?>> queryForPage(List<SearchQuery> queries, List<Class<?>> classes) {
390+
return queryForPage(queries, classes, resultsMapper);
391+
}
392+
393+
@Override
394+
public List<Page<?>> queryForPage(List<SearchQuery> queries, List<Class<?>> classes, SearchResultMapper mapper) {
395+
MultiSearchRequest request = new MultiSearchRequest();
396+
Iterator<Class<?>> it = classes.iterator();
397+
for (SearchQuery query : queries) {
398+
request.add(prepareSearch(prepareSearch(query, it.next()), query));
399+
}
400+
return doMultiSearch(queries, classes, request, mapper);
401+
}
402+
338403
@Override
339404
public <T> T query(SearchQuery query, ResultsExtractor<T> resultsExtractor) {
340405
SearchResponse response = doSearch(prepareSearch(query, Optional.ofNullable(query.getQuery())), query);
@@ -1026,6 +1091,16 @@ public <T> Page<T> moreLikeThis(MoreLikeThisQuery query, Class<T> clazz) {
10261091
}
10271092

10281093
private SearchResponse doSearch(SearchRequest searchRequest, SearchQuery searchQuery) {
1094+
prepareSearch(searchRequest, searchQuery);
1095+
1096+
try {
1097+
return client.search(searchRequest);
1098+
} catch (IOException e) {
1099+
throw new ElasticsearchException("Error for search request with scroll: " + searchRequest.toString(), e);
1100+
}
1101+
}
1102+
1103+
private SearchRequest prepareSearch(SearchRequest searchRequest, SearchQuery searchQuery) {
10291104
if (searchQuery.getFilter() != null) {
10301105
searchRequest.source().postFilter(searchQuery.getFilter());
10311106
}
@@ -1074,12 +1149,7 @@ private SearchResponse doSearch(SearchRequest searchRequest, SearchQuery searchQ
10741149
searchRequest.source().aggregation(aggregatedFacet.getFacet());
10751150
}
10761151
}
1077-
1078-
try {
1079-
return client.search(searchRequest);
1080-
} catch (IOException e) {
1081-
throw new ElasticsearchException("Error for search request with scroll: " + searchRequest.toString(), e);
1082-
}
1152+
return searchRequest;
10831153
}
10841154

10851155
private SearchResponse getSearchResponse(ActionFuture<SearchResponse> response) {

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

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@
5050
import org.elasticsearch.action.get.MultiGetRequestBuilder;
5151
import org.elasticsearch.action.get.MultiGetResponse;
5252
import org.elasticsearch.action.index.IndexRequestBuilder;
53+
import org.elasticsearch.action.search.MultiSearchRequest;
54+
import org.elasticsearch.action.search.MultiSearchResponse;
5355
import org.elasticsearch.action.search.SearchRequestBuilder;
5456
import org.elasticsearch.action.search.SearchResponse;
5557
import org.elasticsearch.action.update.UpdateRequestBuilder;
@@ -134,6 +136,7 @@
134136
* @author Sascha Woo
135137
* @author Ted Liang
136138
* @author Jean-Baptiste Nizet
139+
* @author Zetang Zeng
137140
*/
138141
public class ElasticsearchTemplate implements ElasticsearchOperations, EsClient<Client>, ApplicationContextAware {
139142

@@ -314,6 +317,65 @@ public <T> AggregatedPage<T> queryForPage(SearchQuery query, Class<T> clazz, Sea
314317
return mapper.mapResults(response, clazz, query.getPageable());
315318
}
316319

320+
@Override
321+
public <T> List<Page<T>> queryForPage(List<SearchQuery> queries, Class<T> clazz) {
322+
return queryForPage(queries, clazz, resultsMapper);
323+
}
324+
325+
@Override
326+
public <T> List<Page<T>> queryForPage(List<SearchQuery> queries, Class<T> clazz, SearchResultMapper mapper) {
327+
MultiSearchRequest request = new MultiSearchRequest();
328+
for (SearchQuery query : queries) {
329+
request.add(prepareSearch(prepareSearch(query, clazz), query));
330+
}
331+
return doMultiSearch(queries, clazz, request, mapper);
332+
}
333+
334+
private <T> List<Page<T>> doMultiSearch(List<SearchQuery> queries, Class<T> clazz, MultiSearchRequest request, SearchResultMapper resultsMapper) {
335+
MultiSearchResponse.Item[] items = getMultiSearchResult(request);
336+
List<Page<T>> res = new ArrayList<>(queries.size());
337+
int c = 0;
338+
for (SearchQuery query : queries) {
339+
res.add(resultsMapper.mapResults(items[c++].getResponse(), clazz, query.getPageable()));
340+
}
341+
return res;
342+
}
343+
344+
private List<Page<?>> doMultiSearch(List<SearchQuery> queries, List<Class<?>> classes, MultiSearchRequest request, SearchResultMapper resultsMapper) {
345+
MultiSearchResponse.Item[] items = getMultiSearchResult(request);
346+
List<Page<?>> res = new ArrayList<>(queries.size());
347+
int c = 0;
348+
Iterator<Class<?>> it = classes.iterator();
349+
for (SearchQuery query : queries) {
350+
res.add(resultsMapper.mapResults(items[c++].getResponse(), it.next(), query.getPageable()));
351+
}
352+
return res;
353+
}
354+
355+
private MultiSearchResponse.Item[] getMultiSearchResult(MultiSearchRequest request) {
356+
ActionFuture<MultiSearchResponse> future = client.multiSearch(request);
357+
MultiSearchResponse response = future.actionGet();
358+
MultiSearchResponse.Item[] items = response.getResponses();
359+
Assert.isTrue(items.length == request.requests().size(), "Response should have same length with queries");
360+
return items;
361+
}
362+
363+
@Override
364+
public List<Page<?>> queryForPage(List<SearchQuery> queries, List<Class<?>> classes) {
365+
return queryForPage(queries, classes, resultsMapper);
366+
}
367+
368+
@Override
369+
public List<Page<?>> queryForPage(List<SearchQuery> queries, List<Class<?>> classes, SearchResultMapper mapper) {
370+
Assert.isTrue(queries.size() == classes.size(), "Queries should have same length with classes");
371+
MultiSearchRequest request = new MultiSearchRequest();
372+
Iterator<Class<?>> it = classes.iterator();
373+
for (SearchQuery query : queries) {
374+
request.add(prepareSearch(prepareSearch(query, it.next()), query));
375+
}
376+
return doMultiSearch(queries, classes, request, mapper);
377+
}
378+
317379
@Override
318380
public <T> T query(SearchQuery query, ResultsExtractor<T> resultsExtractor) {
319381
SearchResponse response = doSearch(prepareSearch(query), query);
@@ -887,6 +949,11 @@ public <T> Page<T> moreLikeThis(MoreLikeThisQuery query, Class<T> clazz) {
887949
}
888950

889951
private SearchResponse doSearch(SearchRequestBuilder searchRequest, SearchQuery searchQuery) {
952+
SearchRequestBuilder requestBuilder = prepareSearch(searchRequest, searchQuery);
953+
return getSearchResponse(requestBuilder);
954+
}
955+
956+
private SearchRequestBuilder prepareSearch(SearchRequestBuilder searchRequest, SearchQuery searchQuery) {
890957
if (searchQuery.getFilter() != null) {
891958
searchRequest.setPostFilter(searchQuery.getFilter());
892959
}
@@ -935,7 +1002,7 @@ private SearchResponse doSearch(SearchRequestBuilder searchRequest, SearchQuery
9351002
searchRequest.addAggregation(aggregatedFacet.getFacet());
9361003
}
9371004
}
938-
return getSearchResponse(searchRequest.setQuery(searchQuery.getQuery()));
1005+
return searchRequest.setQuery(searchQuery.getQuery());
9391006
}
9401007

9411008
private SearchResponse getSearchResponse(SearchRequestBuilder requestBuilder) {

0 commit comments

Comments
 (0)