Cwikius-Spring/spring-data-elasticsearch
1
0
Fork 0
mirror of https://github.com/spring-projects/spring-data-elasticsearch.git synced 2025-06-29 23:32:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

DATAES-605 - Make batch size for streamQuery configurable.

Browse Source 
Original PR: #292
This commit is contained in:
rasmusfaber 2019-07-10 22:42:14 +02:00 committed by Peter-Josef Meisch
parent 7fd9986c6e
commit 216fc89fd2
2 changed files with 52 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java
View File

@ -16,7 +16,6 @@
package org.springframework.data.elasticsearch.repository.query; package org.springframework.data.elasticsearch.repository.query;
import org.springframework.data.domain.PageRequest; import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.elasticsearch.core.ElasticsearchOperations; import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty; import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
import org.springframework.data.elasticsearch.core.query.CriteriaQuery; import org.springframework.data.elasticsearch.core.query.CriteriaQuery;
@ -24,9 +23,9 @@ import org.springframework.data.elasticsearch.repository.query.parser.Elasticsea
import org.springframework.data.mapping.context.MappingContext; import org.springframework.data.mapping.context.MappingContext;
import org.springframework.data.repository.query.ParametersParameterAccessor; import org.springframework.data.repository.query.ParametersParameterAccessor;
import org.springframework.data.repository.query.parser.PartTree; import org.springframework.data.repository.query.parser.PartTree;
import org.springframework.util.ClassUtils;
import org.springframework.data.util.CloseableIterator; import org.springframework.data.util.CloseableIterator;
import org.springframework.data.util.StreamUtils; import org.springframework.data.util.StreamUtils;
import org.springframework.util.ClassUtils;
/** /**
* ElasticsearchPartQuery * ElasticsearchPartQuery
@ -35,9 +34,10 @@ import org.springframework.data.util.StreamUtils;
* @author Mohsin Husen * @author Mohsin Husen
* @author Kevin Leturc * @author Kevin Leturc
* @author Mark Paluch * @author Mark Paluch
* @author Rasmus Faber-Espensen
*/ */
public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery { public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery {
private static final int DEFAULT_STREAM_BATCH_SIZE = 500;
private final PartTree tree; private final PartTree tree;
private final MappingContext<?, ElasticsearchPersistentProperty> mappingContext; private final MappingContext<?, ElasticsearchPersistentProperty> mappingContext;
@ -51,7 +51,7 @@ public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery
public Object execute(Object[] parameters) { public Object execute(Object[] parameters) {
ParametersParameterAccessor accessor = new ParametersParameterAccessor(queryMethod.getParameters(), parameters); ParametersParameterAccessor accessor = new ParametersParameterAccessor(queryMethod.getParameters(), parameters);
CriteriaQuery query = createQuery(accessor); CriteriaQuery query = createQuery(accessor);
if(tree.isDelete()) { if (tree.isDelete()) {
Object result = countOrGetDocumentsForDelete(query, accessor); Object result = countOrGetDocumentsForDelete(query, accessor);
elasticsearchOperations.delete(query, queryMethod.getEntityInformation().getJavaType()); elasticsearchOperations.delete(query, queryMethod.getEntityInformation().getJavaType());
return result; return result;
@ -60,20 +60,22 @@ public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery
return elasticsearchOperations.queryForPage(query, queryMethod.getEntityInformation().getJavaType()); return elasticsearchOperations.queryForPage(query, queryMethod.getEntityInformation().getJavaType());
} else if (queryMethod.isStreamQuery()) { } else if (queryMethod.isStreamQuery()) {
Class<?> entityType = queryMethod.getEntityInformation().getJavaType(); Class<?> entityType = queryMethod.getEntityInformation().getJavaType();
if (query.getPageable().isUnpaged()) { if (accessor.getPageable().isUnpaged()) {
int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType()); query.setPageable(PageRequest.of(0, DEFAULT_STREAM_BATCH_SIZE));
query.setPageable(PageRequest.of(0, Math.max(1, itemCount))); } else {
query.setPageable(accessor.getPageable());
} }
return StreamUtils.createStreamFromIterator((CloseableIterator<Object>) elasticsearchOperations.stream(query, entityType)); return StreamUtils
.createStreamFromIterator((CloseableIterator<Object>) elasticsearchOperations.stream(query, entityType));
} else if (queryMethod.isCollectionQuery()) { } else if (queryMethod.isCollectionQuery()) {
if (accessor.getPageable().isUnpaged()) { if (accessor.getPageable().isUnpaged()) {
int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType()); int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());
query.setPageable(PageRequest.of(0, Math.max(1, itemCount))); query.setPageable(PageRequest.of(0, Math.max(1, itemCount)));
} else { } else {
query.setPageable(accessor.getPageable()); query.setPageable(accessor.getPageable());
} }
return elasticsearchOperations.queryForList(query, queryMethod.getEntityInformation().getJavaType()); return elasticsearchOperations.queryForList(query, queryMethod.getEntityInformation().getJavaType());
} else if (tree.isCountProjection()) { } else if (tree.isCountProjection()) {
return elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType()); return elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());

45
src/test/java/org/springframework/data/elasticsearch/repositories/custommethod/CustomMethodRepositoryBaseTests.java
View File

@ -32,7 +32,6 @@ import java.util.List;
import java.util.stream.Stream; import java.util.stream.Stream;
import org.junit.Test; import org.junit.Test;
import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.annotation.Id; import org.springframework.data.annotation.Id;
import org.springframework.data.annotation.Version; import org.springframework.data.annotation.Version;
@ -65,6 +64,8 @@ public abstract class CustomMethodRepositoryBaseTests {
@Autowired private SampleCustomMethodRepository repository; @Autowired private SampleCustomMethodRepository repository;
@Autowired private SampleStreamingCustomMethodRepository streamingRepository;
@Test @Test
public void shouldExecuteCustomMethod() { public void shouldExecuteCustomMethod() {
@ -711,7 +712,7 @@ public abstract class CustomMethodRepositoryBaseTests {
repository.saveAll(entities); repository.saveAll(entities);
// when // when
Stream<SampleEntity> stream = repository.findByType("abc"); Stream<SampleEntity> stream = streamingRepository.findByType("abc");
// then // then
assertThat(stream).isNotNull(); assertThat(stream).isNotNull();
@ -1236,13 +1237,41 @@ public abstract class CustomMethodRepositoryBaseTests {
assertThat(count).isEqualTo(1L); assertThat(count).isEqualTo(1L);
} }
@Test // DATAES-605
public void streamMethodsShouldWorkWithLargeResultSets() {
// given
List<SampleEntity> entities = createSampleEntities("abc", 10001);
repository.saveAll(entities);
// when
Stream<SampleEntity> stream = streamingRepository.findByType("abc");
// then
assertThat(stream).isNotNull();
assertThat(stream.count()).isEqualTo(10001L);
}
@Test // DATAES-605
public void streamMethodsCanHandlePageable() {
// given
List<SampleEntity> entities = createSampleEntities("abc", 10);
repository.saveAll(entities);
// when
Stream<SampleEntity> stream = streamingRepository.findByType("abc", PageRequest.of(0, 2));
// then
assertThat(stream).isNotNull();
assertThat(stream.count()).isEqualTo(10L);
}
private List<SampleEntity> createSampleEntities(String type, int numberOfEntities) { private List<SampleEntity> createSampleEntities(String type, int numberOfEntities) {
List<SampleEntity> entities = new ArrayList<>(); List<SampleEntity> entities = new ArrayList<>();
for (int i = 0; i < numberOfEntities; i++) { for (int i = 0; i < numberOfEntities; i++) {
SampleEntity entity = new SampleEntity(); SampleEntity entity = new SampleEntity();
entity.setId(randomNumeric(numberOfEntities)); entity.setId(randomNumeric(32));
entity.setAvailable(true); entity.setAvailable(true);
entity.setMessage("Message"); entity.setMessage("Message");
entity.setType(type); entity.setType(type);
@ -1328,8 +1357,6 @@ public abstract class CustomMethodRepositoryBaseTests {
Page<SampleEntity> findByLocationNear(GeoPoint point, String distance, Pageable pageable); Page<SampleEntity> findByLocationNear(GeoPoint point, String distance, Pageable pageable);
Stream<SampleEntity> findByType(String type);
long countByType(String type); long countByType(String type);
long countByTypeNot(String type); long countByTypeNot(String type);
@ -1371,4 +1398,12 @@ public abstract class CustomMethodRepositoryBaseTests {
long countByLocationNear(GeoPoint point, String distance); long countByLocationNear(GeoPoint point, String distance);
} }
/**
* @author Rasmus Faber-Espensen
*/
public interface SampleStreamingCustomMethodRepository extends ElasticsearchRepository<SampleEntity, String> {
Stream<SampleEntity> findByType(String type);
Stream<SampleEntity> findByType(String type, Pageable pageable);
}
} }