Spring/spring-data-mongodb
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

DATAMONGO-1637 - Add support for aggregation result streaming.

Browse Source 
We now support aggregation result streaming backed by a MongoDB cursor. Result streaming fetches aggregation results in batches from MongoDB and converts results as they are retrieved through the iterator.

Aggregation aggregation = …

CloseableIterator<TagCount> results = mongoOperations.aggregateStream(aggregation, "inputCollection", TagCount.class);

List<TagCount> tagCount = new ArrayList<TagCount>();
while (results.hasNext()) {
	tagCount.add(results.next());
}

results.close();

Original pull request: #447.
This commit is contained in:
Mainder Singh
2017-03-03 19:45:15 +05:30
committed by Mark Paluch
parent f4f5e02e66
commit 1a65828365
3 changed files with 242 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoOperations.java
View File

@@ -56,6 +56,7 @@ import com.mongodb.client.result.UpdateResult;
* @author Chuong Ngo
* @author Christoph Strobl
* @author Thomas Darimont
* @author maninder
*/
public interface MongoOperations {
@@ -274,7 +275,7 @@ public interface MongoOperations {
/**
* Returns the {@link ScriptOperations} that can be performed on {@link com.mongodb.DB} level.
*
*
* @return
* @since 1.7
*/
@@ -379,6 +380,9 @@ public interface MongoOperations {
*/
<O> AggregationResults<O> aggregate(TypedAggregation<?> aggregation, String collectionName, Class<O> outputType);
<O> CloseableIterator<O> aggregateStream(TypedAggregation<?> aggregation, String inputCollectionName,
Class<O> outputType);
/**
* Execute an aggregation operation. The raw results will be mapped to the given entity class. The name of the
* inputCollection is derived from the inputType of the aggregation.
@@ -391,6 +395,19 @@ public interface MongoOperations {
*/
<O> AggregationResults<O> aggregate(TypedAggregation<?> aggregation, Class<O> outputType);
/**
* Execute an aggregation operation. The raw results will be mapped to the given entity class and are returned as
* stream. The name of the inputCollection is derived from the inputType of the aggregation.
*
* @param aggregation The {@link TypedAggregation} specification holding the aggregation operations, must not be
* {@literal null}.
* @param outputType The parameterized type of the returned list, must not be {@literal null}.
* @return The results of the aggregation operation.
* @since 1.11.0
*/
<O> CloseableIterator<O> aggregateStream(TypedAggregation<?> aggregation, Class<O> outputType);
/**
* Execute an aggregation operation. The raw results will be mapped to the given entity class.
*
@@ -404,6 +421,8 @@ public interface MongoOperations {
*/
<O> AggregationResults<O> aggregate(Aggregation aggregation, Class<?> inputType, Class<O> outputType);
<O> CloseableIterator<O> aggregateStream(Aggregation aggregation, Class<?> inputType, Class<O> outputType);
/**
* Execute an aggregation operation. The raw results will be mapped to the given entity class.
*
@@ -417,6 +436,9 @@ public interface MongoOperations {
*/
<O> AggregationResults<O> aggregate(Aggregation aggregation, String collectionName, Class<O> outputType);
<O> CloseableIterator<O> aggregateStream(Aggregation aggregation, String collectionName, Class<O> outputType);
/**
* Execute a map-reduce operation. The map-reduce operation will be formed with an output type of INLINE
*
@@ -668,7 +690,7 @@ public interface MongoOperations {
* @return
*/
<T> T findAndModify(Query query, Update update, FindAndModifyOptions options, Class<T> entityClass,
String collectionName);
String collectionName);
/**
* Map the results of an ad-hoc query on the collection for the entity type to a single instance of an object of the

69
spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java
View File

@@ -15,6 +15,7 @@
*/
package org.springframework.data.mongodb.core;
import static org.springframework.data.mongodb.core.aggregation.AggregationOptions.*;
import static org.springframework.data.mongodb.core.query.Criteria.*;
import static org.springframework.data.mongodb.core.query.SerializationUtils.*;
@@ -125,6 +126,7 @@ import com.mongodb.client.model.UpdateOptions;
import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;
import com.mongodb.util.JSONParseException;
import com.mongodb.AggregationOptions;
/**
* Primary implementation of {@link MongoOperations}.
@@ -144,6 +146,7 @@ import com.mongodb.util.JSONParseException;
* @author Niko Schmuck
* @author Mark Paluch
* @author Laszlo Csontos
* @author maninder
*/
@SuppressWarnings("deprecation")
public class MongoTemplate implements MongoOperations, ApplicationContextAware, IndexOperationsProvider {
@@ -1529,6 +1532,11 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
return aggregate(aggregation, determineCollectionName(aggregation.getInputType()), outputType);
}
@Override
public <O> CloseableIterator<O> aggregateStream(TypedAggregation<?> aggregation, Class<O> outputType) {
return aggregateStream(aggregation, determineCollectionName(aggregation.getInputType()), outputType);
}
@Override
public <O> AggregationResults<O> aggregate(TypedAggregation<?> aggregation, String inputCollectionName,
Class<O> outputType) {
@@ -1540,6 +1548,17 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
return aggregate(aggregation, inputCollectionName, outputType, context);
}
@Override
public <O> CloseableIterator<O> aggregateStream(TypedAggregation<?> aggregation, String inputCollectionName,
Class<O> outputType) {
Assert.notNull(aggregation, "Aggregation pipeline must not be null!");
AggregationOperationContext context = new TypeBasedAggregationOperationContext(aggregation.getInputType(),
mappingContext, queryMapper);
return aggregateStream(aggregation, inputCollectionName, outputType, context);
}
@Override
public <O> AggregationResults<O> aggregate(Aggregation aggregation, Class<?> inputType, Class<O> outputType) {
@@ -1547,11 +1566,23 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper));
}
@Override
public <O> CloseableIterator<O> aggregateStream(Aggregation aggregation, Class<?> inputType, Class<O> outputType) {
return aggregateStream(aggregation, determineCollectionName(inputType), outputType,
new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper));
}
@Override
public <O> AggregationResults<O> aggregate(Aggregation aggregation, String collectionName, Class<O> outputType) {
return aggregate(aggregation, collectionName, outputType, null);
}
@Override
public <O> CloseableIterator<O> aggregateStream(Aggregation aggregation, String collectionName, Class<O> outputType) {
return aggregateStream(aggregation, collectionName, outputType, null);
}
/*
* (non-Javadoc)
* @see org.springframework.data.mongodb.core.MongoOperations#findAllAndRemove(org.springframework.data.mongodb.core.query.Query, java.lang.String)
@@ -1646,6 +1677,44 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
return mappedResults;
}
protected <O> CloseableIterator<O> aggregateStream(final Aggregation aggregation, final String collectionName,
final Class<O> outputType, AggregationOperationContext context) {
Assert.hasText(collectionName, "Collection name must not be null or empty!");
Assert.notNull(aggregation, "Aggregation pipeline must not be null!");
Assert.notNull(outputType, "Output type must not be null!");
AggregationOperationContext rootContext = context == null ? Aggregation.DEFAULT_CONTEXT : context;
final DBObject command = aggregation.toDbObject(collectionName, rootContext);
Assert.isNull(command.get(CURSOR), "Custom options not allowed while streaming");
Assert.isNull(command.get(EXPLAIN), "Explain option can't be used while streaming");
return execute(collectionName, new CollectionCallback<CloseableIterator<O>>() {
@Override
public CloseableIterator<O> doInCollection(DBCollection collection) throws MongoException, DataAccessException {
List<DBObject> pipeline = (List<DBObject>) command.get("pipeline");
Cursor cursor = collection.aggregate(pipeline, getNativeAggregationOptionsFromCommand(command));
ReadDbObjectCallback<O> readCallback = new ReadDbObjectCallback<O>(mongoConverter, outputType, collectionName);
return new CloseableIterableCursorAdapter<O>(cursor, exceptionTranslator, readCallback);
}
private AggregationOptions getNativeAggregationOptionsFromCommand(DBObject command) {
AggregationOptions.Builder builder = AggregationOptions.builder();
Object allowDiskUse = command.get(ALLOW_DISK_USE);
if (allowDiskUse != null && String.valueOf(allowDiskUse).equals("true")) {
builder.allowDiskUse(true);
}
return builder.build();
}
});
}
protected String replaceWithResourceIfNecessary(String function) {
String func = function;

150
spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java
View File

@@ -66,6 +66,7 @@ import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.NearQuery;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.repository.Person;
import org.springframework.data.util.CloseableIterator;
import org.springframework.data.util.Version;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
@@ -82,6 +83,7 @@ import com.mongodb.client.MongoCollection;
* @author Christoph Strobl
* @author Mark Paluch
* @author Nikolay Bogdanov
* @author maninder
*/
@RunWith(SpringJUnit4ClassRunner.class)
@ContextConfiguration("classpath:infrastructure.xml")
@@ -145,7 +147,7 @@ public class AggregationTests {
/**
* Imports the sample dataset (zips.json) if necessary (e.g. if it doesn't exist yet). The dataset can originally be
* found on the mongodb aggregation framework example website:
*
*
* @see <a href="https://docs.mongodb.org/manual/tutorial/aggregation-examples/">MongoDB Aggregation Examples</a>
*/
private void initSampleDataIfNecessary() {
@@ -229,6 +231,37 @@ public class AggregationTests {
assertTagCount("nosql", 1, tagCount.get(2));
}
@Test // DATAMONGO-1637
public void shouldAggregateAndStream() {
createTagDocuments();
Aggregation agg = newAggregation( //
project("tags"), //
unwind("tags"), //
group("tags") //
.count().as("n"), //
project("n") //
.and("tag").previousOperation(), //
sort(DESC, "n") //
);
CloseableIterator<TagCount> iterator = mongoTemplate.aggregateStream(agg, INPUT_COLLECTION, TagCount.class);
assertThat(iterator, is(notNullValue()));
List<TagCount> tagCount = new ArrayList<TagCount>();
while (iterator.hasNext()) {
tagCount.add(iterator.next());
}
assertThat(tagCount, is(notNullValue()));
assertThat(tagCount.size(), is(3));
assertTagCount("spring", 3, tagCount.get(0));
assertTagCount("mongodb", 2, tagCount.get(1));
assertTagCount("nosql", 1, tagCount.get(2));
}
@Test // DATAMONGO-586
public void shouldAggregateEmptyCollection() {
@@ -252,6 +285,32 @@ public class AggregationTests {
assertThat(tagCount.size(), is(0));
}
@Test // DATAMONGO-1637
public void shouldAggregateEmptyCollectionAndStream() {
Aggregation aggregation = newAggregation(//
project("tags"), //
unwind("tags"), //
group("tags") //
.count().as("n"), //
project("n") //
.and("tag").previousOperation(), //
sort(DESC, "n") //
);
CloseableIterator<TagCount> results = mongoTemplate.aggregateStream(aggregation, INPUT_COLLECTION, TagCount.class);
assertThat(results, is(notNullValue()));
List<TagCount> tagCount = new ArrayList<TagCount>();
while (results.hasNext()) {
tagCount.add(results.next());
}
// assertThat(tagCount, is(notNullValue()));
assertThat(tagCount.size(), is(0));
}
@Test // DATAMONGO-1391
public void shouldUnwindWithIndex() {
@@ -333,6 +392,33 @@ public class AggregationTests {
assertTagCount(null, 0, tagCount.get(1));
}
@Test // DATAMONGO-1637
public void shouldDetectResultMismatchWhileStreaming() {
createTagDocuments();
Aggregation aggregation = newAggregation( //
project("tags"), //
unwind("tags"), //
group("tags") //
.count().as("count"), // count field not present
limit(2) //
);
CloseableIterator<TagCount> results = mongoTemplate.aggregateStream(aggregation, INPUT_COLLECTION, TagCount.class);
assertThat(results, is(notNullValue()));
List<TagCount> tagCount = new ArrayList<TagCount>();
while (results.hasNext()) {
tagCount.add(results.next());
}
// assertThat(tagCount, is(notNullValue()));
assertThat(tagCount.size(), is(2));
assertTagCount(null, 0, tagCount.get(0));
assertTagCount(null, 0, tagCount.get(1));
}
@Test // DATAMONGO-586
public void complexAggregationFrameworkUsageLargestAndSmallestCitiesByState() {
/*
@@ -1189,6 +1275,35 @@ public class AggregationTests {
assertLikeStats(result.getMappedResults().get(4), "e", 3);
}
@Test // DATAMONGO-1637
public void returnFiveMostCommonLikesAggregationFrameworkExampleWithSortOnDiskOptionEnabledWhileStreaming() {
assumeTrue(mongoVersion.isGreaterThanOrEqualTo(TWO_DOT_SIX));
createUserWithLikesDocuments();
TypedAggregation<UserWithLikes> agg = createUsersWithCommonLikesAggregation() //
.withOptions(newAggregationOptions().allowDiskUse(true).build());
assertThat(agg, is(notNullValue()));
assertThat(agg.toString(), is(notNullValue()));
CloseableIterator<LikeStats> iterator = mongoTemplate.aggregateStream(agg, LikeStats.class);
List<LikeStats> result = new ArrayList<LikeStats>();
while (iterator.hasNext()) {
result.add(iterator.next());
}
assertThat(result, is(notNullValue()));
assertThat(result, is(notNullValue()));
assertThat(result.size(), is(5));
assertLikeStats(result.get(0), "a", 4);
assertLikeStats(result.get(1), "b", 2);
assertLikeStats(result.get(2), "c", 4);
assertLikeStats(result.get(3), "d", 2);
assertLikeStats(result.get(4), "e", 3);
}
@Test // DATAMONGO-960
public void returnFiveMostCommonLikesShouldReturnStageExecutionInformationWithExplainOptionEnabled() {
@@ -1443,6 +1558,34 @@ public class AggregationTests {
mongoTemplate.dropCollection(tempOutCollection);
}
@Test // DATAMONGO-1637
public void shouldCreateOutputCollectionWhileStreaming() {
assumeTrue(mongoVersion.isGreaterThanOrEqualTo(TWO_DOT_SIX));
mongoTemplate.save(new Person("Anna", "Ivanova", 21, Person.Sex.FEMALE));
mongoTemplate.save(new Person("Pavel", "Sidorov", 36, Person.Sex.MALE));
mongoTemplate.save(new Person("Anastasia", "Volochkova", 29, Person.Sex.FEMALE));
mongoTemplate.save(new Person("Igor", "Stepanov", 31, Person.Sex.MALE));
mongoTemplate.save(new Person("Leoniv", "Yakubov", 55, Person.Sex.MALE));
String tempOutCollection = "personQueryTemp";
TypedAggregation<Person> agg = newAggregation(Person.class, //
group("sex").count().as("count"), //
sort(DESC, "count"), //
out(tempOutCollection));
CloseableIterator<DBObject> iterator = mongoTemplate.aggregateStream(agg, DBObject.class);
List<DBObject> list = mongoTemplate.findAll(DBObject.class, tempOutCollection);
assertThat(list, hasSize(2));
assertThat(list.get(0), isBsonObject().containing("_id", "MALE").containing("count", 3));
assertThat(list.get(1), isBsonObject().containing("_id", "FEMALE").containing("count", 2));
mongoTemplate.dropCollection(tempOutCollection);
}
@Test(expected = IllegalArgumentException.class) // DATAMONGO-1418
public void outShouldOutBeTheLastOperation() {
@@ -1745,6 +1888,7 @@ public class AggregationTests {
}
static class DATAMONGO753 {
PD[] pd;
DATAMONGO753 withPDs(PD... pds) {
@@ -1754,6 +1898,7 @@ public class AggregationTests {
}
static class PD {
String pDch;
@org.springframework.data.mongodb.core.mapping.Field("alias") int up;
@@ -1830,6 +1975,7 @@ public class AggregationTests {
@SuppressWarnings("unused")
static class Descriptors {
private CarDescriptor carDescriptor;
}
@@ -1846,6 +1992,7 @@ public class AggregationTests {
@SuppressWarnings("unused")
static class Entry {
private String make;
private String model;
private int year;
@@ -1964,3 +2111,4 @@ public class AggregationTests {
double price;
}
}